Merge changes from topics "mutable-stl-state", "stl-punch-hole" into main

* changes:
  Remove Element.lastSharedState and SceneState.lastState
  Remove Modifier.punchHole in SceneScope (1/2)
  Remove SceneTransitionLayoutImpl.readyScenes
  Introduce MutableSceneTransitionLayoutState
  Move onChangeScene and transitions to STLState (1/2)

18 files changed, 455 insertions, 429 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
index 249b3e14ec72..d47527a0a191 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
@@ -33,11 +33,11 @@ import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.SceneTransitionLayout
-import com.android.compose.animation.scene.SceneTransitionLayoutState
 import com.android.compose.animation.scene.Swipe
 import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.observableTransitionState
 import com.android.compose.animation.scene.transitions
+import com.android.compose.animation.scene.updateSceneTransitionLayoutState
 import com.android.systemui.communal.shared.model.CommunalSceneKey
 import com.android.systemui.communal.shared.model.ObservableCommunalTransitionState
 import com.android.systemui.communal.ui.viewmodel.BaseCommunalViewModel
@@ -76,7 +76,13 @@ fun CommunalContainer(
         viewModel.currentScene
             .transform { value -> emit(value.toTransitionSceneKey()) }
             .collectAsState(TransitionSceneKey.Blank)
-    val sceneTransitionLayoutState = remember { SceneTransitionLayoutState(currentScene) }
+    val sceneTransitionLayoutState =
+        updateSceneTransitionLayoutState(
+            currentScene,
+            onChangeScene = { viewModel.onSceneChanged(it.toCommunalSceneKey()) },
+            transitions = sceneTransitions,
+        )
+
     // Don't show hub mode UI if keyguard is present. This is important since we're in the shade,
     // which can be opened from many locations.
     val isKeyguardShowing by viewModel.isKeyguardVisible.collectAsState(initial = false)
@@ -98,12 +104,9 @@ fun CommunalContainer(
 
     Box(modifier = modifier.fillMaxSize()) {
         SceneTransitionLayout(
-            modifier = Modifier.fillMaxSize(),
-            currentScene = currentScene,
-            onChangeScene = { sceneKey -> viewModel.onSceneChanged(sceneKey.toCommunalSceneKey()) },
-            transitions = sceneTransitions,
             state = sceneTransitionLayoutState,
-            edgeDetector = FixedSizeEdgeDetector(ContainerDimensions.EdgeSwipeSize)
+            modifier = Modifier.fillMaxSize(),
+            edgeDetector = FixedSizeEdgeDetector(ContainerDimensions.EdgeSwipeSize),
         ) {
             scene(
                 TransitionSceneKey.Blank,
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
index 4eb9089dc589..c35202cd830a 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
@@ -25,7 +25,6 @@ import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.remember
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
@@ -38,11 +37,11 @@ import com.android.compose.animation.scene.Edge as SceneTransitionEdge
 import com.android.compose.animation.scene.ObservableTransitionState as SceneTransitionObservableTransitionState
 import com.android.compose.animation.scene.SceneKey as SceneTransitionSceneKey
 import com.android.compose.animation.scene.SceneTransitionLayout
-import com.android.compose.animation.scene.SceneTransitionLayoutState
 import com.android.compose.animation.scene.Swipe
 import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.UserAction as SceneTransitionUserAction
 import com.android.compose.animation.scene.observableTransitionState
+import com.android.compose.animation.scene.updateSceneTransitionLayoutState
 import com.android.systemui.ribbon.ui.composable.BottomRightCornerRibbon
 import com.android.systemui.scene.shared.model.Direction
 import com.android.systemui.scene.shared.model.Edge
@@ -82,7 +81,12 @@ fun SceneContainer(
     val currentScene = checkNotNull(sceneByKey[currentSceneKey])
     val currentDestinations: Map<UserAction, SceneModel> by
         currentScene.destinationScenes.collectAsState()
-    val state = remember { SceneTransitionLayoutState(currentSceneKey.toTransitionSceneKey()) }
+    val state =
+        updateSceneTransitionLayoutState(
+            currentSceneKey.toTransitionSceneKey(),
+            onChangeScene = viewModel::onSceneChanged,
+            transitions = SceneContainerTransitions,
+        )
 
     DisposableEffect(viewModel, state) {
         viewModel.setTransitionState(state.observableTransitionState().map { it.toModel() })
@@ -93,9 +97,6 @@ fun SceneContainer(
         modifier = Modifier.fillMaxSize(),
     ) {
         SceneTransitionLayout(
-            currentScene = currentSceneKey.toTransitionSceneKey(),
-            onChangeScene = viewModel::onSceneChanged,
-            transitions = SceneContainerTransitions,
             state = state,
             modifier =
                 modifier
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
index ba6d00e3b7f5..7d3b0fbe1725 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
@@ -28,9 +28,9 @@ import kotlinx.coroutines.launch
  * the currently running transition, if there is one.
  */
 internal fun CoroutineScope.animateToScene(
-    layoutState: SceneTransitionLayoutStateImpl,
+    layoutState: BaseSceneTransitionLayoutState,
     target: SceneKey,
-) {
+): TransitionState.Transition? {
     val transitionState = layoutState.transitionState
     if (transitionState.currentScene == target) {
         // This can happen in 3 different situations, for which there isn't anything else to do:
@@ -41,10 +41,10 @@ internal fun CoroutineScope.animateToScene(
         //     a. didn't release their pointer yet.
         //     b. released their pointer such that the swipe gesture was cancelled and the
         //        transition is currently animating back to [target].
-        return
+        return null
     }
 
-    when (transitionState) {
+    return when (transitionState) {
         is TransitionState.Idle -> animate(layoutState, target)
         is TransitionState.Transition -> {
             // A transition is currently running: first check whether `transition.toScene` or
@@ -62,47 +62,43 @@ internal fun CoroutineScope.animateToScene(
                     // finish the current transition early to make sure that the current state
                     // change is committed.
                     layoutState.finishTransition(transitionState, transitionState.currentScene)
+                    null
                 } else {
                     // The transition is in progress: start the canned animation at the same
                     // progress as it was in.
                     // TODO(b/290184746): Also take the current velocity into account.
                     animate(layoutState, target, startProgress = progress)
                 }
-
-                return
-            }
-
-            if (transitionState.fromScene == target) {
+            } else if (transitionState.fromScene == target) {
                 // There is a transition from [target] to another scene: simply animate the same
                 // transition progress to `0`.
-
                 check(transitionState.toScene == transitionState.currentScene)
+
                 val progress = transitionState.progress
                 if (progress.absoluteValue < ProgressVisibilityThreshold) {
                     // The transition is at progress ~= 0: no need to animate.We finish the current
                     // transition early to make sure that the current state change is committed.
                     layoutState.finishTransition(transitionState, transitionState.currentScene)
+                    null
                 } else {
                     // TODO(b/290184746): Also take the current velocity into account.
                     animate(layoutState, target, startProgress = progress, reversed = true)
                 }
-
-                return
+            } else {
+                // Generic interruption; the current transition is neither from or to [target].
+                // TODO(b/290930950): Better handle interruptions here.
+                animate(layoutState, target)
             }
-
-            // Generic interruption; the current transition is neither from or to [target].
-            // TODO(b/290930950): Better handle interruptions here.
-            animate(layoutState, target)
         }
     }
 }
 
 private fun CoroutineScope.animate(
-    layoutState: SceneTransitionLayoutStateImpl,
+    layoutState: BaseSceneTransitionLayoutState,
     target: SceneKey,
     startProgress: Float = 0f,
     reversed: Boolean = false,
-) {
+): TransitionState.Transition {
     val fromScene = layoutState.transitionState.currentScene
     val isUserInput =
         (layoutState.transitionState as? TransitionState.Transition)?.isInitiatedByUserInput
@@ -143,10 +139,15 @@ private fun CoroutineScope.animate(
         }
 
     // Animate the progress to its target value.
-    launch {
-        animatable.animateTo(targetProgress, animationSpec)
-        layoutState.finishTransition(transition, target)
-    }
+    launch { animatable.animateTo(targetProgress, animationSpec) }
+        .invokeOnCompletion {
+            // Settle the state to Idle(target). Note that this will do nothing if this transition
+            // was replaced/interrupted by another one, and this also runs if this coroutine is
+            // cancelled, i.e. if [this] coroutine scope is cancelled.
+            layoutState.finishTransition(transition, target)
+        }
+
+    return transition
 }
 
 private class OneOffTransition(
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 280fbfb7d3d3..a910bca078e8 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
@@ -20,10 +20,10 @@ import androidx.compose.runtime.Stable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshots.SnapshotStateMap
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.geometry.isSpecified
 import androidx.compose.ui.geometry.isUnspecified
 import androidx.compose.ui.geometry.lerp
 import androidx.compose.ui.graphics.drawscope.ContentDrawScope
@@ -46,41 +46,18 @@ import kotlinx.coroutines.launch
 /** An element on screen, that can be composed in one or more scenes. */
 @Stable
 internal class Element(val key: ElementKey) {
-    /**
-     * The last state of this element, coming from any scene. Note that this state will be unstable
-     * if this element is present in multiple scenes but the shared element animation is disabled,
-     * given that multiple instances of the element with different states will write to this state.
-     * You should prefer using [SceneState.lastState] in the current scene when it is defined.
-     */
-    val lastSharedState = State()
-
     /** The mapping between a scene and the state this element has in that scene, if any. */
-    val sceneStates = mutableMapOf<SceneKey, SceneState>()
+    // TODO(b/316901148): Make this a normal map instead once we can make sure that new transitions
+    // are first seen by composition then layout/drawing code. See 316901148#comment2 for details.
+    val sceneStates = SnapshotStateMap<SceneKey, SceneState>()
 
     override fun toString(): String {
         return "Element(key=$key)"
     }
 
-    /** The state of this element, either in a specific scene or in a shared context. */
-    class State {
-        /** The offset of the element, relative to the SceneTransitionLayout containing it. */
-        var offset = Offset.Unspecified
-
-        /** The size of this element. */
-        var size = SizeUnspecified
-
-        /** The draw scale of this element. */
-        var drawScale = Scale.Default
-
-        /** The alpha of this element. */
-        var alpha = AlphaUnspecified
-    }
-
     /** The last and target state of this element in a given scene. */
     @Stable
     class SceneState(val scene: SceneKey) {
-        val lastState = State()
-
         var targetSize by mutableStateOf(SizeUnspecified)
         var targetOffset by mutableStateOf(Offset.Unspecified)
 
@@ -94,7 +71,6 @@ internal class Element(val key: ElementKey) {
 
     companion object {
         val SizeUnspecified = IntSize(Int.MAX_VALUE, Int.MAX_VALUE)
-        val AlphaUnspecified = Float.MIN_VALUE
     }
 }
 
@@ -219,7 +195,7 @@ internal class ElementNode(
     }
 
     override fun ContentDrawScope.draw() {
-        val drawScale = getDrawScale(layoutImpl, element, scene, sceneState)
+        val drawScale = getDrawScale(layoutImpl, element, scene)
         if (drawScale == Scale.Default) {
             drawContent()
         } else {
@@ -264,7 +240,6 @@ private fun shouldDrawElement(
     // Always draw the element if there is no ongoing transition or if the element is not shared.
     if (
         transition == null ||
-            !layoutImpl.isTransitionReady(transition) ||
             transition.fromScene !in element.sceneStates ||
             transition.toScene !in element.sceneStates
     ) {
@@ -304,7 +279,7 @@ internal fun shouldDrawOrComposeSharedElement(
 }
 
 private fun isSharedElementEnabled(
-    layoutState: SceneTransitionLayoutStateImpl,
+    layoutState: BaseSceneTransitionLayoutState,
     transition: TransitionState.Transition,
     element: ElementKey,
 ): Boolean {
@@ -312,7 +287,7 @@ private fun isSharedElementEnabled(
 }
 
 internal fun sharedElementTransformation(
-    layoutState: SceneTransitionLayoutStateImpl,
+    layoutState: BaseSceneTransitionLayoutState,
     transition: TransitionState.Transition,
     element: ElementKey,
 ): SharedElementTransformation? {
@@ -342,18 +317,9 @@ private fun isElementOpaque(
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
     scene: Scene,
-    sceneState: Element.SceneState,
 ): Boolean {
     val transition = layoutImpl.state.currentTransition ?: return true
 
-    if (!layoutImpl.isTransitionReady(transition)) {
-        val lastValue =
-            sceneState.lastState.alpha.takeIf { it != Element.AlphaUnspecified }
-                ?: element.lastSharedState.alpha.takeIf { it != Element.AlphaUnspecified } ?: 1f
-
-        return lastValue == 1f
-    }
-
     val fromScene = transition.fromScene
     val toScene = transition.toScene
     val fromState = element.sceneStates[fromScene]
@@ -383,7 +349,6 @@ private fun elementAlpha(
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
     scene: Scene,
-    sceneState: Element.SceneState,
 ): Float {
     return computeValue(
             layoutImpl,
@@ -393,10 +358,7 @@ private fun elementAlpha(
             transformation = { it.alpha },
             idleValue = 1f,
             currentValue = { 1f },
-            lastValue = {
-                sceneState.lastState.alpha.takeIf { it != Element.AlphaUnspecified }
-                    ?: element.lastSharedState.alpha.takeIf { it != Element.AlphaUnspecified } ?: 1f
-            },
+            isSpecified = { true },
             ::lerp,
         )
         .coerceIn(0f, 1f)
@@ -434,34 +396,23 @@ private fun IntermediateMeasureScope.measure(
             transformation = { it.size },
             idleValue = lookaheadSize,
             currentValue = { measurable.measure(constraints).also { maybePlaceable = it }.size() },
-            lastValue = {
-                sceneState.lastState.size.takeIf { it != Element.SizeUnspecified }
-                    ?: element.lastSharedState.size.takeIf { it != Element.SizeUnspecified }
-                        ?: measurable.measure(constraints).also { maybePlaceable = it }.size()
-            },
+            isSpecified = { it != Element.SizeUnspecified },
             ::lerp,
         )
 
-    val placeable =
-        maybePlaceable
-            ?: measurable.measure(
-                Constraints.fixed(
-                    targetSize.width.coerceAtLeast(0),
-                    targetSize.height.coerceAtLeast(0),
-                )
+    return maybePlaceable
+        ?: measurable.measure(
+            Constraints.fixed(
+                targetSize.width.coerceAtLeast(0),
+                targetSize.height.coerceAtLeast(0),
             )
-
-    val size = placeable.size()
-    element.lastSharedState.size = size
-    sceneState.lastState.size = size
-    return placeable
+        )
 }
 
 private fun getDrawScale(
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
-    scene: Scene,
-    sceneState: Element.SceneState
+    scene: Scene
 ): Scale {
     return computeValue(
         layoutImpl,
@@ -471,10 +422,7 @@ private fun getDrawScale(
         transformation = { it.drawScale },
         idleValue = Scale.Default,
         currentValue = { Scale.Default },
-        lastValue = {
-            sceneState.lastState.drawScale.takeIf { it != Scale.Default }
-                ?: element.lastSharedState.drawScale
-        },
+        isSpecified = { true },
         ::lerp,
     )
 }
@@ -498,9 +446,12 @@ private fun IntermediateMeasureScope.place(
             sceneState.targetOffset = targetOffsetInScene
         }
 
+        // No need to place the element in this scene if we don't want to draw it anyways.
+        if (!shouldDrawElement(layoutImpl, scene, element)) {
+            return
+        }
+
         val currentOffset = lookaheadScopeCoordinates.localPositionOf(coords, Offset.Zero)
-        val lastSharedState = element.lastSharedState
-        val lastSceneState = sceneState.lastState
         val targetOffset =
             computeValue(
                 layoutImpl,
@@ -510,37 +461,19 @@ private fun IntermediateMeasureScope.place(
                 transformation = { it.offset },
                 idleValue = targetOffsetInScene,
                 currentValue = { currentOffset },
-                lastValue = {
-                    lastSceneState.offset.takeIf { it.isSpecified }
-                        ?: lastSharedState.offset.takeIf { it.isSpecified } ?: currentOffset
-                },
+                isSpecified = { it != Offset.Unspecified },
                 ::lerp,
             )
 
-        lastSharedState.offset = targetOffset
-        lastSceneState.offset = targetOffset
-
-        // No need to place the element in this scene if we don't want to draw it anyways. Note that
-        // it's still important to compute the target offset and update last(Shared|Scene)State,
-        // otherwise they will be out of date.
-        if (!shouldDrawElement(layoutImpl, scene, element)) {
-            return
-        }
-
         val offset = (targetOffset - currentOffset).round()
-        if (isElementOpaque(layoutImpl, element, scene, sceneState)) {
+        if (isElementOpaque(layoutImpl, element, scene)) {
             // TODO(b/291071158): Call placeWithLayer() if offset != IntOffset.Zero and size is not
             // animated once b/305195729 is fixed. Test that drawing is not invalidated in that
             // case.
             placeable.place(offset)
-            lastSharedState.alpha = 1f
-            lastSceneState.alpha = 1f
         } else {
             placeable.placeWithLayer(offset) {
-                val alpha = elementAlpha(layoutImpl, element, scene, sceneState)
-                this.alpha = alpha
-                lastSharedState.alpha = alpha
-                lastSceneState.alpha = alpha
+                this.alpha = elementAlpha(layoutImpl, element, scene)
             }
         }
     }
@@ -563,8 +496,6 @@ private fun IntermediateMeasureScope.place(
  *   different than [idleValue] even if the value is not transformed directly because it could be
  *   impacted by the transformations on other elements, like a parent that is being translated or
  *   resized.
- * @param lastValue the last value that was used. This should be equal to [currentValue] if this is
- *   the first time the value is set.
  * @param lerp the linear interpolation function used to interpolate between two values of this
  *   value type.
  */
@@ -576,7 +507,7 @@ private inline fun <T> computeValue(
     transformation: (ElementTransformations) -> PropertyTransformation<T>?,
     idleValue: T,
     currentValue: () -> T,
-    lastValue: () -> T,
+    isSpecified: (T) -> Boolean,
     lerp: (T, T, Float) -> T,
 ): T {
     val transition =
@@ -587,21 +518,16 @@ private inline fun <T> computeValue(
         // layout phase.
         ?: return currentValue()
 
-    // A transition was started but it's not ready yet (not all elements have been composed/laid
-    // out yet). Use the last value that was set, to make sure elements don't unexpectedly jump.
-    if (!layoutImpl.isTransitionReady(transition)) {
-        return lastValue()
-    }
-
     val fromScene = transition.fromScene
     val toScene = transition.toScene
+
     val fromState = element.sceneStates[fromScene]
     val toState = element.sceneStates[toScene]
 
     if (fromState == null && toState == null) {
         // TODO(b/311600838): Throw an exception instead once layers of disposed elements are not
         // run anymore.
-        return lastValue()
+        return idleValue
     }
 
     // The element is shared: interpolate between the value in fromScene and the value in toScene.
@@ -612,6 +538,11 @@ private inline fun <T> computeValue(
         val start = sceneValue(fromState!!)
         val end = sceneValue(toState!!)
 
+        // TODO(b/316901148): Remove checks to isSpecified() once the lookahead pass runs for all
+        // nodes before the intermediate layout pass.
+        if (!isSpecified(start)) return end
+        if (!isSpecified(end)) return start
+
         // Make sure we don't read progress if values are the same and we don't need to interpolate,
         // so we don't invalidate the phase where this is read.
         return if (start == end) start else lerp(start, end, transition.progress)
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 af3c0999c97b..cdc4778dbf4d 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
@@ -174,22 +174,6 @@ private fun shouldComposeMovableElement(
         // If we are idle, there is only one [scene] that is composed so we can compose our
         // movable content here.
         ?: return true
-    val fromScene = transition.fromScene
-    val toScene = transition.toScene
-
-    val fromReady = layoutImpl.isSceneReady(fromScene)
-    val toReady = layoutImpl.isSceneReady(toScene)
-
-    if (!fromReady && !toReady) {
-        // Neither of the scenes will be drawn, so where we compose it doesn't really matter. Note
-        // that we could have slightly more complicated logic here to optimize for this case, but
-        // it's not worth it given that readyScenes should disappear soon (b/316901148).
-        return scene == toScene
-    }
-
-    // If one of the scenes is not ready, compose it in the other one to make sure it is drawn.
-    if (!fromReady) return scene == toScene
-    if (!toReady) return scene == fromScene
 
     // Always compose movable elements in the scene picked by their scene picker.
     return shouldDrawOrComposeSharedElement(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
index 454c0ecf8ac5..0acc76f8d4ef 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
@@ -16,6 +16,7 @@
 
 package com.android.compose.animation.scene
 
+import androidx.compose.runtime.Stable
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.geometry.Size
@@ -32,80 +33,95 @@ import androidx.compose.ui.graphics.drawscope.DrawScope
 import androidx.compose.ui.graphics.drawscope.drawIntoCanvas
 import androidx.compose.ui.graphics.drawscope.translate
 import androidx.compose.ui.graphics.withSaveLayer
+import androidx.compose.ui.layout.LayoutCoordinates
+import androidx.compose.ui.node.DelegatingNode
 import androidx.compose.ui.node.DrawModifierNode
+import androidx.compose.ui.node.GlobalPositionAwareModifierNode
 import androidx.compose.ui.node.ModifierNodeElement
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.toSize
 
-internal fun Modifier.punchHole(
-    layoutImpl: SceneTransitionLayoutImpl,
-    element: ElementKey,
-    bounds: ElementKey,
-    shape: Shape,
-): Modifier = this.then(PunchHoleElement(layoutImpl, element, bounds, shape))
+/**
+ * Punch a hole in this node with the given [size], [offset] and [shape].
+ *
+ * Punching a hole in an element will "remove" any pixel drawn by that element in the hole area.
+ * This can be used to make content drawn below an opaque element visible. For example, if we have
+ * [this lockscreen scene](http://shortn/_VYySFnJDhN) drawn below
+ * [this shade scene](http://shortn/_fpxGUk0Rg7) and punch a hole in the latter using the big clock
+ * time bounds and a RoundedCornerShape(10dp), [this](http://shortn/_qt80IvORFj) would be the
+ * result.
+ */
+@Stable
+fun Modifier.punchHole(
+    size: () -> Size,
+    offset: () -> Offset,
+    shape: Shape = RectangleShape,
+): Modifier = this.then(PunchHoleElement(size, offset, shape))
+
+/**
+ * Punch a hole in this node using the bounds of [coords] and the given [shape].
+ *
+ * You can use [androidx.compose.ui.layout.onGloballyPositioned] to get the last coordinates of a
+ * node.
+ */
+@Stable
+fun Modifier.punchHole(
+    coords: () -> LayoutCoordinates?,
+    shape: Shape = RectangleShape,
+): Modifier = this.then(PunchHoleWithBoundsElement(coords, shape))
 
 private data class PunchHoleElement(
-    private val layoutImpl: SceneTransitionLayoutImpl,
-    private val element: ElementKey,
-    private val bounds: ElementKey,
+    private val size: () -> Size,
+    private val offset: () -> Offset,
     private val shape: Shape,
 ) : ModifierNodeElement<PunchHoleNode>() {
-    override fun create(): PunchHoleNode = PunchHoleNode(layoutImpl, element, bounds, shape)
+    override fun create(): PunchHoleNode = PunchHoleNode(size, offset, { shape })
 
     override fun update(node: PunchHoleNode) {
-        node.layoutImpl = layoutImpl
-        node.element = element
-        node.bounds = bounds
-        node.shape = shape
+        node.size = size
+        node.offset = offset
+        node.shape = { shape }
     }
 }
 
 private class PunchHoleNode(
-    var layoutImpl: SceneTransitionLayoutImpl,
-    var element: ElementKey,
-    var bounds: ElementKey,
-    var shape: Shape,
+    var size: () -> Size,
+    var offset: () -> Offset,
+    var shape: () -> Shape,
 ) : Modifier.Node(), DrawModifierNode {
     private var lastSize: Size = Size.Unspecified
     private var lastLayoutDirection: LayoutDirection = LayoutDirection.Ltr
     private var lastOutline: Outline? = null
 
     override fun ContentDrawScope.draw() {
-        val bounds = layoutImpl.elements[bounds]
-
-        if (
-            bounds == null ||
-                bounds.lastSharedState.size == Element.SizeUnspecified ||
-                bounds.lastSharedState.offset == Offset.Unspecified
-        ) {
+        val holeSize = size()
+        if (holeSize == Size.Zero) {
             drawContent()
             return
         }
 
-        val element = layoutImpl.elements.getValue(element)
         drawIntoCanvas { canvas ->
             canvas.withSaveLayer(size.toRect(), Paint()) {
                 drawContent()
 
-                val offset = bounds.lastSharedState.offset - element.lastSharedState.offset
-                translate(offset.x, offset.y) { drawHole(bounds) }
+                val offset = offset()
+                translate(offset.x, offset.y) { drawHole(holeSize) }
             }
         }
     }
 
-    private fun DrawScope.drawHole(bounds: Element) {
-        val boundsSize = bounds.lastSharedState.size.toSize()
+    private fun DrawScope.drawHole(size: Size) {
         if (shape == RectangleShape) {
-            drawRect(Color.Black, size = boundsSize, blendMode = BlendMode.DstOut)
+            drawRect(Color.Black, size = size, blendMode = BlendMode.DstOut)
             return
         }
 
         val outline =
-            if (boundsSize == lastSize && layoutDirection == lastLayoutDirection) {
+            if (size == lastSize && layoutDirection == lastLayoutDirection) {
                 lastOutline!!
             } else {
-                val newOutline = shape.createOutline(boundsSize, layoutDirection, this)
-                lastSize = boundsSize
+                val newOutline = shape().createOutline(size, layoutDirection, this)
+                lastSize = size
                 lastLayoutDirection = layoutDirection
                 lastOutline = newOutline
                 newOutline
@@ -118,3 +134,39 @@ private class PunchHoleNode(
         )
     }
 }
+
+private data class PunchHoleWithBoundsElement(
+    private val coords: () -> LayoutCoordinates?,
+    private val shape: Shape,
+) : ModifierNodeElement<PunchHoleWithBoundsNode>() {
+    override fun create(): PunchHoleWithBoundsNode = PunchHoleWithBoundsNode(coords, shape)
+
+    override fun update(node: PunchHoleWithBoundsNode) {
+        node.holeCoords = coords
+        node.shape = shape
+    }
+}
+
+private class PunchHoleWithBoundsNode(
+    var holeCoords: () -> LayoutCoordinates?,
+    var shape: Shape,
+) : DelegatingNode(), DrawModifierNode, GlobalPositionAwareModifierNode {
+    private val delegate = delegate(PunchHoleNode(::holeSize, ::holeOffset, ::shape))
+    private var lastCoords: LayoutCoordinates? = null
+
+    override fun onGloballyPositioned(coordinates: LayoutCoordinates) {
+        this.lastCoords = coordinates
+    }
+
+    override fun ContentDrawScope.draw() = with(delegate) { draw() }
+
+    private fun holeSize(): Size {
+        return holeCoords()?.size?.toSize() ?: Size.Zero
+    }
+
+    private fun holeOffset(): Offset {
+        val holeCoords = holeCoords() ?: return Offset.Zero
+        val lastCoords = lastCoords ?: error("draw() was called before onGloballyPositioned()")
+        return lastCoords.localPositionOf(holeCoords, relativeToSource = Offset.Zero)
+    }
+}
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 3537b7989ed5..f67df54b088c 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
@@ -26,7 +26,6 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.layout.intermediateLayout
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.unit.IntSize
@@ -139,12 +138,6 @@ internal class SceneScopeImpl(
             bottomOrRightBehavior = bottomBehavior,
         )
 
-    override fun Modifier.punchHole(
-        element: ElementKey,
-        bounds: ElementKey,
-        shape: Shape
-    ): Modifier = punchHole(layoutImpl, element, bounds, shape)
-
     override fun Modifier.noResizeDuringTransitions(): Modifier {
         return noResizeDuringTransitions(layoutState = layoutImpl.state)
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
index 338557d0942e..c05591900aa0 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
@@ -312,7 +312,7 @@ internal class SceneGestureHandler(
             // immediately go back B => A.
             if (targetScene != swipeTransition._currentScene) {
                 swipeTransition._currentScene = targetScene
-                layoutImpl.onChangeScene(targetScene.key)
+                with(layoutImpl.state) { coroutineScope.onChangeScene(targetScene.key) }
             }
 
             animateOffset(
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 84fade8937ff..80f8c1c9e987 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
@@ -19,17 +19,48 @@ package com.android.compose.animation.scene
 import androidx.annotation.FloatRange
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.SideEffect
 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
 import androidx.compose.ui.platform.LocalDensity
-import kotlinx.coroutines.channels.Channel
+
+/**
+ * [SceneTransitionLayout] is a container that automatically animates its content whenever its state
+ * changes.
+ *
+ * Note: You should use [androidx.compose.animation.AnimatedContent] instead of
+ * [SceneTransitionLayout] if it fits your need. Use [SceneTransitionLayout] over AnimatedContent if
+ * you need support for swipe gestures, shared elements or transitions defined declaratively outside
+ * UI code.
+ *
+ * @param state the state of this layout.
+ * @param edgeDetector the edge detector used to detect which edge a swipe is started from, if any.
+ * @param transitionInterceptionThreshold used during a scene transition. For the scene to be
+ *   intercepted, the progress value must be above the threshold, and below (1 - threshold).
+ * @param scenes the configuration of the different scenes of this layout.
+ * @see updateSceneTransitionLayoutState
+ */
+@Composable
+fun SceneTransitionLayout(
+    state: SceneTransitionLayoutState,
+    modifier: Modifier = Modifier,
+    edgeDetector: EdgeDetector = DefaultEdgeDetector,
+    @FloatRange(from = 0.0, to = 0.5) transitionInterceptionThreshold: Float = 0f,
+    scenes: SceneTransitionLayoutScope.() -> Unit,
+) {
+    SceneTransitionLayoutForTesting(
+        state,
+        modifier,
+        edgeDetector,
+        transitionInterceptionThreshold,
+        onLayoutImpl = null,
+        scenes,
+    )
+}
 
 /**
  * [SceneTransitionLayout] is a container that automatically animates its content whenever
@@ -45,7 +76,6 @@ import kotlinx.coroutines.channels.Channel
  *   This is called when the user commits a transition to a new scene because of a [UserAction], for
  *   instance by triggering back navigation or by swiping to a new scene.
  * @param transitions the definition of the transitions used to animate a change of scene.
- * @param state the observable state of this layout.
  * @param edgeDetector the edge detector used to detect which edge a swipe is started from, if any.
  * @param transitionInterceptionThreshold used during a scene transition. For the scene to be
  *   intercepted, the progress value must be above the threshold, and below (1 - threshold).
@@ -57,20 +87,16 @@ fun SceneTransitionLayout(
     onChangeScene: (SceneKey) -> Unit,
     transitions: SceneTransitions,
     modifier: Modifier = Modifier,
-    state: SceneTransitionLayoutState = remember { SceneTransitionLayoutState(currentScene) },
     edgeDetector: EdgeDetector = DefaultEdgeDetector,
     @FloatRange(from = 0.0, to = 0.5) transitionInterceptionThreshold: Float = 0f,
     scenes: SceneTransitionLayoutScope.() -> Unit,
 ) {
-    SceneTransitionLayoutForTesting(
-        currentScene,
-        onChangeScene,
-        modifier,
-        transitions,
+    val state = updateSceneTransitionLayoutState(currentScene, onChangeScene, transitions)
+    SceneTransitionLayout(
         state,
+        modifier,
         edgeDetector,
         transitionInterceptionThreshold,
-        onLayoutImpl = null,
         scenes,
     )
 }
@@ -203,18 +229,6 @@ interface BaseSceneScope {
     ): Modifier
 
     /**
-     * Punch a hole in this [element] using the bounds of [bounds] in [scene] and the given [shape].
-     *
-     * Punching a hole in an element will "remove" any pixel drawn by that element in the hole area.
-     * This can be used to make content drawn below an opaque element visible. For example, if we
-     * have [this lockscreen scene](http://shortn/_VYySFnJDhN) drawn below
-     * [this shade scene](http://shortn/_fpxGUk0Rg7) and punch a hole in the latter using the big
-     * clock time bounds and a RoundedCornerShape(10dp), [this](http://shortn/_qt80IvORFj) would be
-     * the result.
-     */
-    fun Modifier.punchHole(element: ElementKey, bounds: ElementKey, shape: Shape): Modifier
-
-    /**
      * Don't resize during transitions. This can for instance be used to make sure that scrollable
      * lists keep a constant size during transitions even if its elements are growing/shrinking.
      */
@@ -346,11 +360,8 @@ enum class SwipeDirection(val orientation: Orientation) {
  */
 @Composable
 internal fun SceneTransitionLayoutForTesting(
-    currentScene: SceneKey,
-    onChangeScene: (SceneKey) -> Unit,
+    state: SceneTransitionLayoutState,
     modifier: Modifier = Modifier,
-    transitions: SceneTransitions = transitions {},
-    state: SceneTransitionLayoutState = remember { SceneTransitionLayoutState(currentScene) },
     edgeDetector: EdgeDetector = DefaultEdgeDetector,
     transitionInterceptionThreshold: Float = 0f,
     onLayoutImpl: ((SceneTransitionLayoutImpl) -> Unit)? = null,
@@ -360,8 +371,7 @@ internal fun SceneTransitionLayoutForTesting(
     val coroutineScope = rememberCoroutineScope()
     val layoutImpl = remember {
         SceneTransitionLayoutImpl(
-                state = state as SceneTransitionLayoutStateImpl,
-                onChangeScene = onChangeScene,
+                state = state as BaseSceneTransitionLayoutState,
                 density = density,
                 edgeDetector = edgeDetector,
                 transitionInterceptionThreshold = transitionInterceptionThreshold,
@@ -375,7 +385,6 @@ internal fun SceneTransitionLayoutForTesting(
     // SnapshotStateMap anymore.
     layoutImpl.updateScenes(scenes)
 
-    val targetSceneChannel = remember { Channel<SceneKey>(Channel.CONFLATED) }
     SideEffect {
         if (state != layoutImpl.state) {
             error(
@@ -384,23 +393,8 @@ internal fun SceneTransitionLayoutForTesting(
             )
         }
 
-        layoutImpl.onChangeScene = onChangeScene
-        (state as SceneTransitionLayoutStateImpl).transitions = transitions
         layoutImpl.density = density
         layoutImpl.edgeDetector = edgeDetector
-
-        state.transitions = transitions
-
-        targetSceneChannel.trySend(currentScene)
-    }
-
-    LaunchedEffect(targetSceneChannel) {
-        for (newKey in targetSceneChannel) {
-            // Inspired by AnimateAsState.kt: let's poll the last value to avoid being one frame
-            // late.
-            val newKey = targetSceneChannel.tryReceive().getOrNull() ?: newKey
-            animateToScene(layoutImpl.state, newKey)
-        }
     }
 
     layoutImpl.Content(modifier)
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 0227aba94b53..7cc9d2623e9c 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
@@ -20,14 +20,11 @@ import androidx.activity.compose.BackHandler
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.DisposableEffect
-import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.key
 import androidx.compose.runtime.snapshots.SnapshotStateMap
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.drawWithContent
 import androidx.compose.ui.layout.LookaheadScope
 import androidx.compose.ui.layout.intermediateLayout
 import androidx.compose.ui.unit.Density
@@ -48,13 +45,12 @@ internal typealias MovableElementContent =
 
 @Stable
 internal class SceneTransitionLayoutImpl(
-    internal val state: SceneTransitionLayoutStateImpl,
-    internal var onChangeScene: (SceneKey) -> Unit,
+    internal val state: BaseSceneTransitionLayoutState,
     internal var density: Density,
     internal var edgeDetector: EdgeDetector,
     internal var transitionInterceptionThreshold: Float,
     builder: SceneTransitionLayoutScope.() -> Unit,
-    coroutineScope: CoroutineScope,
+    private val coroutineScope: CoroutineScope,
 ) {
     /**
      * The map of [Scene]s.
@@ -100,16 +96,6 @@ internal class SceneTransitionLayoutImpl(
                 ?: mutableMapOf<ValueKey, MutableMap<ElementKey?, SnapshotStateMap<SceneKey, *>>>()
                     .also { _sharedValues = it }
 
-    /**
-     * The scenes that are "ready", i.e. they were composed and fully laid-out at least once.
-     *
-     * Note that this map is *read* during composition, so it is a [SnapshotStateMap] to make sure
-     * that we recompose when modifications are made to this map.
-     *
-     * TODO(b/316901148): Remove this map.
-     */
-    private val readyScenes = SnapshotStateMap<SceneKey, Boolean>()
-
     private val horizontalGestureHandler: SceneGestureHandler
     private val verticalGestureHandler: SceneGestureHandler
 
@@ -244,49 +230,19 @@ internal class SceneTransitionLayoutImpl(
                 // TODO(b/290184746): Make sure that this works with SystemUI once we use
                 // SceneTransitionLayout in Flexiglass.
                 scene(state.transitionState.currentScene).userActions[Back]?.let { backScene ->
-                    BackHandler { onChangeScene(backScene) }
+                    BackHandler { with(state) { coroutineScope.onChangeScene(backScene) } }
                 }
 
                 Box {
                     scenesToCompose.fastForEach { scene ->
                         val key = scene.key
-                        key(key) {
-                            // Mark this scene as ready once it has been composed, laid out and
-                            // drawn the first time. We have to do this in a LaunchedEffect here
-                            // because DisposableEffect runs between composition and layout.
-                            LaunchedEffect(key) { readyScenes[key] = true }
-                            DisposableEffect(key) { onDispose { readyScenes.remove(key) } }
-
-                            scene.Content(
-                                Modifier.drawWithContent {
-                                    if (state.currentTransition == null) {
-                                        drawContent()
-                                    } else {
-                                        // Don't draw scenes that are not ready yet.
-                                        if (readyScenes.containsKey(key)) {
-                                            drawContent()
-                                        }
-                                    }
-                                }
-                            )
-                        }
+                        key(key) { scene.Content() }
                     }
                 }
             }
         }
     }
 
-    /**
-     * Return whether [transition] is ready, i.e. the elements of both scenes of the transition were
-     * laid out at least once.
-     */
-    internal fun isTransitionReady(transition: TransitionState.Transition): Boolean {
-        return readyScenes.containsKey(transition.fromScene) &&
-            readyScenes.containsKey(transition.toScene)
-    }
-
-    internal fun isSceneReady(scene: SceneKey): Boolean = readyScenes.containsKey(scene)
-
     internal fun setScenesTargetSizeForTest(size: IntSize) {
         scenes.values.forEach { it.targetSize = size }
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
index 0607aa148157..956e326dc03e 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
@@ -16,12 +16,23 @@
 
 package com.android.compose.animation.scene
 
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.Channel
 
-/** The state of a [SceneTransitionLayout]. */
+/**
+ * The state of a [SceneTransitionLayout].
+ *
+ * @see MutableSceneTransitionLayoutState
+ * @see updateSceneTransitionLayoutState
+ */
 @Stable
 sealed interface SceneTransitionLayoutState {
     /**
@@ -36,6 +47,9 @@ sealed interface SceneTransitionLayoutState {
     val currentTransition: TransitionState.Transition?
         get() = transitionState as? TransitionState.Transition
 
+    /** The [SceneTransitions] used when animating this state. */
+    val transitions: SceneTransitions
+
     /**
      * Whether we are transitioning. If [from] or [to] is empty, we will also check that they match
      * the scenes we are animating from and/or to.
@@ -46,9 +60,68 @@ sealed interface SceneTransitionLayoutState {
     fun isTransitioningBetween(scene: SceneKey, other: SceneKey): Boolean
 }
 
-/** Create a new [SceneTransitionLayoutState] that is currently idle at scene [currentScene]. */
-fun SceneTransitionLayoutState(currentScene: SceneKey): SceneTransitionLayoutState {
-    return SceneTransitionLayoutStateImpl(currentScene, SceneTransitions.Empty)
+/** A [SceneTransitionLayoutState] whose target scene can be imperatively set. */
+sealed interface MutableSceneTransitionLayoutState : SceneTransitionLayoutState {
+    /** The [SceneTransitions] used when animating this state. */
+    override var transitions: SceneTransitions
+
+    /**
+     * Set the target scene of this state to [targetScene].
+     *
+     * If [targetScene] is the same as the [currentScene][TransitionState.currentScene] of
+     * [transitionState], then nothing will happen and this will return `null`. Note that this means
+     * that this will also do nothing if the user is currently swiping from [targetScene] to another
+     * scene, or if we were already animating to [targetScene].
+     *
+     * If [targetScene] is different than the [currentScene][TransitionState.currentScene] of
+     * [transitionState], then this will animate to [targetScene]. The associated
+     * [TransitionState.Transition] will be returned and will be set as the current
+     * [transitionState] of this [MutableSceneTransitionLayoutState].
+     *
+     * Note that because a non-null [TransitionState.Transition] is returned does not mean that the
+     * transition will finish and that we will settle to [targetScene]. The returned transition
+     * might still be interrupted, for instance by another call to [setTargetScene] or by a user
+     * gesture.
+     *
+     * If [this] [CoroutineScope] is cancelled during the transition and that the transition was
+     * still active, then the [transitionState] of this [MutableSceneTransitionLayoutState] will be
+     * set to `TransitionState.Idle(targetScene)`.
+     *
+     * TODO(b/318794193): Add APIs to await() and cancel() any [TransitionState.Transition].
+     */
+    fun setTargetScene(
+        targetScene: SceneKey,
+        coroutineScope: CoroutineScope,
+    ): TransitionState.Transition?
+}
+
+/** Return a [MutableSceneTransitionLayoutState] initially idle at [initialScene]. */
+fun MutableSceneTransitionLayoutState(
+    initialScene: SceneKey,
+    transitions: SceneTransitions = SceneTransitions.Empty,
+): MutableSceneTransitionLayoutState {
+    return MutableSceneTransitionLayoutStateImpl(initialScene, transitions)
+}
+
+/**
+ * Sets up a [SceneTransitionLayoutState] and keeps it synced with [currentScene], [onChangeScene]
+ * and [transitions]. New transitions will automatically be started whenever [currentScene] is
+ * changed.
+ *
+ * @param currentScene the current scene
+ * @param onChangeScene a mutator that should set [currentScene] to the given scene when called.
+ *   This is called when the user commits a transition to a new scene because of a [UserAction], for
+ *   instance by triggering back navigation or by swiping to a new scene.
+ * @param transitions the definition of the transitions used to animate a change of scene.
+ */
+@Composable
+fun updateSceneTransitionLayoutState(
+    currentScene: SceneKey,
+    onChangeScene: (SceneKey) -> Unit,
+    transitions: SceneTransitions = SceneTransitions.Empty,
+): SceneTransitionLayoutState {
+    return remember { HoistedSceneTransitionLayoutScene(currentScene, transitions, onChangeScene) }
+        .apply { update(currentScene, onChangeScene, transitions) }
 }
 
 @Stable
@@ -109,13 +182,11 @@ sealed interface TransitionState {
     }
 }
 
-internal class SceneTransitionLayoutStateImpl(
-    initialScene: SceneKey,
-    internal var transitions: SceneTransitions,
-) : SceneTransitionLayoutState {
+internal abstract class BaseSceneTransitionLayoutState(initialScene: SceneKey) :
+    SceneTransitionLayoutState {
     override var transitionState: TransitionState by
         mutableStateOf(TransitionState.Idle(initialScene))
-        private set
+        protected set
 
     /**
      * The current [transformationSpec] associated to [transitionState]. Accessing this value makes
@@ -123,6 +194,14 @@ internal class SceneTransitionLayoutStateImpl(
      */
     internal var transformationSpec: TransformationSpecImpl = TransformationSpec.Empty
 
+    /**
+     * Called when the [current scene][TransitionState.currentScene] should be changed to [scene].
+     *
+     * When this is called, the source of truth for the current scene should be changed so that
+     * [transitionState] will animate and settle to [scene].
+     */
+    internal abstract fun CoroutineScope.onChangeScene(scene: SceneKey)
+
     override fun isTransitioning(from: SceneKey?, to: SceneKey?): Boolean {
         val transition = currentTransition ?: return false
         return transition.isTransitioning(from, to)
@@ -154,3 +233,62 @@ internal class SceneTransitionLayoutStateImpl(
         }
     }
 }
+
+/**
+ * A [SceneTransitionLayout] whose current scene/source of truth is hoisted (its current value comes
+ * from outside).
+ */
+internal class HoistedSceneTransitionLayoutScene(
+    initialScene: SceneKey,
+    override var transitions: SceneTransitions,
+    private var changeScene: (SceneKey) -> Unit,
+) : BaseSceneTransitionLayoutState(initialScene) {
+    private val targetSceneChannel = Channel<SceneKey>(Channel.CONFLATED)
+
+    override fun CoroutineScope.onChangeScene(scene: SceneKey) = changeScene(scene)
+
+    @Composable
+    fun update(
+        currentScene: SceneKey,
+        onChangeScene: (SceneKey) -> Unit,
+        transitions: SceneTransitions,
+    ) {
+        SideEffect {
+            this.changeScene = onChangeScene
+            this.transitions = transitions
+
+            targetSceneChannel.trySend(currentScene)
+        }
+
+        LaunchedEffect(targetSceneChannel) {
+            for (newKey in targetSceneChannel) {
+                // Inspired by AnimateAsState.kt: let's poll the last value to avoid being one frame
+                // late.
+                val newKey = targetSceneChannel.tryReceive().getOrNull() ?: newKey
+                animateToScene(layoutState = this@HoistedSceneTransitionLayoutScene, newKey)
+            }
+        }
+    }
+}
+
+/** A [MutableSceneTransitionLayoutState] that holds the value for the current scene. */
+internal class MutableSceneTransitionLayoutStateImpl(
+    initialScene: SceneKey,
+    override var transitions: SceneTransitions,
+) : MutableSceneTransitionLayoutState, BaseSceneTransitionLayoutState(initialScene) {
+    override fun setTargetScene(
+        targetScene: SceneKey,
+        coroutineScope: CoroutineScope
+    ): TransitionState.Transition? {
+        return with(this) {
+            coroutineScope.animateToScene(
+                layoutState = this@MutableSceneTransitionLayoutStateImpl,
+                target = targetScene,
+            )
+        }
+    }
+
+    override fun CoroutineScope.onChangeScene(scene: SceneKey) {
+        setTargetScene(scene, coroutineScope = this)
+    }
+}
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
index 35a5054cbd2a..c0de87abbfe8 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
@@ -16,7 +16,6 @@
 
 package com.android.compose.animation.scene
 
-import androidx.compose.animation.core.LinearEasing
 import androidx.compose.animation.core.tween
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.layout.Box
@@ -35,17 +34,11 @@ import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.layout.intermediateLayout
-import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.test.junit4.createComposeRule
-import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.Dp
-import androidx.compose.ui.unit.DpOffset
 import androidx.compose.ui.unit.dp
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.compose.test.subjects.DpOffsetSubject
-import com.android.compose.test.subjects.assertThat
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
@@ -263,8 +256,11 @@ class ElementTest {
 
         rule.setContent {
             SceneTransitionLayoutForTesting(
-                currentScene = currentScene,
-                onChangeScene = { currentScene = it },
+                state =
+                    updateSceneTransitionLayoutState(
+                        currentScene = currentScene,
+                        onChangeScene = { currentScene = it }
+                    ),
                 onLayoutImpl = { nullableLayoutImpl = it },
             ) {
                 scene(TestScenes.SceneA) { /* Nothing */}
@@ -428,8 +424,11 @@ class ElementTest {
 
         rule.setContent {
             SceneTransitionLayoutForTesting(
-                currentScene = TestScenes.SceneA,
-                onChangeScene = {},
+                state =
+                    updateSceneTransitionLayoutState(
+                        currentScene = TestScenes.SceneA,
+                        onChangeScene = {}
+                    ),
                 onLayoutImpl = { nullableLayoutImpl = it },
             ) {
                 scene(TestScenes.SceneA) { Box(Modifier.element(key)) }
@@ -478,8 +477,11 @@ class ElementTest {
             scrollScope = rememberCoroutineScope()
 
             SceneTransitionLayoutForTesting(
-                currentScene = TestScenes.SceneA,
-                onChangeScene = {},
+                state =
+                    updateSceneTransitionLayoutState(
+                        currentScene = TestScenes.SceneA,
+                        onChangeScene = {}
+                    ),
                 onLayoutImpl = { nullableLayoutImpl = it },
             ) {
                 scene(TestScenes.SceneA) {
@@ -565,86 +567,4 @@ class ElementTest {
             after { assertThat(fooCompositions).isEqualTo(1) }
         }
     }
-
-    @Test
-    fun sharedElementOffsetIsUpdatedEvenWhenNotPlaced() {
-        var nullableLayoutImpl: SceneTransitionLayoutImpl? = null
-        var density: Density? = null
-
-        fun layoutImpl() = nullableLayoutImpl ?: error("nullableLayoutImpl was not set")
-
-        fun density() = density ?: error("density was not set")
-
-        fun Offset.toDpOffset() = with(density()) { DpOffset(x.toDp(), y.toDp()) }
-
-        fun foo() = layoutImpl().elements[TestElements.Foo] ?: error("Foo not in elements map")
-
-        fun Element.lastSharedOffset() = lastSharedState.offset.toDpOffset()
-
-        fun Element.lastOffsetIn(scene: SceneKey) =
-            (sceneStates[scene] ?: error("$scene not in sceneValues map"))
-                .lastState
-                .offset
-                .toDpOffset()
-
-        rule.testTransition(
-            from = TestScenes.SceneA,
-            to = TestScenes.SceneB,
-            transitionLayout = { currentScene, onChangeScene ->
-                density = LocalDensity.current
-
-                SceneTransitionLayoutForTesting(
-                    currentScene = currentScene,
-                    onChangeScene = onChangeScene,
-                    onLayoutImpl = { nullableLayoutImpl = it },
-                    transitions =
-                        transitions {
-                            from(TestScenes.SceneA, to = TestScenes.SceneB) {
-                                spec = tween(durationMillis = 4 * 16, easing = LinearEasing)
-                            }
-                        }
-                ) {
-                    scene(TestScenes.SceneA) { Box(Modifier.element(TestElements.Foo)) }
-                    scene(TestScenes.SceneB) {
-                        Box(Modifier.offset(x = 40.dp, y = 80.dp).element(TestElements.Foo))
-                    }
-                }
-            }
-        ) {
-            val tolerance = DpOffsetSubject.DefaultTolerance
-
-            before {
-                val expected = DpOffset(0.dp, 0.dp)
-                assertThat(foo().lastSharedOffset()).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneA)).isWithin(tolerance).of(expected)
-            }
-
-            at(16) {
-                val expected = DpOffset(10.dp, 20.dp)
-                assertThat(foo().lastSharedOffset()).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneA)).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneB)).isWithin(tolerance).of(expected)
-            }
-
-            at(32) {
-                val expected = DpOffset(20.dp, 40.dp)
-                assertThat(foo().lastSharedOffset()).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneA)).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneB)).isWithin(tolerance).of(expected)
-            }
-
-            at(48) {
-                val expected = DpOffset(30.dp, 60.dp)
-                assertThat(foo().lastSharedOffset()).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneA)).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneB)).isWithin(tolerance).of(expected)
-            }
-
-            after {
-                val expected = DpOffset(40.dp, 80.dp)
-                assertThat(foo().lastSharedOffset()).isWithin(tolerance).of(expected)
-                assertThat(foo().lastOffsetIn(TestScenes.SceneB)).isWithin(tolerance).of(expected)
-            }
-        }
-    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
index 04b3f8a1dfe7..0f9b0249b93f 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
@@ -32,7 +32,7 @@ class ObservableTransitionStateTest {
 
     @Test
     fun testObservableTransitionState() = runTest {
-        val state = SceneTransitionLayoutState(TestScenes.SceneA)
+        lateinit var state: SceneTransitionLayoutState
 
         // Collect the current observable state into [observableState].
         // TODO(b/290184746): Use collectValues {} once it is extracted into a library that can be
@@ -58,12 +58,14 @@ class ObservableTransitionStateTest {
             from = TestScenes.SceneA,
             to = TestScenes.SceneB,
             transitionLayout = { currentScene, onChangeScene ->
-                SceneTransitionLayout(
-                    currentScene,
-                    onChangeScene,
-                    EmptyTestTransitions,
-                    state = state,
-                ) {
+                state =
+                    updateSceneTransitionLayoutState(
+                        currentScene,
+                        onChangeScene,
+                        EmptyTestTransitions
+                    )
+
+                SceneTransitionLayout(state = state) {
                     scene(TestScenes.SceneA) {}
                     scene(TestScenes.SceneB) {}
                 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneGestureHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneGestureHandlerTest.kt
index d9ce5191f3d9..066a3e45fb3c 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneGestureHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneGestureHandlerTest.kt
@@ -18,9 +18,6 @@ package com.android.compose.animation.scene
 
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.material3.Text
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.setValue
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
 import androidx.compose.ui.input.nestedscroll.NestedScrollSource
@@ -53,10 +50,8 @@ class SceneGestureHandlerTest {
     private class TestGestureScope(
         val coroutineScope: MonotonicClockTestScope,
     ) {
-        private var internalCurrentScene: SceneKey by mutableStateOf(SceneA)
-
         private val layoutState =
-            SceneTransitionLayoutStateImpl(internalCurrentScene, EmptyTestTransitions)
+            MutableSceneTransitionLayoutStateImpl(SceneA, EmptyTestTransitions)
 
         val mutableUserActionsA: MutableMap<UserAction, SceneKey> =
             mutableMapOf(Swipe.Up to SceneB, Swipe.Down to SceneC)
@@ -94,7 +89,6 @@ class SceneGestureHandlerTest {
         private val layoutImpl =
             SceneTransitionLayoutImpl(
                     state = layoutState,
-                    onChangeScene = { internalCurrentScene = it },
                     density = Density(1f),
                     edgeDetector = DefaultEdgeDetector,
                     transitionInterceptionThreshold = transitionInterceptionThreshold,
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
index 75dee47a91cd..48825fb88096 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
@@ -18,7 +18,11 @@ package com.android.compose.animation.scene
 
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.compose.test.runMonotonicClockTest
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -29,7 +33,7 @@ class SceneTransitionLayoutStateTest {
 
     @Test
     fun isTransitioningTo_idle() {
-        val state = SceneTransitionLayoutState(TestScenes.SceneA)
+        val state = MutableSceneTransitionLayoutStateImpl(TestScenes.SceneA, SceneTransitions.Empty)
 
         assertThat(state.isTransitioning()).isFalse()
         assertThat(state.isTransitioning(from = TestScenes.SceneA)).isFalse()
@@ -40,7 +44,7 @@ class SceneTransitionLayoutStateTest {
 
     @Test
     fun isTransitioningTo_transition() {
-        val state = SceneTransitionLayoutStateImpl(TestScenes.SceneA, SceneTransitions.Empty)
+        val state = MutableSceneTransitionLayoutStateImpl(TestScenes.SceneA, SceneTransitions.Empty)
         state.startTransition(transition(from = TestScenes.SceneA, to = TestScenes.SceneB))
 
         assertThat(state.isTransitioning()).isTrue()
@@ -50,4 +54,56 @@ class SceneTransitionLayoutStateTest {
         assertThat(state.isTransitioning(to = TestScenes.SceneA)).isFalse()
         assertThat(state.isTransitioning(from = TestScenes.SceneA, to = TestScenes.SceneB)).isTrue()
     }
+
+    @Test
+    fun setTargetScene_idleToSameScene() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(TestScenes.SceneA)
+        assertThat(state.setTargetScene(TestScenes.SceneA, coroutineScope = this)).isNull()
+    }
+
+    @Test
+    fun setTargetScene_idleToDifferentScene() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(TestScenes.SceneA)
+        val transition = state.setTargetScene(TestScenes.SceneB, coroutineScope = this)
+        assertThat(transition).isNotNull()
+        assertThat(state.transitionState).isEqualTo(transition)
+
+        testScheduler.advanceUntilIdle()
+        assertThat(state.transitionState).isEqualTo(TransitionState.Idle(TestScenes.SceneB))
+    }
+
+    @Test
+    fun setTargetScene_transitionToSameScene() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(TestScenes.SceneA)
+        assertThat(state.setTargetScene(TestScenes.SceneB, coroutineScope = this)).isNotNull()
+        assertThat(state.setTargetScene(TestScenes.SceneB, coroutineScope = this)).isNull()
+        testScheduler.advanceUntilIdle()
+        assertThat(state.transitionState).isEqualTo(TransitionState.Idle(TestScenes.SceneB))
+    }
+
+    @Test
+    fun setTargetScene_transitionToDifferentScene() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(TestScenes.SceneA)
+        assertThat(state.setTargetScene(TestScenes.SceneB, coroutineScope = this)).isNotNull()
+        assertThat(state.setTargetScene(TestScenes.SceneC, coroutineScope = this)).isNotNull()
+        testScheduler.advanceUntilIdle()
+        assertThat(state.transitionState).isEqualTo(TransitionState.Idle(TestScenes.SceneC))
+    }
+
+    @Test
+    fun setTargetScene_coroutineScopeCancelled() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(TestScenes.SceneA)
+
+        lateinit var transition: TransitionState.Transition
+        val job =
+            launch(start = CoroutineStart.UNDISPATCHED) {
+                transition = state.setTargetScene(TestScenes.SceneB, coroutineScope = this)!!
+            }
+        assertThat(state.transitionState).isEqualTo(transition)
+
+        // Cancelling the scope/job still sets the state to Idle(targetScene).
+        job.cancel()
+        testScheduler.advanceUntilIdle()
+        assertThat(state.transitionState).isEqualTo(TransitionState.Idle(TestScenes.SceneB))
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
index 649e4991434e..efaea71f8d2c 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
@@ -63,7 +63,7 @@ class SceneTransitionLayoutTest {
     }
 
     private var currentScene by mutableStateOf(TestScenes.SceneA)
-    private val layoutState = SceneTransitionLayoutState(currentScene)
+    private lateinit var layoutState: SceneTransitionLayoutState
 
     // We use createAndroidComposeRule() here and not createComposeRule() because we need an
     // activity for testBack().
@@ -72,10 +72,14 @@ class SceneTransitionLayoutTest {
     /** The content under test. */
     @Composable
     private fun TestContent() {
+        layoutState =
+            updateSceneTransitionLayoutState(
+                currentScene,
+                { currentScene = it },
+                EmptyTestTransitions
+            )
+
         SceneTransitionLayout(
-            currentScene,
-            { currentScene = it },
-            EmptyTestTransitions,
             state = layoutState,
             modifier = Modifier.size(LayoutSize),
         ) {
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
index 58d853ef5a00..1ec3c8ba2301 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
@@ -20,9 +20,6 @@ import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.size
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.platform.LocalViewConfiguration
@@ -58,18 +55,15 @@ class SwipeToSceneTest {
             get() = Offset(0f, (LayoutHeight / 2).toPx())
     }
 
-    private var currentScene by mutableStateOf(TestScenes.SceneA)
-    private val layoutState = SceneTransitionLayoutState(currentScene)
-
     @get:Rule val rule = createComposeRule()
 
+    private fun layoutState(initialScene: SceneKey = TestScenes.SceneA) =
+        MutableSceneTransitionLayoutState(initialScene, EmptyTestTransitions)
+
     /** The content under test. */
     @Composable
-    private fun TestContent() {
+    private fun TestContent(layoutState: SceneTransitionLayoutState) {
         SceneTransitionLayout(
-            currentScene,
-            { currentScene = it },
-            EmptyTestTransitions,
             state = layoutState,
             modifier = Modifier.size(LayoutWidth, LayoutHeight).testTag(TestElements.Foo.debugName),
         ) {
@@ -109,9 +103,11 @@ class SwipeToSceneTest {
         // The draggable touch slop, i.e. the min px distance a touch pointer must move before it is
         // detected as a drag event.
         var touchSlop = 0f
+
+        val layoutState = layoutState()
         rule.setContent {
             touchSlop = LocalViewConfiguration.current.touchSlop
-            TestContent()
+            TestContent(layoutState)
         }
 
         assertThat(layoutState.transitionState).isInstanceOf(TransitionState.Idle::class.java)
@@ -195,9 +191,10 @@ class SwipeToSceneTest {
         // The draggable touch slop, i.e. the min px distance a touch pointer must move before it is
         // detected as a drag event.
         var touchSlop = 0f
+        val layoutState = layoutState()
         rule.setContent {
             touchSlop = LocalViewConfiguration.current.touchSlop
-            TestContent()
+            TestContent(layoutState)
         }
 
         assertThat(layoutState.transitionState).isInstanceOf(TransitionState.Idle::class.java)
@@ -260,14 +257,14 @@ class SwipeToSceneTest {
     @Test
     fun multiPointerSwipe() {
         // Start at scene C.
-        currentScene = TestScenes.SceneC
+        val layoutState = layoutState(TestScenes.SceneC)
 
         // The draggable touch slop, i.e. the min px distance a touch pointer must move before it is
         // detected as a drag event.
         var touchSlop = 0f
         rule.setContent {
             touchSlop = LocalViewConfiguration.current.touchSlop
-            TestContent()
+            TestContent(layoutState)
         }
 
         assertThat(layoutState.transitionState).isInstanceOf(TransitionState.Idle::class.java)
@@ -299,14 +296,14 @@ class SwipeToSceneTest {
     @Test
     fun defaultEdgeSwipe() {
         // Start at scene C.
-        currentScene = TestScenes.SceneC
+        val layoutState = layoutState(TestScenes.SceneC)
 
         // The draggable touch slop, i.e. the min px distance a touch pointer must move before it is
         // detected as a drag event.
         var touchSlop = 0f
         rule.setContent {
             touchSlop = LocalViewConfiguration.current.touchSlop
-            TestContent()
+            TestContent(layoutState)
         }
 
         assertThat(layoutState.transitionState).isInstanceOf(TransitionState.Idle::class.java)
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/RunMonotonicClockTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/RunMonotonicClockTest.kt
index cb122dc8e25e..fbcd5b27836e 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/RunMonotonicClockTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/RunMonotonicClockTest.kt
@@ -13,7 +13,7 @@ import kotlinx.coroutines.withContext
  *
  * The [TestCoroutineScheduler] is passed to provide the functionality to wait for idle.
  */
-@ExperimentalTestApi
+@OptIn(ExperimentalTestApi::class)
 fun runMonotonicClockTest(block: suspend MonotonicClockTestScope.() -> Unit) = runTest {
     // We need a CoroutineScope (like a TestScope) to create a TestMonotonicFrameClock.
     withContext(TestMonotonicFrameClock(this)) {