[STL] Add default PredictiveBack transition key

Bug: 350705972
Test: TransitionDslTest
Flag: com.android.systemui.scene_container
Change-Id: I73c7211fd9d9271bd3e116aaa7b68ff54b08d710

5 files changed, 86 insertions, 12 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
index 3f8f5e742079..ced177ccb9a0 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
@@ -153,4 +153,15 @@ class TransitionKey(debugName: String, identity: Any = Object()) : Key(debugName
     override fun toString(): String {
         return "TransitionKey(debugName=$debugName)"
     }
+
+    companion object {
+        /**
+         * A special transition key indicating that the associated transition should be used for
+         * Predictive Back gestures.
+         *
+         * Use this key when defining a transition that you want to be specifically triggered when
+         * the user performs a Predictive Back gesture.
+         */
+        val PredictiveBack = TransitionKey("PredictiveBack")
+    }
 }
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 be4fea10602f..c2fb8fce1d6e 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
@@ -43,7 +43,9 @@ internal fun PredictiveBackHandler(
             createSwipeAnimation(
                 layoutImpl,
                 layoutImpl.coroutineScope,
-                result,
+                result.userActionCopy(
+                    transitionKey = result.transitionKey ?: TransitionKey.PredictiveBack
+                ),
                 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.
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
index b3f74f749a0e..e4534302267a 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -492,6 +492,17 @@ sealed class UserActionResult(
 ) {
     internal abstract fun toContent(currentScene: SceneKey): ContentKey
 
+    internal fun userActionCopy(
+        transitionKey: TransitionKey? = this.transitionKey
+    ): UserActionResult {
+        return when (this) {
+            is ChangeScene -> copy(transitionKey = transitionKey)
+            is ShowOverlay -> copy(transitionKey = transitionKey)
+            is HideOverlay -> copy(transitionKey = transitionKey)
+            is ReplaceByOverlay -> copy(transitionKey = transitionKey)
+        }
+    }
+
     data class ChangeScene
     internal constructor(
         /** The scene we should be transitioning to during the [UserAction]. */
@@ -503,19 +514,19 @@ sealed class UserActionResult(
     }
 
     /** A [UserActionResult] that shows [overlay]. */
-    class ShowOverlay(
+    data class ShowOverlay(
         val overlay: OverlayKey,
-        transitionKey: TransitionKey? = null,
-        requiresFullDistanceSwipe: Boolean = false,
+        override val transitionKey: TransitionKey? = null,
+        override val requiresFullDistanceSwipe: Boolean = false,
     ) : UserActionResult(transitionKey, requiresFullDistanceSwipe) {
         override fun toContent(currentScene: SceneKey): ContentKey = overlay
     }
 
     /** A [UserActionResult] that hides [overlay]. */
-    class HideOverlay(
+    data class HideOverlay(
         val overlay: OverlayKey,
-        transitionKey: TransitionKey? = null,
-        requiresFullDistanceSwipe: Boolean = false,
+        override val transitionKey: TransitionKey? = null,
+        override val requiresFullDistanceSwipe: Boolean = false,
     ) : UserActionResult(transitionKey, requiresFullDistanceSwipe) {
         override fun toContent(currentScene: SceneKey): ContentKey = currentScene
     }
@@ -526,10 +537,10 @@ sealed class UserActionResult(
      * Note: This result can only be used for user actions of overlays and an exception will be
      * thrown if it is used for a scene.
      */
-    class ReplaceByOverlay(
+    data class ReplaceByOverlay(
         val overlay: OverlayKey,
-        transitionKey: TransitionKey? = null,
-        requiresFullDistanceSwipe: Boolean = false,
+        override val transitionKey: TransitionKey? = null,
+        override val requiresFullDistanceSwipe: Boolean = false,
     ) : UserActionResult(transitionKey, requiresFullDistanceSwipe) {
         override fun toContent(currentScene: SceneKey): ContentKey = overlay
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
index cefcff75f13a..e65ed9b7dc97 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
@@ -90,10 +90,19 @@ internal constructor(
             return relaxedSpec
         }
 
-        return transition(from, to, key) {
+        val relaxedReversed =
+            transition(from, to, key) {
                 (it.from == to && it.to == null) || (it.to == from && it.from == null)
             }
-            ?.reversed() ?: defaultTransition(from, to)
+        if (relaxedReversed != null) {
+            return relaxedReversed.reversed()
+        }
+
+        return if (key != null) {
+            findSpec(from, to, null)
+        } else {
+            defaultTransition(from, to)
+        }
     }
 
     private fun transition(
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
index bed6cefa459d..f8068e612f11 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
@@ -232,6 +232,47 @@ class TransitionDslTest {
     }
 
     @Test
+    fun defaultPredictiveBack() {
+        val transitions = transitions {
+            from(
+                TestScenes.SceneA,
+                to = TestScenes.SceneB,
+                preview = { fractionRange(start = 0.1f, end = 0.8f) { fade(TestElements.Foo) } }
+            ) {
+                spec = tween(500)
+                fractionRange(start = 0.1f, end = 0.8f) { fade(TestElements.Foo) }
+                timestampRange(startMillis = 100, endMillis = 300) { fade(TestElements.Foo) }
+            }
+        }
+
+        // Verify that fetching the transitionSpec with the PredictiveBack key defaults to the above
+        // transition despite it not having the PredictiveBack key set.
+        val transitionSpec =
+            transitions.transitionSpec(
+                from = TestScenes.SceneA,
+                to = TestScenes.SceneB,
+                key = TransitionKey.PredictiveBack
+            )
+
+        val transformations = transitionSpec.transformationSpec().transformations
+
+        assertThat(transformations)
+            .comparingElementsUsing(TRANSFORMATION_RANGE)
+            .containsExactly(
+                TransformationRange(start = 0.1f, end = 0.8f),
+                TransformationRange(start = 100 / 500f, end = 300 / 500f),
+            )
+
+        val previewTransformations = transitionSpec.previewTransformationSpec()?.transformations
+
+        assertThat(previewTransformations)
+            .comparingElementsUsing(TRANSFORMATION_RANGE)
+            .containsExactly(
+                TransformationRange(start = 0.1f, end = 0.8f),
+            )
+    }
+
+    @Test
     fun springSpec() {
         val defaultSpec = spring<Float>(stiffness = 1f)
         val specFromAToC = spring<Float>(stiffness = 2f)