Reuse SwipeAnimation in PredictiveBackHandler

This CL makes PredictiveBackHandler reuse SwipeAnimation instead of
implementing its own transition. This makes the back transition work out
of the box with overlays.

Because PredictiveBackHandler used to handle previews but SwipeAnimation
didn't, the preview animation logic from PredictiveBackHandler was moved
to SwipeAnimation, making previews now work with all swipe transitions.

Bug: 353679003
Test: PredictiveBackHandlerTest
Test: OverlayTest
Flag: com.android.systemui.scene_container
Change-Id: Ibef4b8f2a1a2658e8028c41149852b88440e0da0

4 files changed, 145 insertions, 109 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 e7e6b2a257d8..ca379775cc4b 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,120 +18,76 @@ package com.android.compose.animation.scene
 
 import androidx.activity.BackEventCompat
 import androidx.activity.compose.PredictiveBackHandler
-import androidx.compose.animation.core.Animatable
-import androidx.compose.animation.core.AnimationVector1D
+import androidx.compose.animation.core.spring
+import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableFloatStateOf
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.setValue
-import com.android.compose.animation.scene.content.state.TransitionState
+import com.android.compose.animation.scene.content.Content
 import kotlin.coroutines.cancellation.CancellationException
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart
-import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.launch
 
 @Composable
 internal fun PredictiveBackHandler(
-    state: MutableSceneTransitionLayoutStateImpl,
-    coroutineScope: CoroutineScope,
-    targetSceneForBack: SceneKey? = null,
+    layoutImpl: SceneTransitionLayoutImpl,
+    result: UserActionResult?,
 ) {
     PredictiveBackHandler(
-        enabled = targetSceneForBack != null,
+        enabled = result != null,
     ) { progress: Flow<BackEventCompat> ->
-        val fromScene = state.transitionState.currentScene
-        if (targetSceneForBack == null || targetSceneForBack == fromScene) {
+        if (result == null) {
             // Note: We have to collect progress otherwise PredictiveBackHandler will throw.
             progress.first()
             return@PredictiveBackHandler
         }
 
-        val transition =
-            PredictiveBackTransition(state, coroutineScope, fromScene, toScene = targetSceneForBack)
-        state.startTransition(transition)
-        try {
-            progress.collect { backEvent -> transition.dragProgress = backEvent.progress }
-
-            // Back gesture successful.
-            transition.animateTo(targetSceneForBack)
-        } catch (e: CancellationException) {
-            // Back gesture cancelled.
-            transition.animateTo(fromScene)
-        }
+        val animation =
+            createSwipeAnimation(
+                layoutImpl,
+                result,
+                isUpOrLeft = false,
+                // Note that the orientation does not matter here given that it's only used to
+                // compute the distance. In our case the distance is always 1f.
+                orientation = Orientation.Horizontal,
+                distance = 1f,
+            )
+
+        animate(layoutImpl, animation, progress)
     }
 }
 
-private class PredictiveBackTransition(
-    val state: MutableSceneTransitionLayoutStateImpl,
-    val coroutineScope: CoroutineScope,
-    fromScene: SceneKey,
-    toScene: SceneKey,
-) : TransitionState.Transition.ChangeCurrentScene(fromScene, toScene) {
-    override var currentScene by mutableStateOf(fromScene)
-        private set
-
-    /** The animated progress once the gesture was committed or cancelled. */
-    private var progressAnimatable by mutableStateOf<Animatable<Float, AnimationVector1D>?>(null)
-    var dragProgress: Float by mutableFloatStateOf(0f)
-
-    override val previewProgress: Float
-        get() = dragProgress
-
-    override val previewProgressVelocity: Float
-        get() = 0f // Currently, velocity is not exposed by predictive back API
-
-    override val isInPreviewStage: Boolean
-        get() = previewTransformationSpec != null && currentScene == fromScene
-
-    override val progress: Float
-        get() = progressAnimatable?.value ?: previewTransformationSpec?.let { 0f } ?: dragProgress
-
-    override val progressVelocity: Float
-        get() = progressAnimatable?.velocity ?: 0f
-
-    override val isInitiatedByUserInput: Boolean
-        get() = true
-
-    override val isUserInputOngoing: Boolean
-        get() = progressAnimatable == null
-
-    private var animationJob: Job? = null
-
-    override fun finish(): Job = animateTo(currentScene)
-
-    fun animateTo(scene: SceneKey): Job {
-        check(scene == fromScene || scene == toScene)
-        animationJob?.let {
-            return it
+private suspend fun <T : Content> animate(
+    layoutImpl: SceneTransitionLayoutImpl,
+    animation: SwipeAnimation<T>,
+    progress: Flow<BackEventCompat>,
+) {
+    fun animateOffset(targetContent: T) {
+        if (
+            layoutImpl.state.transitionState != animation.contentTransition || animation.isFinishing
+        ) {
+            return
         }
 
-        if (scene != currentScene && state.transitionState == this && state.canChangeScene(scene)) {
-            currentScene = scene
-        }
+        animation.animateOffset(
+            layoutImpl.coroutineScope,
+            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(),
+        )
+    }
 
-        val targetProgress =
-            when (currentScene) {
-                fromScene -> 0f
-                toScene -> 1f
-                else -> error("scene $currentScene should be either $fromScene or $toScene")
-            }
-        val startProgress = if (previewTransformationSpec != null) 0f else dragProgress
-        val animatable = Animatable(startProgress).also { progressAnimatable = it }
+    layoutImpl.state.startTransition(animation.contentTransition)
+    try {
+        progress.collect { backEvent -> animation.dragOffset = backEvent.progress }
 
-        // Important: We start atomically to make sure that we start the coroutine even if it is
-        // cancelled right after it is launched, so that finishTransition() is correctly called.
-        return coroutineScope
-            .launch(start = CoroutineStart.ATOMIC) {
-                try {
-                    animatable.animateTo(targetProgress)
-                } finally {
-                    state.finishTransition(this@PredictiveBackTransition)
-                }
-            }
-            .also { animationJob = it }
+        // Back gesture successful.
+        animateOffset(animation.toContent)
+    } catch (e: CancellationException) {
+        // Back gesture cancelled.
+        animateOffset(animation.fromContent)
     }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
index 258be8122c1d..18e3e4b48aef 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
@@ -353,19 +353,8 @@ internal class SceneTransitionLayoutImpl(
 
     @Composable
     private fun BackHandler() {
-        val targetSceneForBack =
-            when (val result = contentForUserActions().userActions[Back.Resolved]) {
-                null -> null
-                is UserActionResult.ChangeScene -> result.toScene
-                is UserActionResult.ShowOverlay,
-                is UserActionResult.HideOverlay,
-                is UserActionResult.ReplaceByOverlay -> {
-                    // TODO(b/353679003): Support overlay transitions when going back
-                    null
-                }
-            }
-
-        PredictiveBackHandler(state, coroutineScope, targetSceneForBack)
+        val result = contentForUserActions().userActions[Back.Resolved]
+        PredictiveBackHandler(layoutImpl = this, result = result)
     }
 
     @Composable
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 8ca90f18f3e0..ec9574ceb70d 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
@@ -18,6 +18,7 @@ package com.android.compose.animation.scene
 
 import androidx.compose.animation.core.Animatable
 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
@@ -40,6 +41,7 @@ internal fun createSwipeAnimation(
     result: UserActionResult,
     isUpOrLeft: Boolean,
     orientation: Orientation,
+    distance: Float = DistanceUnspecified,
 ): SwipeAnimation<*> {
     fun <T : Content> swipeAnimation(fromContent: T, toContent: T): SwipeAnimation<T> {
         return SwipeAnimation(
@@ -50,6 +52,7 @@ internal fun createSwipeAnimation(
             orientation = orientation,
             isUpOrLeft = isUpOrLeft,
             requiresFullDistanceSwipe = result.requiresFullDistanceSwipe,
+            lastDistance = distance,
         )
     }
 
@@ -147,7 +150,13 @@ internal class SwipeAnimation<T : Content>(
             // Important: If we are going to return early because distance is equal to 0, we should
             // still make sure we read the offset before returning so that the calling code still
             // subscribes to the offset value.
-            val offset = offsetAnimation?.animatable?.value ?: dragOffset
+            val animatable = offsetAnimation?.animatable
+            val offset =
+                when {
+                    animatable != null -> animatable.value
+                    contentTransition.previewTransformationSpec != null -> 0f
+                    else -> dragOffset
+                }
 
             return computeProgress(offset)
         }
@@ -172,6 +181,15 @@ internal class SwipeAnimation<T : Content>(
             return velocityInDistanceUnit / distance.absoluteValue
         }
 
+    val previewProgress: Float
+        get() = computeProgress(dragOffset)
+
+    val previewProgressVelocity: Float
+        get() = 0f
+
+    val isInPreviewStage: Boolean
+        get() = contentTransition.previewTransformationSpec != null && currentContent == fromContent
+
     override var bouncingContent: ContentKey? = null
 
     /** The current offset caused by the drag gesture. */
@@ -266,6 +284,7 @@ internal class SwipeAnimation<T : Content>(
         coroutineScope: CoroutineScope,
         initialVelocity: Float,
         targetContent: T,
+        spec: SpringSpec<Float>? = null,
     ): OffsetAnimation {
         val initialProgress = progress
         // Skip the animation if we have already reached the target content and the overscroll does
@@ -304,7 +323,9 @@ internal class SwipeAnimation<T : Content>(
         }
 
         return startOffsetAnimation {
-            val animatable = Animatable(dragOffset, OffsetVisibilityThreshold)
+            val startProgress =
+                if (contentTransition.previewTransformationSpec != null) 0f else dragOffset
+            val animatable = Animatable(startProgress, OffsetVisibilityThreshold)
             val isTargetGreater = targetOffset > animatable.value
             val startedWhenOvercrollingTargetContent =
                 if (targetContent == fromContent) initialProgress < 0f else initialProgress > 1f
@@ -325,7 +346,8 @@ internal class SwipeAnimation<T : Content>(
 
                         try {
                             val swipeSpec =
-                                contentTransition.transformationSpec.swipeSpec
+                                spec
+                                    ?: contentTransition.transformationSpec.swipeSpec
                                     ?: layoutImpl.state.transitions.defaultSwipeSpec
                             animatable.animateTo(
                                 targetValue = targetOffset,
@@ -471,6 +493,15 @@ private class ChangeCurrentSceneSwipeTransition(
     override val progressVelocity: Float
         get() = swipeAnimation.progressVelocity
 
+    override val previewProgress: Float
+        get() = swipeAnimation.previewProgress
+
+    override val previewProgressVelocity: Float
+        get() = swipeAnimation.previewProgressVelocity
+
+    override val isInPreviewStage: Boolean
+        get() = swipeAnimation.isInPreviewStage
+
     override val isInitiatedByUserInput: Boolean = true
 
     override val isUserInputOngoing: Boolean
@@ -519,6 +550,15 @@ private class ShowOrHideOverlaySwipeTransition(
     override val progressVelocity: Float
         get() = swipeAnimation.progressVelocity
 
+    override val previewProgress: Float
+        get() = swipeAnimation.previewProgress
+
+    override val previewProgressVelocity: Float
+        get() = swipeAnimation.previewProgressVelocity
+
+    override val isInPreviewStage: Boolean
+        get() = swipeAnimation.isInPreviewStage
+
     override val isInitiatedByUserInput: Boolean = true
 
     override val isUserInputOngoing: Boolean
@@ -561,6 +601,15 @@ private class ReplaceOverlaySwipeTransition(
     override val progressVelocity: Float
         get() = swipeAnimation.progressVelocity
 
+    override val previewProgress: Float
+        get() = swipeAnimation.previewProgress
+
+    override val previewProgressVelocity: Float
+        get() = swipeAnimation.previewProgressVelocity
+
+    override val isInPreviewStage: Boolean
+        get() = swipeAnimation.isInPreviewStage
+
     override val isInitiatedByUserInput: Boolean = true
 
     override val isUserInputOngoing: Boolean
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 00c75882a587..c5b6cdf12385 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
@@ -20,10 +20,16 @@ import androidx.activity.BackEventCompat
 import androidx.activity.ComponentActivity
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.size
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.test.assertIsDisplayed
+import androidx.compose.ui.test.hasTestTag
 import androidx.compose.ui.test.junit4.createAndroidComposeRule
+import androidx.compose.ui.unit.dp
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.compose.animation.scene.TestOverlays.OverlayA
+import com.android.compose.animation.scene.TestOverlays.OverlayB
 import com.android.compose.animation.scene.TestScenes.SceneA
 import com.android.compose.animation.scene.TestScenes.SceneB
 import com.android.compose.animation.scene.TestScenes.SceneC
@@ -198,6 +204,42 @@ class PredictiveBackHandlerTest {
         assertThat(canChangeSceneCalled).isFalse()
     }
 
+    @Test
+    fun backDismissesOverlayWithHighestZIndexByDefault() {
+        val state =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(
+                    SceneA,
+                    initialOverlays = setOf(OverlayA, OverlayB)
+                )
+            }
+
+        rule.setContent {
+            SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                scene(SceneA) { Box(Modifier.fillMaxSize()) }
+                overlay(OverlayA) { Box(Modifier.fillMaxSize()) }
+                overlay(OverlayB) { Box(Modifier.fillMaxSize()) }
+            }
+        }
+
+        // Initial state.
+        rule.onNode(hasTestTag(SceneA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayB.testTag)).assertIsDisplayed()
+
+        // Press back. This should hide overlay B because it has a higher zIndex than overlay A.
+        rule.runOnUiThread { rule.activity.onBackPressedDispatcher.onBackPressed() }
+        rule.onNode(hasTestTag(SceneA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayB.testTag)).assertDoesNotExist()
+
+        // Press back again. This should hide overlay A.
+        rule.runOnUiThread { rule.activity.onBackPressedDispatcher.onBackPressed() }
+        rule.onNode(hasTestTag(SceneA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayA.testTag)).assertDoesNotExist()
+        rule.onNode(hasTestTag(OverlayB.testTag)).assertDoesNotExist()
+    }
+
     private fun backEvent(progress: Float = 0f): BackEventCompat {
         return BackEventCompat(
             touchX = 0f,