Merge changes from topics "stl-movable-refactor", "stl-picker-rename", "stl-remove-isbackground", "stl-ucs" into main

* changes:
  Introduce MovableElementScenePicker
  Remove ElementKey.isBackground (1/2)
  Rename SharedElementScenePicker to ElementScenePicker (1/2)
  Make (Movable)ElementScope extend ElementBoxScope
  Work around b/317972419
  Make the Element map a normal Map
  Extract movable contents outside of the Element class
  Rename Element.TargetValues to Element.SceneState
  Extract shared values outside of the Element class
  Replace AnimatedState.valueOrNull by unsafeCompositionState (1/2)
  Refactor animate(Scene|Element)FooAsState (1/2)
  Make animated values normal State objects
  Compose movable elements in the scene picked by the ScenePicker (1/2)
  Split SceneScope.MovableElement with SceneSope.Element (1/2)
  Move scene picker definition inside ElementKey (1/2)
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/AmbientIndicationSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/AmbientIndicationSection.kt
index 0e7ac5e..1e5481e 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/AmbientIndicationSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/AmbientIndicationSection.kt
@@ -35,14 +35,16 @@
             key = AmbientIndicationElementKey,
             modifier = modifier,
         ) {
-            Box(
-                modifier = Modifier.fillMaxWidth().background(Color.Green),
-            ) {
-                Text(
-                    text = "TODO(b/316211368): Ambient indication",
-                    color = Color.White,
-                    modifier = Modifier.align(Alignment.Center),
-                )
+            content {
+                Box(
+                    modifier = Modifier.fillMaxWidth().background(Color.Green),
+                ) {
+                    Text(
+                        text = "TODO(b/316211368): Ambient indication",
+                        color = Color.White,
+                        modifier = Modifier.align(Alignment.Center),
+                    )
+                }
             }
         }
     }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt
index 4f3498e..8bd0d45 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt
@@ -74,20 +74,22 @@
             key = if (isStart) StartButtonElementKey else EndButtonElementKey,
             modifier = modifier,
         ) {
-            Shortcut(
-                viewId = if (isStart) R.id.start_button else R.id.end_button,
-                viewModel = if (isStart) viewModel.startButton else viewModel.endButton,
-                transitionAlpha = viewModel.transitionAlpha,
-                falsingManager = falsingManager,
-                vibratorHelper = vibratorHelper,
-                indicationController = indicationController,
-                modifier =
-                    if (applyPadding) {
-                        Modifier.shortcutPadding()
-                    } else {
-                        Modifier
-                    }
-            )
+            content {
+                Shortcut(
+                    viewId = if (isStart) R.id.start_button else R.id.end_button,
+                    viewModel = if (isStart) viewModel.startButton else viewModel.endButton,
+                    transitionAlpha = viewModel.transitionAlpha,
+                    falsingManager = falsingManager,
+                    vibratorHelper = vibratorHelper,
+                    indicationController = indicationController,
+                    modifier =
+                        if (applyPadding) {
+                            Modifier.shortcutPadding()
+                        } else {
+                            Modifier
+                        }
+                )
+            }
         }
     }
 
@@ -99,11 +101,13 @@
             key = IndicationAreaElementKey,
             modifier = modifier.shortcutPadding(),
         ) {
-            IndicationArea(
-                indicationAreaViewModel = indicationAreaViewModel,
-                alphaViewModel = alphaViewModel,
-                indicationController = indicationController,
-            )
+            content {
+                IndicationArea(
+                    indicationAreaViewModel = indicationAreaViewModel,
+                    alphaViewModel = alphaViewModel,
+                    indicationController = indicationController,
+                )
+            }
         }
     }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/ClockSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/ClockSection.kt
index 0b49922..0f3fc47 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/ClockSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/ClockSection.kt
@@ -49,17 +49,19 @@
             key = ClockElementKey,
             modifier = modifier,
         ) {
-            Box(
-                modifier =
-                    Modifier.fillMaxWidth()
-                        .background(Color.Magenta)
-                        .onTopPlacementChanged(onTopChanged)
-            ) {
-                Text(
-                    text = "TODO(b/316211368): Small clock",
-                    color = Color.White,
-                    modifier = Modifier.align(Alignment.Center),
-                )
+            content {
+                Box(
+                    modifier =
+                        Modifier.fillMaxWidth()
+                            .background(Color.Magenta)
+                            .onTopPlacementChanged(onTopChanged)
+                ) {
+                    Text(
+                        text = "TODO(b/316211368): Small clock",
+                        color = Color.White,
+                        modifier = Modifier.align(Alignment.Center),
+                    )
+                }
             }
         }
     }
@@ -74,14 +76,16 @@
             key = ClockElementKey,
             modifier = modifier,
         ) {
-            Box(
-                modifier = Modifier.fillMaxWidth().background(Color.Blue),
-            ) {
-                Text(
-                    text = "TODO(b/316211368): Large clock",
-                    color = Color.White,
-                    modifier = Modifier.align(Alignment.Center),
-                )
+            content {
+                Box(
+                    modifier = Modifier.fillMaxWidth().background(Color.Blue),
+                ) {
+                    Text(
+                        text = "TODO(b/316211368): Large clock",
+                        color = Color.White,
+                        modifier = Modifier.align(Alignment.Center),
+                    )
+                }
             }
         }
     }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt
index c547e2b..900616f 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt
@@ -22,7 +22,6 @@
 import com.android.systemui.notifications.ui.composable.NotificationStack
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
 import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
 
 class NotificationSection
 @Inject
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/StatusBarSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/StatusBarSection.kt
index 5727e34..ddc12ff 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/StatusBarSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/StatusBarSection.kt
@@ -53,37 +53,43 @@
             key = StatusBarElementKey,
             modifier = modifier,
         ) {
-            AndroidView(
-                factory = {
-                    notificationPanelView.get().findViewById<View>(R.id.keyguard_header)?.let {
-                        (it.parent as ViewGroup).removeView(it)
-                    }
-
-                    val provider =
-                        object : ShadeViewStateProvider {
-                            override val lockscreenShadeDragProgress: Float = 0f
-                            override val panelViewExpandedHeight: Float = 0f
-                            override fun shouldHeadsUpBeVisible(): Boolean {
-                                return false
-                            }
+            content {
+                AndroidView(
+                    factory = {
+                        notificationPanelView.get().findViewById<View>(R.id.keyguard_header)?.let {
+                            (it.parent as ViewGroup).removeView(it)
                         }
 
-                    @SuppressLint("InflateParams")
-                    val view =
-                        LayoutInflater.from(context)
-                            .inflate(
-                                R.layout.keyguard_status_bar,
-                                null,
-                                false,
-                            ) as KeyguardStatusBarView
-                    componentFactory.build(view, provider).keyguardStatusBarViewController.init()
-                    view
-                },
-                modifier =
-                    Modifier.fillMaxWidth().padding(horizontal = 16.dp).height {
-                        Utils.getStatusBarHeaderHeightKeyguard(context)
+                        val provider =
+                            object : ShadeViewStateProvider {
+                                override val lockscreenShadeDragProgress: Float = 0f
+                                override val panelViewExpandedHeight: Float = 0f
+
+                                override fun shouldHeadsUpBeVisible(): Boolean {
+                                    return false
+                                }
+                            }
+
+                        @SuppressLint("InflateParams")
+                        val view =
+                            LayoutInflater.from(context)
+                                .inflate(
+                                    R.layout.keyguard_status_bar,
+                                    null,
+                                    false,
+                                ) as KeyguardStatusBarView
+                        componentFactory
+                            .build(view, provider)
+                            .keyguardStatusBarViewController
+                            .init()
+                        view
                     },
-            )
+                    modifier =
+                        Modifier.fillMaxWidth().padding(horizontal = 16.dp).height {
+                            Utils.getStatusBarHeaderHeightKeyguard(context)
+                        },
+                )
+            }
         }
     }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
index 12f1b30..0eec024 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
@@ -44,7 +44,7 @@
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.ValueKey
-import com.android.compose.animation.scene.animateSharedFloatAsState
+import com.android.compose.animation.scene.animateElementFloatAsState
 import com.android.systemui.notifications.ui.composable.Notifications.Form
 import com.android.systemui.notifications.ui.composable.Notifications.SharedValues.SharedExpansionValue
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
@@ -157,10 +157,10 @@
     modifier: Modifier = Modifier,
 ) {
     val elementKey = Notifications.Elements.NotificationPlaceholder
-    Box(
+    Element(
+        elementKey,
         modifier =
             modifier
-                .element(elementKey)
                 .debugBackground(viewModel)
                 .onSizeChanged { size: IntSize ->
                     debugLog(viewModel) { "STACK onSizeChanged: size=$size" }
@@ -182,19 +182,23 @@
                 }
     ) {
         val animatedExpansion by
-            animateSharedFloatAsState(
+            animateElementFloatAsState(
                 value = if (form == Form.HunFromTop) 0f else 1f,
-                key = SharedExpansionValue,
-                element = elementKey
+                key = SharedExpansionValue
             )
         debugLog(viewModel) { "STACK composed: expansion=$animatedExpansion" }
-        if (viewModel.isPlaceholderTextVisible) {
-            Text(
-                text = "Notifications",
-                style = MaterialTheme.typography.titleLarge,
-                color = MaterialTheme.colorScheme.onSurface,
-                modifier = Modifier.align(Alignment.Center),
-            )
+
+        content {
+            if (viewModel.isPlaceholderTextVisible) {
+                Box(Modifier.fillMaxSize()) {
+                    Text(
+                        text = "Notifications",
+                        style = MaterialTheme.typography.titleLarge,
+                        color = MaterialTheme.colorScheme.onSurface,
+                        modifier = Modifier.align(Alignment.Center),
+                    )
+                }
+            }
         }
     }
 }
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 f3cde53..65a53f5 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
@@ -98,7 +98,7 @@
         key = QuickSettings.Elements.Content,
         modifier = modifier.fillMaxWidth().defaultMinSize(minHeight = 300.dp)
     ) {
-        QuickSettingsContent(qsSceneAdapter = qsSceneAdapter, contentState)
+        content { QuickSettingsContent(qsSceneAdapter = qsSceneAdapter, contentState) }
     }
 }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
index 4bbb78b..99f81ee 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
@@ -50,7 +50,7 @@
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.ValueKey
-import com.android.compose.animation.scene.animateSharedFloatAsState
+import com.android.compose.animation.scene.animateSceneFloatAsState
 import com.android.compose.windowsizeclass.LocalWindowSizeClass
 import com.android.settingslib.Utils
 import com.android.systemui.battery.BatteryMeterView
@@ -69,7 +69,6 @@
 
 object ShadeHeader {
     object Elements {
-        val FormatPlaceholder = ElementKey("ShadeHeaderFormatPlaceholder")
         val ExpandedContent = ElementKey("ShadeHeaderExpandedContent")
         val CollapsedContent = ElementKey("ShadeHeaderCollapsedContent")
     }
@@ -92,14 +91,7 @@
     statusBarIconController: StatusBarIconController,
     modifier: Modifier = Modifier,
 ) {
-    // TODO(b/298153892): Remove this once animateSharedFloatAsState.element can be null.
-    Spacer(Modifier.element(ShadeHeader.Elements.FormatPlaceholder))
-    val formatProgress =
-        animateSharedFloatAsState(
-            0.0f,
-            ShadeHeader.Keys.transitionProgress,
-            ShadeHeader.Elements.FormatPlaceholder
-        )
+    val formatProgress = animateSceneFloatAsState(0.0f, ShadeHeader.Keys.transitionProgress)
 
     val cutoutWidth = LocalDisplayCutout.current.width()
     val cutoutLocation = LocalDisplayCutout.current.location
@@ -217,14 +209,7 @@
     statusBarIconController: StatusBarIconController,
     modifier: Modifier = Modifier,
 ) {
-    // TODO(b/298153892): Remove this once animateSharedFloatAsState.element can be null.
-    Spacer(Modifier.element(ShadeHeader.Elements.FormatPlaceholder))
-    val formatProgress =
-        animateSharedFloatAsState(
-            1.0f,
-            ShadeHeader.Keys.transitionProgress,
-            ShadeHeader.Elements.FormatPlaceholder
-        )
+    val formatProgress = animateSceneFloatAsState(1.0f, ShadeHeader.Keys.transitionProgress)
     val useExpandedFormat by
         remember(formatProgress) { derivedStateOf { formatProgress.value > 0.5f } }
 
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 2944bd9..b26194f 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
@@ -17,10 +17,15 @@
 package com.android.compose.animation.scene
 
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.SideEffect
+import androidx.compose.runtime.Stable
 import androidx.compose.runtime.State
-import androidx.compose.runtime.derivedStateOf
+import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.snapshots.Snapshot
+import androidx.compose.runtime.snapshotFlow
+import androidx.compose.runtime.snapshots.SnapshotStateMap
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.lerp
 import androidx.compose.ui.unit.Dp
@@ -28,180 +33,263 @@
 import com.android.compose.ui.util.lerp
 
 /**
- * Animate a shared Int value.
+ * A [State] whose [value] is animated.
  *
- * @see SceneScope.animateSharedValueAsState
+ * Important: This animated value should always be ready *after* composition, e.g. during layout,
+ * drawing or inside a LaunchedEffect. If you read [value] during composition, it will probably
+ * throw an exception, for 2 important reasons:
+ * 1. You should never read animated values during composition, because this will probably lead to
+ *    bad performance.
+ * 2. Given that this value depends on the target value in different scenes, its current value
+ *    (depending on the current transition state) can only be computed once the full tree has been
+ *    composed.
+ *
+ * If you don't have the choice and *have to* get the value during composition, for instance because
+ * a Modifier or Composable reading this value does not have a lazy/lambda-based API, then you can
+ * access [unsafeCompositionState] and use a fallback value for the first frame where this animated
+ * value can not be computed yet. Note however that doing so will be bad for performance and might
+ * lead to late-by-one-frame flickers.
+ */
+@Stable
+interface AnimatedState<T> : State<T> {
+    /**
+     * Return a [State] that can be read during composition.
+     *
+     * Important: You should avoid using this as much as possible and instead read [value] during
+     * layout/drawing, otherwise you will probably end up with a few frames that have a value that
+     * is not correctly interpolated.
+     */
+    @Composable fun unsafeCompositionState(initialValue: T): State<T>
+}
+
+/**
+ * Animate a scene Int value.
+ *
+ * @see SceneScope.animateSceneValueAsState
  */
 @Composable
-fun SceneScope.animateSharedIntAsState(
+fun SceneScope.animateSceneIntAsState(
     value: Int,
     key: ValueKey,
-    element: ElementKey?,
     canOverflow: Boolean = true,
-): State<Int> {
-    return animateSharedValueAsState(value, key, element, ::lerp, canOverflow)
+): AnimatedState<Int> {
+    return animateSceneValueAsState(value, key, ::lerp, canOverflow)
 }
 
 /**
- * Animate a shared Int value.
+ * Animate a shared element Int value.
  *
- * @see MovableElementScope.animateSharedValueAsState
+ * @see ElementScope.animateElementValueAsState
  */
 @Composable
-fun MovableElementScope.animateSharedIntAsState(
+fun ElementScope<*>.animateElementIntAsState(
     value: Int,
-    debugName: String,
+    key: ValueKey,
     canOverflow: Boolean = true,
-): State<Int> {
-    return animateSharedValueAsState(value, debugName, ::lerp, canOverflow)
+): AnimatedState<Int> {
+    return animateElementValueAsState(value, key, ::lerp, canOverflow)
 }
 
 /**
- * Animate a shared Float value.
+ * Animate a scene Float value.
  *
- * @see SceneScope.animateSharedValueAsState
+ * @see SceneScope.animateSceneValueAsState
  */
 @Composable
-fun SceneScope.animateSharedFloatAsState(
+fun SceneScope.animateSceneFloatAsState(
     value: Float,
     key: ValueKey,
-    element: ElementKey?,
     canOverflow: Boolean = true,
-): State<Float> {
-    return animateSharedValueAsState(value, key, element, ::lerp, canOverflow)
+): AnimatedState<Float> {
+    return animateSceneValueAsState(value, key, ::lerp, canOverflow)
 }
 
 /**
- * Animate a shared Float value.
+ * Animate a shared element Float value.
  *
- * @see MovableElementScope.animateSharedValueAsState
+ * @see ElementScope.animateElementValueAsState
  */
 @Composable
-fun MovableElementScope.animateSharedFloatAsState(
+fun ElementScope<*>.animateElementFloatAsState(
     value: Float,
-    debugName: String,
+    key: ValueKey,
     canOverflow: Boolean = true,
-): State<Float> {
-    return animateSharedValueAsState(value, debugName, ::lerp, canOverflow)
+): AnimatedState<Float> {
+    return animateElementValueAsState(value, key, ::lerp, canOverflow)
 }
 
 /**
- * Animate a shared Dp value.
+ * Animate a scene Dp value.
  *
- * @see SceneScope.animateSharedValueAsState
+ * @see SceneScope.animateSceneValueAsState
  */
 @Composable
-fun SceneScope.animateSharedDpAsState(
+fun SceneScope.animateSceneDpAsState(
     value: Dp,
     key: ValueKey,
-    element: ElementKey?,
     canOverflow: Boolean = true,
-): State<Dp> {
-    return animateSharedValueAsState(value, key, element, ::lerp, canOverflow)
+): AnimatedState<Dp> {
+    return animateSceneValueAsState(value, key, ::lerp, canOverflow)
 }
 
 /**
- * Animate a shared Dp value.
+ * Animate a shared element Dp value.
  *
- * @see MovableElementScope.animateSharedValueAsState
+ * @see ElementScope.animateElementValueAsState
  */
 @Composable
-fun MovableElementScope.animateSharedDpAsState(
+fun ElementScope<*>.animateElementDpAsState(
     value: Dp,
-    debugName: String,
+    key: ValueKey,
     canOverflow: Boolean = true,
-): State<Dp> {
-    return animateSharedValueAsState(value, debugName, ::lerp, canOverflow)
+): AnimatedState<Dp> {
+    return animateElementValueAsState(value, key, ::lerp, canOverflow)
 }
 
 /**
- * Animate a shared Color value.
+ * Animate a scene Color value.
  *
- * @see SceneScope.animateSharedValueAsState
+ * @see SceneScope.animateSceneValueAsState
  */
 @Composable
-fun SceneScope.animateSharedColorAsState(
+fun SceneScope.animateSceneColorAsState(
     value: Color,
     key: ValueKey,
-    element: ElementKey?,
-): State<Color> {
-    return animateSharedValueAsState(value, key, element, ::lerp, canOverflow = false)
+): AnimatedState<Color> {
+    return animateSceneValueAsState(value, key, ::lerp, canOverflow = false)
 }
 
 /**
- * Animate a shared Color value.
+ * Animate a shared element Color value.
  *
- * @see MovableElementScope.animateSharedValueAsState
+ * @see ElementScope.animateElementValueAsState
  */
 @Composable
-fun MovableElementScope.animateSharedColorAsState(
+fun ElementScope<*>.animateElementColorAsState(
     value: Color,
-    debugName: String,
-): State<Color> {
-    return animateSharedValueAsState(value, debugName, ::lerp, canOverflow = false)
+    key: ValueKey,
+): AnimatedState<Color> {
+    return animateElementValueAsState(value, key, ::lerp, canOverflow = false)
 }
 
 @Composable
 internal fun <T> animateSharedValueAsState(
     layoutImpl: SceneTransitionLayoutImpl,
-    scene: Scene,
-    element: Element?,
+    scene: SceneKey,
+    element: ElementKey?,
     key: ValueKey,
     value: T,
     lerp: (T, T, Float) -> T,
     canOverflow: Boolean,
-): State<T> {
-    val sharedValue =
-        Snapshot.withoutReadObservation {
-            val sharedValues =
-                element?.sceneValues?.getValue(scene.key)?.sharedValues ?: scene.sharedValues
-            sharedValues.getOrPut(key) { Element.SharedValue(key, value) } as Element.SharedValue<T>
-        }
+): AnimatedState<T> {
+    DisposableEffect(layoutImpl, scene, element, key) {
+        // Create the associated maps that hold the current value for each (element, scene) pair.
+        val valueMap = layoutImpl.sharedValues.getOrPut(key) { mutableMapOf() }
+        val sceneToValueMap =
+            valueMap.getOrPut(element) { SnapshotStateMap<SceneKey, Any>() }
+                as SnapshotStateMap<SceneKey, T>
+        sceneToValueMap[scene] = value
 
-    if (value != sharedValue.value) {
-        sharedValue.value = value
+        onDispose {
+            // Remove the value associated to the current scene, and eventually remove the maps if
+            // they are empty.
+            sceneToValueMap.remove(scene)
+
+            if (sceneToValueMap.isEmpty() && valueMap[element] === sceneToValueMap) {
+                valueMap.remove(element)
+
+                if (valueMap.isEmpty() && layoutImpl.sharedValues[key] === valueMap) {
+                    layoutImpl.sharedValues.remove(key)
+                }
+            }
+        }
     }
 
-    return remember(layoutImpl, element, sharedValue, lerp, canOverflow) {
-        derivedStateOf { computeValue(layoutImpl, element, sharedValue, lerp, canOverflow) }
+    // Update the current value. Note that side effects run after disposable effects, so we know
+    // that the associated maps were created at this point.
+    SideEffect { sceneToValueMap<T>(layoutImpl, key, element)[scene] = value }
+
+    return remember(layoutImpl, scene, element, lerp, canOverflow) {
+        object : AnimatedState<T> {
+            override val value: T
+                get() = value(layoutImpl, scene, element, key, lerp, canOverflow)
+
+            @Composable
+            override fun unsafeCompositionState(initialValue: T): State<T> {
+                val state = remember { mutableStateOf(initialValue) }
+
+                val animatedState = this
+                LaunchedEffect(animatedState) {
+                    snapshotFlow { animatedState.value }.collect { state.value = it }
+                }
+
+                return state
+            }
+        }
     }
 }
 
-private fun <T> computeValue(
+private fun <T> sceneToValueMap(
     layoutImpl: SceneTransitionLayoutImpl,
-    element: Element?,
-    sharedValue: Element.SharedValue<T>,
+    key: ValueKey,
+    element: ElementKey?
+): MutableMap<SceneKey, T> {
+    return layoutImpl.sharedValues[key]?.get(element)?.let { it as SnapshotStateMap<SceneKey, T> }
+        ?: error(valueReadTooEarlyMessage(key))
+}
+
+private fun valueReadTooEarlyMessage(key: ValueKey) =
+    "Animated value $key was read before its target values were set. This probably " +
+        "means that you are reading it during composition, which you should not do. See the " +
+        "documentation of AnimatedState for more information."
+
+private fun <T> value(
+    layoutImpl: SceneTransitionLayoutImpl,
+    scene: SceneKey,
+    element: ElementKey?,
+    key: ValueKey,
     lerp: (T, T, Float) -> T,
     canOverflow: Boolean,
 ): T {
-    val transition = layoutImpl.state.currentTransition
-    if (transition == null || !layoutImpl.isTransitionReady(transition)) {
-        return sharedValue.value
-    }
+    return valueOrNull(layoutImpl, scene, element, key, lerp, canOverflow)
+        ?: error(valueReadTooEarlyMessage(key))
+}
 
-    fun sceneValue(scene: SceneKey): Element.SharedValue<T>? {
-        val sharedValues =
-            if (element == null) {
-                layoutImpl.scene(scene).sharedValues
-            } else {
-                element.sceneValues[scene]?.sharedValues
-            }
-                ?: return null
-        val value = sharedValues[sharedValue.key] ?: return null
-        return value as Element.SharedValue<T>
-    }
+private fun <T> valueOrNull(
+    layoutImpl: SceneTransitionLayoutImpl,
+    scene: SceneKey,
+    element: ElementKey?,
+    key: ValueKey,
+    lerp: (T, T, Float) -> T,
+    canOverflow: Boolean,
+): T? {
+    val sceneToValueMap = sceneToValueMap<T>(layoutImpl, key, element)
+    fun sceneValue(scene: SceneKey): T? = sceneToValueMap[scene]
 
-    val fromValue = sceneValue(transition.fromScene)
-    val toValue = sceneValue(transition.toScene)
-    return if (fromValue != null && toValue != null) {
-        val progress =
-            if (canOverflow) transition.progress else transition.progress.coerceIn(0f, 1f)
-        lerp(fromValue.value, toValue.value, progress)
-    } else if (fromValue != null) {
-        fromValue.value
-    } else if (toValue != null) {
-        toValue.value
-    } else {
-        sharedValue.value
+    return when (val transition = layoutImpl.state.transitionState) {
+        is TransitionState.Idle -> sceneValue(transition.currentScene)
+        is TransitionState.Transition -> {
+            // Note: no need to check for transition ready here given that all target values are
+            // defined during composition, we should already have the correct values to interpolate
+            // between here.
+            val fromValue = sceneValue(transition.fromScene)
+            val toValue = sceneValue(transition.toScene)
+            if (fromValue != null && toValue != null) {
+                if (fromValue == toValue) {
+                    // Optimization: avoid reading progress if the values are the same, so we don't
+                    // relayout/redraw for nothing.
+                    fromValue
+                } else {
+                    val progress =
+                        if (canOverflow) transition.progress
+                        else transition.progress.coerceIn(0f, 1f)
+                    lerp(fromValue, toValue, progress)
+                }
+            } else fromValue ?: toValue
+        }
     }
+    // TODO(b/311600838): Remove this. We should not have to fallback to the current scene value,
+    // but we have to because code of removed nodes can still run if they are placed with a graphics
+    // layer.
+    ?: sceneValue(scene)
 }
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 a85d9bf..280fbfb 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
@@ -16,15 +16,10 @@
 
 package com.android.compose.animation.scene
 
-import android.graphics.Picture
-import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.movableContentOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.compose.runtime.snapshots.Snapshot
-import androidx.compose.runtime.snapshots.SnapshotStateMap
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
@@ -52,43 +47,22 @@
 @Stable
 internal class Element(val key: ElementKey) {
     /**
-     * The last values of this element, coming from any scene. Note that this value will be unstable
+     * The last state of this element, coming from any scene. Note that this state will be unstable
      * if this element is present in multiple scenes but the shared element animation is disabled,
-     * given that multiple instances of the element with different states will write to these
-     * values. You should prefer using [TargetValues.lastValues] in the current scene if it is
-     * defined.
+     * given that multiple instances of the element with different states will write to this state.
+     * You should prefer using [SceneState.lastState] in the current scene when it is defined.
      */
-    val lastSharedValues = Values()
+    val lastSharedState = State()
 
-    /** The mapping between a scene and the values/state this element has in that scene, if any. */
-    val sceneValues = SnapshotStateMap<SceneKey, TargetValues>()
-
-    /**
-     * The movable content of this element, if this element is composed using
-     * [SceneScope.MovableElement].
-     */
-    private var _movableContent: (@Composable (@Composable () -> Unit) -> Unit)? = null
-    val movableContent: @Composable (@Composable () -> Unit) -> Unit
-        get() =
-            _movableContent
-                ?: movableContentOf { content: @Composable () -> Unit -> content() }
-                    .also { _movableContent = it }
-
-    /**
-     * The [Picture] to which we save the last drawing commands of this element, if it is movable.
-     * This is necessary because the content of this element might not be composed in the scene it
-     * should currently be drawn.
-     */
-    private var _picture: Picture? = null
-    val picture: Picture
-        get() = _picture ?: Picture().also { _picture = it }
+    /** The mapping between a scene and the state this element has in that scene, if any. */
+    val sceneStates = mutableMapOf<SceneKey, SceneState>()
 
     override fun toString(): String {
         return "Element(key=$key)"
     }
 
-    /** The current values of this element, either in a specific scene or in a shared context. */
-    class Values {
+    /** The state of this element, either in a specific scene or in a shared context. */
+    class State {
         /** The offset of the element, relative to the SceneTransitionLayout containing it. */
         var offset = Offset.Unspecified
 
@@ -102,16 +76,14 @@
         var alpha = AlphaUnspecified
     }
 
-    /** The target values of this element in a given scene. */
+    /** The last and target state of this element in a given scene. */
     @Stable
-    class TargetValues(val scene: SceneKey) {
-        val lastValues = Values()
+    class SceneState(val scene: SceneKey) {
+        val lastState = State()
 
         var targetSize by mutableStateOf(SizeUnspecified)
         var targetOffset by mutableStateOf(Offset.Unspecified)
 
-        val sharedValues = SnapshotStateMap<ValueKey, SharedValue<*>>()
-
         /**
          * The attached [ElementNode] a Modifier.element() for a given element and scene. During
          * composition, this set could have 0 to 2 elements. After composition and after all
@@ -120,12 +92,6 @@
         val nodes = mutableSetOf<ElementNode>()
     }
 
-    /** A shared value of this element. */
-    @Stable
-    class SharedValue<T>(val key: ValueKey, initialValue: T) {
-        var value by mutableStateOf(initialValue)
-    }
-
     companion object {
         val SizeUnspecified = IntSize(Int.MAX_VALUE, Int.MAX_VALUE)
         val AlphaUnspecified = Float.MIN_VALUE
@@ -147,27 +113,18 @@
     scene: Scene,
     key: ElementKey,
 ): Modifier {
-    val element: Element
-    val sceneValues: Element.TargetValues
-
-    // Get the element associated to [key] if it was already composed in another scene,
-    // otherwise create it and add it to our Map<ElementKey, Element>. This is done inside a
-    // withoutReadObservation() because there is no need to recompose when that map is mutated.
-    Snapshot.withoutReadObservation {
-        element = layoutImpl.elements[key] ?: Element(key).also { layoutImpl.elements[key] = it }
-        sceneValues =
-            element.sceneValues[scene.key]
-                ?: Element.TargetValues(scene.key).also { element.sceneValues[scene.key] = it }
-    }
-
-    return this.then(ElementModifier(layoutImpl, scene, element, sceneValues))
+    return this.then(ElementModifier(layoutImpl, scene, key))
         // TODO(b/311132415): Move this into ElementNode once we can create a delegate
         // IntermediateLayoutModifierNode.
         .intermediateLayout { measurable, constraints ->
-            val placeable =
-                measure(layoutImpl, scene, element, sceneValues, measurable, constraints)
+            // TODO(b/311132415): No need to fetch the element and sceneState from the map anymore
+            // once this is merged into ElementNode.
+            val element = layoutImpl.elements.getValue(key)
+            val sceneState = element.sceneStates.getValue(scene.key)
+
+            val placeable = measure(layoutImpl, scene, element, sceneState, measurable, constraints)
             layout(placeable.width, placeable.height) {
-                place(layoutImpl, scene, element, sceneValues, placeable, placementScope = this)
+                place(layoutImpl, scene, element, sceneState, placeable, placementScope = this)
             }
         }
         .testTag(key.testTag)
@@ -180,72 +137,89 @@
 private data class ElementModifier(
     private val layoutImpl: SceneTransitionLayoutImpl,
     private val scene: Scene,
-    private val element: Element,
-    private val sceneValues: Element.TargetValues,
+    private val key: ElementKey,
 ) : ModifierNodeElement<ElementNode>() {
-    override fun create(): ElementNode = ElementNode(layoutImpl, scene, element, sceneValues)
+    override fun create(): ElementNode = ElementNode(layoutImpl, scene, key)
 
     override fun update(node: ElementNode) {
-        node.update(layoutImpl, scene, element, sceneValues)
+        node.update(layoutImpl, scene, key)
     }
 }
 
 internal class ElementNode(
     private var layoutImpl: SceneTransitionLayoutImpl,
     private var scene: Scene,
-    private var element: Element,
-    private var sceneValues: Element.TargetValues,
+    private var key: ElementKey,
 ) : Modifier.Node(), DrawModifierNode {
+    private var _element: Element? = null
+    private val element: Element
+        get() = _element!!
+
+    private var _sceneState: Element.SceneState? = null
+    private val sceneState: Element.SceneState
+        get() = _sceneState!!
 
     override fun onAttach() {
         super.onAttach()
-        addNodeToSceneValues()
+        updateElementAndSceneValues()
+        addNodeToSceneState()
     }
 
-    private fun addNodeToSceneValues() {
-        sceneValues.nodes.add(this)
+    private fun updateElementAndSceneValues() {
+        val element =
+            layoutImpl.elements[key] ?: Element(key).also { layoutImpl.elements[key] = it }
+        _element = element
+        _sceneState =
+            element.sceneStates[scene.key]
+                ?: Element.SceneState(scene.key).also { element.sceneStates[scene.key] = it }
+    }
+
+    private fun addNodeToSceneState() {
+        sceneState.nodes.add(this)
 
         coroutineScope.launch {
             // At this point all [CodeLocationNode] have been attached or detached, which means that
-            // [sceneValues.codeLocations] should have exactly 1 element, otherwise this means that
+            // [sceneState.codeLocations] should have exactly 1 element, otherwise this means that
             // this element was composed multiple times in the same scene.
-            val nCodeLocations = sceneValues.nodes.size
-            if (nCodeLocations != 1 || !sceneValues.nodes.contains(this@ElementNode)) {
-                error("${element.key} was composed $nCodeLocations times in ${sceneValues.scene}")
+            val nCodeLocations = sceneState.nodes.size
+            if (nCodeLocations != 1 || !sceneState.nodes.contains(this@ElementNode)) {
+                error("$key was composed $nCodeLocations times in ${sceneState.scene}")
             }
         }
     }
 
     override fun onDetach() {
         super.onDetach()
-        removeNodeFromSceneValues()
-        maybePruneMaps(layoutImpl, element, sceneValues)
+        removeNodeFromSceneState()
+        maybePruneMaps(layoutImpl, element, sceneState)
+
+        _element = null
+        _sceneState = null
     }
 
-    private fun removeNodeFromSceneValues() {
-        sceneValues.nodes.remove(this)
+    private fun removeNodeFromSceneState() {
+        sceneState.nodes.remove(this)
     }
 
     fun update(
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
-        element: Element,
-        sceneValues: Element.TargetValues,
+        key: ElementKey,
     ) {
         check(layoutImpl == this.layoutImpl && scene == this.scene)
-        removeNodeFromSceneValues()
+        removeNodeFromSceneState()
 
         val prevElement = this.element
-        val prevSceneValues = this.sceneValues
-        this.element = element
-        this.sceneValues = sceneValues
+        val prevSceneState = this.sceneState
+        this.key = key
+        updateElementAndSceneValues()
 
-        addNodeToSceneValues()
-        maybePruneMaps(layoutImpl, prevElement, prevSceneValues)
+        addNodeToSceneState()
+        maybePruneMaps(layoutImpl, prevElement, prevSceneState)
     }
 
     override fun ContentDrawScope.draw() {
-        val drawScale = getDrawScale(layoutImpl, element, scene, sceneValues)
+        val drawScale = getDrawScale(layoutImpl, element, scene, sceneState)
         if (drawScale == Scale.Default) {
             drawContent()
         } else {
@@ -263,18 +237,16 @@
         private fun maybePruneMaps(
             layoutImpl: SceneTransitionLayoutImpl,
             element: Element,
-            sceneValues: Element.TargetValues,
+            sceneState: Element.SceneState,
         ) {
             // If element is not composed from this scene anymore, remove the scene values. This
             // works because [onAttach] is called before [onDetach], so if an element is moved from
             // the UI tree we will first add the new code location then remove the old one.
-            if (
-                sceneValues.nodes.isEmpty() && element.sceneValues[sceneValues.scene] == sceneValues
-            ) {
-                element.sceneValues.remove(sceneValues.scene)
+            if (sceneState.nodes.isEmpty() && element.sceneStates[sceneState.scene] == sceneState) {
+                element.sceneStates.remove(sceneState.scene)
 
                 // If the element is not composed in any scene, remove it from the elements map.
-                if (element.sceneValues.isEmpty() && layoutImpl.elements[element.key] == element) {
+                if (element.sceneStates.isEmpty() && layoutImpl.elements[element.key] == element) {
                     layoutImpl.elements.remove(element.key)
                 }
             }
@@ -293,8 +265,8 @@
     if (
         transition == null ||
             !layoutImpl.isTransitionReady(transition) ||
-            transition.fromScene !in element.sceneValues ||
-            transition.toScene !in element.sceneValues
+            transition.fromScene !in element.sceneStates ||
+            transition.toScene !in element.sceneStates
     ) {
         return true
     }
@@ -310,7 +282,6 @@
         transition,
         scene.key,
         element.key,
-        sharedTransformation,
     )
 }
 
@@ -319,17 +290,14 @@
     transition: TransitionState.Transition,
     scene: SceneKey,
     element: ElementKey,
-    sharedTransformation: SharedElementTransformation?
 ): Boolean {
-    val scenePicker = sharedTransformation?.scenePicker ?: DefaultSharedElementScenePicker
+    val scenePicker = element.scenePicker
     val fromScene = transition.fromScene
     val toScene = transition.toScene
 
     return scenePicker.sceneDuringTransition(
         element = element,
-        fromScene = fromScene,
-        toScene = toScene,
-        progress = transition::progress,
+        transition = transition,
         fromSceneZIndex = layoutImpl.scenes.getValue(fromScene).zIndex,
         toSceneZIndex = layoutImpl.scenes.getValue(toScene).zIndex,
     ) == scene
@@ -374,28 +342,28 @@
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
     scene: Scene,
-    sceneValues: Element.TargetValues,
+    sceneState: Element.SceneState,
 ): Boolean {
     val transition = layoutImpl.state.currentTransition ?: return true
 
     if (!layoutImpl.isTransitionReady(transition)) {
         val lastValue =
-            sceneValues.lastValues.alpha.takeIf { it != Element.AlphaUnspecified }
-                ?: element.lastSharedValues.alpha.takeIf { it != Element.AlphaUnspecified } ?: 1f
+            sceneState.lastState.alpha.takeIf { it != Element.AlphaUnspecified }
+                ?: element.lastSharedState.alpha.takeIf { it != Element.AlphaUnspecified } ?: 1f
 
         return lastValue == 1f
     }
 
     val fromScene = transition.fromScene
     val toScene = transition.toScene
-    val fromValues = element.sceneValues[fromScene]
-    val toValues = element.sceneValues[toScene]
+    val fromState = element.sceneStates[fromScene]
+    val toState = element.sceneStates[toScene]
 
-    if (fromValues == null && toValues == null) {
+    if (fromState == null && toState == null) {
         error("This should not happen, element $element is neither in $fromScene or $toScene")
     }
 
-    val isSharedElement = fromValues != null && toValues != null
+    val isSharedElement = fromState != null && toState != null
     if (isSharedElement && isSharedElementEnabled(layoutImpl.state, transition, element.key)) {
         return true
     }
@@ -415,7 +383,7 @@
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
     scene: Scene,
-    sceneValues: Element.TargetValues,
+    sceneState: Element.SceneState,
 ): Float {
     return computeValue(
             layoutImpl,
@@ -426,9 +394,8 @@
             idleValue = 1f,
             currentValue = { 1f },
             lastValue = {
-                sceneValues.lastValues.alpha.takeIf { it != Element.AlphaUnspecified }
-                    ?: element.lastSharedValues.alpha.takeIf { it != Element.AlphaUnspecified }
-                        ?: 1f
+                sceneState.lastState.alpha.takeIf { it != Element.AlphaUnspecified }
+                    ?: element.lastSharedState.alpha.takeIf { it != Element.AlphaUnspecified } ?: 1f
             },
             ::lerp,
         )
@@ -440,15 +407,15 @@
     layoutImpl: SceneTransitionLayoutImpl,
     scene: Scene,
     element: Element,
-    sceneValues: Element.TargetValues,
+    sceneState: Element.SceneState,
     measurable: Measurable,
     constraints: Constraints,
 ): Placeable {
     // Update the size this element has in this scene when idle.
     val targetSizeInScene = lookaheadSize
-    if (targetSizeInScene != sceneValues.targetSize) {
+    if (targetSizeInScene != sceneState.targetSize) {
         // TODO(b/290930950): Better handle when this changes to avoid instant size jumps.
-        sceneValues.targetSize = targetSizeInScene
+        sceneState.targetSize = targetSizeInScene
     }
 
     // Some lambdas called (max once) by computeValue() will need to measure [measurable], in which
@@ -468,8 +435,8 @@
             idleValue = lookaheadSize,
             currentValue = { measurable.measure(constraints).also { maybePlaceable = it }.size() },
             lastValue = {
-                sceneValues.lastValues.size.takeIf { it != Element.SizeUnspecified }
-                    ?: element.lastSharedValues.size.takeIf { it != Element.SizeUnspecified }
+                sceneState.lastState.size.takeIf { it != Element.SizeUnspecified }
+                    ?: element.lastSharedState.size.takeIf { it != Element.SizeUnspecified }
                         ?: measurable.measure(constraints).also { maybePlaceable = it }.size()
             },
             ::lerp,
@@ -485,8 +452,8 @@
             )
 
     val size = placeable.size()
-    element.lastSharedValues.size = size
-    sceneValues.lastValues.size = size
+    element.lastSharedState.size = size
+    sceneState.lastState.size = size
     return placeable
 }
 
@@ -494,7 +461,7 @@
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
     scene: Scene,
-    sceneValues: Element.TargetValues
+    sceneState: Element.SceneState
 ): Scale {
     return computeValue(
         layoutImpl,
@@ -505,8 +472,8 @@
         idleValue = Scale.Default,
         currentValue = { Scale.Default },
         lastValue = {
-            sceneValues.lastValues.drawScale.takeIf { it != Scale.Default }
-                ?: element.lastSharedValues.drawScale
+            sceneState.lastState.drawScale.takeIf { it != Scale.Default }
+                ?: element.lastSharedState.drawScale
         },
         ::lerp,
     )
@@ -517,7 +484,7 @@
     layoutImpl: SceneTransitionLayoutImpl,
     scene: Scene,
     element: Element,
-    sceneValues: Element.TargetValues,
+    sceneState: Element.SceneState,
     placeable: Placeable,
     placementScope: Placeable.PlacementScope,
 ) {
@@ -526,14 +493,14 @@
         // when idle.
         val coords = coordinates ?: error("Element ${element.key} does not have any coordinates")
         val targetOffsetInScene = lookaheadScopeCoordinates.localLookaheadPositionOf(coords)
-        if (targetOffsetInScene != sceneValues.targetOffset) {
+        if (targetOffsetInScene != sceneState.targetOffset) {
             // TODO(b/290930950): Better handle when this changes to avoid instant offset jumps.
-            sceneValues.targetOffset = targetOffsetInScene
+            sceneState.targetOffset = targetOffsetInScene
         }
 
         val currentOffset = lookaheadScopeCoordinates.localPositionOf(coords, Offset.Zero)
-        val lastSharedValues = element.lastSharedValues
-        val lastValues = sceneValues.lastValues
+        val lastSharedState = element.lastSharedState
+        val lastSceneState = sceneState.lastState
         val targetOffset =
             computeValue(
                 layoutImpl,
@@ -544,36 +511,36 @@
                 idleValue = targetOffsetInScene,
                 currentValue = { currentOffset },
                 lastValue = {
-                    lastValues.offset.takeIf { it.isSpecified }
-                        ?: lastSharedValues.offset.takeIf { it.isSpecified } ?: currentOffset
+                    lastSceneState.offset.takeIf { it.isSpecified }
+                        ?: lastSharedState.offset.takeIf { it.isSpecified } ?: currentOffset
                 },
                 ::lerp,
             )
 
-        lastSharedValues.offset = targetOffset
-        lastValues.offset = targetOffset
+        lastSharedState.offset = targetOffset
+        lastSceneState.offset = targetOffset
 
         // No need to place the element in this scene if we don't want to draw it anyways. Note that
-        // it's still important to compute the target offset and update lastValues, otherwise it
-        // will be out of date.
+        // it's still important to compute the target offset and update last(Shared|Scene)State,
+        // otherwise they will be out of date.
         if (!shouldDrawElement(layoutImpl, scene, element)) {
             return
         }
 
         val offset = (targetOffset - currentOffset).round()
-        if (isElementOpaque(layoutImpl, element, scene, sceneValues)) {
+        if (isElementOpaque(layoutImpl, element, scene, sceneState)) {
             // TODO(b/291071158): Call placeWithLayer() if offset != IntOffset.Zero and size is not
             // animated once b/305195729 is fixed. Test that drawing is not invalidated in that
             // case.
             placeable.place(offset)
-            lastSharedValues.alpha = 1f
-            lastValues.alpha = 1f
+            lastSharedState.alpha = 1f
+            lastSceneState.alpha = 1f
         } else {
             placeable.placeWithLayer(offset) {
-                val alpha = elementAlpha(layoutImpl, element, scene, sceneValues)
+                val alpha = elementAlpha(layoutImpl, element, scene, sceneState)
                 this.alpha = alpha
-                lastSharedValues.alpha = alpha
-                lastValues.alpha = alpha
+                lastSharedState.alpha = alpha
+                lastSceneState.alpha = alpha
             }
         }
     }
@@ -605,7 +572,7 @@
     layoutImpl: SceneTransitionLayoutImpl,
     scene: Scene,
     element: Element,
-    sceneValue: (Element.TargetValues) -> T,
+    sceneValue: (Element.SceneState) -> T,
     transformation: (ElementTransformations) -> PropertyTransformation<T>?,
     idleValue: T,
     currentValue: () -> T,
@@ -628,10 +595,10 @@
 
     val fromScene = transition.fromScene
     val toScene = transition.toScene
-    val fromValues = element.sceneValues[fromScene]
-    val toValues = element.sceneValues[toScene]
+    val fromState = element.sceneStates[fromScene]
+    val toState = element.sceneStates[toScene]
 
-    if (fromValues == null && toValues == null) {
+    if (fromState == null && toState == null) {
         // TODO(b/311600838): Throw an exception instead once layers of disposed elements are not
         // run anymore.
         return lastValue()
@@ -640,10 +607,10 @@
     // The element is shared: interpolate between the value in fromScene and the value in toScene.
     // TODO(b/290184746): Support non linear shared paths as well as a way to make sure that shared
     // elements follow the finger direction.
-    val isSharedElement = fromValues != null && toValues != null
+    val isSharedElement = fromState != null && toState != null
     if (isSharedElement && isSharedElementEnabled(layoutImpl.state, transition, element.key)) {
-        val start = sceneValue(fromValues!!)
-        val end = sceneValue(toValues!!)
+        val start = sceneValue(fromState!!)
+        val end = sceneValue(toState!!)
 
         // Make sure we don't read progress if values are the same and we don't need to interpolate,
         // so we don't invalidate the phase where this is read.
@@ -659,12 +626,12 @@
 
     // Get the transformed value, i.e. the target value at the beginning (for entering elements) or
     // end (for leaving elements) of the transition.
-    val sceneValues =
+    val sceneState =
         checkNotNull(
             when {
-                isSharedElement && scene.key == fromScene -> fromValues
-                isSharedElement -> toValues
-                else -> fromValues ?: toValues
+                isSharedElement && scene.key == fromScene -> fromState
+                isSharedElement -> toState
+                else -> fromState ?: toState
             }
         )
 
@@ -673,7 +640,7 @@
             layoutImpl,
             scene,
             element,
-            sceneValues,
+            sceneState,
             transition,
             idleValue,
         )
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 84d3b86..90f46bd 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
@@ -64,10 +64,10 @@
     identity: Any = Object(),
 
     /**
-     * Whether this element is a background and usually drawn below other elements. This should be
-     * set to true to make sure that shared backgrounds are drawn below elements of other scenes.
+     * The [ElementScenePicker] to use when deciding in which scene we should draw shared Elements
+     * or compose MovableElements.
      */
-    val isBackground: Boolean = false,
+    val scenePicker: ElementScenePicker = DefaultElementScenePicker,
 ) : Key(name, identity), ElementMatcher {
     @VisibleForTesting
     // TODO(b/240432457): Make internal once PlatformComposeSceneTransitionLayoutTestsUtils can
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 49df2f6..af3c099 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
@@ -16,27 +16,36 @@
 
 package com.android.compose.animation.scene
 
-import android.util.Log
 import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.BoxScope
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.State
 import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.movableContentOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.snapshots.Snapshot
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.drawBehind
-import androidx.compose.ui.draw.drawWithCache
-import androidx.compose.ui.graphics.Canvas
-import androidx.compose.ui.graphics.drawscope.draw
-import androidx.compose.ui.graphics.drawscope.drawIntoCanvas
-import androidx.compose.ui.graphics.nativeCanvas
-import androidx.compose.ui.layout.layout
-import androidx.compose.ui.unit.Constraints
+import androidx.compose.ui.layout.Layout
 import androidx.compose.ui.unit.IntSize
 
-private const val TAG = "MovableElement"
+@Composable
+internal fun Element(
+    layoutImpl: SceneTransitionLayoutImpl,
+    scene: Scene,
+    key: ElementKey,
+    modifier: Modifier,
+    content: @Composable ElementScope<ElementContentScope>.() -> Unit,
+) {
+    Box(modifier.element(layoutImpl, scene, key)) {
+        val sceneScope = scene.scope
+        val boxScope = this
+        val elementScope =
+            remember(layoutImpl, key, scene, sceneScope, boxScope) {
+                ElementScopeImpl(layoutImpl, key, scene, sceneScope, boxScope)
+            }
+
+        content(elementScope)
+    }
+}
 
 @Composable
 internal fun MovableElement(
@@ -44,72 +53,113 @@
     scene: Scene,
     key: ElementKey,
     modifier: Modifier,
-    content: @Composable MovableElementScope.() -> Unit,
+    content: @Composable ElementScope<MovableElementContentScope>.() -> Unit,
 ) {
     Box(modifier.element(layoutImpl, scene, key)) {
-        // Get the Element from the map. It will always be the same and we don't want to recompose
-        // every time an element is added/removed from SceneTransitionLayoutImpl.elements, so we
-        // disable read observation during the look-up in that map.
-        val element = Snapshot.withoutReadObservation { layoutImpl.elements.getValue(key) }
-        val movableElementScope =
-            remember(layoutImpl, element, scene) {
-                MovableElementScopeImpl(layoutImpl, element, scene)
+        val sceneScope = scene.scope
+        val boxScope = this
+        val elementScope =
+            remember(layoutImpl, key, scene, sceneScope, boxScope) {
+                MovableElementScopeImpl(layoutImpl, key, scene, sceneScope, boxScope)
             }
 
-        // The [Picture] to which we save the last drawing commands of this element. This is
-        // necessary because the content of this element might not be composed in this scene, in
-        // which case we still need to draw it.
-        val picture = element.picture
+        content(elementScope)
+    }
+}
 
+private abstract class BaseElementScope<ContentScope>(
+    private val layoutImpl: SceneTransitionLayoutImpl,
+    private val element: ElementKey,
+    private val scene: Scene,
+) : ElementScope<ContentScope> {
+    @Composable
+    override fun <T> animateElementValueAsState(
+        value: T,
+        key: ValueKey,
+        lerp: (start: T, stop: T, fraction: Float) -> T,
+        canOverflow: Boolean
+    ): AnimatedState<T> {
+        return animateSharedValueAsState(
+            layoutImpl,
+            scene.key,
+            element,
+            key,
+            value,
+            lerp,
+            canOverflow,
+        )
+    }
+}
+
+private class ElementScopeImpl(
+    layoutImpl: SceneTransitionLayoutImpl,
+    element: ElementKey,
+    scene: Scene,
+    private val sceneScope: SceneScope,
+    private val boxScope: BoxScope,
+) : BaseElementScope<ElementContentScope>(layoutImpl, element, scene) {
+    private val contentScope =
+        object : ElementContentScope, SceneScope by sceneScope, BoxScope by boxScope {}
+
+    @Composable
+    override fun content(content: @Composable ElementContentScope.() -> Unit) {
+        contentScope.content()
+    }
+}
+
+private class MovableElementScopeImpl(
+    private val layoutImpl: SceneTransitionLayoutImpl,
+    private val element: ElementKey,
+    private val scene: Scene,
+    private val sceneScope: BaseSceneScope,
+    private val boxScope: BoxScope,
+) : BaseElementScope<MovableElementContentScope>(layoutImpl, element, scene) {
+    private val contentScope =
+        object : MovableElementContentScope, BaseSceneScope by sceneScope, BoxScope by boxScope {}
+
+    @Composable
+    override fun content(content: @Composable MovableElementContentScope.() -> Unit) {
         // Whether we should compose the movable element here. The scene picker logic to know in
         // which scene we should compose/draw a movable element might depend on the current
         // transition progress, so we put this in a derivedStateOf to prevent many recompositions
         // during the transition.
+        // TODO(b/317026105): Use derivedStateOf only if the scene picker reads the progress in its
+        // logic.
         val shouldComposeMovableElement by
             remember(layoutImpl, scene.key, element) {
                 derivedStateOf { shouldComposeMovableElement(layoutImpl, scene.key, element) }
             }
 
         if (shouldComposeMovableElement) {
-            Box(
-                Modifier.drawWithCache {
-                    val width = size.width.toInt()
-                    val height = size.height.toInt()
-
-                    onDrawWithContent {
-                        // Save the draw commands into [picture] for later to draw the last content
-                        // even when this movable content is not composed.
-                        val pictureCanvas = Canvas(picture.beginRecording(width, height))
-                        draw(this, this.layoutDirection, pictureCanvas, this.size) {
-                            this@onDrawWithContent.drawContent()
+            val movableContent: MovableElementContent =
+                layoutImpl.movableContents[element]
+                    ?: movableContentOf {
+                            contentScope: MovableElementContentScope,
+                            content: @Composable MovableElementContentScope.() -> Unit ->
+                            contentScope.content()
                         }
-                        picture.endRecording()
+                        .also { layoutImpl.movableContents[element] = it }
 
-                        // Draw the content.
-                        drawIntoCanvas { canvas -> canvas.nativeCanvas.drawPicture(picture) }
-                    }
-                }
-            ) {
-                element.movableContent { movableElementScope.content() }
-            }
+            // Important: Don't introduce any parent Box or other layout here, because contentScope
+            // delegates its BoxScope implementation to the Box where this content() function is
+            // called, so it's important that this movableContent is composed directly under that
+            // Box.
+            movableContent(contentScope, content)
         } else {
-            // If we are not composed, we draw the previous drawing commands at the same size as the
-            // movable content when it was composed in this scene.
-            val sceneValues = element.sceneValues.getValue(scene.key)
-
-            Spacer(
-                Modifier.layout { measurable, _ ->
-                        val size =
-                            sceneValues.targetSize.takeIf { it != Element.SizeUnspecified }
-                                ?: IntSize.Zero
-                        val placeable =
-                            measurable.measure(Constraints.fixed(size.width, size.height))
-                        layout(size.width, size.height) { placeable.place(0, 0) }
-                    }
-                    .drawBehind {
-                        drawIntoCanvas { canvas -> canvas.nativeCanvas.drawPicture(picture) }
-                    }
-            )
+            // If we are not composed, we still need to lay out an empty space with the same *target
+            // 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.
+            Layout { _, _ ->
+                // No need to measure or place anything.
+                val size =
+                    placeholderContentSize(
+                        layoutImpl,
+                        scene.key,
+                        layoutImpl.elements.getValue(element),
+                    )
+                layout(size.width, size.height) {}
+            }
         }
     }
 }
@@ -117,7 +167,7 @@
 private fun shouldComposeMovableElement(
     layoutImpl: SceneTransitionLayoutImpl,
     scene: SceneKey,
-    element: Element,
+    element: ElementKey,
 ): Boolean {
     val transition =
         layoutImpl.state.currentTransition
@@ -130,72 +180,55 @@
     val fromReady = layoutImpl.isSceneReady(fromScene)
     val toReady = layoutImpl.isSceneReady(toScene)
 
-    val otherScene =
-        when (scene) {
-            fromScene -> toScene
-            toScene -> fromScene
-            else ->
-                error(
-                    "shouldComposeMovableElement(scene=$scene) called with fromScene=$fromScene " +
-                        "and toScene=$toScene"
-                )
-        }
-
-    val isShared = otherScene in element.sceneValues
-
-    if (isShared && !toReady && !fromReady) {
-        // This should usually not happen given that fromScene should be ready, but let's log a
-        // warning here in case it does so it helps debugging flicker issues caused by this part of
-        // the code.
-        Log.w(
-            TAG,
-            "MovableElement $element might have to be composed for the first time in both " +
-                "fromScene=$fromScene and toScene=$toScene. This will probably lead to a flicker " +
-                "where the size of the element will jump from IntSize.Zero to its actual size " +
-                "during the transition."
-        )
-    }
-
-    // Element is not shared in this transition.
-    if (!isShared) {
-        return true
-    }
-
-    // toScene is not ready (because we are composing it for the first time), so we compose it there
-    // first. This is the most common scenario when starting a transition that has a shared movable
-    // element.
-    if (!toReady) {
+    if (!fromReady && !toReady) {
+        // Neither of the scenes will be drawn, so where we compose it doesn't really matter. Note
+        // that we could have slightly more complicated logic here to optimize for this case, but
+        // it's not worth it given that readyScenes should disappear soon (b/316901148).
         return scene == toScene
     }
 
-    // This should usually not happen, but if we are also composing for the first time in fromScene
-    // then we should compose it there only.
-    if (!fromReady) {
-        return scene == fromScene
-    }
+    // If one of the scenes is not ready, compose it in the other one to make sure it is drawn.
+    if (!fromReady) return scene == toScene
+    if (!toReady) return scene == fromScene
 
+    // Always compose movable elements in the scene picked by their scene picker.
     return shouldDrawOrComposeSharedElement(
         layoutImpl,
         transition,
         scene,
-        element.key,
-        sharedElementTransformation(layoutImpl.state, transition, element.key),
+        element,
     )
 }
 
-private class MovableElementScopeImpl(
-    private val layoutImpl: SceneTransitionLayoutImpl,
-    private val element: Element,
-    private val scene: Scene,
-) : MovableElementScope {
-    @Composable
-    override fun <T> animateSharedValueAsState(
-        value: T,
-        debugName: String,
-        lerp: (start: T, stop: T, fraction: Float) -> T,
-        canOverflow: Boolean,
-    ): State<T> {
-        val key = remember { ValueKey(debugName) }
-        return animateSharedValueAsState(layoutImpl, scene, element, key, value, lerp, canOverflow)
+/**
+ * Return the size of the placeholder/space that is composed when the movable content is not
+ * composed in a scene.
+ */
+private fun placeholderContentSize(
+    layoutImpl: SceneTransitionLayoutImpl,
+    scene: SceneKey,
+    element: Element,
+): IntSize {
+    // If the content of the movable element was already composed in this scene before, use that
+    // target size.
+    val targetValueInScene = element.sceneStates.getValue(scene).targetSize
+    if (targetValueInScene != Element.SizeUnspecified) {
+        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
+
+    // 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 == scene) transition.toScene else transition.fromScene
+    val targetValueInOtherScene = element.sceneStates[otherScene]?.targetSize
+    if (targetValueInOtherScene != null && targetValueInOtherScene != Element.SizeUnspecified) {
+        return targetValueInOtherScene
+    }
+
+    return IntSize.Zero
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
index 560e92b..454c0ec 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
@@ -75,8 +75,8 @@
 
         if (
             bounds == null ||
-                bounds.lastSharedValues.size == Element.SizeUnspecified ||
-                bounds.lastSharedValues.offset == Offset.Unspecified
+                bounds.lastSharedState.size == Element.SizeUnspecified ||
+                bounds.lastSharedState.offset == Offset.Unspecified
         ) {
             drawContent()
             return
@@ -87,14 +87,14 @@
             canvas.withSaveLayer(size.toRect(), Paint()) {
                 drawContent()
 
-                val offset = bounds.lastSharedValues.offset - element.lastSharedValues.offset
+                val offset = bounds.lastSharedState.offset - element.lastSharedState.offset
                 translate(offset.x, offset.y) { drawHole(bounds) }
             }
         }
     }
 
     private fun DrawScope.drawHole(bounds: Element) {
-        val boundsSize = bounds.lastSharedValues.size.toSize()
+        val boundsSize = bounds.lastSharedState.size.toSize()
         if (shape == RectangleShape) {
             drawRect(Color.Black, size = boundsSize, blendMode = BlendMode.DstOut)
             return
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
index 30e50a9..3537b79 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
@@ -20,13 +20,10 @@
 import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Stable
-import androidx.compose.runtime.State
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableFloatStateOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.compose.runtime.snapshots.Snapshot
-import androidx.compose.runtime.snapshots.SnapshotStateMap
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Shape
@@ -45,16 +42,13 @@
     actions: Map<UserAction, SceneKey>,
     zIndex: Float,
 ) {
-    private val scope = SceneScopeImpl(layoutImpl, this)
+    internal val scope = SceneScopeImpl(layoutImpl, this)
 
     var content by mutableStateOf(content)
     var userActions by mutableStateOf(actions)
     var zIndex by mutableFloatStateOf(zIndex)
     var targetSize by mutableStateOf(IntSize.Zero)
 
-    /** The shared values in this scene that are not tied to a specific element. */
-    val sharedValues = SnapshotStateMap<ValueKey, Element.SharedValue<*>>()
-
     @Composable
     @OptIn(ExperimentalComposeUiApi::class)
     fun Content(modifier: Modifier = Modifier) {
@@ -77,7 +71,7 @@
     }
 }
 
-private class SceneScopeImpl(
+internal class SceneScopeImpl(
     private val layoutImpl: SceneTransitionLayoutImpl,
     private val scene: Scene,
 ) : SceneScope {
@@ -87,6 +81,42 @@
         return element(layoutImpl, scene, key)
     }
 
+    @Composable
+    override fun Element(
+        key: ElementKey,
+        modifier: Modifier,
+        content: @Composable (ElementScope<ElementContentScope>.() -> Unit)
+    ) {
+        Element(layoutImpl, scene, key, modifier, content)
+    }
+
+    @Composable
+    override fun MovableElement(
+        key: ElementKey,
+        modifier: Modifier,
+        content: @Composable (ElementScope<MovableElementContentScope>.() -> Unit)
+    ) {
+        MovableElement(layoutImpl, scene, key, modifier, content)
+    }
+
+    @Composable
+    override fun <T> animateSceneValueAsState(
+        value: T,
+        key: ValueKey,
+        lerp: (T, T, Float) -> T,
+        canOverflow: Boolean
+    ): AnimatedState<T> {
+        return animateSharedValueAsState(
+            layoutImpl = layoutImpl,
+            scene = scene.key,
+            element = null,
+            key = key,
+            value = value,
+            lerp = lerp,
+            canOverflow = canOverflow,
+        )
+    }
+
     override fun Modifier.horizontalNestedScrollToScene(
         leftBehavior: NestedScrollBehavior,
         rightBehavior: NestedScrollBehavior,
@@ -109,45 +139,6 @@
             bottomOrRightBehavior = bottomBehavior,
         )
 
-    @Composable
-    override fun <T> animateSharedValueAsState(
-        value: T,
-        key: ValueKey,
-        element: ElementKey?,
-        lerp: (T, T, Float) -> T,
-        canOverflow: Boolean
-    ): State<T> {
-        val element =
-            element?.let { key ->
-                Snapshot.withoutReadObservation {
-                    layoutImpl.elements[key]
-                        ?: error(
-                            "Element $key is not composed. Make sure to call " +
-                                "animateSharedXAsState *after* Modifier.element(key)."
-                        )
-                }
-            }
-
-        return animateSharedValueAsState(
-            layoutImpl,
-            scene,
-            element,
-            key,
-            value,
-            lerp,
-            canOverflow,
-        )
-    }
-
-    @Composable
-    override fun MovableElement(
-        key: ElementKey,
-        modifier: Modifier,
-        content: @Composable MovableElementScope.() -> Unit,
-    ) {
-        MovableElement(layoutImpl, scene, key, modifier, content)
-    }
-
     override fun Modifier.punchHole(
         element: ElementKey,
         bounds: ElementKey,
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 5eb339e..84fade89 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -22,9 +22,9 @@
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.Stable
-import androidx.compose.runtime.State
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
@@ -98,9 +98,9 @@
  */
 @DslMarker annotation class ElementDsl
 
-@ElementDsl
 @Stable
-interface SceneScope {
+@ElementDsl
+interface BaseSceneScope {
     /** The state of the [SceneTransitionLayout] in which this scene is contained. */
     val layoutState: SceneTransitionLayoutState
 
@@ -111,21 +111,74 @@
      * that the element can be transformed and animated when the scene transitions in or out.
      *
      * Additionally, this [key] will be used to detect elements that are shared between scenes to
-     * automatically interpolate their size, offset and [shared values][animateSharedValueAsState].
+     * automatically interpolate their size and offset. If you need to animate shared element values
+     * (i.e. values associated to this element that change depending on which scene it is composed
+     * in), use [Element] instead.
      *
      * Note that shared elements tagged using this function will be duplicated in each scene they
      * are part of, so any **internal** state (e.g. state created using `remember {
      * mutableStateOf(...) }`) will be lost. If you need to preserve internal state, you should use
      * [MovableElement] instead.
      *
+     * @see Element
      * @see MovableElement
-     *
-     * TODO(b/291566282): Migrate this to the new Modifier Node API and remove the @Composable
-     *   constraint.
      */
     fun Modifier.element(key: ElementKey): Modifier
 
     /**
+     * Create an element identified by [key].
+     *
+     * Similar to [element], this creates an element that will be automatically shared when present
+     * in multiple scenes and that can be transformed during transitions, the same way that
+     * [element] does.
+     *
+     * The only difference with [element] is that the provided [ElementScope] allows you to
+     * [animate element values][ElementScope.animateElementValueAsState] or specify its
+     * [movable content][Element.movableContent] that will be "moved" and composed only once during
+     * transitions (as opposed to [element] that duplicates shared elements) so that any internal
+     * state is preserved during and after the transition.
+     *
+     * @see element
+     * @see MovableElement
+     */
+    @Composable
+    fun Element(
+        key: ElementKey,
+        modifier: Modifier,
+
+        // TODO(b/317026105): As discussed in http://shortn/_gJVdltF8Si, remove the @Composable
+        // scope here to make sure that callers specify the content in ElementScope.content {} or
+        // ElementScope.movableContent {}.
+        content: @Composable ElementScope<ElementContentScope>.() -> Unit,
+    )
+
+    /**
+     * Create a *movable* element identified by [key].
+     *
+     * Similar to [Element], this creates an element that will be automatically shared when present
+     * in multiple scenes and that can be transformed during transitions, and you can also use the
+     * provided [ElementScope] to [animate element values][ElementScope.animateElementValueAsState].
+     *
+     * The important difference with [element] and [Element] is that this element
+     * [content][ElementScope.content] will be "moved" and composed only once during transitions, as
+     * opposed to [element] and [Element] that duplicates shared elements, so that any internal
+     * state is preserved during and after the transition.
+     *
+     * @see element
+     * @see Element
+     */
+    @Composable
+    fun MovableElement(
+        key: ElementKey,
+        modifier: Modifier,
+
+        // TODO(b/317026105): As discussed in http://shortn/_gJVdltF8Si, remove the @Composable
+        // scope here to make sure that callers specify the content in ElementScope.content {} or
+        // ElementScope.movableContent {}.
+        content: @Composable ElementScope<MovableElementContentScope>.() -> Unit,
+    )
+
+    /**
      * Adds a [NestedScrollConnection] to intercept scroll events not handled by the scrollable
      * component.
      *
@@ -150,51 +203,6 @@
     ): Modifier
 
     /**
-     * Create a *movable* element identified by [key].
-     *
-     * This creates an element that will be automatically shared when present in multiple scenes and
-     * that can be transformed during transitions, the same way that [element] does. The major
-     * difference with [element] is that elements created with [MovableElement] will be "moved" and
-     * composed only once during transitions (as opposed to [element] that duplicates shared
-     * elements) so that any internal state is preserved during and after the transition.
-     *
-     * @see element
-     */
-    @Composable
-    fun MovableElement(
-        key: ElementKey,
-        modifier: Modifier,
-        content: @Composable MovableElementScope.() -> Unit,
-    )
-
-    /**
-     * Animate some value of a shared element.
-     *
-     * @param value the value of this shared value in the current scene.
-     * @param key the key of this shared value.
-     * @param element the element associated with this value. If `null`, this value will be
-     *   associated at the scene level, which means that [key] should be used maximum once in the
-     *   same scene.
-     * @param lerp the *linear* interpolation function that should be used to interpolate between
-     *   two different values. Note that it has to be linear because the [fraction] passed to this
-     *   interpolator is already interpolated.
-     * @param canOverflow whether this value can overflow past the values it is interpolated
-     *   between, for instance because the transition is animated using a bouncy spring.
-     * @see animateSharedIntAsState
-     * @see animateSharedFloatAsState
-     * @see animateSharedDpAsState
-     * @see animateSharedColorAsState
-     */
-    @Composable
-    fun <T> animateSharedValueAsState(
-        value: T,
-        key: ValueKey,
-        element: ElementKey?,
-        lerp: (start: T, stop: T, fraction: Float) -> T,
-        canOverflow: Boolean,
-    ): State<T>
-
-    /**
      * Punch a hole in this [element] using the bounds of [bounds] in [scene] and the given [shape].
      *
      * Punching a hole in an element will "remove" any pixel drawn by that element in the hole area.
@@ -213,19 +221,96 @@
     fun Modifier.noResizeDuringTransitions(): Modifier
 }
 
-// TODO(b/291053742): Add animateSharedValueAsState(targetValue) without any ValueKey and ElementKey
-// arguments to allow sharing values inside a movable element.
+@Stable
 @ElementDsl
-interface MovableElementScope {
+interface SceneScope : BaseSceneScope {
+    /**
+     * Animate some value at the scene level.
+     *
+     * @param value the value of this shared value in the current scene.
+     * @param key the key of this shared value.
+     * @param lerp the *linear* interpolation function that should be used to interpolate between
+     *   two different values. Note that it has to be linear because the [fraction] passed to this
+     *   interpolator is already interpolated.
+     * @param canOverflow whether this value can overflow past the values it is interpolated
+     *   between, for instance because the transition is animated using a bouncy spring.
+     * @see animateSceneIntAsState
+     * @see animateSceneFloatAsState
+     * @see animateSceneDpAsState
+     * @see animateSceneColorAsState
+     */
     @Composable
-    fun <T> animateSharedValueAsState(
+    fun <T> animateSceneValueAsState(
         value: T,
-        debugName: String,
+        key: ValueKey,
         lerp: (start: T, stop: T, fraction: Float) -> T,
         canOverflow: Boolean,
-    ): State<T>
+    ): AnimatedState<T>
 }
 
+@Stable
+@ElementDsl
+interface ElementScope<ContentScope> {
+    /**
+     * Animate some value associated to this element.
+     *
+     * @param value the value of this shared value in the current scene.
+     * @param key the key of this shared value.
+     * @param lerp the *linear* interpolation function that should be used to interpolate between
+     *   two different values. Note that it has to be linear because the [fraction] passed to this
+     *   interpolator is already interpolated.
+     * @param canOverflow whether this value can overflow past the values it is interpolated
+     *   between, for instance because the transition is animated using a bouncy spring.
+     * @see animateElementIntAsState
+     * @see animateElementFloatAsState
+     * @see animateElementDpAsState
+     * @see animateElementColorAsState
+     */
+    @Composable
+    fun <T> animateElementValueAsState(
+        value: T,
+        key: ValueKey,
+        lerp: (start: T, stop: T, fraction: Float) -> T,
+        canOverflow: Boolean,
+    ): AnimatedState<T>
+
+    /**
+     * The content of this element.
+     *
+     * Important: This must be called exactly once, after all calls to [animateElementValueAsState].
+     */
+    @Composable fun content(content: @Composable ContentScope.() -> Unit)
+}
+
+/**
+ * The exact same scope as [androidx.compose.foundation.layout.BoxScope].
+ *
+ * We can't reuse BoxScope directly because of the @LayoutScopeMarker annotation on it, which would
+ * prevent us from calling Modifier.element() and other methods of [SceneScope] inside any Box {} in
+ * the [content][ElementScope.content] of a [SceneScope.Element] or a [SceneScope.MovableElement].
+ */
+@Stable
+@ElementDsl
+interface ElementBoxScope {
+    /** @see [androidx.compose.foundation.layout.BoxScope.align]. */
+    @Stable fun Modifier.align(alignment: Alignment): Modifier
+
+    /** @see [androidx.compose.foundation.layout.BoxScope.matchParentSize]. */
+    @Stable fun Modifier.matchParentSize(): Modifier
+}
+
+/** The scope for "normal" (not movable) elements. */
+@Stable @ElementDsl interface ElementContentScope : SceneScope, ElementBoxScope
+
+/**
+ * The scope for the content of movable elements.
+ *
+ * Note that it extends [BaseSceneScope] and not [SceneScope] because movable elements should not
+ * call [SceneScope.animateSceneValueAsState], given that their content is not composed in all
+ * scenes.
+ */
+@Stable @ElementDsl interface MovableElementContentScope : BaseSceneScope, ElementBoxScope
+
 /** An action performed by the user. */
 sealed interface UserAction
 
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 45e1a0f..0227aba 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
@@ -36,6 +36,16 @@
 import com.android.compose.ui.util.lerp
 import kotlinx.coroutines.CoroutineScope
 
+/**
+ * The type for the content of movable elements.
+ *
+ * TODO(b/317972419): Revert back to make this movable content have a single @Composable lambda
+ *   parameter.
+ */
+internal typealias MovableElementContent =
+    @Composable
+    (MovableElementContentScope, @Composable MovableElementContentScope.() -> Unit) -> Unit
+
 @Stable
 internal class SceneTransitionLayoutImpl(
     internal val state: SceneTransitionLayoutStateImpl,
@@ -56,16 +66,47 @@
     /**
      * The map of [Element]s.
      *
-     * Note that this map is *mutated* directly during composition, so it is a [SnapshotStateMap] to
-     * make sure that mutations are reverted if composition is cancelled.
+     * Important: [Element]s from this map should never be accessed during composition because the
+     * Elements are added when the associated Modifier.element() node is attached to the Modifier
+     * tree, i.e. after composition.
      */
-    internal val elements = SnapshotStateMap<ElementKey, Element>()
+    internal val elements = mutableMapOf<ElementKey, Element>()
+
+    /**
+     * The map of contents of movable elements.
+     *
+     * Note that given that this map is mutated directly during a composition, it has to be a
+     * [SnapshotStateMap] to make sure that mutations are reverted if composition is cancelled.
+     */
+    private var _movableContents: SnapshotStateMap<ElementKey, MovableElementContent>? = null
+    val movableContents: SnapshotStateMap<ElementKey, MovableElementContent>
+        get() =
+            _movableContents
+                ?: SnapshotStateMap<ElementKey, MovableElementContent>().also {
+                    _movableContents = it
+                }
+
+    /**
+     * The different values of a shared value keyed by a a [ValueKey] and the different elements and
+     * scenes it is associated to.
+     */
+    private var _sharedValues:
+        MutableMap<ValueKey, MutableMap<ElementKey?, SnapshotStateMap<SceneKey, *>>>? =
+        null
+    internal val sharedValues:
+        MutableMap<ValueKey, MutableMap<ElementKey?, SnapshotStateMap<SceneKey, *>>>
+        get() =
+            _sharedValues
+                ?: mutableMapOf<ValueKey, MutableMap<ElementKey?, SnapshotStateMap<SceneKey, *>>>()
+                    .also { _sharedValues = it }
 
     /**
      * The scenes that are "ready", i.e. they were composed and fully laid-out at least once.
      *
      * Note that this map is *read* during composition, so it is a [SnapshotStateMap] to make sure
      * that we recompose when modifications are made to this map.
+     *
+     * TODO(b/316901148): Remove this map.
      */
     private val readyScenes = SnapshotStateMap<SceneKey, Boolean>()
 
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 d1ba582..0607aa1 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
@@ -92,6 +92,20 @@
 
         /** Whether user input is currently driving the transition. */
         abstract val isUserInputOngoing: Boolean
+
+        /**
+         * Whether we are transitioning. If [from] or [to] is empty, we will also check that they
+         * match the scenes we are animating from and/or to.
+         */
+        fun isTransitioning(from: SceneKey? = null, to: SceneKey? = null): Boolean {
+            return (from == null || fromScene == from) && (to == null || toScene == to)
+        }
+
+        /** Whether we are transitioning from [scene] to [other], or from [other] to [scene]. */
+        fun isTransitioningBetween(scene: SceneKey, other: SceneKey): Boolean {
+            return isTransitioning(from = scene, to = other) ||
+                isTransitioning(from = other, to = scene)
+        }
     }
 }
 
@@ -111,13 +125,12 @@
 
     override fun isTransitioning(from: SceneKey?, to: SceneKey?): Boolean {
         val transition = currentTransition ?: return false
-        return (from == null || transition.fromScene == from) &&
-            (to == null || transition.toScene == to)
+        return transition.isTransitioning(from, to)
     }
 
     override fun isTransitioningBetween(scene: SceneKey, other: SceneKey): Boolean {
-        return isTransitioning(from = scene, to = other) ||
-            isTransitioning(from = other, to = scene)
+        val transition = currentTransition ?: return false
+        return transition.isTransitioningBetween(scene, other)
     }
 
     /** Start a new [transition], instantly interrupting any ongoing transition if there was one. */
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
index dfa2a9a..dc8505c 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
@@ -119,14 +119,8 @@
      *
      * @param enabled whether the matched element(s) should actually be shared in this transition.
      *   Defaults to true.
-     * @param scenePicker the [SharedElementScenePicker] to use when deciding in which scene we
-     *   should draw or compose this shared element.
      */
-    fun sharedElement(
-        matcher: ElementMatcher,
-        enabled: Boolean = true,
-        scenePicker: SharedElementScenePicker = DefaultSharedElementScenePicker,
-    )
+    fun sharedElement(matcher: ElementMatcher, enabled: Boolean = true)
 
     /**
      * Adds the transformations in [builder] but in reversed order. This allows you to partially
@@ -136,37 +130,65 @@
     fun reversed(builder: TransitionBuilder.() -> Unit)
 }
 
-interface SharedElementScenePicker {
+/**
+ * An interface to decide where we should draw shared Elements or compose MovableElements.
+ *
+ * @see DefaultElementScenePicker
+ * @see HighestZIndexScenePicker
+ * @see LowestZIndexScenePicker
+ * @see MovableElementScenePicker
+ */
+interface ElementScenePicker {
     /**
      * Return the scene in which [element] should be drawn (when using `Modifier.element(key)`) or
-     * composed (when using `MovableElement(key)`) during the transition from [fromScene] to
-     * [toScene].
+     * composed (when using `MovableElement(key)`) during the given [transition].
+     *
+     * Important: For [MovableElements][SceneScope.MovableElement], this scene picker will *always*
+     * be used during transitions to decide whether we should compose that element in a given scene
+     * or not. Therefore, you should make sure that the returned [SceneKey] contains the movable
+     * element, otherwise that element will not be composed in any scene during the transition.
      */
     fun sceneDuringTransition(
         element: ElementKey,
-        fromScene: SceneKey,
-        toScene: SceneKey,
-        progress: () -> Float,
+        transition: TransitionState.Transition,
         fromSceneZIndex: Float,
         toSceneZIndex: Float,
     ): SceneKey
-}
 
-object DefaultSharedElementScenePicker : SharedElementScenePicker {
-    override fun sceneDuringTransition(
+    /**
+     * Return [transition.fromScene] if it is in [scenes] and [transition.toScene] is not, or return
+     * [transition.toScene] if it is in [scenes] and [transition.fromScene] is not, otherwise throw
+     * an exception (i.e. if neither or both of fromScene and toScene are in [scenes]).
+     *
+     * This function can be useful when computing the scene in which a movable element should be
+     * composed.
+     */
+    fun pickSingleSceneIn(
+        scenes: Set<SceneKey>,
+        transition: TransitionState.Transition,
         element: ElementKey,
-        fromScene: SceneKey,
-        toScene: SceneKey,
-        progress: () -> Float,
-        fromSceneZIndex: Float,
-        toSceneZIndex: Float
     ): SceneKey {
-        // By default shared elements are drawn in the highest scene possible, unless it is a
-        // background.
-        return if (
-            (fromSceneZIndex > toSceneZIndex && !element.isBackground) ||
-                (fromSceneZIndex < toSceneZIndex && element.isBackground)
-        ) {
+        val fromScene = transition.fromScene
+        val toScene = transition.toScene
+        val fromSceneInScenes = scenes.contains(fromScene)
+        val toSceneInScenes = scenes.contains(toScene)
+        if (fromSceneInScenes && toSceneInScenes) {
+            error(
+                "Element $element can be in both $fromScene and $toScene. You should add a " +
+                    "special case for this transition before calling pickSingleSceneIn()."
+            )
+        }
+
+        if (!fromSceneInScenes && !toSceneInScenes) {
+            error(
+                "Element $element can be neither in $fromScene and $toScene. This either means " +
+                    "that you should add one of them in the scenes set passed to " +
+                    "pickSingleSceneIn(), or there is an internal error and this element was " +
+                    "composed when it shouldn't be."
+            )
+        }
+
+        return if (fromSceneInScenes) {
             fromScene
         } else {
             toScene
@@ -174,6 +196,66 @@
     }
 }
 
+/** An [ElementScenePicker] that draws/composes elements in the scene with the highest z-order. */
+object HighestZIndexScenePicker : ElementScenePicker {
+    override fun sceneDuringTransition(
+        element: ElementKey,
+        transition: TransitionState.Transition,
+        fromSceneZIndex: Float,
+        toSceneZIndex: Float
+    ): SceneKey {
+        return if (fromSceneZIndex > toSceneZIndex) {
+            transition.fromScene
+        } else {
+            transition.toScene
+        }
+    }
+}
+
+/** An [ElementScenePicker] that draws/composes elements in the scene with the lowest z-order. */
+object LowestZIndexScenePicker : ElementScenePicker {
+    override fun sceneDuringTransition(
+        element: ElementKey,
+        transition: TransitionState.Transition,
+        fromSceneZIndex: Float,
+        toSceneZIndex: Float
+    ): SceneKey {
+        return if (fromSceneZIndex < toSceneZIndex) {
+            transition.fromScene
+        } else {
+            transition.toScene
+        }
+    }
+}
+
+/**
+ * An [ElementScenePicker] that draws/composes elements in the scene we are transitioning to, iff
+ * that scene is in [scenes].
+ *
+ * This picker can be useful for movable elements whose content size depends on its content (because
+ * it wraps it) in at least one scene. That way, the target size of the MovableElement will be
+ * computed in the scene we are going to and, given that this element was probably already composed
+ * in the scene we are going from before starting the transition, the interpolated size of the
+ * movable element during the transition should be correct.
+ *
+ * The downside of this picker is that the zIndex of the element when going from scene A to scene B
+ * is not the same as when going from scene B to scene A, so it's not usable in situations where
+ * z-ordering during the transition matters.
+ */
+class MovableElementScenePicker(private val scenes: Set<SceneKey>) : ElementScenePicker {
+    override fun sceneDuringTransition(
+        element: ElementKey,
+        transition: TransitionState.Transition,
+        fromSceneZIndex: Float,
+        toSceneZIndex: Float,
+    ): SceneKey {
+        return if (scenes.contains(transition.toScene)) transition.toScene else transition.fromScene
+    }
+}
+
+/** The default [ElementScenePicker]. */
+val DefaultElementScenePicker = HighestZIndexScenePicker
+
 @TransitionDsl
 interface PropertyTransformationBuilder {
     /**
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
index 7046866..b96f9be 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
@@ -108,12 +108,8 @@
         range = null
     }
 
-    override fun sharedElement(
-        matcher: ElementMatcher,
-        enabled: Boolean,
-        scenePicker: SharedElementScenePicker,
-    ) {
-        transformations.add(SharedElementTransformation(matcher, enabled, scenePicker))
+    override fun sharedElement(matcher: ElementMatcher, enabled: Boolean) {
+        transformations.add(SharedElementTransformation(matcher, enabled))
     }
 
     override fun timestampRange(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
index 40c814e..124ec29 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
@@ -36,12 +36,12 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: IntSize,
     ): IntSize {
         fun anchorSizeIn(scene: SceneKey): IntSize {
-            val size = layoutImpl.elements[anchor]?.sceneValues?.get(scene)?.targetSize
+            val size = layoutImpl.elements[anchor]?.sceneStates?.get(scene)?.targetSize
             return if (size != null && size != Element.SizeUnspecified) {
                 IntSize(
                     width = if (anchorWidth) size.width else value.width,
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
index a1d6319..7aa702b 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
@@ -35,13 +35,13 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: Offset,
     ): Offset {
         val anchor = layoutImpl.elements[anchor] ?: return value
         fun anchorOffsetIn(scene: SceneKey): Offset? {
-            return anchor.sceneValues[scene]?.targetOffset?.takeIf { it.isSpecified }
+            return anchor.sceneStates[scene]?.targetOffset?.takeIf { it.isSpecified }
         }
 
         // [element] will move the same amount as [anchor] does.
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt
index d1cf8ee..6704a3b 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt
@@ -39,7 +39,7 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: Scale,
     ): Scale {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt
index 70534dd..191a8fb 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt
@@ -34,12 +34,12 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: Offset
     ): Offset {
         val sceneSize = scene.targetSize
-        val elementSize = sceneValues.targetSize
+        val elementSize = sceneState.targetSize
         if (elementSize == Element.SizeUnspecified) {
             return value
         }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt
index 17032dc..41f626e 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt
@@ -30,7 +30,7 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: Float
     ): Float {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt
index 233ae59..f5207dc 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt
@@ -37,7 +37,7 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: IntSize,
     ): IntSize {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt
index 0cd11b9..04254fb 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt
@@ -20,7 +20,6 @@
 import com.android.compose.animation.scene.ElementMatcher
 import com.android.compose.animation.scene.Scene
 import com.android.compose.animation.scene.SceneTransitionLayoutImpl
-import com.android.compose.animation.scene.SharedElementScenePicker
 import com.android.compose.animation.scene.TransitionState
 
 /** A transformation applied to one or more elements during a transition. */
@@ -48,7 +47,6 @@
 internal class SharedElementTransformation(
     override val matcher: ElementMatcher,
     internal val enabled: Boolean,
-    internal val scenePicker: SharedElementScenePicker,
 ) : Transformation
 
 /** A transformation that changes the value of an element property, like its size or offset. */
@@ -62,7 +60,7 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: T,
     ): T
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
index 864b937..04d5914 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
@@ -35,7 +35,7 @@
         layoutImpl: SceneTransitionLayoutImpl,
         scene: Scene,
         element: Element,
-        sceneValues: Element.TargetValues,
+        sceneState: Element.SceneState,
         transition: TransitionState.Transition,
         value: Offset,
     ): Offset {
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 5473186..a116501 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
@@ -18,10 +18,11 @@
 
 import androidx.compose.animation.core.LinearEasing
 import androidx.compose.animation.core.tween
-import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.lerp
@@ -32,6 +33,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.ui.util.lerp
 import com.google.common.truth.Truth.assertThat
+import org.junit.Assert.assertThrows
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -62,17 +64,17 @@
         onCurrentValueChanged: (Values) -> Unit,
     ) {
         val key = TestElements.Foo
-        Box(Modifier.element(key)) {
-            val int by animateSharedIntAsState(targetValues.int, TestValues.Value1, key)
-            val float by animateSharedFloatAsState(targetValues.float, TestValues.Value2, key)
-            val dp by animateSharedDpAsState(targetValues.dp, TestValues.Value3, key)
-            val color by
-                animateSharedColorAsState(targetValues.color, TestValues.Value4, element = null)
+        Element(key, Modifier) {
+            val int by animateElementIntAsState(targetValues.int, key = TestValues.Value1)
+            val float by animateElementFloatAsState(targetValues.float, key = TestValues.Value2)
+            val dp by animateElementDpAsState(targetValues.dp, key = TestValues.Value3)
+            val color by animateElementColorAsState(targetValues.color, key = TestValues.Value4)
 
-            // Make sure we read the values during composition, so that we recompose and call
-            // onCurrentValueChanged() with the latest values.
-            val currentValues = Values(int, float, dp, color)
-            SideEffect { onCurrentValueChanged(currentValues) }
+            content {
+                LaunchedEffect(Unit) {
+                    snapshotFlow { Values(int, float, dp, color) }.collect(onCurrentValueChanged)
+                }
+            }
         }
     }
 
@@ -83,30 +85,34 @@
     ) {
         val key = TestElements.Foo
         MovableElement(key = key, Modifier) {
-            val int by
-                animateSharedIntAsState(targetValues.int, debugName = TestValues.Value1.debugName)
-            val float by
-                animateSharedFloatAsState(
-                    targetValues.float,
-                    debugName = TestValues.Value2.debugName
-                )
-            val dp by
-                animateSharedDpAsState(targetValues.dp, debugName = TestValues.Value3.debugName)
-            val color by
-                animateSharedColorAsState(
-                    targetValues.color,
-                    debugName = TestValues.Value4.debugName
-                )
+            val int by animateElementIntAsState(targetValues.int, key = TestValues.Value1)
+            val float by animateElementFloatAsState(targetValues.float, key = TestValues.Value2)
+            val dp by animateElementDpAsState(targetValues.dp, key = TestValues.Value3)
+            val color by animateElementColorAsState(targetValues.color, key = TestValues.Value4)
 
-            // Make sure we read the values during composition, so that we recompose and call
-            // onCurrentValueChanged() with the latest values.
-            val currentValues = Values(int, float, dp, color)
-            SideEffect { onCurrentValueChanged(currentValues) }
+            LaunchedEffect(Unit) {
+                snapshotFlow { Values(int, float, dp, color) }.collect(onCurrentValueChanged)
+            }
+        }
+    }
+
+    @Composable
+    private fun SceneScope.SceneValues(
+        targetValues: Values,
+        onCurrentValueChanged: (Values) -> Unit,
+    ) {
+        val int by animateSceneIntAsState(targetValues.int, key = TestValues.Value1)
+        val float by animateSceneFloatAsState(targetValues.float, key = TestValues.Value2)
+        val dp by animateSceneDpAsState(targetValues.dp, key = TestValues.Value3)
+        val color by animateSceneColorAsState(targetValues.color, key = TestValues.Value4)
+
+        LaunchedEffect(Unit) {
+            snapshotFlow { Values(int, float, dp, color) }.collect(onCurrentValueChanged)
         }
     }
 
     @Test
-    fun animateSharedValues() {
+    fun animateElementValues() {
         val fromValues = Values(int = 0, float = 0f, dp = 0.dp, color = Color.Red)
         val toValues = Values(int = 100, float = 100f, dp = 100.dp, color = Color.Blue)
 
@@ -194,24 +200,183 @@
             }
 
             at(16) {
-                // Given that we use MovableElement here, animateSharedXAsState is composed only
-                // once, in the highest scene (in this case, in toScene).
-                assertThat(lastValueInFrom).isEqualTo(fromValues)
+                assertThat(lastValueInFrom).isEqualTo(lerp(fromValues, toValues, fraction = 0.25f))
                 assertThat(lastValueInTo).isEqualTo(lerp(fromValues, toValues, fraction = 0.25f))
             }
 
             at(32) {
-                assertThat(lastValueInFrom).isEqualTo(fromValues)
+                assertThat(lastValueInFrom).isEqualTo(lerp(fromValues, toValues, fraction = 0.5f))
                 assertThat(lastValueInTo).isEqualTo(lerp(fromValues, toValues, fraction = 0.5f))
             }
 
             at(48) {
-                assertThat(lastValueInFrom).isEqualTo(fromValues)
+                assertThat(lastValueInFrom).isEqualTo(lerp(fromValues, toValues, fraction = 0.75f))
                 assertThat(lastValueInTo).isEqualTo(lerp(fromValues, toValues, fraction = 0.75f))
             }
 
             after {
+                assertThat(lastValueInFrom).isEqualTo(toValues)
+                assertThat(lastValueInTo).isEqualTo(toValues)
+            }
+        }
+    }
+
+    @Test
+    fun animateSceneValues() {
+        val fromValues = Values(int = 0, float = 0f, dp = 0.dp, color = Color.Red)
+        val toValues = Values(int = 100, float = 100f, dp = 100.dp, color = Color.Blue)
+
+        var lastValueInFrom = fromValues
+        var lastValueInTo = toValues
+
+        rule.testTransition(
+            fromSceneContent = {
+                SceneValues(
+                    targetValues = fromValues,
+                    onCurrentValueChanged = { lastValueInFrom = it }
+                )
+            },
+            toSceneContent = {
+                SceneValues(targetValues = toValues, onCurrentValueChanged = { lastValueInTo = it })
+            },
+            transition = {
+                // The transition lasts 64ms = 4 frames.
+                spec = tween(durationMillis = 16 * 4, easing = LinearEasing)
+            },
+            fromScene = TestScenes.SceneA,
+            toScene = TestScenes.SceneB,
+        ) {
+            before {
                 assertThat(lastValueInFrom).isEqualTo(fromValues)
+
+                // to was not composed yet, so lastValueInTo was not set yet.
+                assertThat(lastValueInTo).isEqualTo(toValues)
+            }
+
+            at(16) {
+                // Given that we use scene values here, animateSceneXAsState is composed in both
+                // scenes and values should be interpolated with the transition fraction.
+                val expectedValues = lerp(fromValues, toValues, fraction = 0.25f)
+                assertThat(lastValueInFrom).isEqualTo(expectedValues)
+                assertThat(lastValueInTo).isEqualTo(expectedValues)
+            }
+
+            at(32) {
+                val expectedValues = lerp(fromValues, toValues, fraction = 0.5f)
+                assertThat(lastValueInFrom).isEqualTo(expectedValues)
+                assertThat(lastValueInTo).isEqualTo(expectedValues)
+            }
+
+            at(48) {
+                val expectedValues = lerp(fromValues, toValues, fraction = 0.75f)
+                assertThat(lastValueInFrom).isEqualTo(expectedValues)
+                assertThat(lastValueInTo).isEqualTo(expectedValues)
+            }
+
+            after {
+                assertThat(lastValueInFrom).isEqualTo(toValues)
+                assertThat(lastValueInTo).isEqualTo(toValues)
+            }
+        }
+    }
+
+    @Test
+    fun readingAnimatedStateValueDuringCompositionThrows() {
+        assertThrows(IllegalStateException::class.java) {
+            rule.testTransition(
+                fromSceneContent = { animateSceneIntAsState(0, TestValues.Value1).value },
+                toSceneContent = {},
+                transition = {},
+            ) {}
+        }
+    }
+
+    @Test
+    fun readingAnimatedStateValueDuringCompositionIsStillPossible() {
+        @Composable
+        fun SceneScope.SceneValuesDuringComposition(
+            targetValues: Values,
+            onCurrentValueChanged: (Values) -> Unit,
+        ) {
+            val int by
+                animateSceneIntAsState(targetValues.int, key = TestValues.Value1)
+                    .unsafeCompositionState(targetValues.int)
+            val float by
+                animateSceneFloatAsState(targetValues.float, key = TestValues.Value2)
+                    .unsafeCompositionState(targetValues.float)
+            val dp by
+                animateSceneDpAsState(targetValues.dp, key = TestValues.Value3)
+                    .unsafeCompositionState(targetValues.dp)
+            val color by
+                animateSceneColorAsState(targetValues.color, key = TestValues.Value4)
+                    .unsafeCompositionState(targetValues.color)
+
+            val values = Values(int, float, dp, color)
+            SideEffect { onCurrentValueChanged(values) }
+        }
+
+        val fromValues = Values(int = 0, float = 0f, dp = 0.dp, color = Color.Red)
+        val toValues = Values(int = 100, float = 100f, dp = 100.dp, color = Color.Blue)
+
+        var lastValueInFrom = fromValues
+        var lastValueInTo = toValues
+
+        rule.testTransition(
+            fromSceneContent = {
+                SceneValuesDuringComposition(
+                    targetValues = fromValues,
+                    onCurrentValueChanged = { lastValueInFrom = it },
+                )
+            },
+            toSceneContent = {
+                SceneValuesDuringComposition(
+                    targetValues = toValues,
+                    onCurrentValueChanged = { lastValueInTo = it },
+                )
+            },
+            transition = {
+                // The transition lasts 64ms = 4 frames.
+                spec = tween(durationMillis = 16 * 4, easing = LinearEasing)
+            },
+        ) {
+            before {
+                assertThat(lastValueInFrom).isEqualTo(fromValues)
+
+                // to was not composed yet, so lastValueInTo was not set yet.
+                assertThat(lastValueInTo).isEqualTo(toValues)
+            }
+
+            at(16) {
+                // Because we are using unsafeCompositionState(), values are one frame behind their
+                // expected progress so at this first frame we are at progress = 0% instead of 25%.
+                val expectedValues = lerp(fromValues, toValues, fraction = 0f)
+                assertThat(lastValueInFrom).isEqualTo(expectedValues)
+                assertThat(lastValueInTo).isEqualTo(expectedValues)
+            }
+
+            at(32) {
+                // One frame behind, so 25% instead of 50%.
+                val expectedValues = lerp(fromValues, toValues, fraction = 0.25f)
+                assertThat(lastValueInFrom).isEqualTo(expectedValues)
+                assertThat(lastValueInTo).isEqualTo(expectedValues)
+            }
+
+            at(48) {
+                // One frame behind, so 50% instead of 75%.
+                val expectedValues = lerp(fromValues, toValues, fraction = 0.5f)
+                assertThat(lastValueInFrom).isEqualTo(expectedValues)
+                assertThat(lastValueInTo).isEqualTo(expectedValues)
+            }
+
+            after {
+                // from should have been last composed at progress = 100% before it is removed from
+                // composition, but given that we are one frame behind the last values are stuck at
+                // 75%.
+                assertThat(lastValueInFrom).isEqualTo(lerp(fromValues, toValues, fraction = 0.75f))
+
+                // The after {} block resumes the clock and will run as many frames as necessary so
+                // that the application is idle, so the toScene settle to the idle state and to the
+                // final values.
                 assertThat(lastValueInTo).isEqualTo(toValues)
             }
         }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementScenePickerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementScenePickerTest.kt
new file mode 100644
index 0000000..3b022e8
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementScenePickerTest.kt
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2023 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.animation.core.LinearEasing
+import androidx.compose.animation.core.tween
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.size
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.test.assertIsDisplayed
+import androidx.compose.ui.test.assertIsNotDisplayed
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.unit.dp
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class ElementScenePickerTest {
+    @get:Rule val rule = createComposeRule()
+
+    @Test
+    fun highestZIndexPicker() {
+        val key = ElementKey("TestElement", scenePicker = HighestZIndexScenePicker)
+        rule.testTransition(
+            fromSceneContent = { Box(Modifier.element(key).size(10.dp)) },
+            toSceneContent = { Box(Modifier.element(key).size(10.dp)) },
+            transition = { spec = tween(4 * 16, easing = LinearEasing) },
+            fromScene = TestScenes.SceneA,
+            toScene = TestScenes.SceneB,
+        ) {
+            before {
+                onElement(key, TestScenes.SceneA).assertIsDisplayed()
+                onElement(key, TestScenes.SceneB).assertDoesNotExist()
+            }
+            at(32) {
+                // Scene B has the highest index, so the element is placed only there.
+                onElement(key, TestScenes.SceneA).assertExists().assertIsNotDisplayed()
+                onElement(key, TestScenes.SceneB).assertIsDisplayed()
+            }
+            after {
+                onElement(key, TestScenes.SceneA).assertDoesNotExist()
+                onElement(key, TestScenes.SceneB).assertIsDisplayed()
+            }
+        }
+    }
+
+    @Test
+    fun lowestZIndexPicker() {
+        val key = ElementKey("TestElement", scenePicker = LowestZIndexScenePicker)
+        rule.testTransition(
+            fromSceneContent = { Box(Modifier.element(key).size(10.dp)) },
+            toSceneContent = { Box(Modifier.element(key).size(10.dp)) },
+            transition = { spec = tween(4 * 16, easing = LinearEasing) },
+            fromScene = TestScenes.SceneA,
+            toScene = TestScenes.SceneB,
+        ) {
+            before {
+                onElement(key, TestScenes.SceneA).assertIsDisplayed()
+                onElement(key, TestScenes.SceneB).assertDoesNotExist()
+            }
+            at(32) {
+                // Scene A has the lowest index, so the element is placed only there.
+                onElement(key, TestScenes.SceneA).assertIsDisplayed()
+                onElement(key, TestScenes.SceneB).assertExists().assertIsNotDisplayed()
+            }
+            after {
+                onElement(key, TestScenes.SceneA).assertDoesNotExist()
+                onElement(key, TestScenes.SceneB).assertIsDisplayed()
+            }
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
index da5a0a0..54c5de7 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
@@ -306,7 +306,7 @@
 
         assertThat(layoutImpl.elements.keys).containsExactly(key)
         val element = layoutImpl.elements.getValue(key)
-        assertThat(element.sceneValues.keys).containsExactly(TestScenes.SceneB)
+        assertThat(element.sceneStates.keys).containsExactly(TestScenes.SceneB)
 
         // Scene C, state 0: the same element is reused.
         currentScene = TestScenes.SceneC
@@ -315,7 +315,7 @@
 
         assertThat(layoutImpl.elements.keys).containsExactly(key)
         assertThat(layoutImpl.elements.getValue(key)).isSameInstanceAs(element)
-        assertThat(element.sceneValues.keys).containsExactly(TestScenes.SceneC)
+        assertThat(element.sceneStates.keys).containsExactly(TestScenes.SceneC)
 
         // Scene C, state 1: the same element is reused.
         sceneCState = 1
@@ -323,7 +323,7 @@
 
         assertThat(layoutImpl.elements.keys).containsExactly(key)
         assertThat(layoutImpl.elements.getValue(key)).isSameInstanceAs(element)
-        assertThat(element.sceneValues.keys).containsExactly(TestScenes.SceneC)
+        assertThat(element.sceneStates.keys).containsExactly(TestScenes.SceneC)
 
         // Scene D, state 0: the same element is reused.
         currentScene = TestScenes.SceneD
@@ -332,7 +332,7 @@
 
         assertThat(layoutImpl.elements.keys).containsExactly(key)
         assertThat(layoutImpl.elements.getValue(key)).isSameInstanceAs(element)
-        assertThat(element.sceneValues.keys).containsExactly(TestScenes.SceneD)
+        assertThat(element.sceneStates.keys).containsExactly(TestScenes.SceneD)
 
         // Scene D, state 1: the same element is reused.
         sceneDState = 1
@@ -340,13 +340,13 @@
 
         assertThat(layoutImpl.elements.keys).containsExactly(key)
         assertThat(layoutImpl.elements.getValue(key)).isSameInstanceAs(element)
-        assertThat(element.sceneValues.keys).containsExactly(TestScenes.SceneD)
+        assertThat(element.sceneStates.keys).containsExactly(TestScenes.SceneD)
 
         // Scene D, state 2: the element is removed from the map.
         sceneDState = 2
         rule.waitForIdle()
 
-        assertThat(element.sceneValues).isEmpty()
+        assertThat(element.sceneStates).isEmpty()
         assertThat(layoutImpl.elements).isEmpty()
     }
 
@@ -442,7 +442,7 @@
         // There is only Foo in the elements map.
         assertThat(layoutImpl.elements.keys).containsExactly(TestElements.Foo)
         val fooElement = layoutImpl.elements.getValue(TestElements.Foo)
-        assertThat(fooElement.sceneValues.keys).containsExactly(TestScenes.SceneA)
+        assertThat(fooElement.sceneStates.keys).containsExactly(TestScenes.SceneA)
 
         key = TestElements.Bar
 
@@ -450,8 +450,8 @@
         rule.waitForIdle()
         assertThat(layoutImpl.elements.keys).containsExactly(TestElements.Bar)
         val barElement = layoutImpl.elements.getValue(TestElements.Bar)
-        assertThat(barElement.sceneValues.keys).containsExactly(TestScenes.SceneA)
-        assertThat(fooElement.sceneValues).isEmpty()
+        assertThat(barElement.sceneStates.keys).containsExactly(TestScenes.SceneA)
+        assertThat(fooElement.sceneStates).isEmpty()
     }
 
     @Test
@@ -505,7 +505,7 @@
         // There is only Foo in the elements map.
         assertThat(layoutImpl.elements.keys).containsExactly(TestElements.Foo)
         val element = layoutImpl.elements.getValue(TestElements.Foo)
-        val sceneValues = element.sceneValues
+        val sceneValues = element.sceneStates
         assertThat(sceneValues.keys).containsExactly(TestScenes.SceneA)
 
         // Get the ElementModifier node that should be reused later on when coming back to this
@@ -528,7 +528,7 @@
 
         assertThat(layoutImpl.elements.keys).containsExactly(TestElements.Foo)
         val newElement = layoutImpl.elements.getValue(TestElements.Foo)
-        val newSceneValues = newElement.sceneValues
+        val newSceneValues = newElement.sceneStates
         assertThat(newElement).isNotEqualTo(element)
         assertThat(newSceneValues).isNotEqualTo(sceneValues)
         assertThat(newSceneValues.keys).containsExactly(TestScenes.SceneA)
@@ -579,11 +579,11 @@
 
         fun foo() = layoutImpl().elements[TestElements.Foo] ?: error("Foo not in elements map")
 
-        fun Element.lastSharedOffset() = lastSharedValues.offset.toDpOffset()
+        fun Element.lastSharedOffset() = lastSharedState.offset.toDpOffset()
 
         fun Element.lastOffsetIn(scene: SceneKey) =
-            (sceneValues[scene] ?: error("$scene not in sceneValues map"))
-                .lastValues
+            (sceneStates[scene] ?: error("$scene not in sceneValues map"))
+                .lastState
                 .offset
                 .toDpOffset()
 
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementScenePickerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementScenePickerTest.kt
new file mode 100644
index 0000000..fb46a34
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementScenePickerTest.kt
@@ -0,0 +1,53 @@
+/*
+ * 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.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class MovableElementScenePickerTest {
+    @Test
+    fun toSceneInScenes() {
+        val picker = MovableElementScenePicker(scenes = setOf(TestScenes.SceneA, TestScenes.SceneB))
+        assertThat(
+                picker.sceneDuringTransition(
+                    TestElements.Foo,
+                    transition(from = TestScenes.SceneA, to = TestScenes.SceneB),
+                    fromSceneZIndex = 0f,
+                    toSceneZIndex = 1f,
+                )
+            )
+            .isEqualTo(TestScenes.SceneB)
+    }
+
+    @Test
+    fun toSceneNotInScenes() {
+        val picker = MovableElementScenePicker(scenes = emptySet())
+        assertThat(
+                picker.sceneDuringTransition(
+                    TestElements.Foo,
+                    transition(from = TestScenes.SceneA, to = TestScenes.SceneB),
+                    fromSceneZIndex = 0f,
+                    toSceneZIndex = 1f,
+                )
+            )
+            .isEqualTo(TestScenes.SceneA)
+    }
+}
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 3cd65cd..35cb691 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
@@ -28,19 +28,24 @@
 import androidx.compose.runtime.mutableIntStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.test.assertIsDisplayed
 import androidx.compose.ui.test.assertIsNotDisplayed
+import androidx.compose.ui.test.assertPositionInRootIsEqualTo
 import androidx.compose.ui.test.hasParent
 import androidx.compose.ui.test.hasText
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onAllNodesWithText
+import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import androidx.compose.ui.unit.dp
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.test.assertSizeIsEqualTo
 import com.google.common.truth.Truth.assertThat
+import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -58,7 +63,7 @@
 
     @Composable
     private fun SceneScope.MovableCounter(key: ElementKey, modifier: Modifier) {
-        MovableElement(key, modifier) { Counter() }
+        MovableElement(key, modifier) { content { Counter() } }
     }
 
     @Test
@@ -142,39 +147,37 @@
 
     @Test
     fun movableElementIsMovedAndComposedOnlyOnce() {
-        rule.testTransition(
-            fromSceneContent = { MovableCounter(TestElements.Foo, Modifier.size(50.dp)) },
-            toSceneContent = { MovableCounter(TestElements.Foo, Modifier.size(100.dp)) },
-            transition = {
-                spec = tween(durationMillis = 16 * 4, easing = LinearEasing)
-                sharedElement(
-                    TestElements.Foo,
-                    scenePicker =
-                        object : SharedElementScenePicker {
-                            override fun sceneDuringTransition(
-                                element: ElementKey,
-                                fromScene: SceneKey,
-                                toScene: SceneKey,
-                                progress: () -> Float,
-                                fromSceneZIndex: Float,
-                                toSceneZIndex: Float
-                            ): SceneKey {
-                                assertThat(fromScene).isEqualTo(TestScenes.SceneA)
-                                assertThat(toScene).isEqualTo(TestScenes.SceneB)
-                                assertThat(fromSceneZIndex).isEqualTo(0)
-                                assertThat(toSceneZIndex).isEqualTo(1)
+        val key =
+            ElementKey(
+                "Foo",
+                scenePicker =
+                    object : ElementScenePicker {
+                        override fun sceneDuringTransition(
+                            element: ElementKey,
+                            transition: TransitionState.Transition,
+                            fromSceneZIndex: Float,
+                            toSceneZIndex: Float
+                        ): SceneKey {
+                            assertThat(transition.fromScene).isEqualTo(TestScenes.SceneA)
+                            assertThat(transition.toScene).isEqualTo(TestScenes.SceneB)
+                            assertThat(fromSceneZIndex).isEqualTo(0)
+                            assertThat(toSceneZIndex).isEqualTo(1)
 
-                                // Compose Foo in Scene A if progress < 0.65f, otherwise compose it
-                                // in Scene B.
-                                return if (progress() < 0.65f) {
-                                    TestScenes.SceneA
-                                } else {
-                                    TestScenes.SceneB
-                                }
+                            // Compose Foo in Scene A if progress < 0.65f, otherwise compose it
+                            // in Scene B.
+                            return if (transition.progress < 0.65f) {
+                                TestScenes.SceneA
+                            } else {
+                                TestScenes.SceneB
                             }
                         }
-                )
-            },
+                    }
+            )
+
+        rule.testTransition(
+            fromSceneContent = { MovableCounter(key, Modifier.size(50.dp)) },
+            toSceneContent = { MovableCounter(key, Modifier.size(100.dp)) },
+            transition = { spec = tween(durationMillis = 16 * 4, easing = LinearEasing) },
             fromScene = TestScenes.SceneA,
             toScene = TestScenes.SceneB,
         ) {
@@ -257,4 +260,73 @@
             }
         }
     }
+
+    @Test
+    @Ignore("b/317972419#comment2")
+    fun movableElementContentIsRecomposedIfContentParametersChange() {
+        @Composable
+        fun SceneScope.MovableFoo(text: String, modifier: Modifier = Modifier) {
+            MovableElement(TestElements.Foo, modifier) { content { Text(text) } }
+        }
+
+        rule.testTransition(
+            fromSceneContent = { MovableFoo(text = "fromScene") },
+            toSceneContent = { MovableFoo(text = "toScene") },
+            transition = { spec = tween(durationMillis = 16 * 4, easing = LinearEasing) },
+            fromScene = TestScenes.SceneA,
+            toScene = TestScenes.SceneB,
+        ) {
+            // Before the transition, only fromScene is composed.
+            before {
+                rule.onNodeWithText("fromScene").assertIsDisplayed()
+                rule.onNodeWithText("toScene").assertDoesNotExist()
+            }
+
+            // During the transition, the element is composed in toScene.
+            at(32) {
+                rule.onNodeWithText("fromScene").assertDoesNotExist()
+                rule.onNodeWithText("toScene").assertIsDisplayed()
+            }
+
+            // At the end of the transition, the element is composed in toScene.
+            after {
+                rule.onNodeWithText("fromScene").assertDoesNotExist()
+                rule.onNodeWithText("toScene").assertIsDisplayed()
+            }
+        }
+    }
+
+    @Test
+    fun elementScopeExtendsBoxScope() {
+        rule.setContent {
+            TestSceneScope {
+                Element(TestElements.Foo, Modifier.size(200.dp)) {
+                    content {
+                        Box(Modifier.testTag("bottomEnd").align(Alignment.BottomEnd))
+                        Box(Modifier.testTag("matchParentSize").matchParentSize())
+                    }
+                }
+            }
+        }
+
+        rule.onNodeWithTag("bottomEnd").assertPositionInRootIsEqualTo(200.dp, 200.dp)
+        rule.onNodeWithTag("matchParentSize").assertSizeIsEqualTo(200.dp, 200.dp)
+    }
+
+    @Test
+    fun movableElementScopeExtendsBoxScope() {
+        rule.setContent {
+            TestSceneScope {
+                MovableElement(TestElements.Foo, Modifier.size(200.dp)) {
+                    content {
+                        Box(Modifier.testTag("bottomEnd").align(Alignment.BottomEnd))
+                        Box(Modifier.testTag("matchParentSize").matchParentSize())
+                    }
+                }
+            }
+        }
+
+        rule.onNodeWithTag("bottomEnd").assertPositionInRootIsEqualTo(200.dp, 200.dp)
+        rule.onNodeWithTag("matchParentSize").assertSizeIsEqualTo(200.dp, 200.dp)
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
index c5b8d9a..75dee47 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
@@ -50,13 +50,4 @@
         assertThat(state.isTransitioning(to = TestScenes.SceneA)).isFalse()
         assertThat(state.isTransitioning(from = TestScenes.SceneA, to = TestScenes.SceneB)).isTrue()
     }
-
-    private fun transition(from: SceneKey, to: SceneKey): TransitionState.Transition {
-        return object : TransitionState.Transition(from, to) {
-            override val currentScene: SceneKey = from
-            override val progress: Float = 0f
-            override val isInitiatedByUserInput: Boolean = false
-            override val isUserInputOngoing: Boolean = false
-        }
-    }
 }
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 ebbd500..649e499 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
@@ -113,25 +113,21 @@
 
     @Composable
     private fun SceneScope.SharedFoo(size: Dp, childOffset: Dp, modifier: Modifier = Modifier) {
-        Box(
-            modifier
-                .size(size)
-                .background(Color.Red)
-                .element(TestElements.Foo)
-                .testTag(TestElements.Foo.debugName)
-        ) {
+        Element(TestElements.Foo, modifier.size(size).background(Color.Red)) {
             // Offset the single child of Foo by some animated shared offset.
-            val offset by animateSharedDpAsState(childOffset, TestValues.Value1, TestElements.Foo)
+            val offset by animateElementDpAsState(childOffset, TestValues.Value1)
 
-            Box(
-                Modifier.offset {
-                        val pxOffset = offset.roundToPx()
-                        IntOffset(pxOffset, pxOffset)
-                    }
-                    .size(30.dp)
-                    .background(Color.Blue)
-                    .testTag(TestElements.Bar.debugName)
-            )
+            content {
+                Box(
+                    Modifier.offset {
+                            val pxOffset = offset.roundToPx()
+                            IntOffset(pxOffset, pxOffset)
+                        }
+                        .size(30.dp)
+                        .background(Color.Blue)
+                        .testTag(TestElements.Bar.debugName)
+                )
+            }
         }
     }
 
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
new file mode 100644
index 0000000..238b21e1
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
@@ -0,0 +1,33 @@
+/*
+ * 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
+
+/** A utility to easily create a [TransitionState.Transition] in tests. */
+fun transition(
+    from: SceneKey,
+    to: SceneKey,
+    progress: () -> Float = { 0f },
+    isInitiatedByUserInput: Boolean = false,
+    isUserInputOngoing: Boolean = false,
+): TransitionState.Transition {
+    return object : TransitionState.Transition(from, to) {
+        override val currentScene: SceneKey = from
+        override val progress: Float = progress()
+        override val isInitiatedByUserInput: Boolean = isInitiatedByUserInput
+        override val isUserInputOngoing: Boolean = isUserInputOngoing
+    }
+}