From cd442a23957b92e9ce7b0fbbce01d6733aaa6ca8 Mon Sep 17 00:00:00 2001
From: Jordan Demeulenaere <jdemeulenaere@google.com>
Date: Fri, 13 Jan 2023 15:26:11 +0100
Subject: [PATCH] Preserve Composable state when expanding it

This CL fixes the last known bug of Expandable.kt by making sure that we
preserve the content state of the Composable being expanded into an
Activity or Dialog, using movableContentOf. This should also help
performance wise, as the Compose runtime will now reuse the nodes that
are being moved at the beginning and end of the animation instead of
recomposing the content from scratch.

Bug: 230830644
Test: In the Gallery app, expand the Composable with a counter into a
      Dialog and notice that its counter value stays the same during and
      after the animation.
Change-Id: Ib6eed637e707df62e59518f8b36f635d32786f11
---
 .../android/compose/animation/Expandable.kt   | 44 ++++++++++---------
 1 file changed, 23 insertions(+), 21 deletions(-)

diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
index 259f0ed5c7a10..4e96ddabfda67 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
@@ -41,6 +41,7 @@ import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.State
 import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.movableContentOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCompositionContext
@@ -73,6 +74,7 @@ import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.ViewTreeLifecycleOwner
 import androidx.lifecycle.ViewTreeViewModelStoreOwner
+import com.android.compose.runtime.movableContentOf
 import com.android.systemui.animation.Expandable
 import com.android.systemui.animation.LaunchAnimator
 import kotlin.math.max
@@ -170,25 +172,25 @@ fun Expandable(
     val contentColor = controller.contentColor
     val shape = controller.shape
 
-    // TODO(b/230830644): Use movableContentOf to preserve the content state instead once the
-    // Compose libraries have been updated and include aosp/2163631.
     val wrappedContent =
-        @Composable { controller: ExpandableController ->
-            CompositionLocalProvider(
-                LocalContentColor provides contentColor,
-            ) {
-                // We make sure that the content itself (wrapped by the background) is at least
-                // 40.dp, which is the same as the M3 buttons. This applies even if onClick is
-                // null, to make it easier to write expandables that are sometimes clickable and
-                // sometimes not. There shouldn't be any Expandable smaller than 40dp because if
-                // the expandable is not clickable directly, then something in its content should
-                // be (and with a size >= 40dp).
-                val minSize = 40.dp
-                Box(
-                    Modifier.defaultMinSize(minWidth = minSize, minHeight = minSize),
-                    contentAlignment = Alignment.Center,
+        remember(content) {
+            movableContentOf { expandable: Expandable ->
+                CompositionLocalProvider(
+                    LocalContentColor provides contentColor,
                 ) {
-                    content(controller.expandable)
+                    // We make sure that the content itself (wrapped by the background) is at least
+                    // 40.dp, which is the same as the M3 buttons. This applies even if onClick is
+                    // null, to make it easier to write expandables that are sometimes clickable and
+                    // sometimes not. There shouldn't be any Expandable smaller than 40dp because if
+                    // the expandable is not clickable directly, then something in its content
+                    // should be (and with a size >= 40dp).
+                    val minSize = 40.dp
+                    Box(
+                        Modifier.defaultMinSize(minWidth = minSize, minHeight = minSize),
+                        contentAlignment = Alignment.Center,
+                    ) {
+                        content(expandable)
+                    }
                 }
             }
         }
@@ -270,7 +272,7 @@ fun Expandable(
                     .onGloballyPositioned {
                         controller.boundsInComposeViewRoot.value = it.boundsInRoot()
                     }
-            ) { wrappedContent(controller) }
+            ) { wrappedContent(controller.expandable) }
         }
         else -> {
             val clickModifier =
@@ -301,7 +303,7 @@ fun Expandable(
                         controller.boundsInComposeViewRoot.value = it.boundsInRoot()
                     },
             ) {
-                wrappedContent(controller)
+                wrappedContent(controller.expandable)
             }
         }
     }
@@ -315,7 +317,7 @@ private fun AnimatedContentInOverlay(
     animatorState: State<LaunchAnimator.State?>,
     overlay: ViewGroupOverlay,
     controller: ExpandableControllerImpl,
-    content: @Composable (ExpandableController) -> Unit,
+    content: @Composable (Expandable) -> Unit,
     composeViewRoot: View,
     onOverlayComposeViewChanged: (View?) -> Unit,
     density: Density,
@@ -370,7 +372,7 @@ private fun AnimatedContentInOverlay(
                             // We center the content in the expanding container.
                             contentAlignment = Alignment.Center,
                         ) {
-                            Box(contentModifier) { content(controller) }
+                            Box(contentModifier) { content(controller.expandable) }
                         }
                     }
                 }