Make (Movable)ElementScope extend ElementBoxScope

This CL exposes the BoxScope of Element and MovableElement so that we
can easily create an Element that wraps its content but that have a
background that is exactly the same size of the element.

Test: MovableElementTest
Bug: 291071158
Flag: N/A
Change-Id: Iee96ece7b792c8e78a3c1dd71726e5317fc111d1
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 964dca8..af3c099 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
@@ -17,6 +17,7 @@
 package com.android.compose.animation.scene
 
 import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.BoxScope
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
@@ -34,20 +35,16 @@
     modifier: Modifier,
     content: @Composable ElementScope<ElementContentScope>.() -> Unit,
 ) {
-    val contentScope = scene.scope
-    val elementScope =
-        remember(layoutImpl, key, scene, contentScope) {
-            ElementScopeImpl(layoutImpl, key, scene, contentScope)
-        }
+    Box(modifier.element(layoutImpl, scene, key)) {
+        val sceneScope = scene.scope
+        val boxScope = this
+        val elementScope =
+            remember(layoutImpl, key, scene, sceneScope, boxScope) {
+                ElementScopeImpl(layoutImpl, key, scene, sceneScope, boxScope)
+            }
 
-    ElementBase(
-        layoutImpl,
-        scene,
-        key,
-        modifier,
-        elementScope,
-        content,
-    )
+        content(elementScope)
+    }
 }
 
 @Composable
@@ -58,32 +55,16 @@
     modifier: Modifier,
     content: @Composable ElementScope<MovableElementContentScope>.() -> Unit,
 ) {
-    val contentScope = scene.scope
-    val elementScope =
-        remember(layoutImpl, key, scene, contentScope) {
-            MovableElementScopeImpl(layoutImpl, key, scene, contentScope)
-        }
+    Box(modifier.element(layoutImpl, scene, key)) {
+        val sceneScope = scene.scope
+        val boxScope = this
+        val elementScope =
+            remember(layoutImpl, key, scene, sceneScope, boxScope) {
+                MovableElementScopeImpl(layoutImpl, key, scene, sceneScope, boxScope)
+            }
 
-    ElementBase(
-        layoutImpl,
-        scene,
-        key,
-        modifier,
-        elementScope,
-        content,
-    )
-}
-
-@Composable
-private inline fun <ContentScope> ElementBase(
-    layoutImpl: SceneTransitionLayoutImpl,
-    scene: Scene,
-    key: ElementKey,
-    modifier: Modifier,
-    elementScope: ElementScope<ContentScope>,
-    content: @Composable (ElementScope<ContentScope>.() -> Unit),
-) {
-    Box(modifier.element(layoutImpl, scene, key)) { elementScope.content() }
+        content(elementScope)
+    }
 }
 
 private abstract class BaseElementScope<ContentScope>(
@@ -114,8 +95,12 @@
     layoutImpl: SceneTransitionLayoutImpl,
     element: ElementKey,
     scene: Scene,
-    private val contentScope: ElementContentScope,
+    private val sceneScope: SceneScope,
+    private val boxScope: BoxScope,
 ) : BaseElementScope<ElementContentScope>(layoutImpl, element, scene) {
+    private val contentScope =
+        object : ElementContentScope, SceneScope by sceneScope, BoxScope by boxScope {}
+
     @Composable
     override fun content(content: @Composable ElementContentScope.() -> Unit) {
         contentScope.content()
@@ -126,8 +111,12 @@
     private val layoutImpl: SceneTransitionLayoutImpl,
     private val element: ElementKey,
     private val scene: Scene,
-    private val contentScope: MovableElementContentScope,
+    private val sceneScope: BaseSceneScope,
+    private val boxScope: BoxScope,
 ) : BaseElementScope<MovableElementContentScope>(layoutImpl, element, scene) {
+    private val contentScope =
+        object : MovableElementContentScope, BaseSceneScope by sceneScope, BoxScope by boxScope {}
+
     @Composable
     override fun content(content: @Composable MovableElementContentScope.() -> Unit) {
         // Whether we should compose the movable element here. The scene picker logic to know in
@@ -151,6 +140,10 @@
                         }
                         .also { layoutImpl.movableContents[element] = it }
 
+            // Important: Don't introduce any parent Box or other layout here, because contentScope
+            // delegates its BoxScope implementation to the Box where this content() function is
+            // called, so it's important that this movableContent is composed directly under that
+            // Box.
             movableContent(contentScope, content)
         } else {
             // If we are not composed, we still need to lay out an empty space with the same *target
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
index 4785716..3537b79 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
@@ -74,7 +74,7 @@
 internal class SceneScopeImpl(
     private val layoutImpl: SceneTransitionLayoutImpl,
     private val scene: Scene,
-) : SceneScope, ElementContentScope, MovableElementContentScope {
+) : SceneScope {
     override val layoutState: SceneTransitionLayoutState = layoutImpl.state
 
     override fun Modifier.element(key: ElementKey): Modifier {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
index 7deaf4d..84fade89 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -24,6 +24,7 @@
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
@@ -281,8 +282,25 @@
     @Composable fun content(content: @Composable ContentScope.() -> Unit)
 }
 
+/**
+ * The exact same scope as [androidx.compose.foundation.layout.BoxScope].
+ *
+ * We can't reuse BoxScope directly because of the @LayoutScopeMarker annotation on it, which would
+ * prevent us from calling Modifier.element() and other methods of [SceneScope] inside any Box {} in
+ * the [content][ElementScope.content] of a [SceneScope.Element] or a [SceneScope.MovableElement].
+ */
+@Stable
+@ElementDsl
+interface ElementBoxScope {
+    /** @see [androidx.compose.foundation.layout.BoxScope.align]. */
+    @Stable fun Modifier.align(alignment: Alignment): Modifier
+
+    /** @see [androidx.compose.foundation.layout.BoxScope.matchParentSize]. */
+    @Stable fun Modifier.matchParentSize(): Modifier
+}
+
 /** The scope for "normal" (not movable) elements. */
-@Stable @ElementDsl interface ElementContentScope : MovableElementContentScope, SceneScope
+@Stable @ElementDsl interface ElementContentScope : SceneScope, ElementBoxScope
 
 /**
  * The scope for the content of movable elements.
@@ -291,7 +309,7 @@
  * call [SceneScope.animateSceneValueAsState], given that their content is not composed in all
  * scenes.
  */
-@Stable @ElementDsl interface MovableElementContentScope : BaseSceneScope
+@Stable @ElementDsl interface MovableElementContentScope : BaseSceneScope, ElementBoxScope
 
 /** An action performed by the user. */
 sealed interface UserAction
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 a80906c..3253289 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
@@ -28,13 +28,17 @@
 import androidx.compose.runtime.mutableIntStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.test.assertIsDisplayed
 import androidx.compose.ui.test.assertIsNotDisplayed
+import androidx.compose.ui.test.assertPositionInRootIsEqualTo
 import androidx.compose.ui.test.hasParent
 import androidx.compose.ui.test.hasText
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onAllNodesWithText
+import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import androidx.compose.ui.unit.dp
@@ -291,4 +295,38 @@
             }
         }
     }
+
+    @Test
+    fun elementScopeExtendsBoxScope() {
+        rule.setContent {
+            TestSceneScope {
+                Element(TestElements.Foo, Modifier.size(200.dp)) {
+                    content {
+                        Box(Modifier.testTag("bottomEnd").align(Alignment.BottomEnd))
+                        Box(Modifier.testTag("matchParentSize").matchParentSize())
+                    }
+                }
+            }
+        }
+
+        rule.onNodeWithTag("bottomEnd").assertPositionInRootIsEqualTo(200.dp, 200.dp)
+        rule.onNodeWithTag("matchParentSize").assertSizeIsEqualTo(200.dp, 200.dp)
+    }
+
+    @Test
+    fun movableElementScopeExtendsBoxScope() {
+        rule.setContent {
+            TestSceneScope {
+                MovableElement(TestElements.Foo, Modifier.size(200.dp)) {
+                    content {
+                        Box(Modifier.testTag("bottomEnd").align(Alignment.BottomEnd))
+                        Box(Modifier.testTag("matchParentSize").matchParentSize())
+                    }
+                }
+            }
+        }
+
+        rule.onNodeWithTag("bottomEnd").assertPositionInRootIsEqualTo(200.dp, 200.dp)
+        rule.onNodeWithTag("matchParentSize").assertSizeIsEqualTo(200.dp, 200.dp)
+    }
 }