[STL] Respect animationSpec for predictive back transition in STL

Bug: 350705972
Test: PredictiveBackHandlerTest
Flag: com.android.systemui.scene_container
Change-Id: I80dcc31e479ab73b3165c82bf1aa072d9bf584f1

3 files changed, 46 insertions, 12 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
index c2fb8fce1d6e..e9300119924e 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
@@ -18,7 +18,7 @@ package com.android.compose.animation.scene
 
 import androidx.activity.BackEventCompat
 import androidx.activity.compose.PredictiveBackHandler
-import androidx.compose.animation.core.spring
+import androidx.compose.animation.core.AnimationSpec
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Composable
 import kotlin.coroutines.cancellation.CancellationException
@@ -62,7 +62,7 @@ private suspend fun <T : ContentKey> animate(
     animation: SwipeAnimation<T>,
     progress: Flow<BackEventCompat>,
 ) {
-    fun animateOffset(targetContent: T) {
+    fun animateOffset(targetContent: T, spec: AnimationSpec<Float>? = null) {
         if (
             layoutImpl.state.transitionState != animation.contentTransition || animation.isFinishing
         ) {
@@ -72,12 +72,7 @@ private suspend fun <T : ContentKey> animate(
         animation.animateOffset(
             initialVelocity = 0f,
             targetContent = targetContent,
-
-            // TODO(b/350705972): Allow to customize or reuse the same customization endpoints as
-            // the normal swipe transitions. We can't just reuse them here because other swipe
-            // transitions animate pixels while this transition animates progress, so the visibility
-            // thresholds will be completely different.
-            spec = spring(),
+            spec = spec,
         )
     }
 
@@ -86,9 +81,15 @@ private suspend fun <T : ContentKey> animate(
         progress.collect { backEvent -> animation.dragOffset = backEvent.progress }
 
         // Back gesture successful.
-        animateOffset(animation.toContent)
+        animateOffset(
+            animation.toContent,
+            animation.contentTransition.transformationSpec.progressSpec
+        )
     } catch (e: CancellationException) {
         // Back gesture cancelled.
+        // If the back gesture is cancelled, the progress is animated back to 0f by the system.
+        // Since the remaining change in progress is usually very small, the progressSpec is omitted
+        // and the default spring spec used instead.
         animateOffset(animation.fromContent)
     }
 }
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 57ff597d7314..be9c5670ceae 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
@@ -17,8 +17,8 @@
 package com.android.compose.animation.scene
 
 import androidx.compose.animation.core.Animatable
+import androidx.compose.animation.core.AnimationSpec
 import androidx.compose.animation.core.AnimationVector1D
-import androidx.compose.animation.core.SpringSpec
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableFloatStateOf
@@ -320,7 +320,7 @@ internal class SwipeAnimation<T : ContentKey>(
     fun animateOffset(
         initialVelocity: Float,
         targetContent: T,
-        spec: SpringSpec<Float>? = null,
+        spec: AnimationSpec<Float>? = null,
     ): OffsetAnimation {
         val initialProgress = progress
         // Skip the animation if we have already reached the target content and the overscroll does
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt
index c5b6cdf12385..9284ffddcee3 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt
@@ -18,6 +18,8 @@ package com.android.compose.animation.scene
 
 import androidx.activity.BackEventCompat
 import androidx.activity.ComponentActivity
+import androidx.compose.animation.core.LinearEasing
+import androidx.compose.animation.core.tween
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.size
@@ -65,7 +67,23 @@ class PredictiveBackHandlerTest {
 
     @Test
     fun testPredictiveBack() {
-        val layoutState = rule.runOnUiThread { MutableSceneTransitionLayoutState(SceneA) }
+        val transitionFrames = 2
+        val layoutState =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(
+                    SceneA,
+                    transitions =
+                        transitions {
+                            from(SceneA, to = SceneB) {
+                                spec =
+                                    tween(
+                                        durationMillis = transitionFrames * 16,
+                                        easing = LinearEasing
+                                    )
+                            }
+                        }
+                )
+            }
         rule.setContent {
             SceneTransitionLayout(layoutState) {
                 scene(SceneA, mapOf(Back to SceneB)) { Box(Modifier.fillMaxSize()) }
@@ -94,12 +112,27 @@ class PredictiveBackHandlerTest {
         assertThat(layoutState.transitionState).hasCurrentScene(SceneA)
         assertThat(layoutState.transitionState).isIdle()
 
+        rule.mainClock.autoAdvance = false
+
         // Start again and commit it.
         rule.runOnUiThread {
             dispatcher.dispatchOnBackStarted(backEvent())
             dispatcher.dispatchOnBackProgressed(backEvent(progress = 0.4f))
             dispatcher.onBackPressed()
         }
+        rule.mainClock.advanceTimeByFrame()
+        rule.waitForIdle()
+        val transition2 = assertThat(layoutState.transitionState).isSceneTransition()
+        // verify that transition picks up progress from preview
+        assertThat(transition2).hasProgress(0.4f, tolerance = 0.0001f)
+
+        rule.mainClock.advanceTimeByFrame()
+        rule.waitForIdle()
+        // verify that transition is half way between preview-end-state (0.4f) and target-state (1f)
+        // after one frame
+        assertThat(transition2).hasProgress(0.7f, tolerance = 0.0001f)
+
+        rule.mainClock.autoAdvance = true
         rule.waitForIdle()
         assertThat(layoutState.transitionState).hasCurrentScene(SceneB)
         assertThat(layoutState.transitionState).isIdle()