STL make DragController.onStop a suspended function

We now have the capability to launch onStop animations within their
respective coroutine scopes.
This improvement will allow us to manage more complex animations in the
future.

Test: Refactor DraggableHandlerTest and MultiPointerDraggableTest
Bug: 378470603
Flag: com.android.systemui.scene_container
Change-Id: I985098f2a5631b3f6d4118e318a52eca158aed8f

6 files changed, 64 insertions, 37 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
index 04c527176cce..8ff51fbebb40 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
@@ -30,8 +30,7 @@ import com.android.compose.nestedscroll.OnStopScope
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
 import com.android.compose.nestedscroll.ScrollController
 import kotlin.math.absoluteValue
-
-internal typealias SuspendedValue<T> = suspend () -> T
+import kotlinx.coroutines.launch
 
 internal interface DraggableHandler {
     /**
@@ -50,6 +49,7 @@ internal interface DragController {
     /**
      * Drag the current scene by [delta] pixels.
      *
+     * @param delta The distance to drag the scene in pixels.
      * @return the consumed [delta]
      */
     fun onDrag(delta: Float): Float
@@ -57,9 +57,18 @@ internal interface DragController {
     /**
      * Stop the current drag with the given [velocity].
      *
+     * @param velocity The velocity of the drag when it stopped.
+     * @param canChangeContent Whether the content can be changed as a result of this drag.
      * @return the consumed [velocity] when the animation complete
      */
-    fun onStop(velocity: Float, canChangeContent: Boolean): SuspendedValue<Float>
+    suspend fun onStop(velocity: Float, canChangeContent: Boolean): Float
+
+    /**
+     * Cancels the current drag.
+     *
+     * @param canChangeContent Whether the content can be changed as a result of this drag.
+     */
+    fun onCancel(canChangeContent: Boolean)
 }
 
 internal class DraggableHandlerImpl(
@@ -350,7 +359,7 @@ private class DragControllerImpl(
         val result = swipes.findUserActionResult(directionOffset = newOffset)
 
         if (result == null) {
-            onStop(velocity = delta, canChangeContent = true)
+            onCancel(canChangeContent = true)
             return 0f
         }
 
@@ -379,11 +388,11 @@ private class DragControllerImpl(
         return consumedDelta
     }
 
-    override fun onStop(velocity: Float, canChangeContent: Boolean): SuspendedValue<Float> {
+    override suspend fun onStop(velocity: Float, canChangeContent: Boolean): Float {
         return onStop(velocity, canChangeContent, swipeAnimation)
     }
 
-    private fun <T : ContentKey> onStop(
+    private suspend fun <T : ContentKey> onStop(
         velocity: Float,
         canChangeContent: Boolean,
 
@@ -392,14 +401,14 @@ private class DragControllerImpl(
         // callbacks (like onAnimationCompleted()) might incorrectly finish a new transition that
         // replaced this one.
         swipeAnimation: SwipeAnimation<T>,
-    ): SuspendedValue<Float> {
+    ): Float {
         // The state was changed since the drag started; don't do anything.
         if (!isDrivingTransition || swipeAnimation.isAnimatingOffset()) {
-            return { 0f }
+            return 0f
         }
 
         val fromContent = swipeAnimation.fromContent
-        val consumedVelocity: SuspendedValue<Float>
+        val consumedVelocity: Float
         if (canChangeContent) {
             // If we are halfway between two contents, we check what the target will be based on the
             // velocity and offset of the transition, then we launch the animation.
@@ -478,6 +487,12 @@ private class DragControllerImpl(
                 isCloserToTarget()
         }
     }
+
+    override fun onCancel(canChangeContent: Boolean) {
+        swipeAnimation.contentTransition.coroutineScope.launch {
+            onStop(velocity = 0f, canChangeContent = canChangeContent)
+        }
+    }
 }
 
 /** The [Swipe] associated to a given fromScene, startedPosition and pointersDown. */
@@ -701,13 +716,14 @@ private fun scrollController(
         }
 
         override suspend fun OnStopScope.onStop(initialVelocity: Float): Float {
-            return dragController
-                .onStop(velocity = initialVelocity, canChangeContent = canChangeScene)
-                .invoke()
+            return dragController.onStop(
+                velocity = initialVelocity,
+                canChangeContent = canChangeScene,
+            )
         }
 
         override fun onCancel() {
-            dragController.onStop(velocity = 0f, canChangeContent = canChangeScene)
+            dragController.onCancel(canChangeScene)
         }
 
         /**
@@ -731,5 +747,9 @@ internal const val OffsetVisibilityThreshold = 0.5f
 private object NoOpDragController : DragController {
     override fun onDrag(delta: Float) = 0f
 
-    override fun onStop(velocity: Float, canChangeContent: Boolean) = suspend { 0f }
+    override suspend fun onStop(velocity: Float, canChangeContent: Boolean) = 0f
+
+    override fun onCancel(canChangeContent: Boolean) {
+        /* do nothing */
+    }
 }
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 160326792f81..e49709375e1e 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
@@ -318,17 +318,13 @@ internal class MultiPointerDraggableNode(
                                             velocityTracker.calculateVelocity(maxVelocity)
                                         }
                                         .toFloat(),
-                                onFling = {
-                                    controller.onStop(it, canChangeContent = true).invoke()
-                                },
+                                onFling = { controller.onStop(it, canChangeContent = true) },
                             )
                         },
                         onDragCancel = { controller ->
                             startFlingGesture(
                                 initialVelocity = 0f,
-                                onFling = {
-                                    controller.onStop(it, canChangeContent = true).invoke()
-                                },
+                                onFling = { controller.onStop(it, canChangeContent = true) },
                             )
                         },
                         swipeDetector = swipeDetector,
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
index f0043e1e89b0..412fdcb8a03c 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
@@ -29,6 +29,7 @@ import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.HasOverscrollProperties.Companion.DistanceUnspecified
 import kotlin.math.absoluteValue
 import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.launch
 
 internal fun createSwipeAnimation(
     layoutState: MutableSceneTransitionLayoutStateImpl,
@@ -317,11 +318,11 @@ internal class SwipeAnimation<T : ContentKey>(
      *
      * @return the velocity consumed
      */
-    fun animateOffset(
+    suspend fun animateOffset(
         initialVelocity: Float,
         targetContent: T,
         spec: AnimationSpec<Float>? = null,
-    ): SuspendedValue<Float> {
+    ): Float {
         check(!isAnimatingOffset()) { "SwipeAnimation.animateOffset() can only be called once" }
 
         val initialProgress = progress
@@ -379,7 +380,7 @@ internal class SwipeAnimation<T : ContentKey>(
         if (skipAnimation) {
             // Unblock the job.
             offsetAnimationRunnable.complete(null)
-            return { 0f }
+            return 0f
         }
 
         val isTargetGreater = targetOffset > animatable.value
@@ -440,7 +441,7 @@ internal class SwipeAnimation<T : ContentKey>(
             }
         }
 
-        return { velocityConsumed.await() }
+        return velocityConsumed.await()
     }
 
     /** An exception thrown during the animation to stop it immediately. */
@@ -469,7 +470,9 @@ internal class SwipeAnimation<T : ContentKey>(
     fun freezeAndAnimateToCurrentState() {
         if (isAnimatingOffset()) return
 
-        animateOffset(initialVelocity = 0f, targetContent = currentContent)
+        contentTransition.coroutineScope.launch {
+            animateOffset(initialVelocity = 0f, targetContent = currentContent)
+        }
     }
 }
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/Seek.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/Seek.kt
index 2b33224022fc..c6912d5e4ad2 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/Seek.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/Seek.kt
@@ -145,7 +145,7 @@ internal suspend fun <T : ContentKey> animateProgress(
     cancelSpec: AnimationSpec<Float>?,
     animationScope: CoroutineScope? = null,
 ) {
-    fun animateOffset(targetContent: T, spec: AnimationSpec<Float>?) {
+    suspend fun animateOffset(targetContent: T, spec: AnimationSpec<Float>?) {
         if (state.transitionState != animation.contentTransition || animation.isAnimatingOffset()) {
             return
         }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
index 61e9bb0db0f7..8d32065e0bc3 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
@@ -45,6 +45,8 @@ import com.android.compose.test.runMonotonicClockTest
 import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.async
 import kotlinx.coroutines.launch
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -266,7 +268,7 @@ class DraggableHandlerTest {
         ) {
             val velocityConsumed = onDragStoppedAnimateLater(velocity, canChangeScene)
             onAnimationStart()
-            onAnimationEnd(velocityConsumed.invoke())
+            onAnimationEnd(velocityConsumed.await())
         }
 
         suspend fun DragController.onDragStoppedAnimateNow(
@@ -285,8 +287,10 @@ class DraggableHandlerTest {
         fun DragController.onDragStoppedAnimateLater(
             velocity: Float,
             canChangeScene: Boolean = true,
-        ): SuspendedValue<Float> {
-            return onStop(velocity, canChangeScene)
+        ): Deferred<Float> {
+            val velocityConsumed = testScope.async { onStop(velocity, canChangeScene) }
+            testScope.testScheduler.runCurrent()
+            return velocityConsumed
         }
 
         fun NestedScrollConnection.scroll(
@@ -1112,6 +1116,7 @@ class DraggableHandlerTest {
         // Freeze the transition.
         val transition = transitionState as Transition
         transition.freezeAndAnimateToCurrentState()
+        runCurrent()
         assertTransition(isUserInputOngoing = false)
         advanceUntilIdle()
         assertIdle(SceneC)
@@ -1279,14 +1284,13 @@ class DraggableHandlerTest {
         // Release the finger.
         dragController.onDragStoppedAnimateNow(
             velocity = -velocityThreshold,
-            onAnimationStart = { assertTransition(fromScene = SceneA, toScene = SceneB) },
+            onAnimationStart = {
+                // Given that we are at progress >= 100% and that the overscroll on scene B is doing
+                // nothing, we are already idle.
+                assertIdle(SceneB)
+            },
             expectedConsumedVelocity = 0f,
         )
-
-        // Exhaust all coroutines *without advancing the clock*. Given that we are at progress >=
-        // 100% and that the overscroll on scene B is doing nothing, we are already idle.
-        runCurrent()
-        assertIdle(SceneB)
     }
 
     @Test
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
index 5ec74f8d2260..8072862acbba 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
@@ -72,9 +72,13 @@ class MultiPointerDraggableTest {
             return delta
         }
 
-        override fun onStop(velocity: Float, canChangeContent: Boolean): SuspendedValue<Float> {
+        override suspend fun onStop(velocity: Float, canChangeContent: Boolean): Float {
             onStop.invoke(velocity)
-            return { velocity }
+            return velocity
+        }
+
+        override fun onCancel(canChangeContent: Boolean) {
+            error("MultiPointerDraggable never calls onCancel()")
         }
     }