Merge "Introduce overlays in SceneTransitionLayout (1/2)" into main
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
index fc4a8a5..1921624 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
@@ -121,6 +121,8 @@
                         )
                 }
             }
+        is TransitionState.Transition.OverlayTransition ->
+            TODO("b/359173565: Handle overlay transitions")
     }
 }
 
@@ -212,7 +214,8 @@
                             addView(view)
                         }
                     },
-                    // When the view changes (e.g. due to a theme change), this will be recomposed
+                    // When the view changes (e.g. due to a theme change), this will be
+                    // recomposed
                     // if needed and the new view will be attached to the FrameLayout here.
                     update = {
                         qsSceneAdapter.setState(state())
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt
index ea708a5..7eef5d6 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt
@@ -393,7 +393,8 @@
         transition: TransitionState.Transition?,
     ): T? {
         if (transition == null) {
-            return sharedValue[layoutImpl.state.transitionState.currentScene]
+            return sharedValue[content]
+                ?: sharedValue[layoutImpl.state.transitionState.currentScene]
         }
 
         val fromValue = sharedValue[transition.fromContent]
@@ -424,10 +425,12 @@
         val targetValues = sharedValue.targetValues
         val transition =
             if (element != null) {
-                layoutImpl.elements[element]?.stateByContent?.let { sceneStates ->
-                    layoutImpl.state.currentTransitions.fastLastOrNull { transition ->
-                        transition.fromContent in sceneStates || transition.toContent in sceneStates
-                    }
+                layoutImpl.elements[element]?.let { element ->
+                    elementState(
+                        layoutImpl.state.transitionStates,
+                        isInContent = { it in element.stateByContent },
+                    )
+                        as? TransitionState.Transition
                 }
             } else {
                 layoutImpl.state.currentTransitions.fastLastOrNull { transition ->
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 f2c2a36..8aa0690 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
@@ -43,12 +43,15 @@
     }
 
     return when (transitionState) {
-        is TransitionState.Idle -> {
+        is TransitionState.Idle,
+        is TransitionState.Transition.ShowOrHideOverlay,
+        is TransitionState.Transition.ReplaceOverlay -> {
             animateToScene(
                 layoutState,
                 target,
                 transitionKey,
                 isInitiatedByUserInput = false,
+                fromScene = transitionState.currentScene,
                 replacedTransition = null,
             )
         }
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 7dac2e4..6ea0285 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
@@ -46,7 +46,7 @@
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.round
 import androidx.compose.ui.util.fastCoerceIn
-import androidx.compose.ui.util.fastLastOrNull
+import androidx.compose.ui.util.fastForEachReversed
 import androidx.compose.ui.util.lerp
 import com.android.compose.animation.scene.content.Content
 import com.android.compose.animation.scene.content.state.TransitionState
@@ -145,8 +145,9 @@
     // layout/drawing.
     // TODO(b/341072461): Revert this and read the current transitions in ElementNode directly once
     // we can ensure that SceneTransitionLayoutImpl will compose new contents first.
-    val currentTransitions = layoutImpl.state.currentTransitions
-    return then(ElementModifier(layoutImpl, currentTransitions, content, key)).testTag(key.testTag)
+    val currentTransitionStates = layoutImpl.state.transitionStates
+    return then(ElementModifier(layoutImpl, currentTransitionStates, content, key))
+        .testTag(key.testTag)
 }
 
 /**
@@ -155,20 +156,21 @@
  */
 private data class ElementModifier(
     private val layoutImpl: SceneTransitionLayoutImpl,
-    private val currentTransitions: List<TransitionState.Transition>,
+    private val currentTransitionStates: List<TransitionState>,
     private val content: Content,
     private val key: ElementKey,
 ) : ModifierNodeElement<ElementNode>() {
-    override fun create(): ElementNode = ElementNode(layoutImpl, currentTransitions, content, key)
+    override fun create(): ElementNode =
+        ElementNode(layoutImpl, currentTransitionStates, content, key)
 
     override fun update(node: ElementNode) {
-        node.update(layoutImpl, currentTransitions, content, key)
+        node.update(layoutImpl, currentTransitionStates, content, key)
     }
 }
 
 internal class ElementNode(
     private var layoutImpl: SceneTransitionLayoutImpl,
-    private var currentTransitions: List<TransitionState.Transition>,
+    private var currentTransitionStates: List<TransitionState>,
     private var content: Content,
     private var key: ElementKey,
 ) : Modifier.Node(), DrawModifierNode, ApproachLayoutModifierNode, TraversableNode {
@@ -226,12 +228,12 @@
 
     fun update(
         layoutImpl: SceneTransitionLayoutImpl,
-        currentTransitions: List<TransitionState.Transition>,
+        currentTransitionStates: List<TransitionState>,
         content: Content,
         key: ElementKey,
     ) {
         check(layoutImpl == this.layoutImpl && content == this.content)
-        this.currentTransitions = currentTransitions
+        this.currentTransitionStates = currentTransitionStates
 
         removeNodeFromContentState()
 
@@ -287,31 +289,72 @@
         measurable: Measurable,
         constraints: Constraints,
     ): MeasureResult {
-        val transitions = currentTransitions
-        val transition = elementTransition(layoutImpl, element, transitions)
+        val elementState = elementState(layoutImpl, element, currentTransitionStates)
+        if (elementState == null) {
+            // If the element is not part of any transition, place it normally in its idle scene.
+            val currentState = currentTransitionStates.last()
+            val placeInThisContent =
+                elementContentWhenIdle(
+                    layoutImpl,
+                    currentState.currentScene,
+                    currentState.currentOverlays,
+                    isInContent = { it in element.stateByContent },
+                ) == content.key
 
-        // If this element is not supposed to be laid out now, either because it is not part of any
-        // ongoing transition or the other content of its transition is overscrolling, then lay out
-        // the element normally and don't place it.
+            return if (placeInThisContent) {
+                placeNormally(measurable, constraints)
+            } else {
+                doNotPlace(measurable, constraints)
+            }
+        }
+
+        val transition = elementState as? TransitionState.Transition
+
+        // If this element is not supposed to be laid out now because the other content of its
+        // transition is overscrolling, then lay out the element normally and don't place it.
         val overscrollScene = transition?.currentOverscrollSpec?.scene
         val isOtherSceneOverscrolling = overscrollScene != null && overscrollScene != content.key
-        val isNotPartOfAnyOngoingTransitions = transitions.isNotEmpty() && transition == null
-        if (isNotPartOfAnyOngoingTransitions || isOtherSceneOverscrolling) {
-            recursivelyClearPlacementValues()
-            stateInContent.lastSize = Element.SizeUnspecified
-
-            val placeable = measurable.measure(constraints)
-            return layout(placeable.width, placeable.height) { /* Do not place */ }
+        if (isOtherSceneOverscrolling) {
+            return doNotPlace(measurable, constraints)
         }
 
         val placeable =
             measure(layoutImpl, element, transition, stateInContent, measurable, constraints)
         stateInContent.lastSize = placeable.size()
-        return layout(placeable.width, placeable.height) { place(transition, placeable) }
+        return layout(placeable.width, placeable.height) { place(elementState, placeable) }
+    }
+
+    private fun ApproachMeasureScope.doNotPlace(
+        measurable: Measurable,
+        constraints: Constraints
+    ): MeasureResult {
+        recursivelyClearPlacementValues()
+        stateInContent.lastSize = Element.SizeUnspecified
+
+        val placeable = measurable.measure(constraints)
+        return layout(placeable.width, placeable.height) { /* Do not place */ }
+    }
+
+    private fun ApproachMeasureScope.placeNormally(
+        measurable: Measurable,
+        constraints: Constraints
+    ): MeasureResult {
+        val placeable = measurable.measure(constraints)
+        stateInContent.lastSize = placeable.size()
+        return layout(placeable.width, placeable.height) {
+            coordinates?.let {
+                with(layoutImpl.lookaheadScope) {
+                    stateInContent.lastOffset =
+                        lookaheadScopeCoordinates.localPositionOf(it, Offset.Zero)
+                }
+            }
+
+            placeable.place(0, 0)
+        }
     }
 
     private fun Placeable.PlacementScope.place(
-        transition: TransitionState.Transition?,
+        elementState: TransitionState,
         placeable: Placeable,
     ) {
         with(layoutImpl.lookaheadScope) {
@@ -321,11 +364,12 @@
                 coordinates ?: error("Element ${element.key} does not have any coordinates")
 
             // No need to place the element in this content if we don't want to draw it anyways.
-            if (!shouldPlaceElement(layoutImpl, content.key, element, transition)) {
+            if (!shouldPlaceElement(layoutImpl, content.key, element, elementState)) {
                 recursivelyClearPlacementValues()
                 return
             }
 
+            val transition = elementState as? TransitionState.Transition
             val currentOffset = lookaheadScopeCoordinates.localPositionOf(coords, Offset.Zero)
             val targetOffset =
                 computeValue(
@@ -391,11 +435,15 @@
                         return@placeWithLayer
                     }
 
-                    val transition = elementTransition(layoutImpl, element, currentTransitions)
-                    if (!shouldPlaceElement(layoutImpl, content.key, element, transition)) {
+                    val elementState = elementState(layoutImpl, element, currentTransitionStates)
+                    if (
+                        elementState == null ||
+                            !shouldPlaceElement(layoutImpl, content.key, element, elementState)
+                    ) {
                         return@placeWithLayer
                     }
 
+                    val transition = elementState as? TransitionState.Transition
                     alpha = elementAlpha(layoutImpl, element, transition, stateInContent)
                     compositingStrategy = CompositingStrategy.ModulateAlpha
                 }
@@ -425,7 +473,9 @@
     override fun ContentDrawScope.draw() {
         element.wasDrawnInAnyContent = true
 
-        val transition = elementTransition(layoutImpl, element, currentTransitions)
+        val transition =
+            elementState(layoutImpl, element, currentTransitionStates)
+                as? TransitionState.Transition
         val drawScale = getDrawScale(layoutImpl, element, transition, stateInContent)
         if (drawScale == Scale.Default) {
             drawContent()
@@ -468,21 +518,15 @@
     }
 }
 
-/**
- * The transition that we should consider for [element]. This is the last transition where one of
- * its contents contains the element.
- */
-private fun elementTransition(
+/** The [TransitionState] that we should consider for [element]. */
+private fun elementState(
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
-    transitions: List<TransitionState.Transition>,
-): TransitionState.Transition? {
-    val transition =
-        transitions.fastLastOrNull { transition ->
-            transition.fromContent in element.stateByContent ||
-                transition.toContent in element.stateByContent
-        }
+    transitionStates: List<TransitionState>,
+): TransitionState? {
+    val state = elementState(transitionStates, isInContent = { it in element.stateByContent })
 
+    val transition = state as? TransitionState.Transition
     val previousTransition = element.lastTransition
     element.lastTransition = transition
 
@@ -497,7 +541,66 @@
         }
     }
 
-    return transition
+    return state
+}
+
+internal inline fun elementState(
+    transitionStates: List<TransitionState>,
+    isInContent: (ContentKey) -> Boolean,
+): TransitionState? {
+    val lastState = transitionStates.last()
+    if (lastState is TransitionState.Idle) {
+        check(transitionStates.size == 1)
+        return lastState
+    }
+
+    // Find the last transition with a content that contains the element.
+    transitionStates.fastForEachReversed { state ->
+        val transition = state as TransitionState.Transition
+        if (isInContent(transition.fromContent) || isInContent(transition.toContent)) {
+            return transition
+        }
+    }
+
+    return null
+}
+
+internal inline fun elementContentWhenIdle(
+    layoutImpl: SceneTransitionLayoutImpl,
+    idle: TransitionState.Idle,
+    isInContent: (ContentKey) -> Boolean,
+): ContentKey {
+    val currentScene = idle.currentScene
+    val overlays = idle.currentOverlays
+    return elementContentWhenIdle(layoutImpl, currentScene, overlays, isInContent)
+}
+
+private inline fun elementContentWhenIdle(
+    layoutImpl: SceneTransitionLayoutImpl,
+    currentScene: SceneKey,
+    overlays: Set<OverlayKey>,
+    isInContent: (ContentKey) -> Boolean,
+): ContentKey {
+    if (overlays.isEmpty()) {
+        return currentScene
+    }
+
+    // Find the overlay with highest zIndex that contains the element.
+    // TODO(b/353679003): Should we cache enabledOverlays into a List<> to avoid a lot of
+    // allocations here?
+    var currentOverlay: OverlayKey? = null
+    for (overlay in overlays) {
+        if (
+            isInContent(overlay) &&
+                (currentOverlay == null ||
+                    (layoutImpl.overlay(overlay).zIndex >
+                        layoutImpl.overlay(currentOverlay).zIndex))
+        ) {
+            currentOverlay = overlay
+        }
+    }
+
+    return currentOverlay ?: currentScene
 }
 
 private fun prepareInterruption(
@@ -693,12 +796,20 @@
     layoutImpl: SceneTransitionLayoutImpl,
     content: ContentKey,
     element: Element,
-    transition: TransitionState.Transition?,
+    elementState: TransitionState,
 ): Boolean {
-    // Always place the element if we are idle.
-    if (transition == null) {
-        return true
-    }
+    val transition =
+        when (elementState) {
+            is TransitionState.Idle -> {
+                return content ==
+                    elementContentWhenIdle(
+                        layoutImpl,
+                        elementState,
+                        isInContent = { it in element.stateByContent },
+                    )
+            }
+            is TransitionState.Transition -> elementState
+        }
 
     // Don't place the element in this content if this content is not part of the current element
     // transition.
@@ -741,16 +852,12 @@
 
     val scenePicker = element.contentPicker
     val pickedScene =
-        when (transition) {
-            is TransitionState.Transition.ChangeCurrentScene -> {
-                scenePicker.contentDuringTransition(
-                    element = element,
-                    transition = transition,
-                    fromContentZIndex = layoutImpl.scene(transition.fromScene).zIndex,
-                    toContentZIndex = layoutImpl.scene(transition.toScene).zIndex,
-                )
-            }
-        }
+        scenePicker.contentDuringTransition(
+            element = element,
+            transition = transition,
+            fromContentZIndex = layoutImpl.content(transition.fromContent).zIndex,
+            toContentZIndex = layoutImpl.content(transition.toContent).zIndex,
+        )
 
     return pickedScene == content
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
index acb436e..3f8f5e7 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
@@ -63,6 +63,18 @@
     }
 }
 
+/** Key for an overlay. */
+class OverlayKey(
+    debugName: String,
+    identity: Any = Object(),
+) : ContentKey(debugName, identity) {
+    override val testTag: String = "overlay:$debugName"
+
+    override fun toString(): String {
+        return "OverlayKey(debugName=$debugName)"
+    }
+}
+
 /** Key for an element. */
 open class ElementKey(
     debugName: String,
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 63d51f9..715222c 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
@@ -26,7 +26,6 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.layout.Layout
 import androidx.compose.ui.unit.IntSize
-import androidx.compose.ui.util.fastLastOrNull
 import com.android.compose.animation.scene.content.Content
 import com.android.compose.animation.scene.content.state.TransitionState
 
@@ -58,6 +57,13 @@
     modifier: Modifier,
     content: @Composable ElementScope<MovableElementContentScope>.() -> Unit,
 ) {
+    check(key.contentPicker.contents.contains(sceneOrOverlay.key)) {
+        val elementName = key.debugName
+        val contentName = sceneOrOverlay.key.debugName
+        "MovableElement $elementName was composed in content $contentName but the " +
+            "MovableElementKey($elementName).contentPicker.contents does not contain $contentName"
+    }
+
     Box(modifier.element(layoutImpl, sceneOrOverlay, key)) {
         val contentScope = sceneOrOverlay.scope
         val boxScope = this
@@ -153,13 +159,20 @@
             // size* as its movable content, i.e. the same *size when idle*. During transitions,
             // this size will be used to interpolate the transition size, during the intermediate
             // layout pass.
+            //
+            // Important: Like in Modifier.element(), we read the transition states during
+            // composition then pass them to Layout to make sure that composition sees new states
+            // before layout and drawing.
+            val transitionStates = layoutImpl.state.transitionStates
             Layout { _, _ ->
                 // No need to measure or place anything.
                 val size =
                     placeholderContentSize(
-                        layoutImpl,
-                        contentKey,
-                        layoutImpl.elements.getValue(element),
+                        layoutImpl = layoutImpl,
+                        content = contentKey,
+                        element = layoutImpl.elements.getValue(element),
+                        elementKey = element,
+                        transitionStates = transitionStates,
                     )
                 layout(size.width, size.height) {}
             }
@@ -172,28 +185,43 @@
     content: ContentKey,
     element: MovableElementKey,
 ): Boolean {
-    val transitions = layoutImpl.state.currentTransitions
-    if (transitions.isEmpty()) {
-        // If we are idle, there is only one [scene] that is composed so we can compose our
-        // movable content here. We still check that [scene] is equal to the current idle scene, to
-        // make sure we only compose it there.
-        return layoutImpl.state.transitionState.currentScene == content
+    return when (
+        val elementState = movableElementState(element, layoutImpl.state.transitionStates)
+    ) {
+        null -> false
+        is TransitionState.Idle ->
+            movableElementContentWhenIdle(layoutImpl, element, elementState) == content
+        is TransitionState.Transition -> {
+            // During transitions, always compose movable elements in the scene picked by their
+            // content picker.
+            shouldPlaceOrComposeSharedElement(
+                layoutImpl,
+                content,
+                element,
+                elementState,
+            )
+        }
     }
+}
 
-    // The current transition for this element is the last transition in which either fromScene or
-    // toScene contains the element.
+private fun movableElementState(
+    element: MovableElementKey,
+    transitionStates: List<TransitionState>,
+): TransitionState? {
+    val content = element.contentPicker.contents
+    return elementState(transitionStates, isInContent = { content.contains(it) })
+}
+
+private fun movableElementContentWhenIdle(
+    layoutImpl: SceneTransitionLayoutImpl,
+    element: MovableElementKey,
+    elementState: TransitionState.Idle,
+): ContentKey {
     val contents = element.contentPicker.contents
-    val transition =
-        transitions.fastLastOrNull { transition ->
-            transition.fromContent in contents || transition.toContent in contents
-        } ?: return false
-
-    // Always compose movable elements in the scene picked by their scene picker.
-    return shouldPlaceOrComposeSharedElement(
+    return elementContentWhenIdle(
         layoutImpl,
-        content,
-        element,
-        transition,
+        elementState,
+        isInContent = { contents.contains(it) },
     )
 }
 
@@ -205,6 +233,8 @@
     layoutImpl: SceneTransitionLayoutImpl,
     content: ContentKey,
     element: Element,
+    elementKey: MovableElementKey,
+    transitionStates: List<TransitionState>,
 ): IntSize {
     // If the content of the movable element was already composed in this scene before, use that
     // target size.
@@ -213,20 +243,21 @@
         return targetValueInScene
     }
 
-    // This code is only run during transitions (otherwise the content would be composed and the
-    // placeholder would not), so it's ok to cast the state into a Transition directly.
-    val transition =
-        layoutImpl.state.transitionState as TransitionState.Transition.ChangeCurrentScene
+    // If the element content was already composed in the other overlay/scene, we use that
+    // target size assuming it doesn't change between scenes.
+    // TODO(b/317026105): Provide a way to give a hint size/content for cases where this is
+    // not true.
+    val otherContent =
+        when (val state = movableElementState(elementKey, transitionStates)) {
+            null -> return IntSize.Zero
+            is TransitionState.Idle -> movableElementContentWhenIdle(layoutImpl, elementKey, state)
+            is TransitionState.Transition ->
+                if (state.fromContent == content) state.toContent else state.fromContent
+        }
 
-    // If the content was already composed in the other scene, we use that target size assuming it
-    // doesn't change between scenes.
-    // TODO(b/317026105): Provide a way to give a hint size/content for cases where this is not
-    // true.
-    val otherScene =
-        if (transition.fromScene == content) transition.toScene else transition.fromScene
-    val targetValueInOtherScene = element.stateByContent[otherScene]?.targetSize
-    if (targetValueInOtherScene != null && targetValueInOtherScene != Element.SizeUnspecified) {
-        return targetValueInOtherScene
+    val targetValueInOtherContent = element.stateByContent[otherContent]?.targetSize
+    if (targetValueInOtherContent != null && targetValueInOtherContent != Element.SizeUnspecified) {
+        return targetValueInOtherContent
     }
 
     return IntSize.Zero
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
index 5071a7f..236e202 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
@@ -43,7 +43,9 @@
     fun currentScene(): Flow<SceneKey> {
         return when (this) {
             is Idle -> flowOf(currentScene)
-            is Transition -> currentScene
+            is Transition.ChangeCurrentScene -> currentScene
+            is Transition.ShowOrHideOverlay -> flowOf(currentScene)
+            is Transition.ReplaceOverlay -> flowOf(currentScene)
         }
     }
 
@@ -51,10 +53,11 @@
     data class Idle(val currentScene: SceneKey) : ObservableTransitionState
 
     /** There is a transition animating between two scenes. */
-    class Transition(
-        val fromScene: SceneKey,
-        val toScene: SceneKey,
-        val currentScene: Flow<SceneKey>,
+    sealed class Transition(
+        // TODO(b/353679003): Rename these to fromContent and toContent.
+        open val fromScene: ContentKey,
+        open val toScene: ContentKey,
+        val currentOverlays: Flow<Set<OverlayKey>>,
         val progress: Flow<Float>,
 
         /**
@@ -76,10 +79,10 @@
         val isUserInputOngoing: Flow<Boolean>,
 
         /** Current progress of the preview part of the transition */
-        val previewProgress: Flow<Float> = flowOf(0f),
+        val previewProgress: Flow<Float>,
 
         /** Whether the transition is currently in the preview stage or not */
-        val isInPreviewStage: Flow<Boolean> = flowOf(false),
+        val isInPreviewStage: Flow<Boolean>,
     ) : ObservableTransitionState {
         override fun toString(): String =
             """Transition
@@ -89,13 +92,109 @@
                 | isUserInputOngoing=$isUserInputOngoing
                 |)"""
                 .trimMargin()
+
+        /** A transition animating between [fromScene] and [toScene]. */
+        class ChangeCurrentScene(
+            override val fromScene: SceneKey,
+            override val toScene: SceneKey,
+            val currentScene: Flow<SceneKey>,
+            currentOverlays: Flow<Set<OverlayKey>>,
+            progress: Flow<Float>,
+            isInitiatedByUserInput: Boolean,
+            isUserInputOngoing: Flow<Boolean>,
+            previewProgress: Flow<Float>,
+            isInPreviewStage: Flow<Boolean>,
+        ) :
+            Transition(
+                fromScene,
+                toScene,
+                currentOverlays,
+                progress,
+                isInitiatedByUserInput,
+                isUserInputOngoing,
+                previewProgress,
+                isInPreviewStage,
+            )
+
+        /** The [overlay] is either showing from [currentScene] or hiding into [currentScene]. */
+        class ShowOrHideOverlay(
+            val overlay: OverlayKey,
+            fromContent: ContentKey,
+            toContent: ContentKey,
+            val currentScene: SceneKey,
+            currentOverlays: Flow<Set<OverlayKey>>,
+            progress: Flow<Float>,
+            isInitiatedByUserInput: Boolean,
+            isUserInputOngoing: Flow<Boolean>,
+            previewProgress: Flow<Float>,
+            isInPreviewStage: Flow<Boolean>,
+        ) :
+            Transition(
+                fromContent,
+                toContent,
+                currentOverlays,
+                progress,
+                isInitiatedByUserInput,
+                isUserInputOngoing,
+                previewProgress,
+                isInPreviewStage,
+            )
+
+        /** We are transitioning from [fromOverlay] to [toOverlay]. */
+        class ReplaceOverlay(
+            val fromOverlay: OverlayKey,
+            val toOverlay: OverlayKey,
+            val currentScene: SceneKey,
+            currentOverlays: Flow<Set<OverlayKey>>,
+            progress: Flow<Float>,
+            isInitiatedByUserInput: Boolean,
+            isUserInputOngoing: Flow<Boolean>,
+            previewProgress: Flow<Float>,
+            isInPreviewStage: Flow<Boolean>,
+        ) :
+            Transition(
+                fromOverlay,
+                toOverlay,
+                currentOverlays,
+                progress,
+                isInitiatedByUserInput,
+                isUserInputOngoing,
+                previewProgress,
+                isInPreviewStage,
+            )
+
+        companion object {
+            operator fun invoke(
+                fromScene: SceneKey,
+                toScene: SceneKey,
+                currentScene: Flow<SceneKey>,
+                progress: Flow<Float>,
+                isInitiatedByUserInput: Boolean,
+                isUserInputOngoing: Flow<Boolean>,
+                previewProgress: Flow<Float> = flowOf(0f),
+                isInPreviewStage: Flow<Boolean> = flowOf(false),
+                currentOverlays: Flow<Set<OverlayKey>> = flowOf(emptySet()),
+            ): ChangeCurrentScene {
+                return ChangeCurrentScene(
+                    fromScene,
+                    toScene,
+                    currentScene,
+                    currentOverlays,
+                    progress,
+                    isInitiatedByUserInput,
+                    isUserInputOngoing,
+                    previewProgress,
+                    isInPreviewStage,
+                )
+            }
+        }
     }
 
     fun isIdle(scene: SceneKey?): Boolean {
         return this is Idle && (scene == null || this.currentScene == scene)
     }
 
-    fun isTransitioning(from: SceneKey? = null, to: SceneKey? = null): Boolean {
+    fun isTransitioning(from: ContentKey? = null, to: ContentKey? = null): Boolean {
         return this is Transition &&
             (from == null || this.fromScene == from) &&
             (to == null || this.toScene == to)
@@ -112,15 +211,44 @@
             when (val state = transitionState) {
                 is TransitionState.Idle -> ObservableTransitionState.Idle(state.currentScene)
                 is TransitionState.Transition.ChangeCurrentScene -> {
-                    ObservableTransitionState.Transition(
+                    ObservableTransitionState.Transition.ChangeCurrentScene(
                         fromScene = state.fromScene,
                         toScene = state.toScene,
                         currentScene = snapshotFlow { state.currentScene },
+                        currentOverlays = flowOf(state.currentOverlays),
                         progress = snapshotFlow { state.progress },
                         isInitiatedByUserInput = state.isInitiatedByUserInput,
                         isUserInputOngoing = snapshotFlow { state.isUserInputOngoing },
                         previewProgress = snapshotFlow { state.previewProgress },
-                        isInPreviewStage = snapshotFlow { state.isInPreviewStage }
+                        isInPreviewStage = snapshotFlow { state.isInPreviewStage },
+                    )
+                }
+                is TransitionState.Transition.ShowOrHideOverlay -> {
+                    check(state.fromOrToScene == state.currentScene)
+                    ObservableTransitionState.Transition.ShowOrHideOverlay(
+                        overlay = state.overlay,
+                        fromContent = state.fromContent,
+                        toContent = state.toContent,
+                        currentScene = state.currentScene,
+                        currentOverlays = snapshotFlow { state.currentOverlays },
+                        progress = snapshotFlow { state.progress },
+                        isInitiatedByUserInput = state.isInitiatedByUserInput,
+                        isUserInputOngoing = snapshotFlow { state.isUserInputOngoing },
+                        previewProgress = snapshotFlow { state.previewProgress },
+                        isInPreviewStage = snapshotFlow { state.isInPreviewStage },
+                    )
+                }
+                is TransitionState.Transition.ReplaceOverlay -> {
+                    ObservableTransitionState.Transition.ReplaceOverlay(
+                        fromOverlay = state.fromOverlay,
+                        toOverlay = state.toOverlay,
+                        currentScene = state.currentScene,
+                        currentOverlays = snapshotFlow { state.currentOverlays },
+                        progress = snapshotFlow { state.progress },
+                        isInitiatedByUserInput = state.isInitiatedByUserInput,
+                        isUserInputOngoing = snapshotFlow { state.isUserInputOngoing },
+                        previewProgress = snapshotFlow { state.previewProgress },
+                        isInPreviewStage = snapshotFlow { state.isInPreviewStage },
                     )
                 }
             }
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 65a7367..aaa2546 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
@@ -49,7 +49,7 @@
  *   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.
+ * @param builder the configuration of the different scenes and overlays of this layout.
  */
 @Composable
 fun SceneTransitionLayout(
@@ -58,7 +58,7 @@
     swipeSourceDetector: SwipeSourceDetector = DefaultEdgeDetector,
     swipeDetector: SwipeDetector = DefaultSwipeDetector,
     @FloatRange(from = 0.0, to = 0.5) transitionInterceptionThreshold: Float = 0.05f,
-    scenes: SceneTransitionLayoutScope.() -> Unit,
+    builder: SceneTransitionLayoutScope.() -> Unit,
 ) {
     SceneTransitionLayoutForTesting(
         state,
@@ -67,7 +67,7 @@
         swipeDetector,
         transitionInterceptionThreshold,
         onLayoutImpl = null,
-        scenes,
+        builder,
     )
 }
 
@@ -86,6 +86,31 @@
         userActions: Map<UserAction, UserActionResult> = emptyMap(),
         content: @Composable ContentScope.() -> Unit,
     )
+
+    /**
+     * Add an overlay to this layout, identified by [key].
+     *
+     * Overlays are displayed above scenes and can be toggled using
+     * [MutableSceneTransitionLayoutState.showOverlay] and
+     * [MutableSceneTransitionLayoutState.hideOverlay].
+     *
+     * Overlays will have a maximum size that is the size of the layout without overlays, i.e. an
+     * overlay can be fillMaxSize() to match the layout size but it won't make the layout bigger.
+     *
+     * By default overlays are centered in their layout but they can be aligned differently using
+     * [alignment].
+     *
+     * Important: overlays must be defined after all scenes. Overlay order along the z-axis follows
+     * call order. Calling overlay(A) followed by overlay(B) will mean that overlay B renders
+     * after/above overlay A.
+     */
+    // TODO(b/353679003): Allow to specify user actions. When overlays are shown, the user actions
+    // of the top-most overlay in currentOverlays will be used.
+    fun overlay(
+        key: OverlayKey,
+        alignment: Alignment = Alignment.Center,
+        content: @Composable ContentScope.() -> Unit,
+    )
 }
 
 /**
@@ -239,7 +264,7 @@
     /**
      * Animate some value at the content level.
      *
-     * @param value the value of this shared value in the current scene.
+     * @param value the value of this shared value in the current content.
      * @param key the key of this shared value.
      * @param type the [SharedValueType] of this animated value.
      * @param canOverflow whether this value can overflow past the values it is interpolated
@@ -292,7 +317,7 @@
     /**
      * Animate some value associated to this element.
      *
-     * @param value the value of this shared value in the current scene.
+     * @param value the value of this shared value in the current content.
      * @param key the key of this shared value.
      * @param type the [SharedValueType] of this animated value.
      * @param canOverflow whether this value can overflow past the values it is interpolated
@@ -509,7 +534,7 @@
     swipeDetector: SwipeDetector = DefaultSwipeDetector,
     transitionInterceptionThreshold: Float = 0f,
     onLayoutImpl: ((SceneTransitionLayoutImpl) -> Unit)? = null,
-    scenes: SceneTransitionLayoutScope.() -> Unit,
+    builder: SceneTransitionLayoutScope.() -> Unit,
 ) {
     val density = LocalDensity.current
     val layoutDirection = LocalLayoutDirection.current
@@ -521,7 +546,7 @@
                 layoutDirection = layoutDirection,
                 swipeSourceDetector = swipeSourceDetector,
                 transitionInterceptionThreshold = transitionInterceptionThreshold,
-                builder = scenes,
+                builder = builder,
                 coroutineScope = coroutineScope,
             )
             .also { onLayoutImpl?.invoke(it) }
@@ -529,7 +554,7 @@
 
     // TODO(b/317014852): Move this into the SideEffect {} again once STLImpl.scenes is not a
     // SnapshotStateMap anymore.
-    layoutImpl.updateScenes(scenes, layoutDirection)
+    layoutImpl.updateContents(builder, layoutDirection)
 
     SideEffect {
         if (state != layoutImpl.state) {
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 392ff7e..21f11e4 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
@@ -18,10 +18,12 @@
 
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.BoxScope
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.key
 import androidx.compose.runtime.snapshots.SnapshotStateMap
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.layout.ApproachLayoutModifierNode
@@ -36,7 +38,9 @@
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.util.fastForEach
 import androidx.compose.ui.util.fastForEachReversed
+import androidx.compose.ui.zIndex
 import com.android.compose.animation.scene.content.Content
+import com.android.compose.animation.scene.content.Overlay
 import com.android.compose.animation.scene.content.Scene
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.ui.util.lerp
@@ -60,7 +64,17 @@
      *
      * TODO(b/317014852): Make this a normal MutableMap instead.
      */
-    internal val scenes = SnapshotStateMap<SceneKey, Scene>()
+    private val scenes = SnapshotStateMap<SceneKey, Scene>()
+
+    /**
+     * The map of [Overlays].
+     *
+     * Note: We lazily create this map to avoid instantiation an expensive SnapshotStateMap in the
+     * common case where there is no overlay in this layout.
+     */
+    private var _overlays: MutableMap<OverlayKey, Overlay>? = null
+    private val overlays
+        get() = _overlays ?: SnapshotStateMap<OverlayKey, Overlay>().also { _overlays = it }
 
     /**
      * The map of [Element]s.
@@ -119,7 +133,7 @@
         private set
 
     init {
-        updateScenes(builder, layoutDirection)
+        updateContents(builder, layoutDirection)
 
         // DraggableHandlerImpl must wait for the scenes to be initialized, in order to access the
         // current scene (required for SwipeTransition).
@@ -152,22 +166,32 @@
         return scenes[key] ?: error("Scene $key is not configured")
     }
 
+    internal fun sceneOrNull(key: SceneKey): Scene? = scenes[key]
+
+    internal fun overlay(key: OverlayKey): Overlay {
+        return overlays[key] ?: error("Overlay $key is not configured")
+    }
+
     internal fun content(key: ContentKey): Content {
         return when (key) {
             is SceneKey -> scene(key)
+            is OverlayKey -> overlay(key)
         }
     }
 
-    internal fun updateScenes(
+    internal fun updateContents(
         builder: SceneTransitionLayoutScope.() -> Unit,
         layoutDirection: LayoutDirection,
     ) {
-        // Keep a reference of the current scenes. After processing [builder], the scenes that were
-        // not configured will be removed.
+        // Keep a reference of the current contents. After processing [builder], the contents that
+        // were not configured will be removed.
         val scenesToRemove = scenes.keys.toMutableSet()
+        val overlaysToRemove =
+            if (_overlays == null) mutableSetOf() else overlays.keys.toMutableSet()
 
         // The incrementing zIndex of each scene.
         var zIndex = 0f
+        var overlaysDefined = false
 
         object : SceneTransitionLayoutScope {
                 override fun scene(
@@ -175,6 +199,8 @@
                     userActions: Map<UserAction, UserActionResult>,
                     content: @Composable ContentScope.() -> Unit,
                 ) {
+                    require(!overlaysDefined) { "all scenes must be defined before overlays" }
+
                     scenesToRemove.remove(key)
 
                     val resolvedUserActions =
@@ -199,10 +225,42 @@
 
                     zIndex++
                 }
+
+                override fun overlay(
+                    key: OverlayKey,
+                    alignment: Alignment,
+                    content: @Composable (ContentScope.() -> Unit)
+                ) {
+                    overlaysDefined = true
+                    overlaysToRemove.remove(key)
+
+                    val overlay = overlays[key]
+                    if (overlay != null) {
+                        // Update an existing overlay.
+                        overlay.content = content
+                        overlay.zIndex = zIndex
+                        overlay.alignment = alignment
+                    } else {
+                        // New overlay.
+                        overlays[key] =
+                            Overlay(
+                                key,
+                                this@SceneTransitionLayoutImpl,
+                                content,
+                                // TODO(b/353679003): Allow to specify user actions
+                                actions = emptyMap(),
+                                zIndex,
+                                alignment,
+                            )
+                    }
+
+                    zIndex++
+                }
             }
             .builder()
 
         scenesToRemove.forEach { scenes.remove(it) }
+        overlaysToRemove.forEach { overlays.remove(it) }
     }
 
     @Composable
@@ -220,8 +278,8 @@
                 lookaheadScope = this
 
                 BackHandler()
-
-                scenesToCompose().fastForEach { scene -> key(scene.key) { scene.Content() } }
+                Scenes()
+                Overlays()
             }
         }
     }
@@ -233,6 +291,11 @@
         PredictiveBackHandler(state, coroutineScope, targetSceneForBack)
     }
 
+    @Composable
+    private fun Scenes() {
+        scenesToCompose().fastForEach { scene -> key(scene.key) { scene.Content() } }
+    }
+
     private fun scenesToCompose(): List<Scene> {
         val transitions = state.currentTransitions
         return if (transitions.isEmpty()) {
@@ -253,15 +316,74 @@
                             maybeAdd(transition.toScene)
                             maybeAdd(transition.fromScene)
                         }
+                        is TransitionState.Transition.ShowOrHideOverlay ->
+                            maybeAdd(transition.fromOrToScene)
+                        is TransitionState.Transition.ReplaceOverlay -> {}
                     }
                 }
+
+                // Make sure that the current scene is always composed.
+                maybeAdd(transitions.last().currentScene)
             }
         }
     }
 
+    @Composable
+    private fun BoxScope.Overlays() {
+        val overlaysOrderedByZIndex = overlaysToComposeOrderedByZIndex()
+        if (overlaysOrderedByZIndex.isEmpty()) {
+            return
+        }
+
+        // We put the overlays inside a Box that is matching the layout size so that overlays are
+        // measured after all scenes and that their max size is the size of the layout without the
+        // overlays.
+        Box(Modifier.matchParentSize().zIndex(overlaysOrderedByZIndex.first().zIndex)) {
+            overlaysOrderedByZIndex.fastForEach { overlay ->
+                key(overlay.key) { overlay.Content(Modifier.align(overlay.alignment)) }
+            }
+        }
+    }
+
+    private fun overlaysToComposeOrderedByZIndex(): List<Overlay> {
+        if (_overlays == null) return emptyList()
+
+        val transitions = state.currentTransitions
+        return if (transitions.isEmpty()) {
+                state.transitionState.currentOverlays.map { overlay(it) }
+            } else {
+                buildList {
+                    val visited = mutableSetOf<OverlayKey>()
+                    fun maybeAdd(key: OverlayKey) {
+                        if (visited.add(key)) {
+                            add(overlay(key))
+                        }
+                    }
+
+                    transitions.fastForEach { transition ->
+                        when (transition) {
+                            is TransitionState.Transition.ChangeCurrentScene -> {}
+                            is TransitionState.Transition.ShowOrHideOverlay ->
+                                maybeAdd(transition.overlay)
+                            is TransitionState.Transition.ReplaceOverlay -> {
+                                maybeAdd(transition.fromOverlay)
+                                maybeAdd(transition.toOverlay)
+                            }
+                        }
+                    }
+
+                    // Make sure that all current overlays are composed.
+                    transitions.last().currentOverlays.forEach { maybeAdd(it) }
+                }
+            }
+            .sortedBy { it.zIndex }
+    }
+
     internal fun setScenesTargetSizeForTest(size: IntSize) {
         scenes.values.forEach { it.targetSize = size }
     }
+
+    internal fun overlaysOrNullForTest(): Map<OverlayKey, Overlay>? = _overlays
 }
 
 private data class LayoutElement(private val layoutImpl: SceneTransitionLayoutImpl) :
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 f37ded0..74cd136 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
@@ -39,6 +39,21 @@
 @Stable
 sealed interface SceneTransitionLayoutState {
     /**
+     * The current effective scene. If a new transition is triggered, it will start from this scene.
+     */
+    val currentScene: SceneKey
+
+    /**
+     * The current set of overlays. This represents the set of overlays that will be visible on
+     * screen once all [currentTransitions] are finished.
+     *
+     * @see MutableSceneTransitionLayoutState.showOverlay
+     * @see MutableSceneTransitionLayoutState.hideOverlay
+     * @see MutableSceneTransitionLayoutState.replaceOverlay
+     */
+    val currentOverlays: Set<OverlayKey>
+
+    /**
      * The current [TransitionState]. All values read here are backed by the Snapshot system.
      *
      * To observe those values outside of Compose/the Snapshot system, use
@@ -110,7 +125,50 @@
     ): TransitionState.Transition?
 
     /** Immediately snap to the given [scene]. */
-    fun snapToScene(scene: SceneKey)
+    fun snapToScene(
+        scene: SceneKey,
+        currentOverlays: Set<OverlayKey> = transitionState.currentOverlays,
+    )
+
+    /**
+     * Request to show [overlay] so that it animates in from [currentScene] and ends up being
+     * visible on screen.
+     *
+     * After this returns, this overlay will be included in [currentOverlays]. This does nothing if
+     * [overlay] is already in [currentOverlays].
+     */
+    fun showOverlay(
+        overlay: OverlayKey,
+        animationScope: CoroutineScope,
+        transitionKey: TransitionKey? = null,
+    )
+
+    /**
+     * Request to hide [overlay] so that it animates out to [currentScene] and ends up *not* being
+     * visible on screen.
+     *
+     * After this returns, this overlay will not be included in [currentOverlays]. This does nothing
+     * if [overlay] is not in [currentOverlays].
+     */
+    fun hideOverlay(
+        overlay: OverlayKey,
+        animationScope: CoroutineScope,
+        transitionKey: TransitionKey? = null,
+    )
+
+    /**
+     * Replace [from] by [to] so that [from] ends up not being visible on screen and [to] ends up
+     * being visible.
+     *
+     * This throws if [from] is not currently in [currentOverlays] or if [to] is already in
+     * [currentOverlays].
+     */
+    fun replaceOverlay(
+        from: OverlayKey,
+        to: OverlayKey,
+        animationScope: CoroutineScope,
+        transitionKey: TransitionKey? = null,
+    )
 }
 
 /**
@@ -128,6 +186,7 @@
 fun MutableSceneTransitionLayoutState(
     initialScene: SceneKey,
     transitions: SceneTransitions = SceneTransitions.Empty,
+    initialOverlays: Set<OverlayKey> = emptySet(),
     canChangeScene: (SceneKey) -> Boolean = { true },
     stateLinks: List<StateLink> = emptyList(),
     enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
@@ -135,6 +194,7 @@
     return MutableSceneTransitionLayoutStateImpl(
         initialScene,
         transitions,
+        initialOverlays,
         canChangeScene,
         stateLinks,
         enableInterruptions,
@@ -145,6 +205,7 @@
 internal class MutableSceneTransitionLayoutStateImpl(
     initialScene: SceneKey,
     override var transitions: SceneTransitions = transitions {},
+    initialOverlays: Set<OverlayKey> = emptySet(),
     internal val canChangeScene: (SceneKey) -> Boolean = { true },
     private val stateLinks: List<StateLink> = emptyList(),
 
@@ -158,13 +219,18 @@
      * 1. A list with a single [TransitionState.Idle] element, when we are idle.
      * 2. A list with one or more [TransitionState.Transition], when we are transitioning.
      */
-    @VisibleForTesting
     internal var transitionStates: List<TransitionState> by
-        mutableStateOf(listOf(TransitionState.Idle(initialScene)))
+        mutableStateOf(listOf(TransitionState.Idle(initialScene, initialOverlays)))
         private set
 
+    override val currentScene: SceneKey
+        get() = transitionState.currentScene
+
+    override val currentOverlays: Set<OverlayKey>
+        get() = transitionState.currentOverlays
+
     override val transitionState: TransitionState
-        get() = transitionStates.last()
+        get() = transitionStates[transitionStates.lastIndex]
 
     override val currentTransitions: List<TransitionState.Transition>
         get() {
@@ -233,6 +299,11 @@
     ) {
         checkThread()
 
+        // Set the current scene and overlays on the transition.
+        val currentState = transitionState
+        transition.currentSceneWhenTransitionStarted = currentState.currentScene
+        transition.currentOverlaysWhenTransitionStarted = currentState.currentOverlays
+
         // Compute the [TransformationSpec] when the transition starts.
         val fromScene = transition.fromScene
         val toScene = transition.toScene
@@ -356,6 +427,7 @@
                     transition.activeTransitionLinks[stateLink] = linkedTransition
                 }
             }
+            else -> error("transition links are not supported with overlays yet")
         }
     }
 
@@ -408,23 +480,28 @@
         // If all transitions are finished, we are idle.
         if (i == nStates) {
             check(finishedTransitions.isEmpty())
-            this.transitionStates = listOf(TransitionState.Idle(lastTransition.currentScene))
+            this.transitionStates =
+                listOf(
+                    TransitionState.Idle(
+                        lastTransition.currentScene,
+                        lastTransition.currentOverlays,
+                    )
+                )
         } else if (i > 0) {
             this.transitionStates = transitionStates.subList(fromIndex = i, toIndex = nStates)
         }
     }
 
-    override fun snapToScene(scene: SceneKey) {
+    override fun snapToScene(scene: SceneKey, currentOverlays: Set<OverlayKey>) {
         checkThread()
 
         // Force finish all transitions.
         while (currentTransitions.isNotEmpty()) {
-            val transition = transitionStates[0] as TransitionState.Transition
-            finishTransition(transition)
+            finishTransition(transitionStates[0] as TransitionState.Transition)
         }
 
         check(transitionStates.size == 1)
-        transitionStates = listOf(TransitionState.Idle(scene))
+        transitionStates = listOf(TransitionState.Idle(scene, currentOverlays))
     }
 
     private fun finishActiveTransitionLinks(transition: TransitionState.Transition) {
@@ -466,6 +543,57 @@
             false
         }
     }
+
+    override fun showOverlay(
+        overlay: OverlayKey,
+        animationScope: CoroutineScope,
+        transitionKey: TransitionKey?
+    ) {
+        checkThread()
+
+        // Overlay is already shown, do nothing.
+        if (overlay in transitionState.currentOverlays) {
+            return
+        }
+
+        // TODO(b/353679003): Animate the overlay instead of instantly snapping to an Idle state.
+        snapToScene(transitionState.currentScene, transitionState.currentOverlays + overlay)
+    }
+
+    override fun hideOverlay(
+        overlay: OverlayKey,
+        animationScope: CoroutineScope,
+        transitionKey: TransitionKey?
+    ) {
+        checkThread()
+
+        // Overlay is not shown, do nothing.
+        if (!transitionState.currentOverlays.contains(overlay)) {
+            return
+        }
+
+        // TODO(b/353679003): Animate the overlay instead of instantly snapping to an Idle state.
+        snapToScene(transitionState.currentScene, transitionState.currentOverlays - overlay)
+    }
+
+    override fun replaceOverlay(
+        from: OverlayKey,
+        to: OverlayKey,
+        animationScope: CoroutineScope,
+        transitionKey: TransitionKey?
+    ) {
+        checkThread()
+        require(from in currentOverlays) {
+            "Overlay ${from.debugName} is not shown so it can't be replaced by ${to.debugName}"
+        }
+        require(to !in currentOverlays) {
+            "Overlay ${to.debugName} is already shown so it can't replace ${from.debugName}"
+        }
+
+        // TODO(b/353679003): Animate from into to instead of hiding/showing the overlays
+        // separately.
+        snapToScene(transitionState.currentScene, transitionState.currentOverlays - from + to)
+    }
 }
 
 private const val TAG = "SceneTransitionLayoutState"
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
index 0f66804..9851b32 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
@@ -35,7 +35,7 @@
     }
 
     override fun SceneKey.targetSize(): IntSize? {
-        return layoutImpl.scenes[this]?.targetSize.takeIf { it != IntSize.Zero }
+        return layoutImpl.sceneOrNull(this)?.targetSize.takeIf { it != IntSize.Zero }
     }
 }
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Overlay.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Overlay.kt
new file mode 100644
index 0000000..ccec9e8
--- /dev/null
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Overlay.kt
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.compose.animation.scene.content
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.Stable
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
+import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.OverlayKey
+import com.android.compose.animation.scene.SceneTransitionLayoutImpl
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
+
+/** An overlay defined in a [SceneTransitionLayout]. */
+@Stable
+internal class Overlay(
+    override val key: OverlayKey,
+    layoutImpl: SceneTransitionLayoutImpl,
+    content: @Composable ContentScope.() -> Unit,
+    actions: Map<UserAction.Resolved, UserActionResult>,
+    zIndex: Float,
+    alignment: Alignment,
+) : Content(key, layoutImpl, content, actions, zIndex) {
+    var alignment by mutableStateOf(alignment)
+
+    override fun toString(): String {
+        return "Overlay(key=$key)"
+    }
+}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
index 22df34b..fdb019f 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
@@ -21,7 +21,11 @@
 import androidx.compose.animation.core.spring
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Stable
+import androidx.compose.runtime.derivedStateOf
+import androidx.compose.runtime.getValue
 import com.android.compose.animation.scene.ContentKey
+import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.OverscrollScope
 import com.android.compose.animation.scene.OverscrollSpecImpl
 import com.android.compose.animation.scene.ProgressVisibilityThreshold
@@ -49,9 +53,20 @@
      */
     val currentScene: SceneKey
 
+    /**
+     * The current set of overlays. This represents the set of overlays that will be visible on
+     * screen once all transitions are finished.
+     *
+     * @see MutableSceneTransitionLayoutState.showOverlay
+     * @see MutableSceneTransitionLayoutState.hideOverlay
+     * @see MutableSceneTransitionLayoutState.replaceOverlay
+     */
+    val currentOverlays: Set<OverlayKey>
+
     /** The scene [currentScene] is idle. */
     data class Idle(
         override val currentScene: SceneKey,
+        override val currentOverlays: Set<OverlayKey> = emptySet(),
     ) : TransitionState
 
     sealed class Transition(
@@ -69,7 +84,125 @@
 
             /** The transition that `this` transition is replacing, if any. */
             replacedTransition: Transition? = null,
-        ) : Transition(fromScene, toScene, replacedTransition)
+        ) : Transition(fromScene, toScene, replacedTransition) {
+            final override val currentOverlays: Set<OverlayKey>
+                get() {
+                    // The set of overlays does not change in a [ChangeCurrentScene] transition.
+                    return currentOverlaysWhenTransitionStarted
+                }
+        }
+
+        /**
+         * A transition that is animating one or more overlays and for which [currentOverlays] will
+         * change over the course of the transition.
+         */
+        sealed class OverlayTransition(
+            fromContent: ContentKey,
+            toContent: ContentKey,
+            replacedTransition: Transition?,
+        ) : Transition(fromContent, toContent, replacedTransition) {
+            final override val currentScene: SceneKey
+                get() {
+                    // The current scene does not change during overlay transitions.
+                    return currentSceneWhenTransitionStarted
+                }
+
+            // Note: We use deriveStateOf() so that the computed set is cached and reused when the
+            // inputs of the computations don't change, to avoid recomputing and allocating a new
+            // set every time currentOverlays is called (which is every frame and for each element).
+            final override val currentOverlays: Set<OverlayKey> by derivedStateOf {
+                computeCurrentOverlays()
+            }
+
+            protected abstract fun computeCurrentOverlays(): Set<OverlayKey>
+        }
+
+        /** The [overlay] is either showing from [fromOrToScene] or hiding into [fromOrToScene]. */
+        abstract class ShowOrHideOverlay(
+            val overlay: OverlayKey,
+            val fromOrToScene: SceneKey,
+            fromContent: ContentKey,
+            toContent: ContentKey,
+            replacedTransition: Transition? = null,
+        ) : OverlayTransition(fromContent, toContent, replacedTransition) {
+            /**
+             * Whether [overlay] is effectively shown. For instance, this will be `false` when
+             * starting a swipe transition to show [overlay] and will be `true` only once the swipe
+             * transition is committed.
+             */
+            protected abstract val isEffectivelyShown: Boolean
+
+            init {
+                check(
+                    (fromContent == fromOrToScene && toContent == overlay) ||
+                        (fromContent == overlay && toContent == fromOrToScene)
+                )
+            }
+
+            final override fun computeCurrentOverlays(): Set<OverlayKey> {
+                return if (isEffectivelyShown) {
+                    currentOverlaysWhenTransitionStarted + overlay
+                } else {
+                    currentOverlaysWhenTransitionStarted - overlay
+                }
+            }
+        }
+
+        /** We are transitioning from [fromOverlay] to [toOverlay]. */
+        abstract class ReplaceOverlay(
+            val fromOverlay: OverlayKey,
+            val toOverlay: OverlayKey,
+            replacedTransition: Transition? = null,
+        ) :
+            OverlayTransition(
+                fromContent = fromOverlay,
+                toContent = toOverlay,
+                replacedTransition,
+            ) {
+            /**
+             * The current effective overlay, either [fromOverlay] or [toOverlay]. For instance,
+             * this will be [fromOverlay] when starting a swipe transition that replaces
+             * [fromOverlay] by [toOverlay] and will [toOverlay] once the swipe transition is
+             * committed.
+             */
+            protected abstract val effectivelyShownOverlay: OverlayKey
+
+            init {
+                check(fromOverlay != toOverlay)
+            }
+
+            final override fun computeCurrentOverlays(): Set<OverlayKey> {
+                return when (effectivelyShownOverlay) {
+                    fromOverlay ->
+                        computeCurrentOverlays(include = fromOverlay, exclude = toOverlay)
+                    toOverlay -> computeCurrentOverlays(include = toOverlay, exclude = fromOverlay)
+                    else ->
+                        error(
+                            "effectivelyShownOverlay=$effectivelyShownOverlay, should be " +
+                                "equal to fromOverlay=$fromOverlay or toOverlay=$toOverlay"
+                        )
+                }
+            }
+
+            private fun computeCurrentOverlays(
+                include: OverlayKey,
+                exclude: OverlayKey
+            ): Set<OverlayKey> {
+                return buildSet {
+                    addAll(currentOverlaysWhenTransitionStarted)
+                    remove(exclude)
+                    add(include)
+                }
+            }
+        }
+
+        /**
+         * The current scene and overlays observed right when this transition started. These are set
+         * when this transition is started in
+         * [com.android.compose.animation.scene.MutableSceneTransitionLayoutStateImpl.startTransition].
+         */
+        internal lateinit var currentSceneWhenTransitionStarted: SceneKey
+        internal lateinit var currentOverlaysWhenTransitionStarted: Set<OverlayKey>
 
         /**
          * The key of this transition. This should usually be null, but it can be specified to use a
@@ -163,6 +296,11 @@
                 isTransitioning(from = other, to = content)
         }
 
+        /** Whether we are transitioning from or to [content]. */
+        fun isTransitioningFromOrTo(content: ContentKey): Boolean {
+            return fromContent == content || toContent == content
+        }
+
         /**
          * Force this transition to finish and animate to an [Idle] state.
          *
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt
index 01895c9..8ebb42a 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt
@@ -168,10 +168,7 @@
                 assertThat(lastValueInTo).isEqualTo(expectedValues)
             }
 
-            after {
-                assertThat(lastValueInFrom).isEqualTo(toValues)
-                assertThat(lastValueInTo).isEqualTo(toValues)
-            }
+            after { assertThat(lastValueInTo).isEqualTo(toValues) }
         }
     }
 
@@ -229,10 +226,7 @@
                 assertThat(lastValueInTo).isEqualTo(lerp(fromValues, toValues, fraction = 0.75f))
             }
 
-            after {
-                assertThat(lastValueInFrom).isEqualTo(toValues)
-                assertThat(lastValueInTo).isEqualTo(toValues)
-            }
+            after { assertThat(lastValueInTo).isEqualTo(toValues) }
         }
     }
 
@@ -288,10 +282,7 @@
                 assertThat(lastValueInTo).isEqualTo(expectedValues)
             }
 
-            after {
-                assertThat(lastValueInFrom).isEqualTo(toValues)
-                assertThat(lastValueInTo).isEqualTo(toValues)
-            }
+            after { assertThat(lastValueInTo).isEqualTo(toValues) }
         }
     }
 
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
index 72a16b7..25be3f9 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
@@ -65,19 +65,19 @@
         var layoutDirection = LayoutDirection.Rtl
             set(value) {
                 field = value
-                layoutImpl.updateScenes(scenesBuilder, layoutDirection)
+                layoutImpl.updateContents(scenesBuilder, layoutDirection)
             }
 
         var mutableUserActionsA = mapOf(Swipe.Up to SceneB, Swipe.Down to SceneC)
             set(value) {
                 field = value
-                layoutImpl.updateScenes(scenesBuilder, layoutDirection)
+                layoutImpl.updateContents(scenesBuilder, layoutDirection)
             }
 
         var mutableUserActionsB = mapOf(Swipe.Up to SceneC, Swipe.Down to SceneA)
             set(value) {
                 field = value
-                layoutImpl.updateScenes(scenesBuilder, layoutDirection)
+                layoutImpl.updateContents(scenesBuilder, layoutDirection)
             }
 
         private val scenesBuilder: SceneTransitionLayoutScope.() -> Unit = {
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
index b7f50fd..a549d03 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
@@ -106,7 +106,7 @@
                 rule
                     .onNode(
                         hasText("count: 3") and
-                            hasParent(isElement(TestElements.Foo, scene = SceneA))
+                            hasParent(isElement(TestElements.Foo, content = SceneA))
                     )
                     .assertExists()
                     .assertIsNotDisplayed()
@@ -114,7 +114,7 @@
                 rule
                     .onNode(
                         hasText("count: 0") and
-                            hasParent(isElement(TestElements.Foo, scene = SceneB))
+                            hasParent(isElement(TestElements.Foo, content = SceneB))
                     )
                     .assertIsDisplayed()
                     .assertSizeIsEqualTo(75.dp, 75.dp)
@@ -213,7 +213,7 @@
                 rule
                     .onNode(
                         hasText("count: 3") and
-                            hasParent(isElement(TestElements.Foo, scene = SceneA))
+                            hasParent(isElement(TestElements.Foo, content = SceneA))
                     )
                     .assertIsDisplayed()
                     .assertSizeIsEqualTo(75.dp, 75.dp)
@@ -234,7 +234,7 @@
                 rule
                     .onNode(
                         hasText("count: 3") and
-                            hasParent(isElement(TestElements.Foo, scene = SceneB))
+                            hasParent(isElement(TestElements.Foo, content = SceneB))
                     )
                     .assertIsDisplayed()
 
@@ -324,7 +324,7 @@
     fun movableElementScopeExtendsBoxScope() {
         val key = MovableElementKey("Foo", contents = setOf(SceneA))
         rule.setContent {
-            TestContentScope {
+            TestContentScope(currentScene = SceneA) {
                 MovableElement(key, Modifier.size(200.dp)) {
                     content {
                         Box(Modifier.testTag("bottomEnd").align(Alignment.BottomEnd))
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
new file mode 100644
index 0000000..d4391e0
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
@@ -0,0 +1,319 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.compose.animation.scene
+
+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.rememberCoroutineScope
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.platform.testTag
+import androidx.compose.ui.test.assertIsDisplayed
+import androidx.compose.ui.test.assertIsNotDisplayed
+import androidx.compose.ui.test.assertPositionInRootIsEqualTo
+import androidx.compose.ui.test.hasTestTag
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.onNodeWithTag
+import androidx.compose.ui.unit.dp
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.compose.animation.scene.TestOverlays.OverlayA
+import com.android.compose.animation.scene.TestOverlays.OverlayB
+import com.android.compose.animation.scene.TestScenes.SceneA
+import com.android.compose.test.assertSizeIsEqualTo
+import kotlinx.coroutines.CoroutineScope
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class OverlayTest {
+    @get:Rule val rule = createComposeRule()
+
+    @Composable
+    private fun ContentScope.Foo() {
+        Box(Modifier.element(TestElements.Foo).size(100.dp))
+    }
+
+    @Test
+    fun showThenHideOverlay() {
+        val state = rule.runOnUiThread { MutableSceneTransitionLayoutState(SceneA) }
+        lateinit var coroutineScope: CoroutineScope
+        rule.setContent {
+            coroutineScope = rememberCoroutineScope()
+            SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                scene(SceneA) { Box(Modifier.fillMaxSize()) { Foo() } }
+                overlay(OverlayA) { Foo() }
+            }
+        }
+
+        // Initial state: overlay A is not shown, so Foo is displayed at the top left in scene A.
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertIsDisplayed()
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+        rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+
+        // Show overlay A: Foo is now centered on screen and placed in overlay A. It is not placed
+        // in scene A.
+        rule.runOnUiThread { state.showOverlay(OverlayA, coroutineScope) }
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertExists()
+            .assertIsNotDisplayed()
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+
+        // Hide overlay A: back to initial state, top-left in scene A.
+        rule.runOnUiThread { state.hideOverlay(OverlayA, coroutineScope) }
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertIsDisplayed()
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+        rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+    }
+
+    @Test
+    fun multipleOverlays() {
+        val state = rule.runOnUiThread { MutableSceneTransitionLayoutState(SceneA) }
+        lateinit var coroutineScope: CoroutineScope
+        rule.setContent {
+            coroutineScope = rememberCoroutineScope()
+            SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                scene(SceneA) { Box(Modifier.fillMaxSize()) { Foo() } }
+                overlay(OverlayA) { Foo() }
+                overlay(OverlayB) { Foo() }
+            }
+        }
+
+        // Initial state.
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertIsDisplayed()
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+        rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+        rule.onNode(isElement(TestElements.Foo, content = OverlayB)).assertDoesNotExist()
+
+        // Show overlay A.
+        rule.runOnUiThread { state.showOverlay(OverlayA, coroutineScope) }
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertExists()
+            .assertIsNotDisplayed()
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+        rule.onNode(isElement(TestElements.Foo, content = OverlayB)).assertDoesNotExist()
+
+        // Replace overlay A by overlay B.
+        rule.runOnUiThread { state.replaceOverlay(OverlayA, OverlayB, coroutineScope) }
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertExists()
+            .assertIsNotDisplayed()
+        rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayB))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+
+        // Show overlay A: Foo is still placed in B because it has a higher zIndex, but it now
+        // exists in A as well.
+        rule.runOnUiThread { state.showOverlay(OverlayA, coroutineScope) }
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertExists()
+            .assertIsNotDisplayed()
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayA))
+            .assertExists()
+            .assertIsNotDisplayed()
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayB))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+
+        // Hide overlay B.
+        rule.runOnUiThread { state.hideOverlay(OverlayB, coroutineScope) }
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertExists()
+            .assertIsNotDisplayed()
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+        rule.onNode(isElement(TestElements.Foo, content = OverlayB)).assertDoesNotExist()
+
+        // Hide overlay A.
+        rule.runOnUiThread { state.hideOverlay(OverlayA, coroutineScope) }
+        rule
+            .onNode(isElement(TestElements.Foo, content = SceneA))
+            .assertIsDisplayed()
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+        rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+        rule.onNode(isElement(TestElements.Foo, content = OverlayB)).assertDoesNotExist()
+    }
+
+    @Test
+    fun movableElement() {
+        val key = MovableElementKey("MovableBar", contents = setOf(SceneA, OverlayA, OverlayB))
+        val elementChildTag = "elementChildTag"
+
+        fun elementChild(content: ContentKey) = hasTestTag(elementChildTag) and inContent(content)
+
+        @Composable
+        fun ContentScope.MovableBar() {
+            MovableElement(key, Modifier) {
+                content { Box(Modifier.testTag(elementChildTag).size(100.dp)) }
+            }
+        }
+
+        val state = rule.runOnUiThread { MutableSceneTransitionLayoutState(SceneA) }
+        lateinit var coroutineScope: CoroutineScope
+        rule.setContent {
+            coroutineScope = rememberCoroutineScope()
+            SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                scene(SceneA) { Box(Modifier.fillMaxSize()) { MovableBar() } }
+                overlay(OverlayA) { MovableBar() }
+                overlay(OverlayB) { MovableBar() }
+            }
+        }
+
+        // Initial state.
+        rule
+            .onNode(elementChild(content = SceneA))
+            .assertIsDisplayed()
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+        rule.onNode(elementChild(content = OverlayA)).assertDoesNotExist()
+        rule.onNode(elementChild(content = OverlayB)).assertDoesNotExist()
+
+        // Show overlay A: movable element child only exists (is only composed) in overlay A.
+        rule.runOnUiThread { state.showOverlay(OverlayA, coroutineScope) }
+        rule.onNode(elementChild(content = SceneA)).assertDoesNotExist()
+        rule
+            .onNode(elementChild(content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+        rule.onNode(elementChild(content = OverlayB)).assertDoesNotExist()
+
+        // Replace overlay A by overlay B: element child is only in overlay B.
+        rule.runOnUiThread { state.replaceOverlay(OverlayA, OverlayB, coroutineScope) }
+        rule.onNode(elementChild(content = SceneA)).assertDoesNotExist()
+        rule.onNode(elementChild(content = OverlayA)).assertDoesNotExist()
+        rule
+            .onNode(elementChild(content = OverlayB))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+
+        // Show overlay A: element child still only exists in overlay B because it has a higher
+        // zIndex.
+        rule.runOnUiThread { state.showOverlay(OverlayA, coroutineScope) }
+        rule.onNode(elementChild(content = SceneA)).assertDoesNotExist()
+        rule.onNode(elementChild(content = OverlayA)).assertDoesNotExist()
+        rule
+            .onNode(elementChild(content = OverlayB))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+
+        // Hide overlay B: element child is in overlay A.
+        rule.runOnUiThread { state.hideOverlay(OverlayB, coroutineScope) }
+        rule.onNode(elementChild(content = SceneA)).assertDoesNotExist()
+        rule
+            .onNode(elementChild(content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+        rule.onNode(elementChild(content = OverlayB)).assertDoesNotExist()
+
+        // Hide overlay A: element child is in scene A.
+        rule.runOnUiThread { state.hideOverlay(OverlayA, coroutineScope) }
+        rule
+            .onNode(elementChild(content = SceneA))
+            .assertIsDisplayed()
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+        rule.onNode(elementChild(content = OverlayA)).assertDoesNotExist()
+        rule.onNode(elementChild(content = OverlayB)).assertDoesNotExist()
+    }
+
+    @Test
+    fun overlayAlignment() {
+        val state =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(SceneA, initialOverlays = setOf(OverlayA))
+            }
+        var alignment by mutableStateOf(Alignment.Center)
+        rule.setContent {
+            SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                scene(SceneA) { Box(Modifier.fillMaxSize()) { Foo() } }
+                overlay(OverlayA, alignment) { Foo() }
+            }
+        }
+
+        // Initial state: 100x100dp centered in 200x200dp layout.
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(50.dp, 50.dp)
+
+        // BottomStart.
+        alignment = Alignment.BottomStart
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(0.dp, 100.dp)
+
+        // TopEnd.
+        alignment = Alignment.TopEnd
+        rule
+            .onNode(isElement(TestElements.Foo, content = OverlayA))
+            .assertSizeIsEqualTo(100.dp)
+            .assertPositionInRootIsEqualTo(100.dp, 0.dp)
+    }
+
+    @Test
+    fun overlayMaxSizeIsCurrentSceneSize() {
+        val state =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(SceneA, initialOverlays = setOf(OverlayA))
+            }
+
+        val contentTag = "overlayContent"
+        rule.setContent {
+            SceneTransitionLayout(state) {
+                scene(SceneA) { Box(Modifier.size(100.dp)) { Foo() } }
+                overlay(OverlayA) { Box(Modifier.testTag(contentTag).fillMaxSize()) }
+            }
+        }
+
+        // Max overlay size is the size of the layout without overlays, not the (max) possible size
+        // of the layout.
+        rule.onNodeWithTag(contentTag).assertSizeIsEqualTo(100.dp)
+    }
+}
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 e97c27e..b8e13da 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
@@ -500,4 +500,19 @@
         assertThat(keyInB).isEqualTo(SceneB)
         assertThat(keyInC).isEqualTo(SceneC)
     }
+
+    @Test
+    fun overlaysMapIsNotAllocatedWhenNoOverlayIsDefined() {
+        lateinit var layoutImpl: SceneTransitionLayoutImpl
+        rule.setContent {
+            SceneTransitionLayoutForTesting(
+                remember { MutableSceneTransitionLayoutState(SceneA) },
+                onLayoutImpl = { layoutImpl = it },
+            ) {
+                scene(SceneA) { Box(Modifier.fillMaxSize()) }
+            }
+        }
+
+        assertThat(layoutImpl.overlaysOrNullForTest()).isNull()
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestContentScope.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestContentScope.kt
index 00adefb..5cccfb1 100644
--- a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestContentScope.kt
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestContentScope.kt
@@ -26,9 +26,9 @@
 @Composable
 fun TestContentScope(
     modifier: Modifier = Modifier,
+    currentScene: SceneKey = remember { SceneKey("current") },
     content: @Composable ContentScope.() -> Unit,
 ) {
-    val currentScene = remember { SceneKey("current") }
     val state = remember { MutableSceneTransitionLayoutState(currentScene) }
     SceneTransitionLayout(state, modifier) { scene(currentScene, content = content) }
 }
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestMatchers.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestMatchers.kt
index 6d063a0..22450d3 100644
--- a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestMatchers.kt
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestMatchers.kt
@@ -20,11 +20,16 @@
 import androidx.compose.ui.test.hasAnyAncestor
 import androidx.compose.ui.test.hasTestTag
 
-/** A [SemanticsMatcher] that matches [element], optionally restricted to scene [scene]. */
-fun isElement(element: ElementKey, scene: SceneKey? = null): SemanticsMatcher {
-    return if (scene == null) {
+/** A [SemanticsMatcher] that matches [element], optionally restricted to content [content]. */
+fun isElement(element: ElementKey, content: ContentKey? = null): SemanticsMatcher {
+    return if (content == null) {
         hasTestTag(element.testTag)
     } else {
-        hasTestTag(element.testTag) and hasAnyAncestor(hasTestTag(scene.testTag))
+        hasTestTag(element.testTag) and inContent(content)
     }
 }
+
+/** A [SemanticsMatcher] that matches anything inside [content]. */
+fun inContent(content: ContentKey): SemanticsMatcher {
+    return hasAnyAncestor(hasTestTag(content.testTag))
+}
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestValues.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestValues.kt
index b83705a..f39dd67 100644
--- a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestValues.kt
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestValues.kt
@@ -21,7 +21,7 @@
 import androidx.compose.animation.core.snap
 import androidx.compose.animation.core.tween
 
-/** Scenes keys that can be reused by tests. */
+/** Scene keys that can be reused by tests. */
 object TestScenes {
     val SceneA = SceneKey("SceneA")
     val SceneB = SceneKey("SceneB")
@@ -29,6 +29,12 @@
     val SceneD = SceneKey("SceneD")
 }
 
+/** Overlay keys that can be reused by tests. */
+object TestOverlays {
+    val OverlayA = OverlayKey("OverlayA")
+    val OverlayB = OverlayKey("OverlayB")
+}
+
 /** Element keys that can be reused by tests. */
 object TestElements {
     val Foo = ElementKey("Foo")
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
index 163b9b0..c633816 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -495,7 +495,9 @@
     private fun getCurrentSceneInUi(): SceneKey {
         return when (val state = transitionState.value) {
             is ObservableTransitionState.Idle -> state.currentScene
-            is ObservableTransitionState.Transition -> state.fromScene
+            is ObservableTransitionState.Transition.ChangeCurrentScene -> state.fromScene
+            is ObservableTransitionState.Transition.ShowOrHideOverlay -> state.currentScene
+            is ObservableTransitionState.Transition.ReplaceOverlay -> state.currentScene
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
index 2d510e1..ea61bd3 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
@@ -118,8 +118,11 @@
         get() =
             when (this) {
                 is ObservableTransitionState.Idle -> currentScene.canBeOccluded
-                is ObservableTransitionState.Transition ->
+                is ObservableTransitionState.Transition.ChangeCurrentScene ->
                     fromScene.canBeOccluded && toScene.canBeOccluded
+                is ObservableTransitionState.Transition.ReplaceOverlay,
+                is ObservableTransitionState.Transition.ShowOrHideOverlay ->
+                    TODO("b/359173565: Handle overlay transitions")
             }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
index 75cb017d..1b9c346 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
@@ -109,7 +109,15 @@
      */
     val transitioningTo: StateFlow<SceneKey?> =
         transitionState
-            .map { state -> (state as? ObservableTransitionState.Transition)?.toScene }
+            .map { state ->
+                when (state) {
+                    is ObservableTransitionState.Idle -> null
+                    is ObservableTransitionState.Transition.ChangeCurrentScene -> state.toScene
+                    is ObservableTransitionState.Transition.ShowOrHideOverlay,
+                    is ObservableTransitionState.Transition.ReplaceOverlay ->
+                        TODO("b/359173565: Handle overlay transitions")
+                }
+            }
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.WhileSubscribed(),
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt
index c6f51b3..ec743ba 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt
@@ -112,7 +112,7 @@
                 // It
                 // happens only when unlocking or when dismissing a dismissible lockscreen.
                 val isTransitioningAwayFromKeyguard =
-                    transitionState is ObservableTransitionState.Transition &&
+                    transitionState is ObservableTransitionState.Transition.ChangeCurrentScene &&
                         transitionState.fromScene.isKeyguard() &&
                         transitionState.toScene == Scenes.Gone
 
@@ -120,7 +120,7 @@
                 val isCurrentSceneShade = currentScene.isShade()
                 // This is true when moving into one of the shade scenes when a non-shade scene.
                 val isTransitioningToShade =
-                    transitionState is ObservableTransitionState.Transition &&
+                    transitionState is ObservableTransitionState.Transition.ChangeCurrentScene &&
                         !transitionState.fromScene.isShade() &&
                         transitionState.toScene.isShade()
 
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
index 8006e94..7d67121 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
@@ -64,7 +64,7 @@
                             0f
                         }
                     )
-                is ObservableTransitionState.Transition ->
+                is ObservableTransitionState.Transition.ChangeCurrentScene ->
                     when {
                         state.fromScene == Scenes.Gone ->
                             if (state.toScene.isExpandable()) {
@@ -88,6 +88,9 @@
                             }
                         else -> flowOf(1f)
                     }
+                is ObservableTransitionState.Transition.ShowOrHideOverlay,
+                is ObservableTransitionState.Transition.ReplaceOverlay ->
+                    TODO("b/359173565: Handle overlay transitions")
             }
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
index 399b8d0..0e984cf 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
@@ -74,7 +74,7 @@
         }
 
     private fun expandFractionForTransition(
-        state: ObservableTransitionState.Transition,
+        state: ObservableTransitionState.Transition.ChangeCurrentScene,
         shadeExpansion: Float,
         shadeMode: ShadeMode,
         qsExpansion: Float,
@@ -113,7 +113,7 @@
                 when (transitionState) {
                     is ObservableTransitionState.Idle ->
                         expandFractionForScene(transitionState.currentScene, shadeExpansion)
-                    is ObservableTransitionState.Transition ->
+                    is ObservableTransitionState.Transition.ChangeCurrentScene ->
                         expandFractionForTransition(
                             transitionState,
                             shadeExpansion,
@@ -121,6 +121,9 @@
                             qsExpansion,
                             quickSettingsScene
                         )
+                    is ObservableTransitionState.Transition.ShowOrHideOverlay,
+                    is ObservableTransitionState.Transition.ReplaceOverlay ->
+                        TODO("b/359173565: Handle overlay transitions")
                 }
             }
             .distinctUntilChanged()
@@ -238,7 +241,7 @@
     }
 }
 
-private fun ObservableTransitionState.Transition.isBetween(
+private fun ObservableTransitionState.Transition.ChangeCurrentScene.isBetween(
     a: (SceneKey) -> Boolean,
     b: (SceneKey) -> Boolean
 ): Boolean = (a(fromScene) && b(toScene)) || (b(fromScene) && a(toScene))