Introduce MovableElementScenePicker

Bug: 317026105
Test: MovableElementScenePickerTest
Flag: N/A
Change-Id: Ie2d051259ef1ce4aa5b222d980dec2e089204cd0
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 f6e2d79..dc8505c 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
@@ -136,6 +136,7 @@
  * @see DefaultElementScenePicker
  * @see HighestZIndexScenePicker
  * @see LowestZIndexScenePicker
+ * @see MovableElementScenePicker
  */
 interface ElementScenePicker {
     /**
@@ -227,6 +228,31 @@
     }
 }
 
+/**
+ * An [ElementScenePicker] that draws/composes elements in the scene we are transitioning to, iff
+ * that scene is in [scenes].
+ *
+ * This picker can be useful for movable elements whose content size depends on its content (because
+ * it wraps it) in at least one scene. That way, the target size of the MovableElement will be
+ * computed in the scene we are going to and, given that this element was probably already composed
+ * in the scene we are going from before starting the transition, the interpolated size of the
+ * movable element during the transition should be correct.
+ *
+ * The downside of this picker is that the zIndex of the element when going from scene A to scene B
+ * is not the same as when going from scene B to scene A, so it's not usable in situations where
+ * z-ordering during the transition matters.
+ */
+class MovableElementScenePicker(private val scenes: Set<SceneKey>) : ElementScenePicker {
+    override fun sceneDuringTransition(
+        element: ElementKey,
+        transition: TransitionState.Transition,
+        fromSceneZIndex: Float,
+        toSceneZIndex: Float,
+    ): SceneKey {
+        return if (scenes.contains(transition.toScene)) transition.toScene else transition.fromScene
+    }
+}
+
 /** The default [ElementScenePicker]. */
 val DefaultElementScenePicker = HighestZIndexScenePicker
 
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementScenePickerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementScenePickerTest.kt
new file mode 100644
index 0000000..fb46a34
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementScenePickerTest.kt
@@ -0,0 +1,53 @@
+/*
+ * 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.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class MovableElementScenePickerTest {
+    @Test
+    fun toSceneInScenes() {
+        val picker = MovableElementScenePicker(scenes = setOf(TestScenes.SceneA, TestScenes.SceneB))
+        assertThat(
+                picker.sceneDuringTransition(
+                    TestElements.Foo,
+                    transition(from = TestScenes.SceneA, to = TestScenes.SceneB),
+                    fromSceneZIndex = 0f,
+                    toSceneZIndex = 1f,
+                )
+            )
+            .isEqualTo(TestScenes.SceneB)
+    }
+
+    @Test
+    fun toSceneNotInScenes() {
+        val picker = MovableElementScenePicker(scenes = emptySet())
+        assertThat(
+                picker.sceneDuringTransition(
+                    TestElements.Foo,
+                    transition(from = TestScenes.SceneA, to = TestScenes.SceneB),
+                    fromSceneZIndex = 0f,
+                    toSceneZIndex = 1f,
+                )
+            )
+            .isEqualTo(TestScenes.SceneA)
+    }
+}
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
index c5b8d9a..75dee47 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
@@ -50,13 +50,4 @@
         assertThat(state.isTransitioning(to = TestScenes.SceneA)).isFalse()
         assertThat(state.isTransitioning(from = TestScenes.SceneA, to = TestScenes.SceneB)).isTrue()
     }
-
-    private fun transition(from: SceneKey, to: SceneKey): TransitionState.Transition {
-        return object : TransitionState.Transition(from, to) {
-            override val currentScene: SceneKey = from
-            override val progress: Float = 0f
-            override val isInitiatedByUserInput: Boolean = false
-            override val isUserInputOngoing: Boolean = false
-        }
-    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
new file mode 100644
index 0000000..238b21e1
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
@@ -0,0 +1,33 @@
+/*
+ * 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
+
+/** A utility to easily create a [TransitionState.Transition] in tests. */
+fun transition(
+    from: SceneKey,
+    to: SceneKey,
+    progress: () -> Float = { 0f },
+    isInitiatedByUserInput: Boolean = false,
+    isUserInputOngoing: Boolean = false,
+): TransitionState.Transition {
+    return object : TransitionState.Transition(from, to) {
+        override val currentScene: SceneKey = from
+        override val progress: Float = progress()
+        override val isInitiatedByUserInput: Boolean = isInitiatedByUserInput
+        override val isUserInputOngoing: Boolean = isUserInputOngoing
+    }
+}