PriorityNestedScrollConnection: make onStop suspendable and add onCancel

The onStop lambda is now a suspend function. This allows for performing
long-running operations, such as animations, within the onStop callback.

The onCancel lambda has been added, this lambda will be invoked when the
nested scrolling operation is canceled, for example, if the user cannot
scroll anymore.

Test: atest PriorityNestedScrollConnectionTest
Bug: 370949877
Flag: com.android.systemui.scene_container
Change-Id: I4cc652fe4f6baac800e1484e7bddfbf5510de2b1

9 files changed, 57 insertions, 24 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
index 779934d8d211..e4c60e166fd5 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
@@ -73,13 +73,19 @@ fun NotificationScrimNestedScrollConnection(
             snapScrimOffset(currentHeight + amountConsumed)
             amountConsumed
         },
-        // Don't consume the velocity on pre/post fling
         onStop = { velocityAvailable ->
             onStop(velocityAvailable)
             if (scrimOffset() < minScrimOffset()) {
                 animateScrimOffset(minScrimOffset())
             }
-            { 0f }
+            // Don't consume the velocity on pre/post fling
+            0f
+        },
+        onCancel = {
+            onStop(0f)
+            if (scrimOffset() < minScrimOffset()) {
+                animateScrimOffset(minScrimOffset())
+            }
         },
     )
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
index 42dd1be51d44..edb05ebd77d1 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
@@ -104,7 +104,8 @@ fun NotificationStackNestedScrollConnection(
         },
         onStop = { velocityAvailable ->
             onStop(velocityAvailable)
-            suspend { velocityAvailable }
+            velocityAvailable
         },
+        onCancel = { onStop(0f) },
     )
 }
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 3fadd1a62495..7e288ddd3a4c 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
@@ -27,9 +27,10 @@ import com.android.compose.animation.scene.content.Content
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.HasOverscrollProperties.Companion.DistanceUnspecified
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
-import com.android.compose.nestedscroll.SuspendedValue
 import kotlin.math.absoluteValue
 
+internal typealias SuspendedValue<T> = suspend () -> T
+
 internal interface DraggableHandler {
     /**
      * Start a drag in the given [startedPosition], with the given [overSlop] and number of
@@ -712,10 +713,18 @@ internal class NestedScrollHandlerImpl(
             },
             onStop = { velocityAvailable ->
                 val controller = dragController ?: error("Should be called after onStart")
-
-                controller
-                    .onStop(velocity = velocityAvailable, canChangeContent = canChangeScene)
-                    .also { dragController = null }
+                try {
+                    controller
+                        .onStop(velocity = velocityAvailable, canChangeContent = canChangeScene)
+                        .invoke()
+                } finally {
+                    dragController = null
+                }
+            },
+            onCancel = {
+                val controller = dragController ?: error("Should be called after onStart")
+                controller.onStop(velocity = 0f, canChangeContent = canChangeScene)
+                dragController = null
             },
         )
     }
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 205267da151a..f0043e1e89b0 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
@@ -27,7 +27,6 @@ import androidx.compose.runtime.setValue
 import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.HasOverscrollProperties.Companion.DistanceUnspecified
-import com.android.compose.nestedscroll.SuspendedValue
 import kotlin.math.absoluteValue
 import kotlinx.coroutines.CompletableDeferred
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
index b05f02224327..ecf64b771d1f 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
@@ -70,6 +70,7 @@ fun LargeTopAppBarNestedScrollConnection(
             amountConsumed
         },
         // Don't consume the velocity on pre/post fling
-        onStop = { { 0f } },
+        onStop = { 0f },
+        onCancel = { /* do nothing */ },
     )
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
index e3cddc1a8f1d..636c55799ff2 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
@@ -16,15 +16,21 @@
 
 package com.android.compose.nestedscroll
 
+import androidx.compose.animation.core.AnimationState
+import androidx.compose.animation.core.DecayAnimationSpec
+import androidx.compose.animation.core.animateDecay
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
 import androidx.compose.ui.input.nestedscroll.NestedScrollSource
 import androidx.compose.ui.unit.Velocity
 import com.android.compose.ui.util.SpaceVectorConverter
+import kotlin.math.abs
 import kotlin.math.sign
-
-internal typealias SuspendedValue<T> = suspend () -> T
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.async
+import kotlinx.coroutines.coroutineScope
 
 /**
  * A [NestedScrollConnection] that intercepts scroll events in priority mode.
@@ -47,6 +53,7 @@ internal typealias SuspendedValue<T> = suspend () -> T
  *   consumed amount.
  * @param onStop lambda that is called when the connection stops consuming scroll events and returns
  *   the consumed velocity.
+ * @param onCancel lambda that is called when the connection is cancelled.
  * @sample LargeTopAppBarNestedScrollConnection
  * @sample com.android.compose.animation.scene.NestedScrollHandlerImpl.nestedScrollConnection
  */
@@ -60,7 +67,8 @@ class PriorityNestedScrollConnection(
     private val canStopOnPreFling: () -> Boolean,
     private val onStart: (offsetAvailable: Float) -> Unit,
     private val onScroll: (offsetAvailable: Float, source: NestedScrollSource) -> Float,
-    private val onStop: (velocityAvailable: Float) -> SuspendedValue<Float>,
+    private val onStop: suspend (velocityAvailable: Float) -> Float,
+    private val onCancel: () -> Unit,
 ) : NestedScrollConnection, SpaceVectorConverter by SpaceVectorConverter(orientation) {
 
     /** In priority mode [onPreScroll] events are first consumed by the parent, via [onScroll]. */
@@ -68,6 +76,9 @@ class PriorityNestedScrollConnection(
 
     private var offsetScrolledBeforePriorityMode = 0f
 
+    /** This job allows us to interrupt the onStop animation */
+    private var onStopJob: Deferred<Float> = CompletableDeferred(0f)
+
     override fun onPostScroll(
         consumed: Offset,
         available: Offset,
@@ -148,7 +159,7 @@ class PriorityNestedScrollConnection(
      */
     fun reset() {
         if (isPriorityMode) {
-            // Step 3c: To ensure that an onStop is always called for every onStart.
+            // Step 3c: To ensure that an onStop (or onCancel) is always called for every onStart.
             cancel()
         } else {
             resetOffsetTracker()
@@ -172,6 +183,8 @@ class PriorityNestedScrollConnection(
         // available offset following a scroll event.
         isPriorityMode = true
 
+        onStopJob.cancel()
+
         // Note: onStop will be called if we cannot continue to scroll (step 3a), or the finger is
         // lifted (step 3b), or this object has been destroyed (step 3c).
         onStart(availableOffset)
@@ -204,13 +217,17 @@ class PriorityNestedScrollConnection(
         check(isPriorityMode) { "This should never happen, stop() was called before start()" }
         isPriorityMode = false
         resetOffsetTracker()
-        return onStop(velocityAvailable).invoke()
+
+        return coroutineScope {
+            onStopJob = async { onStop(velocityAvailable) }
+            onStopJob.await()
+        }
     }
 
     private fun cancel() {
         check(isPriorityMode) { "This should never happen, cancel() was called before start()" }
         isPriorityMode = false
         resetOffsetTracker()
-        onStop(0f)
+        onCancel()
     }
 }
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 ecef6be49df8..57b9423e85d1 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
@@ -39,7 +39,6 @@ import com.android.compose.animation.scene.TestScenes.SceneC
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.Transition
 import com.android.compose.animation.scene.subjects.assertThat
-import com.android.compose.nestedscroll.SuspendedValue
 import com.android.compose.test.MonotonicClockTestScope
 import com.android.compose.test.runMonotonicClockTest
 import com.android.compose.test.transition
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 c8f6e6d99933..3df608717414 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
@@ -46,7 +46,6 @@ import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.Velocity
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.modifiers.thenIf
-import com.android.compose.nestedscroll.SuspendedValue
 import com.google.common.truth.Truth.assertThat
 import kotlin.properties.Delegates
 import kotlinx.coroutines.coroutineScope
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
index 9363ecbf63dc..1a3b86b936df 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
@@ -39,6 +39,7 @@ class PriorityNestedScrollConnectionTest {
     private var lastScroll: Float? = null
     private var consumeScroll = true
     private var lastStop: Float? = null
+    private var isCancelled: Boolean = false
     private var consumeStop = true
 
     private val scrollConnection =
@@ -55,8 +56,9 @@ class PriorityNestedScrollConnectionTest {
             },
             onStop = {
                 lastStop = it
-                { if (consumeStop) it else 0f }
+                if (consumeStop) it else 0f
             },
+            onCancel = { isCancelled = true },
         )
 
     @Test
@@ -147,13 +149,13 @@ class PriorityNestedScrollConnectionTest {
     }
 
     @Test
-    fun step3a_onPriorityMode_shouldStopIfCannotContinue() {
+    fun step3a_onPriorityMode_shouldCancelIfCannotContinue() {
         startPriorityModePostScroll()
         consumeScroll = false
 
-        scrollConnection.onPreScroll(available = Offset.Zero, source = UserInput)
+        scrollConnection.onPreScroll(available = Offset(0f, 1f), source = UserInput)
 
-        assertThat(lastStop).isNotNull()
+        assertThat(isCancelled).isTrue()
     }
 
     @Test
@@ -178,12 +180,12 @@ class PriorityNestedScrollConnectionTest {
     }
 
     @Test
-    fun step3c_onPriorityMode_shouldStopOnReset() {
+    fun step3c_onPriorityMode_shouldCancelOnReset() {
         startPriorityModePostScroll()
 
         scrollConnection.reset()
 
-        assertThat(lastStop).isNotNull()
+        assertThat(isCancelled).isTrue()
     }
 
     @Test