Merge changes from topic "stl-no-modifier-transformation" into main

* changes:
  Migrate Modifier.multiPointerDraggable to the Node Modifier API
  Annotate Stable classes/interfaces to optimize compositions
  Move Element drawing logic in ElementNode
  Remove Modifier transformations (1/2)
  Move PunchHole.kt to the animation/scene/ directory

14 files changed, 252 insertions, 152 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToShadeTransition.kt
index fadbdce80cbf..22dc0aee45b4 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToShadeTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToShadeTransition.kt
@@ -10,7 +10,6 @@ import com.android.systemui.shade.ui.composable.Shade
 fun TransitionBuilder.lockscreenToShadeTransition() {
     spec = tween(durationMillis = 500)
 
-    punchHole(Shade.Elements.QuickSettings, bounds = Shade.Elements.Scrim, Shade.Shapes.Scrim)
     translate(Shade.Elements.Scrim, Edge.Top, startsOutsideLayoutBounds = false)
     fractionRange(end = 0.5f) {
         fade(Shade.Elements.ScrimBackground)
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 3b999e304491..2b1195229c76 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
@@ -17,6 +17,7 @@
 package com.android.compose.animation.scene
 
 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
@@ -25,16 +26,17 @@ 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.draw.drawWithContent
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.geometry.isSpecified
 import androidx.compose.ui.geometry.isUnspecified
 import androidx.compose.ui.geometry.lerp
+import androidx.compose.ui.graphics.drawscope.ContentDrawScope
 import androidx.compose.ui.graphics.drawscope.scale
 import androidx.compose.ui.layout.IntermediateMeasureScope
 import androidx.compose.ui.layout.Measurable
 import androidx.compose.ui.layout.Placeable
 import androidx.compose.ui.layout.intermediateLayout
+import androidx.compose.ui.node.DrawModifierNode
 import androidx.compose.ui.node.ModifierNodeElement
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.unit.Constraints
@@ -46,6 +48,7 @@ import com.android.compose.ui.util.lerp
 import kotlinx.coroutines.launch
 
 /** An element on screen, that can be composed in one or more scenes. */
+@Stable
 internal class Element(val key: ElementKey) {
     /**
      * The last values of this element, coming from any scene. Note that this value will be unstable
@@ -90,6 +93,7 @@ internal class Element(val key: ElementKey) {
     }
 
     /** The target values of this element in a given scene. */
+    @Stable
     class TargetValues(val scene: SceneKey) {
         val lastValues = Values()
 
@@ -107,6 +111,7 @@ internal class Element(val key: ElementKey) {
     }
 
     /** A shared value of this element. */
+    @Stable
     class SharedValue<T>(val key: ValueKey, initialValue: T) {
         var value by mutableStateOf(initialValue)
     }
@@ -126,6 +131,7 @@ data class Scale(val scaleX: Float, val scaleY: Float, val pivot: Offset = Offse
 
 /** The implementation of [SceneScope.element]. */
 @OptIn(ExperimentalComposeUiApi::class)
+@Stable
 internal fun Modifier.element(
     layoutImpl: SceneTransitionLayoutImpl,
     scene: Scene,
@@ -144,24 +150,9 @@ internal fun Modifier.element(
                 ?: Element.TargetValues(scene.key).also { element.sceneValues[scene.key] = it }
     }
 
-    return this.then(ElementModifier(layoutImpl, element, sceneValues))
-        .drawWithContent {
-            if (shouldDrawElement(layoutImpl, scene, element)) {
-                val drawScale = getDrawScale(layoutImpl, element, scene, sceneValues)
-                if (drawScale == Scale.Default) {
-                    drawContent()
-                } else {
-                    scale(
-                        drawScale.scaleX,
-                        drawScale.scaleY,
-                        if (drawScale.pivot.isUnspecified) center else drawScale.pivot,
-                    ) {
-                        this@drawWithContent.drawContent()
-                    }
-                }
-            }
-        }
-        .modifierTransformations(layoutImpl, scene, element, sceneValues)
+    return this.then(ElementModifier(layoutImpl, scene, element, sceneValues))
+        // 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)
@@ -178,22 +169,25 @@ internal fun Modifier.element(
  */
 private data class ElementModifier(
     private val layoutImpl: SceneTransitionLayoutImpl,
+    private val scene: Scene,
     private val element: Element,
     private val sceneValues: Element.TargetValues,
 ) : ModifierNodeElement<ElementNode>() {
-    override fun create(): ElementNode = ElementNode(layoutImpl, element, sceneValues)
+    override fun create(): ElementNode = ElementNode(layoutImpl, scene, element, sceneValues)
 
     override fun update(node: ElementNode) {
-        node.update(layoutImpl, element, sceneValues)
+        node.update(layoutImpl, scene, element, sceneValues)
     }
 }
 
 internal class ElementNode(
     layoutImpl: SceneTransitionLayoutImpl,
+    scene: Scene,
     element: Element,
     sceneValues: Element.TargetValues,
-) : Modifier.Node() {
+) : Modifier.Node(), DrawModifierNode {
     private var layoutImpl: SceneTransitionLayoutImpl = layoutImpl
+    private var scene: Scene = scene
     private var element: Element = element
     private var sceneValues: Element.TargetValues = sceneValues
 
@@ -239,15 +233,34 @@ internal class ElementNode(
 
     fun update(
         layoutImpl: SceneTransitionLayoutImpl,
+        scene: Scene,
         element: Element,
         sceneValues: Element.TargetValues,
     ) {
         removeNodeFromSceneValues()
         this.layoutImpl = layoutImpl
+        this.scene = scene
         this.element = element
         this.sceneValues = sceneValues
         addNodeToSceneValues()
     }
+
+    override fun ContentDrawScope.draw() {
+        if (shouldDrawElement(layoutImpl, scene, element)) {
+            val drawScale = getDrawScale(layoutImpl, element, scene, sceneValues)
+            if (drawScale == Scale.Default) {
+                drawContent()
+            } else {
+                scale(
+                    drawScale.scaleX,
+                    drawScale.scaleY,
+                    if (drawScale.pivot.isUnspecified) center else drawScale.pivot,
+                ) {
+                    this@draw.drawContent()
+                }
+            }
+        }
+    }
 }
 
 private fun shouldDrawElement(
@@ -332,39 +345,6 @@ internal fun sharedElementTransformation(
 }
 
 /**
- * Chain the [com.android.compose.animation.scene.transformation.ModifierTransformation] applied
- * throughout the current transition, if any.
- */
-private fun Modifier.modifierTransformations(
-    layoutImpl: SceneTransitionLayoutImpl,
-    scene: Scene,
-    element: Element,
-    sceneValues: Element.TargetValues,
-): Modifier {
-    when (val state = layoutImpl.state.transitionState) {
-        is TransitionState.Idle -> return this
-        is TransitionState.Transition -> {
-            val fromScene = state.fromScene
-            val toScene = state.toScene
-            if (fromScene == toScene) {
-                // Same as idle.
-                return this
-            }
-
-            return layoutImpl.transitions
-                .transitionSpec(fromScene, state.toScene)
-                .transformations(element.key, scene.key)
-                .modifier
-                .fold(this) { modifier, transformation ->
-                    with(transformation) {
-                        modifier.transform(layoutImpl, scene, element, sceneValues)
-                    }
-                }
-        }
-    }
-}
-
-/**
  * Whether the element is opaque or not.
  *
  * Important: The logic here should closely match the logic in [elementAlpha]. Note that we don't
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 5b752eb4e900..84d3b8647d6c 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
@@ -17,11 +17,13 @@
 package com.android.compose.animation.scene
 
 import androidx.annotation.VisibleForTesting
+import androidx.compose.runtime.Stable
 
 /**
  * A base class to create unique keys, associated to an [identity] that is used to check the
  * equality of two key instances.
  */
+@Stable
 sealed class Key(val debugName: String, val identity: Any) {
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
index d48781a4529b..a0fba8076517 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
@@ -23,20 +23,25 @@ import androidx.compose.foundation.gestures.awaitHorizontalTouchSlopOrCancellati
 import androidx.compose.foundation.gestures.awaitVerticalTouchSlopOrCancellation
 import androidx.compose.foundation.gestures.horizontalDrag
 import androidx.compose.foundation.gestures.verticalDrag
+import androidx.compose.runtime.Stable
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.remember
-import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.composed
 import androidx.compose.ui.geometry.Offset
+import androidx.compose.ui.input.pointer.PointerEvent
 import androidx.compose.ui.input.pointer.PointerEventPass
 import androidx.compose.ui.input.pointer.PointerId
 import androidx.compose.ui.input.pointer.PointerInputChange
 import androidx.compose.ui.input.pointer.PointerInputScope
+import androidx.compose.ui.input.pointer.SuspendingPointerInputModifierNode
 import androidx.compose.ui.input.pointer.pointerInput
 import androidx.compose.ui.input.pointer.positionChange
 import androidx.compose.ui.input.pointer.util.VelocityTracker
 import androidx.compose.ui.input.pointer.util.addPointerInputChange
+import androidx.compose.ui.node.CompositionLocalConsumerModifierNode
+import androidx.compose.ui.node.DelegatingNode
+import androidx.compose.ui.node.ModifierNodeElement
+import androidx.compose.ui.node.PointerInputModifierNode
+import androidx.compose.ui.node.currentValueOf
 import androidx.compose.ui.platform.LocalViewConfiguration
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.Velocity
@@ -56,7 +61,7 @@ import androidx.compose.ui.util.fastForEach
  * dragged) and a second pointer is down and dragged. This is an implementation detail that might
  * change in the future.
  */
-// TODO(b/291055080): Migrate to the Modifier.Node API.
+@Stable
 internal fun Modifier.multiPointerDraggable(
     orientation: Orientation,
     enabled: Boolean,
@@ -64,22 +69,88 @@ internal fun Modifier.multiPointerDraggable(
     onDragStarted: (layoutSize: IntSize, startedPosition: Offset, pointersDown: Int) -> Unit,
     onDragDelta: (Float) -> Unit,
     onDragStopped: (velocity: Float) -> Unit,
-): Modifier = composed {
-    val onDragStarted by rememberUpdatedState(onDragStarted)
-    val onDragStopped by rememberUpdatedState(onDragStopped)
-    val onDragDelta by rememberUpdatedState(onDragDelta)
-    val startDragImmediately by rememberUpdatedState(startDragImmediately)
-
-    val velocityTracker = remember { VelocityTracker() }
-    val maxFlingVelocity =
-        LocalViewConfiguration.current.maximumFlingVelocity.let { max ->
-            val maxF = max.toFloat()
-            Velocity(maxF, maxF)
+): Modifier =
+    this.then(
+        MultiPointerDraggableElement(
+            orientation,
+            enabled,
+            startDragImmediately,
+            onDragStarted,
+            onDragDelta,
+            onDragStopped,
+        )
+    )
+
+private data class MultiPointerDraggableElement(
+    private val orientation: Orientation,
+    private val enabled: Boolean,
+    private val startDragImmediately: Boolean,
+    private val onDragStarted:
+        (layoutSize: IntSize, startedPosition: Offset, pointersDown: Int) -> Unit,
+    private val onDragDelta: (Float) -> Unit,
+    private val onDragStopped: (velocity: Float) -> Unit,
+) : ModifierNodeElement<MultiPointerDraggableNode>() {
+    override fun create(): MultiPointerDraggableNode =
+        MultiPointerDraggableNode(
+            orientation = orientation,
+            enabled = enabled,
+            startDragImmediately = startDragImmediately,
+            onDragStarted = onDragStarted,
+            onDragDelta = onDragDelta,
+            onDragStopped = onDragStopped,
+        )
+
+    override fun update(node: MultiPointerDraggableNode) {
+        node.orientation = orientation
+        node.enabled = enabled
+        node.startDragImmediately = startDragImmediately
+        node.onDragStarted = onDragStarted
+        node.onDragDelta = onDragDelta
+        node.onDragStopped = onDragStopped
+    }
+}
+
+private class MultiPointerDraggableNode(
+    orientation: Orientation,
+    enabled: Boolean,
+    var startDragImmediately: Boolean,
+    var onDragStarted: (layoutSize: IntSize, startedPosition: Offset, pointersDown: Int) -> Unit,
+    var onDragDelta: (Float) -> Unit,
+    var onDragStopped: (velocity: Float) -> Unit,
+) : PointerInputModifierNode, DelegatingNode(), CompositionLocalConsumerModifierNode {
+    private val pointerInputHandler: suspend PointerInputScope.() -> Unit = { pointerInput() }
+    private val delegate = delegate(SuspendingPointerInputModifierNode(pointerInputHandler))
+    private val velocityTracker = VelocityTracker()
+
+    var enabled: Boolean = enabled
+        set(value) {
+            // Reset the pointer input whenever enabled changed.
+            if (value != field) {
+                field = value
+                delegate.resetPointerInputHandler()
+            }
+        }
+
+    var orientation: Orientation = orientation
+        set(value) {
+            // Reset the pointer input whenever enabled orientation.
+            if (value != field) {
+                field = value
+                delegate.resetPointerInputHandler()
+            }
         }
 
-    pointerInput(enabled, orientation, maxFlingVelocity) {
+    override fun onCancelPointerInput() = delegate.onCancelPointerInput()
+
+    override fun onPointerEvent(
+        pointerEvent: PointerEvent,
+        pass: PointerEventPass,
+        bounds: IntSize
+    ) = delegate.onPointerEvent(pointerEvent, pass, bounds)
+
+    private suspend fun PointerInputScope.pointerInput() {
         if (!enabled) {
-            return@pointerInput
+            return
         }
 
         val onDragStart: (Offset, Int) -> Unit = { startedPosition, pointersDown ->
@@ -90,6 +161,12 @@ internal fun Modifier.multiPointerDraggable(
         val onDragCancel: () -> Unit = { onDragStopped(/* velocity= */ 0f) }
 
         val onDragEnd: () -> Unit = {
+            val maxFlingVelocity =
+                currentValueOf(LocalViewConfiguration).maximumFlingVelocity.let { max ->
+                    val maxF = max.toFloat()
+                    Velocity(maxF, maxF)
+                }
+
             val velocity = velocityTracker.calculateVelocity(maxFlingVelocity)
             onDragStopped(
                 when (orientation) {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/PunchHole.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
index 984086b7792a..560e92becba5 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/PunchHole.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 The Android Open Source Project
+ * 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.
@@ -14,10 +14,9 @@
  * limitations under the License.
  */
 
-package com.android.compose.animation.scene.transformation
+package com.android.compose.animation.scene
 
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.drawWithContent
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.geometry.Size
 import androidx.compose.ui.geometry.toRect
@@ -28,52 +27,68 @@ import androidx.compose.ui.graphics.Paint
 import androidx.compose.ui.graphics.RectangleShape
 import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.graphics.drawOutline
+import androidx.compose.ui.graphics.drawscope.ContentDrawScope
 import androidx.compose.ui.graphics.drawscope.DrawScope
 import androidx.compose.ui.graphics.drawscope.drawIntoCanvas
 import androidx.compose.ui.graphics.drawscope.translate
 import androidx.compose.ui.graphics.withSaveLayer
+import androidx.compose.ui.node.DrawModifierNode
+import androidx.compose.ui.node.ModifierNodeElement
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.toSize
-import com.android.compose.animation.scene.Element
-import com.android.compose.animation.scene.ElementKey
-import com.android.compose.animation.scene.ElementMatcher
-import com.android.compose.animation.scene.Scene
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 
-/** Punch a hole in an element using the bounds of another element and a given [shape]. */
-internal class PunchHole(
-    override val matcher: ElementMatcher,
+internal fun Modifier.punchHole(
+    layoutImpl: SceneTransitionLayoutImpl,
+    element: ElementKey,
+    bounds: ElementKey,
+    shape: Shape,
+): Modifier = this.then(PunchHoleElement(layoutImpl, element, bounds, shape))
+
+private data class PunchHoleElement(
+    private val layoutImpl: SceneTransitionLayoutImpl,
+    private val element: ElementKey,
     private val bounds: ElementKey,
     private val shape: Shape,
-) : ModifierTransformation {
+) : ModifierNodeElement<PunchHoleNode>() {
+    override fun create(): PunchHoleNode = PunchHoleNode(layoutImpl, element, bounds, shape)
 
+    override fun update(node: PunchHoleNode) {
+        node.layoutImpl = layoutImpl
+        node.element = element
+        node.bounds = bounds
+        node.shape = shape
+    }
+}
+
+private class PunchHoleNode(
+    var layoutImpl: SceneTransitionLayoutImpl,
+    var element: ElementKey,
+    var bounds: ElementKey,
+    var shape: Shape,
+) : Modifier.Node(), DrawModifierNode {
     private var lastSize: Size = Size.Unspecified
     private var lastLayoutDirection: LayoutDirection = LayoutDirection.Ltr
     private var lastOutline: Outline? = null
 
-    override fun Modifier.transform(
-        layoutImpl: SceneTransitionLayoutImpl,
-        scene: Scene,
-        element: Element,
-        sceneValues: Element.TargetValues,
-    ): Modifier {
-        return drawWithContent {
-            val bounds = layoutImpl.elements[bounds]
-            if (
-                bounds == null ||
-                    bounds.lastSharedValues.size == Element.SizeUnspecified ||
-                    bounds.lastSharedValues.offset == Offset.Unspecified
-            ) {
+    override fun ContentDrawScope.draw() {
+        val bounds = layoutImpl.elements[bounds]
+
+        if (
+            bounds == null ||
+                bounds.lastSharedValues.size == Element.SizeUnspecified ||
+                bounds.lastSharedValues.offset == Offset.Unspecified
+        ) {
+            drawContent()
+            return
+        }
+
+        val element = layoutImpl.elements.getValue(element)
+        drawIntoCanvas { canvas ->
+            canvas.withSaveLayer(size.toRect(), Paint()) {
                 drawContent()
-                return@drawWithContent
-            }
-            drawIntoCanvas { canvas ->
-                canvas.withSaveLayer(size.toRect(), Paint()) {
-                    drawContent()
 
-                    val offset = bounds.lastSharedValues.offset - element.lastSharedValues.offset
-                    translate(offset.x, offset.y) { drawHole(bounds) }
-                }
+                val offset = bounds.lastSharedValues.offset - element.lastSharedValues.offset
+                translate(offset.x, offset.y) { drawHole(bounds) }
             }
         }
     }
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 857a596a1404..f5561cb404b6 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
@@ -19,20 +19,24 @@ package com.android.compose.animation.scene
 import androidx.compose.foundation.gestures.Orientation
 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
 import androidx.compose.ui.layout.intermediateLayout
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.zIndex
 
 /** A scene in a [SceneTransitionLayout]. */
+@Stable
 internal class Scene(
     val key: SceneKey,
     layoutImpl: SceneTransitionLayoutImpl,
@@ -104,11 +108,13 @@ private class SceneScopeImpl(
     ): State<T> {
         val element =
             element?.let { key ->
-                layoutImpl.elements[key]
-                    ?: error(
-                        "Element $key is not composed. Make sure to call animateSharedXAsState " +
-                            "*after* Modifier.element(key)."
-                    )
+                Snapshot.withoutReadObservation {
+                    layoutImpl.elements[key]
+                        ?: error(
+                            "Element $key is not composed. Make sure to call " +
+                                "animateSharedXAsState *after* Modifier.element(key)."
+                        )
+                }
             }
 
         return animateSharedValueAsState(
@@ -130,4 +136,10 @@ private class SceneScopeImpl(
     ) {
         MovableElement(layoutImpl, scene, key, modifier, content)
     }
+
+    override fun Modifier.punchHole(
+        element: ElementKey,
+        bounds: ElementKey,
+        shape: Shape
+    ): Modifier = punchHole(layoutImpl, element, bounds, shape)
 }
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 30d13dfa3f70..07add77eccd4 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -19,10 +19,12 @@ package com.android.compose.animation.scene
 import androidx.annotation.FloatRange
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.Stable
 import androidx.compose.runtime.State
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
 import androidx.compose.ui.platform.LocalDensity
 
@@ -94,6 +96,7 @@ interface SceneTransitionLayoutScope {
 @DslMarker annotation class ElementDsl
 
 @ElementDsl
+@Stable
 interface SceneScope {
     /** The state of the [SceneTransitionLayout] in which this scene is contained. */
     val layoutState: SceneTransitionLayoutState
@@ -177,6 +180,18 @@ interface SceneScope {
         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.
+     * This can be used to make content drawn below an opaque element visible. For example, if we
+     * have [this lockscreen scene](http://shortn/_VYySFnJDhN) drawn below
+     * [this shade scene](http://shortn/_fpxGUk0Rg7) and punch a hole in the latter using the big
+     * clock time bounds and a RoundedCornerShape(10dp), [this](http://shortn/_qt80IvORFj) would be
+     * the result.
+     */
+    fun Modifier.punchHole(element: ElementKey, bounds: ElementKey, shape: Shape): Modifier
 }
 
 // TODO(b/291053742): Add animateSharedValueAsState(targetValue) without any ValueKey and ElementKey
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 60f385aede02..02ddccbc051b 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
@@ -23,6 +23,7 @@ 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.getValue
 import androidx.compose.runtime.key
 import androidx.compose.runtime.mutableStateOf
@@ -41,6 +42,7 @@ import com.android.compose.ui.util.lerp
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.Channel
 
+@Stable
 internal class SceneTransitionLayoutImpl(
     onChangeScene: (SceneKey) -> Unit,
     builder: SceneTransitionLayoutScope.() -> Unit,
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 64c9775a386e..f48e9147eef4 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
@@ -16,11 +16,13 @@
 
 package com.android.compose.animation.scene
 
+import androidx.compose.runtime.Stable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 
 /** The state of a [SceneTransitionLayout]. */
+@Stable
 class SceneTransitionLayoutState(initialScene: SceneKey) {
     /**
      * The current [TransitionState]. All values read here are backed by the Snapshot system.
@@ -29,7 +31,6 @@ class SceneTransitionLayoutState(initialScene: SceneKey) {
      * [SceneTransitionLayoutState.observableTransitionState] instead.
      */
     var transitionState: TransitionState by mutableStateOf(TransitionState.Idle(initialScene))
-        internal set
 
     /**
      * Whether we are transitioning, optionally restricting the check to the transition between
@@ -46,8 +47,15 @@ class SceneTransitionLayoutState(initialScene: SceneKey) {
         return (from == null || transition.fromScene == from) &&
             (to == null || transition.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)
+    }
 }
 
+@Stable
 sealed interface TransitionState {
     /**
      * The current effective scene. If a new transition was triggered, it would start from this
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
index 2172ed300a33..f91895bb0e05 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
@@ -18,6 +18,7 @@ package com.android.compose.animation.scene
 
 import androidx.compose.animation.core.AnimationSpec
 import androidx.compose.animation.core.snap
+import androidx.compose.runtime.Stable
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.util.fastForEach
@@ -27,7 +28,6 @@ import com.android.compose.animation.scene.transformation.AnchoredTranslate
 import com.android.compose.animation.scene.transformation.DrawScale
 import com.android.compose.animation.scene.transformation.EdgeTranslate
 import com.android.compose.animation.scene.transformation.Fade
-import com.android.compose.animation.scene.transformation.ModifierTransformation
 import com.android.compose.animation.scene.transformation.PropertyTransformation
 import com.android.compose.animation.scene.transformation.RangedPropertyTransformation
 import com.android.compose.animation.scene.transformation.ScaleSize
@@ -93,6 +93,7 @@ class SceneTransitions(
 }
 
 /** The definition of a transition between [from] and [to]. */
+@Stable
 data class TransitionSpec(
     val from: SceneKey?,
     val to: SceneKey?,
@@ -122,7 +123,6 @@ data class TransitionSpec(
         scene: SceneKey,
     ): ElementTransformations {
         var shared: SharedElementTransformation? = null
-        val modifier = mutableListOf<ModifierTransformation>()
         var offset: PropertyTransformation<Offset>? = null
         var size: PropertyTransformation<IntSize>? = null
         var drawScale: PropertyTransformation<Scale>? = null
@@ -166,12 +166,11 @@ data class TransitionSpec(
                     throwIfNotNull(shared, element, name = "shared")
                     shared = transformation
                 }
-                is ModifierTransformation -> modifier.add(transformation)
                 is PropertyTransformation<*> -> onPropertyTransformation(transformation)
             }
         }
 
-        return ElementTransformations(shared, modifier, offset, size, drawScale, alpha)
+        return ElementTransformations(shared, offset, size, drawScale, alpha)
     }
 
     private fun throwIfNotNull(
@@ -188,7 +187,6 @@ data class TransitionSpec(
 /** The transformations of an element during a transition. */
 internal class ElementTransformations(
     val shared: SharedElementTransformation?,
-    val modifier: List<ModifierTransformation>,
     val offset: PropertyTransformation<Offset>?,
     val size: PropertyTransformation<IntSize>?,
     val drawScale: PropertyTransformation<Scale>?,
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 ca66dff5e231..f820074ec3d1 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
@@ -18,8 +18,6 @@ package com.android.compose.animation.scene
 
 import androidx.compose.animation.core.AnimationSpec
 import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.graphics.RectangleShape
-import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.dp
 
@@ -131,19 +129,6 @@ interface TransitionBuilder : PropertyTransformationBuilder {
     )
 
     /**
-     * Punch a hole in the element(s) matching [matcher] that has the same bounds as [bounds] and
-     * using the given [shape].
-     *
-     * Punching a hole in an element will "remove" any pixel drawn by that element in the hole area.
-     * This can be used to make content drawn below an opaque element visible. For example, if we
-     * have [this lockscreen scene](http://shortn/_VYySFnJDhN) drawn below
-     * [this shade scene](http://shortn/_fpxGUk0Rg7) and punch a hole in the latter using the big
-     * clock time bounds and a RoundedCornerShape(10dp), [this](http://shortn/_qt80IvORFj) would be
-     * the result.
-     */
-    fun punchHole(matcher: ElementMatcher, bounds: ElementKey, shape: Shape = RectangleShape)
-
-    /**
      * Adds the transformations in [builder] but in reversed order. This allows you to partially
      * reuse the definition of the transition from scene `Foo` to scene `Bar` inside the definition
      * of the transition from scene `Bar` to scene `Foo`.
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 d4909892f492..8c0a5a394331 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
@@ -22,7 +22,6 @@ import androidx.compose.animation.core.Spring
 import androidx.compose.animation.core.VectorConverter
 import androidx.compose.animation.core.spring
 import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.unit.Dp
 import com.android.compose.animation.scene.transformation.AnchoredSize
 import com.android.compose.animation.scene.transformation.AnchoredTranslate
@@ -30,7 +29,6 @@ import com.android.compose.animation.scene.transformation.DrawScale
 import com.android.compose.animation.scene.transformation.EdgeTranslate
 import com.android.compose.animation.scene.transformation.Fade
 import com.android.compose.animation.scene.transformation.PropertyTransformation
-import com.android.compose.animation.scene.transformation.PunchHole
 import com.android.compose.animation.scene.transformation.RangedPropertyTransformation
 import com.android.compose.animation.scene.transformation.ScaleSize
 import com.android.compose.animation.scene.transformation.SharedElementTransformation
@@ -93,10 +91,6 @@ internal class TransitionBuilderImpl : TransitionBuilder {
         spec.vectorize(Float.VectorConverter).durationMillis
     }
 
-    override fun punchHole(matcher: ElementMatcher, bounds: ElementKey, shape: Shape) {
-        transformations.add(PunchHole(matcher, bounds, shape))
-    }
-
     override fun reversed(builder: TransitionBuilder.() -> Unit) {
         reversed = true
         builder()
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 0db8469466ef..206935558179 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
@@ -16,7 +16,6 @@
 
 package com.android.compose.animation.scene.transformation
 
-import androidx.compose.ui.Modifier
 import com.android.compose.animation.scene.Element
 import com.android.compose.animation.scene.ElementMatcher
 import com.android.compose.animation.scene.Scene
@@ -52,19 +51,6 @@ internal class SharedElementTransformation(
     internal val scenePicker: SharedElementScenePicker,
 ) : Transformation
 
-/** A transformation that is applied on the element during the whole transition. */
-internal interface ModifierTransformation : Transformation {
-    /** Apply the transformation to [element]. */
-    // TODO(b/290184746): Figure out a public API for custom transformations that don't have access
-    // to these internal classes.
-    fun Modifier.transform(
-        layoutImpl: SceneTransitionLayoutImpl,
-        scene: Scene,
-        element: Element,
-        sceneValues: Element.TargetValues,
-    ): Modifier
-}
-
 /** A transformation that changes the value of an element property, like its size or offset. */
 internal sealed interface PropertyTransformation<T> : Transformation {
     /**
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 cc7a0b8e33b2..ce3e1db2c3d0 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
@@ -23,6 +23,7 @@ import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.offset
 import androidx.compose.foundation.layout.size
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
@@ -427,4 +428,30 @@ class ElementTest {
         assertThat(barElement.sceneValues.keys).containsExactly(TestScenes.SceneA)
         assertThat(fooElement.sceneValues).isEmpty()
     }
+
+    @Test
+    fun existingElementsDontRecomposeWhenTransitionStateChanges() {
+        var fooCompositions = 0
+
+        rule.testTransition(
+            fromSceneContent = {
+                SideEffect { fooCompositions++ }
+                Box(Modifier.element(TestElements.Foo))
+            },
+            toSceneContent = {},
+            transition = {
+                spec = tween(4 * 16)
+
+                scaleSize(TestElements.Foo, width = 2f, height = 0.5f)
+                translate(TestElements.Foo, x = 10.dp, y = 10.dp)
+                fade(TestElements.Foo)
+            }
+        ) {
+            before { assertThat(fooCompositions).isEqualTo(1) }
+            at(16) { assertThat(fooCompositions).isEqualTo(1) }
+            at(32) { assertThat(fooCompositions).isEqualTo(1) }
+            at(48) { assertThat(fooCompositions).isEqualTo(1) }
+            after { assertThat(fooCompositions).isEqualTo(1) }
+        }
+    }
 }