Add awaitFling lambda to Controller.onStop()

This CL adds a way for NestedDraggable.Controller.onStop() to wait for
the overscroll (wrapping the call to onStop()) to be finished. This will
be used to ensure that a STL transition is marked as finished when its
overscroll effect is finished.

Bug: 378470603
Test: atest NestedDraggableTest
Flag: EXEMPT NestedDraggable is not used yet
Change-Id: I5ec36d85187561ebc663a3aba77649e41dfaf031

2 files changed, 97 insertions, 10 deletions

diff --git a/packages/SystemUI/compose/core/src/com/android/compose/gesture/NestedDraggable.kt b/packages/SystemUI/compose/core/src/com/android/compose/gesture/NestedDraggable.kt
index f47ae4377cd9..e02e8b483543 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/gesture/NestedDraggable.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/gesture/NestedDraggable.kt
@@ -52,10 +52,10 @@ import androidx.compose.ui.node.currentValueOf
 import androidx.compose.ui.platform.LocalViewConfiguration
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.Velocity
-import androidx.compose.ui.util.fastAny
 import androidx.compose.ui.util.fastSumBy
 import com.android.compose.modifiers.thenIf
 import kotlin.math.sign
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.async
 import kotlinx.coroutines.launch
 
@@ -107,12 +107,16 @@ interface NestedDraggable {
         /**
          * Stop the current drag with the given [velocity].
          *
+         * @param velocity the velocity of the drag when it stopped.
+         * @param awaitFling a lambda that can be used to wait for the end of the full fling, i.e.
+         *   wait for the end of the nested scroll fling or overscroll fling performed with the
+         *   unconsumed velocity *after* this call to [onDragStopped] returned.
          * @return the consumed [velocity]. Any non-consumed velocity will be dispatched to the next
          *   nested scroll connection to be consumed by any composable above in the hierarchy. If
          *   the drag was performed on this draggable directly (instead of on a nested scrollable),
          *   any remaining velocity will be used to animate the overscroll of this draggable.
          */
-        suspend fun onDragStopped(velocity: Float): Float
+        suspend fun onDragStopped(velocity: Float, awaitFling: suspend () -> Unit): Float
     }
 }
 
@@ -356,10 +360,20 @@ private class NestedDraggableNode(
         // We launch in the scope of the dispatcher so that the fling is not cancelled if this node
         // is removed right after onDragStopped() is called.
         nestedScrollDispatcher.coroutineScope.launch {
-            flingWithOverscroll(velocity) { velocityFromOverscroll ->
-                flingWithNestedScroll(velocityFromOverscroll) { velocityFromNestedScroll ->
-                    controller.onDragStopped(velocityFromNestedScroll.toFloat()).toVelocity()
+            val flingCompletable = CompletableDeferred<Unit>()
+            try {
+                flingWithOverscroll(velocity) { velocityFromOverscroll ->
+                    flingWithNestedScroll(velocityFromOverscroll) { velocityFromNestedScroll ->
+                        controller
+                            .onDragStopped(
+                                velocityFromNestedScroll.toFloat(),
+                                awaitFling = { flingCompletable.await() },
+                            )
+                            .toVelocity()
+                    }
                 }
+            } finally {
+                flingCompletable.complete(Unit)
             }
         }
     }
@@ -514,8 +528,15 @@ private class NestedDraggableNode(
         }
 
         suspend fun flingWithOverscroll(velocity: Velocity): Velocity {
-            return flingWithOverscroll(overscrollEffect, velocity) {
-                controller.onDragStopped(it.toFloat()).toVelocity()
+            val flingCompletable = CompletableDeferred<Unit>()
+            return try {
+                flingWithOverscroll(overscrollEffect, velocity) {
+                    controller
+                        .onDragStopped(it.toFloat(), awaitFling = { flingCompletable.await() })
+                        .toVelocity()
+                }
+            } finally {
+                flingCompletable.complete(Unit)
             }
         }
     }
diff --git a/packages/SystemUI/compose/core/tests/src/com/android/compose/gesture/NestedDraggableTest.kt b/packages/SystemUI/compose/core/tests/src/com/android/compose/gesture/NestedDraggableTest.kt
index f98b090e77fd..9c49090916e3 100644
--- a/packages/SystemUI/compose/core/tests/src/com/android/compose/gesture/NestedDraggableTest.kt
+++ b/packages/SystemUI/compose/core/tests/src/com/android/compose/gesture/NestedDraggableTest.kt
@@ -45,6 +45,9 @@ import com.google.common.truth.Truth.assertThat
 import kotlin.math.ceil
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
@@ -593,6 +596,63 @@ class NestedDraggableTest(override val orientation: Orientation) : OrientationAw
         assertThat(effect.applyToFlingDone).isTrue()
     }
 
+    @Test
+    fun awaitFling() = runTest {
+        var flingIsDone = false
+        val draggable =
+            TestDraggable(
+                onDragStopped = { _, awaitFling ->
+                    // Start a coroutine in the background that waits for the fling to be finished.
+                    launch {
+                        awaitFling()
+                        flingIsDone = true
+                    }
+
+                    0f
+                }
+            )
+
+        val effectPostFlingCompletable = CompletableDeferred<Unit>()
+        val effect =
+            TestOverscrollEffect(
+                orientation,
+                onPostScroll = { 0f },
+                onPostFling = {
+                    effectPostFlingCompletable.await()
+                    it
+                },
+            )
+
+        val touchSlop =
+            rule.setContentWithTouchSlop {
+                Box(
+                    Modifier.fillMaxSize()
+                        .nestedDraggable(draggable, orientation, overscrollEffect = effect)
+                )
+            }
+
+        assertThat(draggable.onDragStartedCalled).isFalse()
+
+        rule.onRoot().performTouchInput {
+            down(center)
+            moveBy(touchSlop.toOffset())
+            up()
+        }
+
+        // The drag was started and stopped, but the fling is not finished yet as the overscroll
+        // effect is stuck on the effectPostFlingCompletable.
+        runCurrent()
+        rule.waitForIdle()
+        assertThat(draggable.onDragStartedCalled).isTrue()
+        assertThat(draggable.onDragStoppedCalled).isTrue()
+        assertThat(flingIsDone).isFalse()
+
+        effectPostFlingCompletable.complete(Unit)
+        runCurrent()
+        rule.waitForIdle()
+        assertThat(flingIsDone).isTrue()
+    }
+
     private fun ComposeContentTestRule.setContentWithTouchSlop(
         content: @Composable () -> Unit
     ): Float {
@@ -614,7 +674,10 @@ class NestedDraggableTest(override val orientation: Orientation) : OrientationAw
     private class TestDraggable(
         private val onDragStarted: (Offset, Float) -> Unit = { _, _ -> },
         private val onDrag: (Float) -> Float = { it },
-        private val onDragStopped: suspend (Float) -> Float = { it },
+        private val onDragStopped: suspend (Float, awaitFling: suspend () -> Unit) -> Float =
+            { velocity, _ ->
+                velocity
+            },
         private val shouldConsumeNestedScroll: (Float) -> Boolean = { true },
     ) : NestedDraggable {
         var shouldStartDrag = true
@@ -648,9 +711,12 @@ class NestedDraggableTest(override val orientation: Orientation) : OrientationAw
                     return onDrag.invoke(delta)
                 }
 
-                override suspend fun onDragStopped(velocity: Float): Float {
+                override suspend fun onDragStopped(
+                    velocity: Float,
+                    awaitFling: suspend () -> Unit,
+                ): Float {
                     onDragStoppedCalled = true
-                    return onDragStopped.invoke(velocity)
+                    return onDragStopped.invoke(velocity, awaitFling)
                 }
             }
         }