Don't cast ApproachMeasureScope as LookaheadScope

Bug: 290930950
Test: atest SceneTransitionLayoutTest
Flag: com.android.systemui.scene_container
Change-Id: I79abd2a7fe1bc5273c9a04613c54da61d3f6c27a

2 files changed, 20 insertions, 16 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
index f0fb9f62fdad..f1f84dca6b8e 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
@@ -33,7 +33,6 @@ import androidx.compose.ui.graphics.drawscope.scale
 import androidx.compose.ui.layout.ApproachLayoutModifierNode
 import androidx.compose.ui.layout.ApproachMeasureScope
 import androidx.compose.ui.layout.LayoutCoordinates
-import androidx.compose.ui.layout.LookaheadScope
 import androidx.compose.ui.layout.Measurable
 import androidx.compose.ui.layout.MeasureResult
 import androidx.compose.ui.layout.Placeable
@@ -272,13 +271,14 @@ internal class ElementNode(
             val placeable = measurable.measure(constraints)
             sceneState.lastSize = placeable.size()
 
-            this as LookaheadScope
             return layout(placeable.width, placeable.height) {
                 // Update the offset (relative to the SceneTransitionLayout) this element has in
                 // this scene when idle.
                 coordinates?.let { coords ->
-                    sceneState.targetOffset =
-                        lookaheadScopeCoordinates.localLookaheadPositionOf(coords)
+                    with(layoutImpl.lookaheadScope) {
+                        sceneState.targetOffset =
+                            lookaheadScopeCoordinates.localLookaheadPositionOf(coords)
+                    }
                 }
             }
         }
@@ -294,7 +294,6 @@ internal class ElementNode(
                 transition,
                 sceneState,
                 placeable,
-                placementScope = this,
             )
         }
     }
@@ -541,8 +540,7 @@ internal fun shouldDrawOrComposeSharedElement(
             transition = transition,
             fromSceneZIndex = layoutImpl.scenes.getValue(fromScene).zIndex,
             toSceneZIndex = layoutImpl.scenes.getValue(toScene).zIndex,
-        )
-            ?: return false
+        ) ?: return false
 
     return pickedScene == scene || transition.currentOverscrollSpec?.scene == scene
 }
@@ -797,18 +795,15 @@ private fun ContentDrawScope.getDrawScale(
 }
 
 @OptIn(ExperimentalComposeUiApi::class)
-private fun ApproachMeasureScope.place(
+private fun Placeable.PlacementScope.place(
     layoutImpl: SceneTransitionLayoutImpl,
     scene: Scene,
     element: Element,
     transition: TransitionState.Transition?,
     sceneState: Element.SceneState,
     placeable: Placeable,
-    placementScope: Placeable.PlacementScope,
 ) {
-    this as LookaheadScope
-
-    with(placementScope) {
+    with(layoutImpl.lookaheadScope) {
         // Update the offset (relative to the SceneTransitionLayout) this element has in this scene
         // when idle.
         val coords = coordinates ?: error("Element ${element.key} does not have any coordinates")
@@ -985,10 +980,10 @@ private inline fun <T> computeValue(
 
     val transformation =
         transformation(transition.transformationSpec.transformations(element.key, scene.key))
-        // If there is no transformation explicitly associated to this element value, let's use
-        // the value given by the system (like the current position and size given by the layout
-        // pass).
-        ?: return currentValue()
+            // If there is no transformation explicitly associated to this element value, let's use
+            // the value given by the system (like the current position and size given by the layout
+            // pass).
+            ?: return currentValue()
 
     // Get the transformed value, i.e. the target value at the beginning (for entering elements) or
     // end (for leaving elements) of the transition.
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 5fa7c874c879..f32720c4716d 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
@@ -107,6 +107,13 @@ internal class SceneTransitionLayoutImpl(
                     _userActionDistanceScope = it
                 }
 
+    /**
+     * The [LookaheadScope] of this layout, that can be used to compute offsets relative to the
+     * layout.
+     */
+    internal lateinit var lookaheadScope: LookaheadScope
+        private set
+
     init {
         updateScenes(builder)
 
@@ -195,6 +202,8 @@ internal class SceneTransitionLayoutImpl(
                 .then(LayoutElement(layoutImpl = this))
         ) {
             LookaheadScope {
+                lookaheadScope = this
+
                 BackHandler()
 
                 scenesToCompose().fastForEach { scene -> key(scene.key) { scene.Content() } }