Merge changes I591dac9b,I79abd2a7 into main

* changes:
  Set the target state during lookahead pass
  Don't cast ApproachMeasureScope as LookaheadScope
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 f0fb9f6..11d3841 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,9 +33,9 @@
 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.MeasureScope
 import androidx.compose.ui.layout.Placeable
 import androidx.compose.ui.node.DrawModifierNode
 import androidx.compose.ui.node.ModifierNodeElement
@@ -248,13 +248,34 @@
     }
 
     @ExperimentalComposeUiApi
+    override fun MeasureScope.measure(
+        measurable: Measurable,
+        constraints: Constraints
+    ): MeasureResult {
+        check(isLookingAhead)
+
+        return measurable.measure(constraints).run {
+            // Update the size this element has in this scene when idle.
+            sceneState.targetSize = size()
+
+            layout(width, height) {
+                // Update the offset (relative to the SceneTransitionLayout) this element has in
+                // this scene when idle.
+                coordinates?.let { coords ->
+                    with(layoutImpl.lookaheadScope) {
+                        sceneState.targetOffset =
+                            lookaheadScopeCoordinates.localLookaheadPositionOf(coords)
+                    }
+                }
+                place(0, 0)
+            }
+        }
+    }
+
     override fun ApproachMeasureScope.approachMeasure(
         measurable: Measurable,
         constraints: Constraints,
     ): MeasureResult {
-        // Update the size this element has in this scene when idle.
-        sceneState.targetSize = lookaheadSize
-
         val transitions = currentTransitions
         val transition = elementTransition(element, transitions)
 
@@ -272,15 +293,7 @@
             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)
-                }
-            }
+            return layout(placeable.width, placeable.height) { /* Do not place */ }
         }
 
         val placeable =
@@ -294,7 +307,6 @@
                 transition,
                 sceneState,
                 placeable,
-                placementScope = this,
             )
         }
     }
@@ -541,8 +553,7 @@
             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,23 +808,19 @@
 }
 
 @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")
         val targetOffsetInScene = lookaheadScopeCoordinates.localLookaheadPositionOf(coords)
-        sceneState.targetOffset = targetOffsetInScene
 
         // No need to place the element in this scene if we don't want to draw it anyways.
         if (!shouldPlaceElement(layoutImpl, scene, element, transition)) {
@@ -985,10 +992,10 @@
 
     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 5fa7c87..f32720c 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 @@
                     _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 @@
                 .then(LayoutElement(layoutImpl = this))
         ) {
             LookaheadScope {
+                lookaheadScope = this
+
                 BackHandler()
 
                 scenesToCompose().fastForEach { scene -> key(scene.key) { scene.Content() } }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
index 124ec29..b54afae 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
@@ -41,15 +41,20 @@
         value: IntSize,
     ): IntSize {
         fun anchorSizeIn(scene: SceneKey): IntSize {
-            val size = layoutImpl.elements[anchor]?.sceneStates?.get(scene)?.targetSize
-            return if (size != null && size != Element.SizeUnspecified) {
-                IntSize(
-                    width = if (anchorWidth) size.width else value.width,
-                    height = if (anchorHeight) size.height else value.height,
-                )
-            } else {
-                value
-            }
+            val size =
+                layoutImpl.elements[anchor]?.sceneStates?.get(scene)?.targetSize?.takeIf {
+                    it != Element.SizeUnspecified
+                }
+                    ?: throwMissingAnchorException(
+                        transformation = "AnchoredSize",
+                        anchor = anchor,
+                        scene = scene,
+                    )
+
+            return IntSize(
+                width = if (anchorWidth) size.width else value.width,
+                height = if (anchorHeight) size.height else value.height,
+            )
         }
 
         // This simple implementation assumes that the size of [element] is the same as the size of
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
index 7aa702b..2bab4f8 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
@@ -39,7 +39,15 @@
         transition: TransitionState.Transition,
         value: Offset,
     ): Offset {
-        val anchor = layoutImpl.elements[anchor] ?: return value
+        fun throwException(scene: SceneKey?): Nothing {
+            throwMissingAnchorException(
+                transformation = "AnchoredTranslate",
+                anchor = anchor,
+                scene = scene,
+            )
+        }
+
+        val anchor = layoutImpl.elements[anchor] ?: throwException(scene = null)
         fun anchorOffsetIn(scene: SceneKey): Offset? {
             return anchor.sceneStates[scene]?.targetOffset?.takeIf { it.isSpecified }
         }
@@ -47,8 +55,10 @@
         // [element] will move the same amount as [anchor] does.
         // TODO(b/290184746): Also support anchors that are not shared but translated because of
         // other transformations, like an edge translation.
-        val anchorFromOffset = anchorOffsetIn(transition.fromScene) ?: return value
-        val anchorToOffset = anchorOffsetIn(transition.toScene) ?: return value
+        val anchorFromOffset =
+            anchorOffsetIn(transition.fromScene) ?: throwException(transition.fromScene)
+        val anchorToOffset =
+            anchorOffsetIn(transition.toScene) ?: throwException(transition.toScene)
         val offset = anchorToOffset - anchorFromOffset
 
         return if (scene.key == transition.toScene) {
@@ -64,3 +74,20 @@
         }
     }
 }
+
+internal fun throwMissingAnchorException(
+    transformation: String,
+    anchor: ElementKey,
+    scene: SceneKey?,
+): Nothing {
+    error(
+        """
+        Anchor ${anchor.debugName} does not have a target state in scene ${scene?.debugName}.
+        This either means that it was not composed at all during the transition or that it was
+        composed too late, for instance during layout/subcomposition. To avoid flickers in
+        $transformation, you should make sure that the composition and layout of anchor is *not*
+        deferred, for instance by moving it out of lazy layouts.
+    """
+            .trimIndent()
+    )
+}
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt
index d1205e7..46075c3 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt
@@ -27,6 +27,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.animation.scene.TestElements
 import com.android.compose.animation.scene.testTransition
+import com.android.compose.animation.scene.transition
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -83,4 +84,28 @@
             after { onElement(TestElements.Bar).assertPositionInRootIsEqualTo(20.dp, 40.dp) }
         }
     }
+
+    @Test
+    fun anchorPlacedAfterAnchoredElement() {
+        rule.testTransition(
+            fromSceneContent = { Box(Modifier.offset(10.dp, 50.dp).element(TestElements.Foo)) },
+            toSceneContent = {
+                Box(Modifier.offset(20.dp, 40.dp).element(TestElements.Bar))
+                Box(Modifier.offset(30.dp, 10.dp).element(TestElements.Foo))
+            },
+            transition = {
+                spec = tween(16 * 4, easing = LinearEasing)
+                anchoredTranslate(TestElements.Bar, TestElements.Foo)
+            },
+        ) {
+            // No exception is thrown even if Bar is placed before the anchor in toScene.
+            before { onElement(TestElements.Bar).assertDoesNotExist() }
+            at(0) { onElement(TestElements.Bar).assertPositionInRootIsEqualTo(0.dp, 80.dp) }
+            at(16) { onElement(TestElements.Bar).assertPositionInRootIsEqualTo(5.dp, 70.dp) }
+            at(32) { onElement(TestElements.Bar).assertPositionInRootIsEqualTo(10.dp, 60.dp) }
+            at(48) { onElement(TestElements.Bar).assertPositionInRootIsEqualTo(15.dp, 50.dp) }
+            at(64) { onElement(TestElements.Bar).assertPositionInRootIsEqualTo(20.dp, 40.dp) }
+            after { onElement(TestElements.Bar).assertPositionInRootIsEqualTo(20.dp, 40.dp) }
+        }
+    }
 }