Remove Modifier.punchHole in SceneScope (1/2)

This CL removes the Modifier.punchHole in SceneScope and replaces it by
a more generic Modifier that takes any offset and size as a parameter.
The previous punchHole modifier was the last user of
Element.lastSharedState and SceneState.lastState, which will be removed
in the following CL.

This CL also adds an overload of Modifier.punchHole to easily punch a
hole given the LayoutCoordinates of any node.

Bug: 316901148
Test: PunchHoleTest
Test: Manual, swipe down from lockscreen to go to shade scene
Flag: N/A
Change-Id: I0d9adeead15b0fe211bfc634f5e2de0d01597398

3 files changed, 86 insertions, 54 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
index 454c0ecf8ac5..0acc76f8d4ef 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
@@ -16,6 +16,7 @@
 
 package com.android.compose.animation.scene
 
+import androidx.compose.runtime.Stable
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.geometry.Size
@@ -32,80 +33,95 @@ import androidx.compose.ui.graphics.drawscope.DrawScope
 import androidx.compose.ui.graphics.drawscope.drawIntoCanvas
 import androidx.compose.ui.graphics.drawscope.translate
 import androidx.compose.ui.graphics.withSaveLayer
+import androidx.compose.ui.layout.LayoutCoordinates
+import androidx.compose.ui.node.DelegatingNode
 import androidx.compose.ui.node.DrawModifierNode
+import androidx.compose.ui.node.GlobalPositionAwareModifierNode
 import androidx.compose.ui.node.ModifierNodeElement
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.toSize
 
-internal fun Modifier.punchHole(
-    layoutImpl: SceneTransitionLayoutImpl,
-    element: ElementKey,
-    bounds: ElementKey,
-    shape: Shape,
-): Modifier = this.then(PunchHoleElement(layoutImpl, element, bounds, shape))
+/**
+ * Punch a hole in this node with the given [size], [offset] and [shape].
+ *
+ * Punching a hole in an element will "remove" any pixel drawn by that element in the hole area.
+ * This can be used to make content drawn below an opaque element visible. For example, if we have
+ * [this lockscreen scene](http://shortn/_VYySFnJDhN) drawn below
+ * [this shade scene](http://shortn/_fpxGUk0Rg7) and punch a hole in the latter using the big clock
+ * time bounds and a RoundedCornerShape(10dp), [this](http://shortn/_qt80IvORFj) would be the
+ * result.
+ */
+@Stable
+fun Modifier.punchHole(
+    size: () -> Size,
+    offset: () -> Offset,
+    shape: Shape = RectangleShape,
+): Modifier = this.then(PunchHoleElement(size, offset, shape))
+
+/**
+ * Punch a hole in this node using the bounds of [coords] and the given [shape].
+ *
+ * You can use [androidx.compose.ui.layout.onGloballyPositioned] to get the last coordinates of a
+ * node.
+ */
+@Stable
+fun Modifier.punchHole(
+    coords: () -> LayoutCoordinates?,
+    shape: Shape = RectangleShape,
+): Modifier = this.then(PunchHoleWithBoundsElement(coords, shape))
 
 private data class PunchHoleElement(
-    private val layoutImpl: SceneTransitionLayoutImpl,
-    private val element: ElementKey,
-    private val bounds: ElementKey,
+    private val size: () -> Size,
+    private val offset: () -> Offset,
     private val shape: Shape,
 ) : ModifierNodeElement<PunchHoleNode>() {
-    override fun create(): PunchHoleNode = PunchHoleNode(layoutImpl, element, bounds, shape)
+    override fun create(): PunchHoleNode = PunchHoleNode(size, offset, { shape })
 
     override fun update(node: PunchHoleNode) {
-        node.layoutImpl = layoutImpl
-        node.element = element
-        node.bounds = bounds
-        node.shape = shape
+        node.size = size
+        node.offset = offset
+        node.shape = { shape }
     }
 }
 
 private class PunchHoleNode(
-    var layoutImpl: SceneTransitionLayoutImpl,
-    var element: ElementKey,
-    var bounds: ElementKey,
-    var shape: Shape,
+    var size: () -> Size,
+    var offset: () -> Offset,
+    var shape: () -> Shape,
 ) : Modifier.Node(), DrawModifierNode {
     private var lastSize: Size = Size.Unspecified
     private var lastLayoutDirection: LayoutDirection = LayoutDirection.Ltr
     private var lastOutline: Outline? = null
 
     override fun ContentDrawScope.draw() {
-        val bounds = layoutImpl.elements[bounds]
-
-        if (
-            bounds == null ||
-                bounds.lastSharedState.size == Element.SizeUnspecified ||
-                bounds.lastSharedState.offset == Offset.Unspecified
-        ) {
+        val holeSize = size()
+        if (holeSize == Size.Zero) {
             drawContent()
             return
         }
 
-        val element = layoutImpl.elements.getValue(element)
         drawIntoCanvas { canvas ->
             canvas.withSaveLayer(size.toRect(), Paint()) {
                 drawContent()
 
-                val offset = bounds.lastSharedState.offset - element.lastSharedState.offset
-                translate(offset.x, offset.y) { drawHole(bounds) }
+                val offset = offset()
+                translate(offset.x, offset.y) { drawHole(holeSize) }
             }
         }
     }
 
-    private fun DrawScope.drawHole(bounds: Element) {
-        val boundsSize = bounds.lastSharedState.size.toSize()
+    private fun DrawScope.drawHole(size: Size) {
         if (shape == RectangleShape) {
-            drawRect(Color.Black, size = boundsSize, blendMode = BlendMode.DstOut)
+            drawRect(Color.Black, size = size, blendMode = BlendMode.DstOut)
             return
         }
 
         val outline =
-            if (boundsSize == lastSize && layoutDirection == lastLayoutDirection) {
+            if (size == lastSize && layoutDirection == lastLayoutDirection) {
                 lastOutline!!
             } else {
-                val newOutline = shape.createOutline(boundsSize, layoutDirection, this)
-                lastSize = boundsSize
+                val newOutline = shape().createOutline(size, layoutDirection, this)
+                lastSize = size
                 lastLayoutDirection = layoutDirection
                 lastOutline = newOutline
                 newOutline
@@ -118,3 +134,39 @@ private class PunchHoleNode(
         )
     }
 }
+
+private data class PunchHoleWithBoundsElement(
+    private val coords: () -> LayoutCoordinates?,
+    private val shape: Shape,
+) : ModifierNodeElement<PunchHoleWithBoundsNode>() {
+    override fun create(): PunchHoleWithBoundsNode = PunchHoleWithBoundsNode(coords, shape)
+
+    override fun update(node: PunchHoleWithBoundsNode) {
+        node.holeCoords = coords
+        node.shape = shape
+    }
+}
+
+private class PunchHoleWithBoundsNode(
+    var holeCoords: () -> LayoutCoordinates?,
+    var shape: Shape,
+) : DelegatingNode(), DrawModifierNode, GlobalPositionAwareModifierNode {
+    private val delegate = delegate(PunchHoleNode(::holeSize, ::holeOffset, ::shape))
+    private var lastCoords: LayoutCoordinates? = null
+
+    override fun onGloballyPositioned(coordinates: LayoutCoordinates) {
+        this.lastCoords = coordinates
+    }
+
+    override fun ContentDrawScope.draw() = with(delegate) { draw() }
+
+    private fun holeSize(): Size {
+        return holeCoords()?.size?.toSize() ?: Size.Zero
+    }
+
+    private fun holeOffset(): Offset {
+        val holeCoords = holeCoords() ?: return Offset.Zero
+        val lastCoords = lastCoords ?: error("draw() was called before onGloballyPositioned()")
+        return lastCoords.localPositionOf(holeCoords, relativeToSource = Offset.Zero)
+    }
+}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
index 3537b7989ed5..f67df54b088c 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Scene.kt
@@ -26,7 +26,6 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.layout.intermediateLayout
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.unit.IntSize
@@ -139,12 +138,6 @@ internal class SceneScopeImpl(
             bottomOrRightBehavior = bottomBehavior,
         )
 
-    override fun Modifier.punchHole(
-        element: ElementKey,
-        bounds: ElementKey,
-        shape: Shape
-    ): Modifier = punchHole(layoutImpl, element, bounds, shape)
-
     override fun Modifier.noResizeDuringTransitions(): Modifier {
         return noResizeDuringTransitions(layoutState = layoutImpl.state)
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
index 0ada23a76afa..80f8c1c9e987 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -25,7 +25,6 @@ import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
 import androidx.compose.ui.platform.LocalDensity
 
@@ -230,18 +229,6 @@ interface BaseSceneScope {
     ): Modifier
 
     /**
-     * Punch a hole in this [element] using the bounds of [bounds] in [scene] and the given [shape].
-     *
-     * Punching a hole in an element will "remove" any pixel drawn by that element in the hole area.
-     * This can be used to make content drawn below an opaque element visible. For example, if we
-     * have [this lockscreen scene](http://shortn/_VYySFnJDhN) drawn below
-     * [this shade scene](http://shortn/_fpxGUk0Rg7) and punch a hole in the latter using the big
-     * clock time bounds and a RoundedCornerShape(10dp), [this](http://shortn/_qt80IvORFj) would be
-     * the result.
-     */
-    fun Modifier.punchHole(element: ElementKey, bounds: ElementKey, shape: Shape): Modifier
-
-    /**
      * Don't resize during transitions. This can for instance be used to make sure that scrollable
      * lists keep a constant size during transitions even if its elements are growing/shrinking.
      */