Merge "Don't remove finished transitions early when other transitions run" into main

3 files changed, 46 insertions, 40 deletions

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 568a358e4a7e..158256d14d1a 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
@@ -411,9 +411,7 @@ internal class MutableSceneTransitionLayoutStateImpl(
                     if (tooManyTransitions) logTooManyTransitions()
 
                     // Force finish all transitions.
-                    while (currentTransitions.isNotEmpty()) {
-                        finishTransition(transitionStates[0] as TransitionState.Transition)
-                    }
+                    currentTransitions.fastForEach { finishTransition(it) }
 
                     // We finished all transitions, so we are now idle. We remove this state so that
                     // we end up only with the new transition after appending it.
@@ -475,46 +473,36 @@ internal class MutableSceneTransitionLayoutStateImpl(
         // Mark this transition as finished.
         finishedTransitions.add(transition)
 
-        // Keep a reference to the last transition, in case we remove all transitions and should
-        // settle to Idle.
+        if (finishedTransitions.size != transitionStates.size) {
+            // Some transitions were not finished, so we won't settle to idle.
+            return
+        }
+
+        // Keep a reference to the last transition, in case all transitions are finished and we
+        // should settle to Idle.
         val lastTransition = transitionStates.last()
 
-        // Remove all first n finished transitions.
-        var i = 0
-        val nStates = transitionStates.size
-        while (i < nStates) {
-            val t = transitionStates[i]
-            if (!finishedTransitions.contains(t)) {
-                // Stop here.
-                break
+        transitionStates.fastForEach { state ->
+            if (!finishedTransitions.contains(state)) {
+                // Some transitions were not finished, so we won't settle to idle.
+                return
             }
-
-            // Remove the transition from the set of finished transitions.
-            finishedTransitions.remove(t)
-            i++
         }
 
-        // If all transitions are finished, we are idle.
-        if (i == nStates) {
-            check(finishedTransitions.isEmpty())
-            val idle =
-                TransitionState.Idle(lastTransition.currentScene, lastTransition.currentOverlays)
-            Log.i(TAG, "all transitions finished. idle=$idle")
-            this.transitionStates = listOf(idle)
-        } else if (i > 0) {
-            this.transitionStates = transitionStates.subList(fromIndex = i, toIndex = nStates)
-        }
+        val idle = TransitionState.Idle(lastTransition.currentScene, lastTransition.currentOverlays)
+        Log.i(TAG, "all transitions finished. idle=$idle")
+        finishedTransitions.clear()
+        this.transitionStates = listOf(idle)
     }
 
     override fun snapToScene(scene: SceneKey, currentOverlays: Set<OverlayKey>) {
         checkThread()
 
         // Force finish all transitions.
-        while (currentTransitions.isNotEmpty()) {
-            finishTransition(transitionStates[0] as TransitionState.Transition)
-        }
+        currentTransitions.fastForEach { finishTransition(it) }
 
         check(transitionStates.size == 1)
+        check(currentTransitions.isEmpty())
         transitionStates = listOf(TransitionState.Idle(scene, currentOverlays))
     }
 
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 d1bd52b56ddd..5074cd5211ce 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
@@ -198,23 +198,30 @@ class SceneTransitionLayoutStateTest {
         assertThat(state.currentTransitions).containsExactly(aToB, bToC).inOrder()
 
         // C => A. This should automatically call freezeAndAnimateToCurrentState() on bToC.
-        state.startTransitionImmediately(animationScope = backgroundScope, cToA)
+        val cToAJob = state.startTransitionImmediately(animationScope = backgroundScope, cToA)
         assertThat(frozenTransitions).containsExactly(aToB, bToC)
         assertThat(state.finishedTransitions).isEmpty()
         assertThat(state.currentTransitions).containsExactly(aToB, bToC, cToA).inOrder()
 
-        // Mark bToC as finished. The list of current transitions does not change because aToB is
-        // still not marked as finished.
+        // Mark aToB and bToC as finished. The list of current transitions does not change because
+        // cToA is still running.
+        aToB.finish()
+        aToBJob.join()
+        assertThat(state.finishedTransitions).containsExactly(aToB)
+        assertThat(state.currentTransitions).containsExactly(aToB, bToC, cToA).inOrder()
+
         bToC.finish()
         bToCJob.join()
-        assertThat(state.finishedTransitions).containsExactly(bToC)
+        assertThat(state.finishedTransitions).containsExactly(aToB, bToC)
         assertThat(state.currentTransitions).containsExactly(aToB, bToC, cToA).inOrder()
 
-        // Mark aToB as finished. This will remove both aToB and bToC from the list of transitions.
-        aToB.finish()
-        aToBJob.join()
+        // Mark cToA as finished. This should clear all transitions and settle to idle.
+        cToA.finish()
+        cToAJob.join()
         assertThat(state.finishedTransitions).isEmpty()
-        assertThat(state.currentTransitions).containsExactly(cToA).inOrder()
+        assertThat(state.currentTransitions).isEmpty()
+        assertThat(state.transitionState).isIdle()
+        assertThat(state.transitionState).hasCurrentScene(SceneA)
     }
 
     @Test
@@ -473,4 +480,15 @@ class SceneTransitionLayoutStateTest {
                     "SceneKey(debugName=SceneB)"
             )
     }
+
+    @Test
+    fun snapToScene_multipleTransitions() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(SceneA)
+        state.startTransitionImmediately(this, transition(SceneA, SceneB))
+        state.startTransitionImmediately(this, transition(SceneB, SceneC))
+        state.snapToScene(SceneC)
+
+        assertThat(state.transitionState).isIdle()
+        assertThat(state.transitionState).hasCurrentScene(SceneC)
+    }
 }
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 fdbd0f63292a..7c8c6e5f6c12 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
@@ -379,8 +379,8 @@ class SceneTransitionLayoutTest {
         assertThat(transition).hasProgress(0.5f)
         rule.waitForIdle()
 
-        // B and C are composed.
-        rule.onNodeWithTag("aRoot").assertDoesNotExist()
+        // A, B and C are still composed given that B => C is not finished yet.
+        rule.onNodeWithTag("aRoot").assertExists()
         rule.onNodeWithTag("bRoot").assertExists()
         rule.onNodeWithTag("cRoot").assertExists()