Move UserActionDistance to TransitionSpec

This CL moves UserActionDistance from UserActionResult to
TransitionSpec. I realized while adding custom distances to the demo
that this made more sense: usually the distance is the same when going
from A => B and B => A, and the previous API would require us to specify
the distance twice. Moreover, it was impossible to assign a distance to
a generic transition (for example from any scene to Shade or from Shade
to any scene).

Bug: 308961608
Test: SwipeToSceneTest
Flag: N/A
Change-Id: I24f9cf346f37d81dc19f738665b872e2efcd94d9
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposeAwareExtensions.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposeAwareExtensions.kt
index 52900e6..0de4650 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposeAwareExtensions.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposeAwareExtensions.kt
@@ -16,8 +16,6 @@
 
 package com.android.systemui.scene.ui.composable
 
-import androidx.compose.foundation.gestures.Orientation
-import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.Back
 import com.android.compose.animation.scene.Edge as ComposeAwareEdge
 import com.android.compose.animation.scene.SceneKey as ComposeAwareSceneKey
@@ -25,15 +23,12 @@
 import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.TransitionKey as ComposeAwareTransitionKey
 import com.android.compose.animation.scene.UserAction as ComposeAwareUserAction
-import com.android.compose.animation.scene.UserActionDistance as ComposeAwareUserActionDistance
-import com.android.compose.animation.scene.UserActionDistanceScope
 import com.android.compose.animation.scene.UserActionResult as ComposeAwareUserActionResult
 import com.android.systemui.scene.shared.model.Direction
 import com.android.systemui.scene.shared.model.Edge
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.TransitionKey
 import com.android.systemui.scene.shared.model.UserAction
-import com.android.systemui.scene.shared.model.UserActionDistance
 import com.android.systemui.scene.shared.model.UserActionResult
 
 // TODO(b/293899074): remove this file once we can use the types from SceneTransitionLayout.
@@ -82,22 +77,5 @@
     return ComposeAwareUserActionResult(
         toScene = composeUnaware.toScene.asComposeAware(),
         transitionKey = composeUnaware.transitionKey?.asComposeAware(),
-        distance = composeUnaware.distance?.asComposeAware(),
     )
 }
-
-fun UserActionDistance.asComposeAware(): ComposeAwareUserActionDistance {
-    val composeUnware = this
-    return object : ComposeAwareUserActionDistance {
-        override fun UserActionDistanceScope.absoluteDistance(
-            fromSceneSize: IntSize,
-            orientation: Orientation,
-        ): Float {
-            return composeUnware.absoluteDistance(
-                fromSceneWidth = fromSceneSize.width,
-                fromSceneHeight = fromSceneSize.height,
-                isHorizontal = orientation == Orientation.Horizontal,
-            )
-        }
-    }
-}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
index 8edf636..c408560 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
@@ -55,14 +55,17 @@
         if (isDrivingTransition || force) {
             layoutState.startTransition(newTransition, newTransition.key)
 
-            // Initialize SwipeTransition.swipeSpec. Note that this must be called right after
-            // layoutState.startTransition() is called, because it computes the
-            // layoutState.transformationSpec().
+            // Initialize SwipeTransition.transformationSpec and .swipeSpec. Note that this must be
+            // called right after layoutState.startTransition() is called, because it computes the
+            // current layoutState.transformationSpec().
+            val transformationSpec = layoutState.transformationSpec
+            newTransition.transformationSpec = transformationSpec
             newTransition.swipeSpec =
-                layoutState.transformationSpec.swipeSpec ?: layoutState.transitions.defaultSwipeSpec
+                transformationSpec.swipeSpec ?: layoutState.transitions.defaultSwipeSpec
         } else {
-            // We were not driving the transition and we don't force the update, so the spec won't
-            // be used and it doesn't matter which one we set here.
+            // We were not driving the transition and we don't force the update, so the specs won't
+            // be used and it doesn't matter which ones we set here.
+            newTransition.transformationSpec = TransformationSpec.Empty
             newTransition.swipeSpec = SceneTransitions.DefaultSwipeSpec
         }
 
@@ -472,43 +475,20 @@
 ): SwipeTransition {
     val upOrLeftResult = swipes.upOrLeftResult
     val downOrRightResult = swipes.downOrRightResult
-    val userActionDistance = result.distance ?: DefaultSwipeDistance
-
-    // The absolute distance of the gesture. Note that the UserActionDistance might return 0f or a
-    // negative value at first if it needs the size or offset of an element that is not composed yet
-    // when computing the distance. We call UserActionDistance.absoluteDistance() until it returns a
-    // value different than 0.
-    var lastAbsoluteDistance = 0f
-    val absoluteDistance: () -> Float = {
-        if (lastAbsoluteDistance > 0f) {
-            lastAbsoluteDistance
-        } else {
-            with(userActionDistance) {
-                    layoutImpl.userActionDistanceScope.absoluteDistance(
-                        fromScene.targetSize,
-                        orientation,
-                    )
-                }
-                .also { lastAbsoluteDistance = it }
-        }
-    }
-
-    // The signed distance of the gesture.
-    val distance: () -> Float = {
-        val absoluteDistance = absoluteDistance()
-        when {
-            absoluteDistance <= 0f -> SwipeTransition.DistanceUnspecified
-            result == upOrLeftResult -> -absoluteDistance
-            result == downOrRightResult -> absoluteDistance
+    val isUpOrLeft =
+        when (result) {
+            upOrLeftResult -> true
+            downOrRightResult -> false
             else -> error("Unknown result $result ($upOrLeftResult $downOrRightResult)")
         }
-    }
 
     return SwipeTransition(
         key = result.transitionKey,
         _fromScene = fromScene,
         _toScene = layoutImpl.scene(result.toScene),
-        distance = distance,
+        userActionDistanceScope = layoutImpl.userActionDistanceScope,
+        orientation = orientation,
+        isUpOrLeft = isUpOrLeft,
     )
 }
 
@@ -516,16 +496,9 @@
     val key: TransitionKey?,
     val _fromScene: Scene,
     val _toScene: Scene,
-
-    /**
-     * The signed distance between [fromScene] and [toScene]. It is negative if [fromScene] is above
-     * or to the left of [toScene].
-     *
-     * Note that this distance can be equal to [DistanceUnspecified] during the first frame of a
-     * transition when the distance depends on the size or position of an element that is composed
-     * in the scene we are going to.
-     */
-    val distance: () -> Float,
+    private val userActionDistanceScope: UserActionDistanceScope,
+    private val orientation: Orientation,
+    private val isUpOrLeft: Boolean,
 ) : TransitionState.Transition(_fromScene.key, _toScene.key) {
     var _currentScene by mutableStateOf(_fromScene)
     override val currentScene: SceneKey
@@ -566,9 +539,50 @@
     /** Job to check that there is at most one offset animation in progress. */
     private var offsetAnimationJob: Job? = null
 
+    /**
+     * The [TransformationSpecImpl] associated to this transition.
+     *
+     * Note: This is lateinit because this [SwipeTransition] is needed by
+     * [BaseSceneTransitionLayoutState] to compute the [TransitionSpec], and it will be set right
+     * after [BaseSceneTransitionLayoutState.startTransition] is called with this transition.
+     */
+    lateinit var transformationSpec: TransformationSpecImpl
+
     /** The spec to use when animating this transition to either [fromScene] or [toScene]. */
     lateinit var swipeSpec: SpringSpec<Float>
 
+    private var lastDistance = DistanceUnspecified
+
+    /**
+     * The signed distance between [fromScene] and [toScene]. It is negative if [fromScene] is above
+     * or to the left of [toScene].
+     *
+     * Note that this distance can be equal to [DistanceUnspecified] during the first frame of a
+     * transition when the distance depends on the size or position of an element that is composed
+     * in the scene we are going to.
+     */
+    fun distance(): Float {
+        if (lastDistance != DistanceUnspecified) {
+            return lastDistance
+        }
+
+        val absoluteDistance =
+            with(transformationSpec.distance ?: DefaultSwipeDistance) {
+                userActionDistanceScope.absoluteDistance(
+                    _fromScene.targetSize,
+                    orientation,
+                )
+            }
+
+        if (absoluteDistance <= 0f) {
+            return DistanceUnspecified
+        }
+
+        val distance = if (isUpOrLeft) -absoluteDistance else absoluteDistance
+        lastDistance = distance
+        return distance
+    }
+
     /** Ends any previous [offsetAnimationJob] and runs the new [job]. */
     private fun startOffsetAnimation(job: () -> Job) {
         cancelOffsetAnimation()
@@ -611,6 +625,7 @@
         }
         isAnimatingOffset = true
 
+        val animationSpec = transformationSpec
         offsetAnimatable.animateTo(
             targetValue = targetOffset,
             animationSpec = swipeSpec,
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 11085d9..1e3842a 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
@@ -395,22 +395,9 @@
     /** The scene we should be transitioning to during the [UserAction]. */
     val toScene: SceneKey,
 
-    /**
-     * The distance the action takes to animate from 0% to 100%.
-     *
-     * If `null`, a default distance will be used that depends on the [UserAction] performed.
-     */
-    val distance: UserActionDistance? = null,
-
     /** The key of the transition that should be used. */
     val transitionKey: TransitionKey? = null,
-) {
-    constructor(
-        toScene: SceneKey,
-        distance: Dp,
-        transitionKey: TransitionKey? = null,
-    ) : this(toScene, FixedDistance(distance), transitionKey)
-}
+)
 
 interface UserActionDistance {
     /**
@@ -449,7 +436,7 @@
 }
 
 /** The user action has a fixed [absoluteDistance]. */
-private class FixedDistance(private val distance: Dp) : UserActionDistance {
+class FixedDistance(private val distance: Dp) : UserActionDistance {
     override fun UserActionDistanceScope.absoluteDistance(
         fromSceneSize: IntSize,
         orientation: Orientation,
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 b8f9359..8ee23b6 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
@@ -163,6 +163,14 @@
      */
     val swipeSpec: SpringSpec<Float>?
 
+    /**
+     * The distance it takes for this transition to animate from 0% to 100% when it is driven by a
+     * [UserAction].
+     *
+     * If `null`, a default distance will be used that depends on the [UserAction] performed.
+     */
+    val distance: UserActionDistance?
+
     /** The list of [Transformation] applied to elements during this transition. */
     val transformations: List<Transformation>
 
@@ -171,6 +179,7 @@
             TransformationSpecImpl(
                 progressSpec = snap(),
                 swipeSpec = null,
+                distance = null,
                 transformations = emptyList(),
             )
         internal val EmptyProvider = { Empty }
@@ -193,6 +202,7 @@
                 TransformationSpecImpl(
                     progressSpec = reverse.progressSpec,
                     swipeSpec = reverse.swipeSpec,
+                    distance = reverse.distance,
                     transformations = reverse.transformations.map { it.reversed() }
                 )
             }
@@ -209,6 +219,7 @@
 internal class TransformationSpecImpl(
     override val progressSpec: AnimationSpec<Float>,
     override val swipeSpec: SpringSpec<Float>?,
+    override val distance: UserActionDistance?,
     override val transformations: List<Transformation>,
 ) : TransformationSpec {
     private val cache = mutableMapOf<ElementKey, MutableMap<SceneKey, ElementTransformations>>()
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
index d93911d..8a09b00 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
@@ -91,6 +91,14 @@
     var swipeSpec: SpringSpec<Float>?
 
     /**
+     * The distance it takes for this transition to animate from 0% to 100% when it is driven by a
+     * [UserAction].
+     *
+     * If `null`, a default distance will be used that depends on the [UserAction] performed.
+     */
+    var distance: UserActionDistance?
+
+    /**
      * Define a progress-based range for the transformations inside [builder].
      *
      * For instance, the following will fade `Foo` during the first half of the transition then it
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
index 9b16d46..7828999 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
@@ -77,6 +77,7 @@
             return TransformationSpecImpl(
                 progressSpec = impl.spec,
                 swipeSpec = impl.swipeSpec,
+                distance = impl.distance,
                 transformations = impl.transformations,
             )
         }
@@ -91,6 +92,7 @@
     val transformations = mutableListOf<Transformation>()
     override var spec: AnimationSpec<Float> = spring(stiffness = Spring.StiffnessLow)
     override var swipeSpec: SpringSpec<Float>? = null
+    override var distance: UserActionDistance? = null
 
     private var range: TransformationRange? = null
     private var reversed = false
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
index 44b5d7f..99372a5 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
@@ -64,8 +64,10 @@
 
     @get:Rule val rule = createComposeRule()
 
-    private fun layoutState(initialScene: SceneKey = TestScenes.SceneA) =
-        MutableSceneTransitionLayoutState(initialScene, EmptyTestTransitions)
+    private fun layoutState(
+        initialScene: SceneKey = TestScenes.SceneA,
+        transitions: SceneTransitions = EmptyTestTransitions,
+    ) = MutableSceneTransitionLayoutState(initialScene, transitions)
 
     /** The content under test. */
     @Composable
@@ -373,8 +375,16 @@
         // detected as a drag event.
         var touchSlop = 0f
 
-        val layoutState = layoutState()
         val verticalSwipeDistance = 50.dp
+        val layoutState =
+            layoutState(
+                transitions =
+                    transitions {
+                        from(TestScenes.SceneA, to = TestScenes.SceneB) {
+                            distance = FixedDistance(verticalSwipeDistance)
+                        }
+                    }
+            )
         assertThat(verticalSwipeDistance).isNotEqualTo(LayoutHeight)
 
         rule.setContent {
@@ -386,14 +396,7 @@
             ) {
                 scene(
                     TestScenes.SceneA,
-                    userActions =
-                        mapOf(
-                            Swipe.Down to
-                                UserActionResult(
-                                    toScene = TestScenes.SceneB,
-                                    distance = verticalSwipeDistance,
-                                )
-                        ),
+                    userActions = mapOf(Swipe.Down to TestScenes.SceneB),
                 ) {
                     Spacer(Modifier.fillMaxSize())
                 }
@@ -554,7 +557,6 @@
 
     @Test
     fun dynamicSwipeDistance() {
-        val state = MutableSceneTransitionLayoutState(TestScenes.SceneA)
         val swipeDistance =
             object : UserActionDistance {
                 override fun UserActionDistanceScope.absoluteDistance(
@@ -572,6 +574,14 @@
                 }
             }
 
+        val state =
+            MutableSceneTransitionLayoutState(
+                TestScenes.SceneA,
+                transitions {
+                    from(TestScenes.SceneA, to = TestScenes.SceneB) { distance = swipeDistance }
+                }
+            )
+
         val layoutSize = 200.dp
         val fooYOffset = 50.dp
         val fooSize = 25.dp
@@ -581,14 +591,7 @@
             touchSlop = LocalViewConfiguration.current.touchSlop
 
             SceneTransitionLayout(state, Modifier.size(layoutSize)) {
-                scene(
-                    TestScenes.SceneA,
-                    userActions =
-                        mapOf(
-                            Swipe.Up to
-                                UserActionResult(TestScenes.SceneB, distance = swipeDistance)
-                        )
-                ) {
+                scene(TestScenes.SceneA, userActions = mapOf(Swipe.Up to TestScenes.SceneB)) {
                     Box(Modifier.fillMaxSize())
                 }
                 scene(TestScenes.SceneB) {
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/UserActionResult.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/UserActionResult.kt
index e1b96e4..c6ae215 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/UserActionResult.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/UserActionResult.kt
@@ -22,13 +22,6 @@
     val toScene: SceneKey,
 
     /**
-     * The distance the action takes to animate from 0% to 100%.
-     *
-     * If `null`, a default distance will be used depending on the [UserAction] performed.
-     */
-    val distance: UserActionDistance? = null,
-
-    /**
      * The key of the transition that should be used, if a specific one should be used.
      *
      * If `null`, the transition used will be the corresponding transition from the collection