Introduce InterruptionHandler to STL

This CL introduces a new InterruptionHandler interface that allows to
specify what to do when a transition is interrupted and that we should
now animate to a new target scene.

When such an interuptions occurs, the caller can now specify 2
parameters of interruption:
 - animateFrom: the scene from which we animate to the new target scene.
   It must be either the fromScene or toScene of the interrupted
   transition.
 - chain: a boolean indicating whether the new transition is added to
   the list of current transitions (the default) or if all current
   transitions should be removed and that the new transition is the only
   remaining one in the list of ongoing transitions.

Bug: 290930950
Flag: N/A
Test: InterruptionHandlerTest
Change-Id: Ifebe37e5a35e304792aa602eb06212af3e6f1441

7 files changed, 336 insertions, 40 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
index 6b289f3c66a3..79ef22c41063 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
@@ -47,8 +47,11 @@ internal fun CoroutineScope.animateToScene(
     }
 
     return when (transitionState) {
-        is TransitionState.Idle -> animate(layoutState, target, transitionKey)
+        is TransitionState.Idle ->
+            animate(layoutState, target, transitionKey, isInitiatedByUserInput = false)
         is TransitionState.Transition -> {
+            val isInitiatedByUserInput = transitionState.isInitiatedByUserInput
+
             // A transition is currently running: first check whether `transition.toScene` or
             // `transition.fromScene` is the same as our target scene, in which case the transition
             // can be accelerated or reversed to end up in the target state.
@@ -69,7 +72,13 @@ internal fun CoroutineScope.animateToScene(
                     // The transition is in progress: start the canned animation at the same
                     // progress as it was in.
                     // TODO(b/290184746): Also take the current velocity into account.
-                    animate(layoutState, target, transitionKey, startProgress = progress)
+                    animate(
+                        layoutState,
+                        target,
+                        transitionKey,
+                        isInitiatedByUserInput,
+                        startProgress = progress,
+                    )
                 }
             } else if (transitionState.fromScene == target) {
                 // There is a transition from [target] to another scene: simply animate the same
@@ -88,14 +97,47 @@ internal fun CoroutineScope.animateToScene(
                         layoutState,
                         target,
                         transitionKey,
+                        isInitiatedByUserInput,
                         startProgress = progress,
                         reversed = true,
                     )
                 }
             } else {
                 // Generic interruption; the current transition is neither from or to [target].
-                // TODO(b/290930950): Better handle interruptions here.
-                animate(layoutState, target, transitionKey)
+                val interruptionResult =
+                    layoutState.transitions.interruptionHandler.onInterruption(
+                        transitionState,
+                        target,
+                    )
+                        ?: DefaultInterruptionHandler.onInterruption(transitionState, target)
+
+                val animateFrom = interruptionResult.animateFrom
+                if (
+                    animateFrom != transitionState.toScene &&
+                        animateFrom != transitionState.fromScene
+                ) {
+                    error(
+                        "InterruptionResult.animateFrom must be either the fromScene " +
+                            "(${transitionState.fromScene.debugName}) or the toScene " +
+                            "(${transitionState.toScene.debugName}) of the interrupted transition."
+                    )
+                }
+
+                // If we were A => B and that we are now animating A => C, add a transition B => A
+                // to the list of transitions so that B "disappears back to A".
+                val chain = interruptionResult.chain
+                if (chain && animateFrom != transitionState.currentScene) {
+                    animateToScene(layoutState, animateFrom, transitionKey = null)
+                }
+
+                animate(
+                    layoutState,
+                    target,
+                    transitionKey,
+                    isInitiatedByUserInput,
+                    fromScene = animateFrom,
+                    chain = chain,
+                )
             }
         }
     }
@@ -103,32 +145,30 @@ internal fun CoroutineScope.animateToScene(
 
 private fun CoroutineScope.animate(
     layoutState: BaseSceneTransitionLayoutState,
-    target: SceneKey,
+    targetScene: SceneKey,
     transitionKey: TransitionKey?,
+    isInitiatedByUserInput: Boolean,
     startProgress: Float = 0f,
     reversed: Boolean = false,
+    fromScene: SceneKey = layoutState.transitionState.currentScene,
+    chain: Boolean = true,
 ): TransitionState.Transition {
-    val fromScene = layoutState.transitionState.currentScene
-    val isUserInput =
-        (layoutState.transitionState as? TransitionState.Transition)?.isInitiatedByUserInput
-            ?: false
-
     val targetProgress = if (reversed) 0f else 1f
     val transition =
         if (reversed) {
             OneOffTransition(
-                fromScene = target,
+                fromScene = targetScene,
                 toScene = fromScene,
-                currentScene = target,
-                isInitiatedByUserInput = isUserInput,
+                currentScene = targetScene,
+                isInitiatedByUserInput = isInitiatedByUserInput,
                 isUserInputOngoing = false,
             )
         } else {
             OneOffTransition(
                 fromScene = fromScene,
-                toScene = target,
-                currentScene = target,
-                isInitiatedByUserInput = isUserInput,
+                toScene = targetScene,
+                currentScene = targetScene,
+                isInitiatedByUserInput = isInitiatedByUserInput,
                 isUserInputOngoing = false,
             )
         }
@@ -136,7 +176,7 @@ private fun CoroutineScope.animate(
     // Change the current layout state to start this new transition. This will compute the
     // TransformationSpec associated to this transition, which we need to initialize the Animatable
     // that will actually animate it.
-    layoutState.startTransition(transition, transitionKey)
+    layoutState.startTransition(transition, transitionKey, chain)
 
     // The transition now contains the transformation spec that we should use to instantiate the
     // Animatable.
@@ -156,7 +196,7 @@ private fun CoroutineScope.animate(
                     // Settle the state to Idle(target). Note that this will do nothing if this
                     // transition was replaced/interrupted by another one, and this also runs if
                     // this coroutine is cancelled, i.e. if [this] coroutine scope is cancelled.
-                    layoutState.finishTransition(transition, target)
+                    layoutState.finishTransition(transition, targetScene)
                 }
             }
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/InterruptionHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/InterruptionHandler.kt
new file mode 100644
index 000000000000..54c64fd721ec
--- /dev/null
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/InterruptionHandler.kt
@@ -0,0 +1,85 @@
+/*
+ * 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 handler to specify how a transition should be interrupted.
+ *
+ * @see DefaultInterruptionHandler
+ * @see SceneTransitionsBuilder.interruptionHandler
+ */
+interface InterruptionHandler {
+    /**
+     * This function is called when [interrupted] is interrupted: it is currently animating between
+     * [interrupted.fromScene] and [interrupted.toScene], and we will now animate to
+     * [newTargetScene].
+     *
+     * If this returns `null`, then the [default behavior][DefaultInterruptionHandler] will be used:
+     * we will animate from [interrupted.currentScene] and chaining will be enabled (see
+     * [InterruptionResult] for more information about chaining).
+     *
+     * @see InterruptionResult
+     */
+    fun onInterruption(
+        interrupted: TransitionState.Transition,
+        newTargetScene: SceneKey,
+    ): InterruptionResult?
+}
+
+/**
+ * The result of an interruption that specifies how we should handle a transition A => B now that we
+ * have to animate to C.
+ *
+ * For instance, if the interrupted transition was A => B and currentScene = B:
+ * - animateFrom = B && chain = true => there will be 2 transitions running in parallel, A => B and
+ *   B => C.
+ * - animateFrom = A && chain = true => there will be 2 transitions running in parallel, B => A and
+ *   A => C.
+ * - animateFrom = B && chain = false => there will be 1 transition running, B => C.
+ * - animateFrom = A && chain = false => there will be 1 transition running, A => C.
+ */
+class InterruptionResult(
+    /**
+     * The scene we should animate from when transitioning to C.
+     *
+     * Important: This **must** be either [TransitionState.Transition.fromScene] or
+     * [TransitionState.Transition.toScene] of the transition that was interrupted.
+     */
+    val animateFrom: SceneKey,
+
+    /**
+     * Whether chaining is enabled, i.e. if the new transition to C should run in parallel with the
+     * previous one(s) or if it should be the only remaining transition that is running.
+     */
+    val chain: Boolean = true,
+)
+
+/**
+ * The default interruption handler: we animate from [TransitionState.Transition.currentScene] and
+ * chaining is enabled.
+ */
+object DefaultInterruptionHandler : InterruptionHandler {
+    override fun onInterruption(
+        interrupted: TransitionState.Transition,
+        newTargetScene: SceneKey,
+    ): InterruptionResult {
+        return InterruptionResult(
+            animateFrom = interrupted.currentScene,
+            chain = true,
+        )
+    }
+}
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 5fda77a3e0ae..d5526f32fbd2 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
@@ -422,13 +422,18 @@ internal abstract class BaseSceneTransitionLayoutState(
     }
 
     /**
-     * Start a new [transition], instantly interrupting any ongoing transition if there was one.
+     * Start a new [transition].
+     *
+     * If [chain] is `true`, then the transitions will simply be added to [currentTransitions] and
+     * will run in parallel to the current transitions. If [chain] is `false`, then the list of
+     * [currentTransitions] will be cleared and [transition] will be the only running transition.
      *
      * Important: you *must* call [finishTransition] once the transition is finished.
      */
     internal fun startTransition(
         transition: TransitionState.Transition,
         transitionKey: TransitionKey?,
+        chain: Boolean = true,
     ) {
         // Compute the [TransformationSpec] when the transition starts.
         val fromScene = transition.fromScene
@@ -471,26 +476,10 @@ internal abstract class BaseSceneTransitionLayoutState(
                     finishTransition(currentState, currentState.currentScene)
                 }
 
-                // Check that we don't have too many concurrent transitions.
-                if (transitionStates.size >= MAX_CONCURRENT_TRANSITIONS) {
-                    Log.wtf(
-                        TAG,
-                        buildString {
-                            appendLine("Potential leak detected in SceneTransitionLayoutState!")
-                            appendLine(
-                                "  Some transition(s) never called STLState.finishTransition()."
-                            )
-                            appendLine("  Transitions (size=${transitionStates.size}):")
-                            transitionStates.fastForEach { state ->
-                                val transition = state as TransitionState.Transition
-                                val from = transition.fromScene
-                                val to = transition.toScene
-                                val indicator =
-                                    if (finishedTransitions.contains(transition)) "x" else " "
-                                appendLine("  [$indicator] $from => $to ($transition)")
-                            }
-                        }
-                    )
+                val tooManyTransitions = transitionStates.size >= MAX_CONCURRENT_TRANSITIONS
+                val clearCurrentTransitions = !chain || tooManyTransitions
+                if (clearCurrentTransitions) {
+                    if (tooManyTransitions) logTooManyTransitions()
 
                     // Force finish all transitions.
                     while (currentTransitions.isNotEmpty()) {
@@ -511,6 +500,24 @@ internal abstract class BaseSceneTransitionLayoutState(
         }
     }
 
+    private fun logTooManyTransitions() {
+        Log.wtf(
+            TAG,
+            buildString {
+                appendLine("Potential leak detected in SceneTransitionLayoutState!")
+                appendLine("  Some transition(s) never called STLState.finishTransition().")
+                appendLine("  Transitions (size=${transitionStates.size}):")
+                transitionStates.fastForEach { state ->
+                    val transition = state as TransitionState.Transition
+                    val from = transition.fromScene
+                    val to = transition.toScene
+                    val indicator = if (finishedTransitions.contains(transition)) "x" else " "
+                    appendLine("  [$indicator] $from => $to ($transition)")
+                }
+            }
+        )
+    }
+
     private fun cancelActiveTransitionLinks() {
         for ((link, linkedTransition) in activeTransitionLinks) {
             link.target.finishTransition(linkedTransition, linkedTransition.currentScene)
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 b46614397ff4..0f6a1d276578 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
@@ -44,6 +44,7 @@ internal constructor(
     internal val defaultSwipeSpec: SpringSpec<Float>,
     internal val transitionSpecs: List<TransitionSpecImpl>,
     internal val overscrollSpecs: List<OverscrollSpecImpl>,
+    internal val interruptionHandler: InterruptionHandler,
 ) {
     private val transitionCache =
         mutableMapOf<
@@ -145,6 +146,7 @@ internal constructor(
                 defaultSwipeSpec = DefaultSwipeSpec,
                 transitionSpecs = emptyList(),
                 overscrollSpecs = emptyList(),
+                interruptionHandler = DefaultInterruptionHandler,
             )
     }
 }
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 6bc397e86cfa..a4682ff2a885 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
@@ -40,6 +40,12 @@ interface SceneTransitionsBuilder {
     var defaultSwipeSpec: SpringSpec<Float>
 
     /**
+     * The [InterruptionHandler] used when transitions are interrupted. Defaults to
+     * [DefaultInterruptionHandler].
+     */
+    var interruptionHandler: InterruptionHandler
+
+    /**
      * Define the default animation to be played when transitioning [to] the specified scene, from
      * any scene. For the animation specification to apply only when transitioning between two
      * specific scenes, use [from] instead.
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 1c9080fa085d..802ab1f2eebb 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
@@ -47,12 +47,14 @@ internal fun transitionsImpl(
     return SceneTransitions(
         impl.defaultSwipeSpec,
         impl.transitionSpecs,
-        impl.transitionOverscrollSpecs
+        impl.transitionOverscrollSpecs,
+        impl.interruptionHandler,
     )
 }
 
 private class SceneTransitionsBuilderImpl : SceneTransitionsBuilder {
     override var defaultSwipeSpec: SpringSpec<Float> = SceneTransitions.DefaultSwipeSpec
+    override var interruptionHandler: InterruptionHandler = DefaultInterruptionHandler
 
     val transitionSpecs = mutableListOf<TransitionSpecImpl>()
     val transitionOverscrollSpecs = mutableListOf<OverscrollSpecImpl>()
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/InterruptionHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/InterruptionHandlerTest.kt
new file mode 100644
index 000000000000..bbd058b5a999
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/InterruptionHandlerTest.kt
@@ -0,0 +1,154 @@
+/*
+ * 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.compose.animation.core.tween
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.compose.animation.scene.TestScenes.SceneA
+import com.android.compose.animation.scene.TestScenes.SceneB
+import com.android.compose.animation.scene.TestScenes.SceneC
+import com.android.compose.test.runMonotonicClockTest
+import com.google.common.truth.Correspondence
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class InterruptionHandlerTest {
+    @get:Rule val rule = createComposeRule()
+
+    @Test
+    fun default() = runMonotonicClockTest {
+        val state =
+            MutableSceneTransitionLayoutState(
+                SceneA,
+                transitions { /* default interruption handler */},
+            )
+
+        state.setTargetScene(SceneB, coroutineScope = this)
+        state.setTargetScene(SceneC, coroutineScope = this)
+
+        assertThat(state.currentTransitions)
+            .comparingElementsUsing(FromToCurrentTriple)
+            .containsExactly(
+                // A to B.
+                Triple(SceneA, SceneB, SceneB),
+
+                // B to C.
+                Triple(SceneB, SceneC, SceneC),
+            )
+            .inOrder()
+    }
+
+    @Test
+    fun chainingDisabled() = runMonotonicClockTest {
+        val state =
+            MutableSceneTransitionLayoutState(
+                SceneA,
+                transitions {
+                    // Handler that animates from currentScene (default) but disables chaining.
+                    interruptionHandler =
+                        object : InterruptionHandler {
+                            override fun onInterruption(
+                                interrupted: TransitionState.Transition,
+                                newTargetScene: SceneKey
+                            ): InterruptionResult {
+                                return InterruptionResult(
+                                    animateFrom = interrupted.currentScene,
+                                    chain = false,
+                                )
+                            }
+                        }
+                },
+            )
+
+        state.setTargetScene(SceneB, coroutineScope = this)
+        state.setTargetScene(SceneC, coroutineScope = this)
+
+        assertThat(state.currentTransitions)
+            .comparingElementsUsing(FromToCurrentTriple)
+            .containsExactly(
+                // B to C.
+                Triple(SceneB, SceneC, SceneC),
+            )
+            .inOrder()
+    }
+
+    @Test
+    fun animateFromOtherScene() = runMonotonicClockTest {
+        val duration = 500
+        val state =
+            MutableSceneTransitionLayoutState(
+                SceneA,
+                transitions {
+                    // Handler that animates from the scene that is not currentScene.
+                    interruptionHandler =
+                        object : InterruptionHandler {
+                            override fun onInterruption(
+                                interrupted: TransitionState.Transition,
+                                newTargetScene: SceneKey
+                            ): InterruptionResult {
+                                return InterruptionResult(
+                                    animateFrom =
+                                        if (interrupted.currentScene == interrupted.toScene) {
+                                            interrupted.fromScene
+                                        } else {
+                                            interrupted.toScene
+                                        }
+                                )
+                            }
+                        }
+
+                    from(SceneA, to = SceneB) { spec = tween(duration) }
+                },
+            )
+
+        // Animate to B and advance the transition a little bit so that progress > visibility
+        // threshold and that reversing from B back to A won't immediately snap to A.
+        state.setTargetScene(SceneB, coroutineScope = this)
+        testScheduler.advanceTimeBy(duration / 2L)
+
+        state.setTargetScene(SceneC, coroutineScope = this)
+
+        assertThat(state.currentTransitions)
+            .comparingElementsUsing(FromToCurrentTriple)
+            .containsExactly(
+                // Initial transition A to B. This transition will never be consumed by anyone given
+                // that it has the same (from, to) pair as the next transition.
+                Triple(SceneA, SceneB, SceneB),
+
+                // Initial transition reversed, B back to A.
+                Triple(SceneA, SceneB, SceneA),
+
+                // A to C.
+                Triple(SceneA, SceneC, SceneC),
+            )
+            .inOrder()
+    }
+
+    companion object {
+        val FromToCurrentTriple =
+            Correspondence.transforming(
+                { transition: TransitionState.Transition? ->
+                    Triple(transition?.fromScene, transition?.toScene, transition?.currentScene)
+                },
+                "(from, to, current) triple"
+            )
+    }
+}