Work around b/317972419

This CL works around b/317972419 by making the parameters of the movable
content of movable elements more explicit instead of relying on new
lambda creation.

Test: Manual. I added a test but it still fails, probably because of a
 bug in movableContentOf or in the testing libraries (see
 b/317972419#comment2 for details). I left the test in this CL but
 @Ignored it.
Bug: 317972419
Flag: N/A
Change-Id: Id54c237ea0b4f0e974a9cb9adcffb6426bf8e8c3
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
index 9402d84..964dca8 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
@@ -34,14 +34,19 @@
     modifier: Modifier,
     content: @Composable ElementScope<ElementContentScope>.() -> Unit,
 ) {
+    val contentScope = scene.scope
+    val elementScope =
+        remember(layoutImpl, key, scene, contentScope) {
+            ElementScopeImpl(layoutImpl, key, scene, contentScope)
+        }
+
     ElementBase(
         layoutImpl,
         scene,
         key,
         modifier,
+        elementScope,
         content,
-        scene.scope,
-        isMovable = false,
     )
 }
 
@@ -53,14 +58,19 @@
     modifier: Modifier,
     content: @Composable ElementScope<MovableElementContentScope>.() -> Unit,
 ) {
+    val contentScope = scene.scope
+    val elementScope =
+        remember(layoutImpl, key, scene, contentScope) {
+            MovableElementScopeImpl(layoutImpl, key, scene, contentScope)
+        }
+
     ElementBase(
         layoutImpl,
         scene,
         key,
         modifier,
+        elementScope,
         content,
-        scene.scope,
-        isMovable = true,
     )
 }
 
@@ -70,26 +80,16 @@
     scene: Scene,
     key: ElementKey,
     modifier: Modifier,
-    content: @Composable ElementScope<ContentScope>.() -> Unit,
-    contentScope: ContentScope,
-    isMovable: Boolean,
+    elementScope: ElementScope<ContentScope>,
+    content: @Composable (ElementScope<ContentScope>.() -> Unit),
 ) {
-    Box(modifier.element(layoutImpl, scene, key)) {
-        val elementScope =
-            remember(layoutImpl, key, scene, contentScope, isMovable) {
-                ElementScopeImpl(layoutImpl, key, scene, contentScope, isMovable)
-            }
-
-        elementScope.content()
-    }
+    Box(modifier.element(layoutImpl, scene, key)) { elementScope.content() }
 }
 
-private class ElementScopeImpl<ContentScope>(
+private abstract class BaseElementScope<ContentScope>(
     private val layoutImpl: SceneTransitionLayoutImpl,
     private val element: ElementKey,
     private val scene: Scene,
-    private val contentScope: ContentScope,
-    private val isMovable: Boolean,
 ) : ElementScope<ContentScope> {
     @Composable
     override fun <T> animateElementValueAsState(
@@ -108,14 +108,28 @@
             canOverflow,
         )
     }
+}
 
+private class ElementScopeImpl(
+    layoutImpl: SceneTransitionLayoutImpl,
+    element: ElementKey,
+    scene: Scene,
+    private val contentScope: ElementContentScope,
+) : BaseElementScope<ElementContentScope>(layoutImpl, element, scene) {
     @Composable
-    override fun content(content: @Composable ContentScope.() -> Unit) {
-        if (!isMovable) {
-            contentScope.content()
-            return
-        }
+    override fun content(content: @Composable ElementContentScope.() -> Unit) {
+        contentScope.content()
+    }
+}
 
+private class MovableElementScopeImpl(
+    private val layoutImpl: SceneTransitionLayoutImpl,
+    private val element: ElementKey,
+    private val scene: Scene,
+    private val contentScope: MovableElementContentScope,
+) : BaseElementScope<MovableElementContentScope>(layoutImpl, element, scene) {
+    @Composable
+    override fun content(content: @Composable MovableElementContentScope.() -> Unit) {
         // Whether we should compose the movable element here. The scene picker logic to know in
         // which scene we should compose/draw a movable element might depend on the current
         // transition progress, so we put this in a derivedStateOf to prevent many recompositions
@@ -130,10 +144,14 @@
         if (shouldComposeMovableElement) {
             val movableContent: MovableElementContent =
                 layoutImpl.movableContents[element]
-                    ?: movableContentOf { content: @Composable () -> Unit -> content() }
+                    ?: movableContentOf {
+                            contentScope: MovableElementContentScope,
+                            content: @Composable MovableElementContentScope.() -> Unit ->
+                            contentScope.content()
+                        }
                         .also { layoutImpl.movableContents[element] = it }
 
-            movableContent { contentScope.content() }
+            movableContent(contentScope, content)
         } else {
             // If we are not composed, we still need to lay out an empty space with the same *target
             // size* as its movable content, i.e. the same *size when idle*. During transitions,
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
index a930309..0227aba 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
@@ -36,7 +36,15 @@
 import com.android.compose.ui.util.lerp
 import kotlinx.coroutines.CoroutineScope
 
-internal typealias MovableElementContent = @Composable (@Composable () -> Unit) -> Unit
+/**
+ * The type for the content of movable elements.
+ *
+ * TODO(b/317972419): Revert back to make this movable content have a single @Composable lambda
+ *   parameter.
+ */
+internal typealias MovableElementContent =
+    @Composable
+    (MovableElementContentScope, @Composable MovableElementContentScope.() -> Unit) -> Unit
 
 @Stable
 internal class SceneTransitionLayoutImpl(
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
index 3824cf5..a80906c 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
@@ -41,6 +41,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.test.assertSizeIsEqualTo
 import com.google.common.truth.Truth.assertThat
+import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -255,4 +256,39 @@
             }
         }
     }
+
+    @Test
+    @Ignore("b/317972419#comment2")
+    fun movableElementContentIsRecomposedIfContentParametersChange() {
+        @Composable
+        fun SceneScope.MovableFoo(text: String, modifier: Modifier = Modifier) {
+            MovableElement(TestElements.Foo, modifier) { content { Text(text) } }
+        }
+
+        rule.testTransition(
+            fromSceneContent = { MovableFoo(text = "fromScene") },
+            toSceneContent = { MovableFoo(text = "toScene") },
+            transition = { spec = tween(durationMillis = 16 * 4, easing = LinearEasing) },
+            fromScene = TestScenes.SceneA,
+            toScene = TestScenes.SceneB,
+        ) {
+            // Before the transition, only fromScene is composed.
+            before {
+                rule.onNodeWithText("fromScene").assertIsDisplayed()
+                rule.onNodeWithText("toScene").assertDoesNotExist()
+            }
+
+            // During the transition, the element is composed in toScene.
+            at(32) {
+                rule.onNodeWithText("fromScene").assertDoesNotExist()
+                rule.onNodeWithText("toScene").assertIsDisplayed()
+            }
+
+            // At the end of the transition, the element is composed in toScene.
+            after {
+                rule.onNodeWithText("fromScene").assertDoesNotExist()
+                rule.onNodeWithText("toScene").assertIsDisplayed()
+            }
+        }
+    }
 }