Fix interruption + overscroll bug

This CL fixes a tricky interruption + overscroll bug: because the scene
in which an element is placed can change when overscrolling, the
computation of the interruption deltas for placement-related values
should also be saved on the scene we are not placing the element in if
the element is shared in both scenes. This avoids jumpcuts that can
happen when an interruption is ongoing and that we start overscrolling,
placing the element in another scene.

Bug: 290930950
Test: atest ElementTest
Flag: com.android.systemui.scene_container
Change-Id: Id9ba18e70013ca4561ad8e35e25be5d8d69d6789

2 files changed, 193 insertions, 22 deletions

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 edf8943509df..f074bdda9212 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
@@ -280,7 +280,7 @@ internal class ElementNode(
         constraints: Constraints,
     ): MeasureResult {
         val transitions = currentTransitions
-        val transition = elementTransition(element, transitions)
+        val transition = elementTransition(layoutImpl, element, transitions)
 
         // If this element is not supposed to be laid out now, either because it is not part of any
         // ongoing transition or the other scene of its transition is overscrolling, then lay out
@@ -318,13 +318,10 @@ internal class ElementNode(
             val targetOffsetInScene = lookaheadScopeCoordinates.localLookaheadPositionOf(coords)
 
             // No need to place the element in this scene if we don't want to draw it anyways.
-            if (!shouldPlaceElement(layoutImpl, scene, element, transition)) {
+            if (!shouldPlaceElement(layoutImpl, scene.key, element, transition)) {
                 sceneState.lastOffset = Offset.Unspecified
                 sceneState.lastScale = Scale.Unspecified
                 sceneState.lastAlpha = Element.AlphaUnspecified
-
-                sceneState.clearValuesBeforeInterruption()
-                sceneState.clearInterruptionDeltas()
                 return
             }
 
@@ -353,7 +350,17 @@ internal class ElementNode(
                     getValueBeforeInterruption = { sceneState.offsetBeforeInterruption },
                     setValueBeforeInterruption = { sceneState.offsetBeforeInterruption = it },
                     getInterruptionDelta = { sceneState.offsetInterruptionDelta },
-                    setInterruptionDelta = { sceneState.offsetInterruptionDelta = it },
+                    setInterruptionDelta = { delta ->
+                        setPlacementInterruptionDelta(
+                            element = element,
+                            sceneState = sceneState,
+                            transition = transition,
+                            delta = delta,
+                            setter = { sceneState, delta ->
+                                sceneState.offsetInterruptionDelta = delta
+                            },
+                        )
+                    },
                     diff = { a, b -> a - b },
                     add = { a, b, bProgress -> a + b * bProgress },
                 )
@@ -363,7 +370,7 @@ internal class ElementNode(
             val offset = (interruptedOffset - currentOffset).round()
             if (
                 isElementOpaque(scene, element, transition) &&
-                    interruptedAlpha(layoutImpl, transition, sceneState, alpha = 1f) == 1f
+                    interruptedAlpha(layoutImpl, element, transition, sceneState, alpha = 1f) == 1f
             ) {
                 sceneState.lastAlpha = 1f
 
@@ -379,8 +386,8 @@ internal class ElementNode(
                     // also need to recompute the current transition to make sure that we are using
                     // the current transition and not a reference to an old one. See b/343138966 for
                     // details.
-                    val transition = elementTransition(element, currentTransitions)
-                    if (!shouldPlaceElement(layoutImpl, scene, element, transition)) {
+                    val transition = elementTransition(layoutImpl, element, currentTransitions)
+                    if (!shouldPlaceElement(layoutImpl, scene.key, element, transition)) {
                         return@placeWithLayer
                     }
 
@@ -394,7 +401,7 @@ internal class ElementNode(
     override fun ContentDrawScope.draw() {
         element.wasDrawnInAnyScene = true
 
-        val transition = elementTransition(element, currentTransitions)
+        val transition = elementTransition(layoutImpl, element, currentTransitions)
         val drawScale = getDrawScale(layoutImpl, scene, element, transition, sceneState)
         if (drawScale == Scale.Default) {
             drawContent()
@@ -435,6 +442,7 @@ internal class ElementNode(
  * its scenes contains the element.
  */
 private fun elementTransition(
+    layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
     transitions: List<TransitionState.Transition>,
 ): TransitionState.Transition? {
@@ -448,7 +456,7 @@ private fun elementTransition(
 
     if (transition != previousTransition && transition != null && previousTransition != null) {
         // The previous transition was interrupted by another transition.
-        prepareInterruption(element, transition, previousTransition)
+        prepareInterruption(layoutImpl, element, transition, previousTransition)
     } else if (transition == null && previousTransition != null) {
         // The transition was just finished.
         element.sceneStates.values.forEach {
@@ -461,18 +469,43 @@ private fun elementTransition(
 }
 
 private fun prepareInterruption(
+    layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
     transition: TransitionState.Transition,
     previousTransition: TransitionState.Transition,
 ) {
     val sceneStates = element.sceneStates
-    sceneStates[previousTransition.fromScene]?.selfUpdateValuesBeforeInterruption()
-    sceneStates[previousTransition.toScene]?.selfUpdateValuesBeforeInterruption()
-    sceneStates[transition.fromScene]?.selfUpdateValuesBeforeInterruption()
-    sceneStates[transition.toScene]?.selfUpdateValuesBeforeInterruption()
+    fun updatedSceneState(key: SceneKey): Element.SceneState? {
+        return sceneStates[key]?.also { it.selfUpdateValuesBeforeInterruption() }
+    }
+
+    val previousFromState = updatedSceneState(previousTransition.fromScene)
+    val previousToState = updatedSceneState(previousTransition.toScene)
+    val fromState = updatedSceneState(transition.fromScene)
+    val toState = updatedSceneState(transition.toScene)
 
     reconcileStates(element, previousTransition)
     reconcileStates(element, transition)
+
+    // Remove the interruption values to all scenes but the scene(s) where the element will be
+    // placed, to make sure that interruption deltas are computed only right after this interruption
+    // is prepared.
+    fun maybeCleanPlacementValuesBeforeInterruption(sceneState: Element.SceneState) {
+        if (!shouldPlaceElement(layoutImpl, sceneState.scene, element, transition)) {
+            sceneState.offsetBeforeInterruption = Offset.Unspecified
+            sceneState.alphaBeforeInterruption = Element.AlphaUnspecified
+            sceneState.scaleBeforeInterruption = Scale.Unspecified
+
+            sceneState.offsetInterruptionDelta = Offset.Zero
+            sceneState.alphaInterruptionDelta = 0f
+            sceneState.scaleInterruptionDelta = Scale.Zero
+        }
+    }
+
+    previousFromState?.let { maybeCleanPlacementValuesBeforeInterruption(it) }
+    previousToState?.let { maybeCleanPlacementValuesBeforeInterruption(it) }
+    fromState?.let { maybeCleanPlacementValuesBeforeInterruption(it) }
+    toState?.let { maybeCleanPlacementValuesBeforeInterruption(it) }
 }
 
 /**
@@ -579,9 +612,38 @@ private inline fun <T> computeInterruptedValue(
     }
 }
 
+/**
+ * Set the interruption delta of a *placement/drawing*-related value (offset, alpha, scale). This
+ * ensures that the delta is also set on the other scene in the transition for shared elements, so
+ * that there is no jump cut if the scene where the element is placed has changed.
+ */
+private inline fun <T> setPlacementInterruptionDelta(
+    element: Element,
+    sceneState: Element.SceneState,
+    transition: TransitionState.Transition?,
+    delta: T,
+    setter: (Element.SceneState, T) -> Unit,
+) {
+    // Set the interruption delta on the current scene.
+    setter(sceneState, delta)
+
+    if (transition == null) {
+        return
+    }
+
+    // If the element is shared, also set the delta on the other scene so that it is used by that
+    // scene if we start overscrolling it and change the scene where the element is placed.
+    val otherScene =
+        if (sceneState.scene == transition.fromScene) transition.toScene else transition.fromScene
+    val otherSceneState = element.sceneStates[otherScene] ?: return
+    if (isSharedElementEnabled(element.key, transition)) {
+        setter(otherSceneState, delta)
+    }
+}
+
 private fun shouldPlaceElement(
     layoutImpl: SceneTransitionLayoutImpl,
-    scene: Scene,
+    scene: SceneKey,
     element: Element,
     transition: TransitionState.Transition?,
 ): Boolean {
@@ -592,7 +654,7 @@ private fun shouldPlaceElement(
 
     // Don't place the element in this scene if this scene is not part of the current element
     // transition.
-    if (scene.key != transition.fromScene && scene.key != transition.toScene) {
+    if (scene != transition.fromScene && scene != transition.toScene) {
         return false
     }
 
@@ -610,7 +672,7 @@ private fun shouldPlaceElement(
 
     return shouldPlaceOrComposeSharedElement(
         layoutImpl,
-        scene.key,
+        scene,
         element.key,
         transition,
     )
@@ -740,13 +802,14 @@ private fun elementAlpha(
         element.sceneStates.forEach { it.value.alphaBeforeInterruption = 0f }
     }
 
-    val interruptedAlpha = interruptedAlpha(layoutImpl, transition, sceneState, alpha)
+    val interruptedAlpha = interruptedAlpha(layoutImpl, element, transition, sceneState, alpha)
     sceneState.lastAlpha = interruptedAlpha
     return interruptedAlpha
 }
 
 private fun interruptedAlpha(
     layoutImpl: SceneTransitionLayoutImpl,
+    element: Element,
     transition: TransitionState.Transition?,
     sceneState: Element.SceneState,
     alpha: Float,
@@ -760,7 +823,15 @@ private fun interruptedAlpha(
         getValueBeforeInterruption = { sceneState.alphaBeforeInterruption },
         setValueBeforeInterruption = { sceneState.alphaBeforeInterruption = it },
         getInterruptionDelta = { sceneState.alphaInterruptionDelta },
-        setInterruptionDelta = { sceneState.alphaInterruptionDelta = it },
+        setInterruptionDelta = { delta ->
+            setPlacementInterruptionDelta(
+                element = element,
+                sceneState = sceneState,
+                transition = transition,
+                delta = delta,
+                setter = { sceneState, delta -> sceneState.alphaInterruptionDelta = delta },
+            )
+        },
         diff = { a, b -> a - b },
         add = { a, b, bProgress -> a + b * bProgress },
     )
@@ -867,7 +938,15 @@ private fun ContentDrawScope.getDrawScale(
             getValueBeforeInterruption = { sceneState.scaleBeforeInterruption },
             setValueBeforeInterruption = { sceneState.scaleBeforeInterruption = it },
             getInterruptionDelta = { sceneState.scaleInterruptionDelta },
-            setInterruptionDelta = { sceneState.scaleInterruptionDelta = it },
+            setInterruptionDelta = { delta ->
+                setPlacementInterruptionDelta(
+                    element = element,
+                    sceneState = sceneState,
+                    transition = transition,
+                    delta = delta,
+                    setter = { sceneState, delta -> sceneState.scaleInterruptionDelta = delta },
+                )
+            },
             diff = { a, b ->
                 Scale(
                     scaleX = a.scaleX - b.scaleX,
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 47c9b9cbfd10..7827d85bcfb0 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
@@ -47,7 +47,6 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.layout.approachLayout
 import androidx.compose.ui.platform.LocalViewConfiguration
-import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.test.assertIsDisplayed
 import androidx.compose.ui.test.assertIsNotDisplayed
 import androidx.compose.ui.test.assertPositionInRootIsEqualTo
@@ -1632,4 +1631,97 @@ class ElementTest {
         rule.onNode(hasText(fooInA)).assertIsDisplayed()
         rule.onNode(hasText(fooInB)).assertDoesNotExist()
     }
+
+    @Test
+    fun interruptionThenOverscroll() = runTest {
+        val state =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutStateImpl(
+                    SceneA,
+                    transitions {
+                        overscroll(SceneB, Orientation.Vertical) {
+                            translate(TestElements.Foo, y = 15.dp)
+                        }
+                    }
+                )
+            }
+
+        @Composable
+        fun SceneScope.SceneWithFoo(offset: DpOffset, modifier: Modifier = Modifier) {
+            Box(modifier.fillMaxSize()) {
+                Box(Modifier.offset(offset.x, offset.y).element(TestElements.Foo).size(100.dp))
+            }
+        }
+
+        rule.setContent {
+            SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                scene(SceneA) { SceneWithFoo(offset = DpOffset.Zero) }
+                scene(SceneB) { SceneWithFoo(offset = DpOffset(x = 40.dp, y = 0.dp)) }
+                scene(SceneC) { SceneWithFoo(offset = DpOffset(x = 40.dp, y = 40.dp)) }
+            }
+        }
+
+        // Start A => B at 75%.
+        rule.runOnUiThread {
+            state.startTransition(
+                transition(
+                    from = SceneA,
+                    to = SceneB,
+                    progress = { 0.75f },
+                    onFinish = neverFinish(),
+                )
+            )
+        }
+
+        // Foo should be at offset (30dp, 0dp) and placed in scene B.
+        rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsNotDisplayed()
+        rule.onNode(isElement(TestElements.Foo, SceneB)).assertPositionInRootIsEqualTo(30.dp, 0.dp)
+        rule.onNode(isElement(TestElements.Foo, SceneC)).assertIsNotDisplayed()
+
+        // Interrupt A => B with B => C at 0%.
+        var progress by mutableStateOf(0f)
+        var interruptionProgress by mutableStateOf(1f)
+        rule.runOnUiThread {
+            state.startTransition(
+                transition(
+                    from = SceneB,
+                    to = SceneC,
+                    progress = { progress },
+                    interruptionProgress = { interruptionProgress },
+                    orientation = Orientation.Vertical,
+                    onFinish = neverFinish(),
+                )
+            )
+        }
+
+        // Because interruption progress is at 100M, Foo should still be at offset (30dp, 0dp) but
+        // placed in scene C.
+        rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsNotDisplayed()
+        rule.onNode(isElement(TestElements.Foo, SceneB)).assertIsNotDisplayed()
+        rule.onNode(isElement(TestElements.Foo, SceneC)).assertPositionInRootIsEqualTo(30.dp, 0.dp)
+
+        // Overscroll B => C on scene B at -100%. Because overscrolling on B => C translates Foo
+        // vertically by -15dp and that interruptionProgress is still 100%, we should now be at
+        // (30dp, -15dp)
+        progress = -1f
+        rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsNotDisplayed()
+        rule
+            .onNode(isElement(TestElements.Foo, SceneB))
+            .assertPositionInRootIsEqualTo(30.dp, -15.dp)
+        rule.onNode(isElement(TestElements.Foo, SceneC)).assertIsNotDisplayed()
+
+        // Finish the interruption, we should now be at (40dp, -15dp), still on scene B.
+        interruptionProgress = 0f
+        rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsNotDisplayed()
+        rule
+            .onNode(isElement(TestElements.Foo, SceneB))
+            .assertPositionInRootIsEqualTo(40.dp, -15.dp)
+        rule.onNode(isElement(TestElements.Foo, SceneC)).assertIsNotDisplayed()
+
+        // Finish the transition, we should be at the final position (40dp, 40dp) on scene C.
+        progress = 1f
+        rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsNotDisplayed()
+        rule.onNode(isElement(TestElements.Foo, SceneB)).assertIsNotDisplayed()
+        rule.onNode(isElement(TestElements.Foo, SceneC)).assertPositionInRootIsEqualTo(40.dp, 40.dp)
+    }
 }