Make it possible to compute the swipe distance lazily

This CL makes it possible to compute the swipe distance of swipes
lazily, so that the distance can depend on the size or position of
elements from the scene we are transitioning to.

The way it works is pretty simple: UserActionDistance.absoluteDistance
can return 0f until the swipe distance can be computed. This CL also
exposes a UserActionDistanceScope to get the target size and offset of
an element inside a scene, which are usually useful to compute swipe
distances that depend on an element target state.

Note that some of the code in this CL is going to be moved in
ag/26316632.

Bug: 308961608
Test: SwipeToSceneTest
Flag: N/A
Change-Id: If70c48de0aba7a793942badcb5a24993277302b1

6 files changed, 230 insertions, 30 deletions

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 a7de1eede1f4..52900e643a1f 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
@@ -17,7 +17,6 @@
 package com.android.systemui.scene.ui.composable
 
 import androidx.compose.foundation.gestures.Orientation
-import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.Back
 import com.android.compose.animation.scene.Edge as ComposeAwareEdge
@@ -27,6 +26,7 @@ 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
@@ -89,7 +89,7 @@ fun UserActionResult.asComposeAware(): ComposeAwareUserActionResult {
 fun UserActionDistance.asComposeAware(): ComposeAwareUserActionDistance {
     val composeUnware = this
     return object : ComposeAwareUserActionDistance {
-        override fun Density.absoluteDistance(
+        override fun UserActionDistanceScope.absoluteDistance(
             fromSceneSize: IntSize,
             orientation: Orientation,
         ): Float {
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 76e7c95f274a..8edf636a2dcb 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
@@ -27,7 +27,6 @@ import androidx.compose.runtime.mutableFloatStateOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.unit.round
@@ -285,16 +284,21 @@ internal class SceneGestureHandler(
     ): Pair<Scene, Float> {
         val toScene = swipeTransition._toScene
         val fromScene = swipeTransition._fromScene
-        val absoluteDistance = swipeTransition.distance.absoluteValue
+        val distance = swipeTransition.distance()
 
-        // If the swipe was not committed, don't do anything.
-        if (swipeTransition._currentScene != toScene) {
+        // If the swipe was not committed or if the swipe distance is not computed yet, don't do
+        // anything.
+        if (
+            swipeTransition._currentScene != toScene ||
+                distance == SwipeTransition.DistanceUnspecified
+        ) {
             return fromScene to 0f
         }
 
         // If the offset is past the distance then let's change fromScene so that the user can swipe
         // to the next screen or go back to the previous one.
         val offset = swipeTransition.dragOffset
+        val absoluteDistance = distance.absoluteValue
         return if (offset <= -absoluteDistance && swipes!!.upOrLeftResult?.toScene == toScene.key) {
             toScene to absoluteDistance
         } else if (
@@ -347,16 +351,17 @@ internal class SceneGestureHandler(
 
             // Compute the destination scene (and therefore offset) to settle in.
             val offset = swipeTransition.dragOffset
-            val distance = swipeTransition.distance
+            val distance = swipeTransition.distance()
             var targetScene: Scene
             var targetOffset: Float
             if (
-                shouldCommitSwipe(
-                    offset,
-                    distance,
-                    velocity,
-                    wasCommitted = swipeTransition._currentScene == toScene,
-                )
+                distance != SwipeTransition.DistanceUnspecified &&
+                    shouldCommitSwipe(
+                        offset,
+                        distance,
+                        velocity,
+                        wasCommitted = swipeTransition._currentScene == toScene,
+                    )
             ) {
                 targetScene = toScene
                 targetOffset = distance
@@ -372,7 +377,15 @@ internal class SceneGestureHandler(
                 // We wanted to change to a new scene but we are not allowed to, so we animate back
                 // to the current scene.
                 targetScene = swipeTransition._currentScene
-                targetOffset = if (targetScene == fromScene) 0f else distance
+                targetOffset =
+                    if (targetScene == fromScene) {
+                        0f
+                    } else {
+                        check(distance != SwipeTransition.DistanceUnspecified) {
+                            "distance is equal to ${SwipeTransition.DistanceUnspecified}"
+                        }
+                        distance
+                    }
             }
 
             animateTo(targetScene = targetScene, targetOffset = targetOffset)
@@ -460,21 +473,42 @@ private fun SwipeTransition(
     val upOrLeftResult = swipes.upOrLeftResult
     val downOrRightResult = swipes.downOrRightResult
     val userActionDistance = result.distance ?: DefaultSwipeDistance
-    val absoluteDistance =
-        with(userActionDistance) {
-            layoutImpl.density.absoluteDistance(fromScene.targetSize, orientation)
+
+    // 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
+            else -> error("Unknown result $result ($upOrLeftResult $downOrRightResult)")
+        }
+    }
 
     return SwipeTransition(
         key = result.transitionKey,
         _fromScene = fromScene,
         _toScene = layoutImpl.scene(result.toScene),
-        distance =
-            when (result) {
-                upOrLeftResult -> -absoluteDistance
-                downOrRightResult -> absoluteDistance
-                else -> error("Unknown result $result ($upOrLeftResult $downOrRightResult)")
-            },
+        distance = distance,
     )
 }
 
@@ -482,11 +516,16 @@ private class SwipeTransition(
     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]
+     * 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,
+    val distance: () -> Float,
 ) : TransitionState.Transition(_fromScene.key, _toScene.key) {
     var _currentScene by mutableStateOf(_fromScene)
     override val currentScene: SceneKey
@@ -494,7 +533,16 @@ private class SwipeTransition(
 
     override val progress: Float
         get() {
+            // 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 = if (isAnimatingOffset) offsetAnimatable.value else dragOffset
+
+            val distance = distance()
+            if (distance == DistanceUnspecified) {
+                return 0f
+            }
+
             return offset / distance
         }
 
@@ -571,10 +619,14 @@ private class SwipeTransition(
 
         finishOffsetAnimation()
     }
+
+    companion object {
+        const val DistanceUnspecified = 0f
+    }
 }
 
 private object DefaultSwipeDistance : UserActionDistance {
-    override fun Density.absoluteDistance(
+    override fun UserActionDistanceScope.absoluteDistance(
         fromSceneSize: IntSize,
         orientation: Orientation,
     ): Float {
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 e1f8a0959f6f..11085d9c95e8 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
@@ -25,6 +25,7 @@ import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.unit.Density
@@ -415,15 +416,44 @@ interface UserActionDistance {
     /**
      * Return the **absolute** distance of the user action given the size of the scene we are
      * animating from and the [orientation].
+     *
+     * Note: This function will be called for each drag event until it returns a value > 0f. This
+     * for instance allows you to return 0f or a negative value until the first layout pass of a
+     * scene, so that you can use the size and position of elements in the scene we are
+     * transitioning to when computing this absolute distance.
      */
-    fun Density.absoluteDistance(fromSceneSize: IntSize, orientation: Orientation): Float
+    fun UserActionDistanceScope.absoluteDistance(
+        fromSceneSize: IntSize,
+        orientation: Orientation
+    ): Float
+}
+
+interface UserActionDistanceScope : Density {
+    /**
+     * Return the *target* size of [this] element in the given [scene], i.e. the size of the element
+     * when idle, or `null` if the element is not composed and measured in that scene (yet).
+     */
+    fun ElementKey.targetSize(scene: SceneKey): IntSize?
+
+    /**
+     * Return the *target* offset of [this] element in the given [scene], i.e. the size of the
+     * element when idle, or `null` if the element is not composed and placed in that scene (yet).
+     */
+    fun ElementKey.targetOffset(scene: SceneKey): Offset?
+
+    /**
+     * Return the *target* size of [this] scene, i.e. the size of the scene when idle, or `null` if
+     * the scene was never composed.
+     */
+    fun SceneKey.targetSize(): IntSize?
 }
 
 /** The user action has a fixed [absoluteDistance]. */
 private class FixedDistance(private val distance: Dp) : UserActionDistance {
-    override fun Density.absoluteDistance(fromSceneSize: IntSize, orientation: Orientation): Float {
-        return distance.toPx()
-    }
+    override fun UserActionDistanceScope.absoluteDistance(
+        fromSceneSize: IntSize,
+        orientation: Orientation,
+    ): Float = distance.toPx()
 }
 
 /**
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 08399ff03f63..039a5b0c9523 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
@@ -96,9 +96,18 @@ internal class SceneTransitionLayoutImpl(
                 ?: mutableMapOf<ValueKey, MutableMap<ElementKey?, SnapshotStateMap<SceneKey, *>>>()
                     .also { _sharedValues = it }
 
+    // TODO(b/317958526): Lazily allocate scene gesture handlers the first time they are needed.
     private val horizontalGestureHandler: SceneGestureHandler
     private val verticalGestureHandler: SceneGestureHandler
 
+    private var _userActionDistanceScope: UserActionDistanceScope? = null
+    internal val userActionDistanceScope: UserActionDistanceScope
+        get() =
+            _userActionDistanceScope
+                ?: UserActionDistanceScopeImpl(layoutImpl = this).also {
+                    _userActionDistanceScope = it
+                }
+
     init {
         updateScenes(builder)
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
new file mode 100644
index 000000000000..228d19f09cff
--- /dev/null
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.compose.animation.scene
+
+import androidx.compose.ui.geometry.Offset
+import androidx.compose.ui.unit.IntSize
+
+internal class UserActionDistanceScopeImpl(
+    private val layoutImpl: SceneTransitionLayoutImpl,
+) : UserActionDistanceScope {
+    override val density: Float
+        get() = layoutImpl.density.density
+
+    override val fontScale: Float
+        get() = layoutImpl.density.fontScale
+
+    override fun ElementKey.targetSize(scene: SceneKey): IntSize? {
+        return layoutImpl.elements[this]?.sceneStates?.get(scene)?.targetSize.takeIf {
+            it != Element.SizeUnspecified
+        }
+    }
+
+    override fun ElementKey.targetOffset(scene: SceneKey): Offset? {
+        return layoutImpl.elements[this]?.sceneStates?.get(scene)?.targetOffset.takeIf {
+            it != Offset.Unspecified
+        }
+    }
+
+    override fun SceneKey.targetSize(): IntSize? {
+        return layoutImpl.scenes[this]?.targetSize.takeIf { it != IntSize.Zero }
+    }
+}
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 543ed0482430..44b5d7f3a5fd 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
@@ -16,9 +16,11 @@
 
 package com.android.compose.animation.scene
 
+import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.offset
 import androidx.compose.foundation.layout.size
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
@@ -33,6 +35,7 @@ import androidx.compose.ui.test.onRoot
 import androidx.compose.ui.test.performTouchInput
 import androidx.compose.ui.test.swipeWithVelocity
 import androidx.compose.ui.unit.Density
+import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.dp
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.google.common.truth.Truth.assertThat
@@ -548,4 +551,64 @@ class SwipeToSceneTest {
         assertThat(state.isTransitioning(from = TestScenes.SceneA, to = TestScenes.SceneB)).isTrue()
         assertThat(state.transformationSpec.transformations).hasSize(2)
     }
+
+    @Test
+    fun dynamicSwipeDistance() {
+        val state = MutableSceneTransitionLayoutState(TestScenes.SceneA)
+        val swipeDistance =
+            object : UserActionDistance {
+                override fun UserActionDistanceScope.absoluteDistance(
+                    fromSceneSize: IntSize,
+                    orientation: Orientation,
+                ): Float {
+                    // Foo is going to have a vertical offset of 50dp. Let's make the swipe distance
+                    // the difference between the bottom of the scene and the bottom of the element,
+                    // so that we use the offset and size of the element as well as the size of the
+                    // scene.
+                    val fooSize = TestElements.Foo.targetSize(TestScenes.SceneB) ?: return 0f
+                    val fooOffset = TestElements.Foo.targetOffset(TestScenes.SceneB) ?: return 0f
+                    val sceneSize = TestScenes.SceneB.targetSize() ?: return 0f
+                    return sceneSize.height - fooOffset.y - fooSize.height
+                }
+            }
+
+        val layoutSize = 200.dp
+        val fooYOffset = 50.dp
+        val fooSize = 25.dp
+
+        var touchSlop = 0f
+        rule.setContent {
+            touchSlop = LocalViewConfiguration.current.touchSlop
+
+            SceneTransitionLayout(state, Modifier.size(layoutSize)) {
+                scene(
+                    TestScenes.SceneA,
+                    userActions =
+                        mapOf(
+                            Swipe.Up to
+                                UserActionResult(TestScenes.SceneB, distance = swipeDistance)
+                        )
+                ) {
+                    Box(Modifier.fillMaxSize())
+                }
+                scene(TestScenes.SceneB) {
+                    Box(Modifier.fillMaxSize()) {
+                        Box(Modifier.offset(y = fooYOffset).element(TestElements.Foo).size(fooSize))
+                    }
+                }
+            }
+        }
+
+        // Swipe up by half the expected distance to get to 50% progress.
+        val expectedDistance = layoutSize - fooYOffset - fooSize
+        rule.onRoot().performTouchInput {
+            val middle = (layoutSize / 2).toPx()
+            down(Offset(middle, middle))
+            moveBy(Offset(0f, -touchSlop - (expectedDistance / 2f).toPx()), delayMillis = 1_000)
+        }
+
+        rule.waitForIdle()
+        assertThat(state.isTransitioning(from = TestScenes.SceneA, to = TestScenes.SceneB)).isTrue()
+        assertThat(state.currentTransition!!.progress).isWithin(0.01f).of(0.5f)
+    }
 }