Merge "Fix MovableElements during overscroll on overlays" into main

6 files changed, 185 insertions, 16 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
index 4c0feb883c84..56c08b9d4789 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
@@ -853,19 +853,32 @@ private fun shouldPlaceElement(
         content,
         element.key,
         transition,
+        isInContent = { it in element.stateByContent },
     )
 }
 
-internal fun shouldPlaceOrComposeSharedElement(
+internal inline fun shouldPlaceOrComposeSharedElement(
     layoutImpl: SceneTransitionLayoutImpl,
     content: ContentKey,
     element: ElementKey,
     transition: TransitionState.Transition,
+    isInContent: (ContentKey) -> Boolean,
 ): Boolean {
-    // If we are overscrolling, only place/compose the element in the overscrolling scene.
-    val overscrollScene = transition.currentOverscrollSpec?.content
-    if (overscrollScene != null) {
-        return content == overscrollScene
+    val overscrollContent = transition.currentOverscrollSpec?.content
+    if (overscrollContent != null) {
+        return when (transition) {
+            // If we are overscrolling between scenes, only place/compose the element in the
+            // overscrolling scene.
+            is TransitionState.Transition.ChangeScene -> content == overscrollContent
+
+            // If we are overscrolling an overlay, place/compose the element if [content] is the
+            // overscrolling content or if [content] is the current scene and the overscrolling
+            // overlay does not contain the element.
+            is TransitionState.Transition.ReplaceOverlay,
+            is TransitionState.Transition.ShowOrHideOverlay ->
+                content == overscrollContent ||
+                    (content == transition.currentScene && !isInContent(overscrollContent))
+        }
     }
 
     val scenePicker = element.contentPicker
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
index 715222cfd9da..471ad3fee9fb 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MovableElement.kt
@@ -194,11 +194,13 @@ private fun shouldComposeMovableElement(
         is TransitionState.Transition -> {
             // During transitions, always compose movable elements in the scene picked by their
             // content picker.
+            val contents = element.contentPicker.contents
             shouldPlaceOrComposeSharedElement(
                 layoutImpl,
                 content,
                 element,
                 elementState,
+                isInContent = { contents.contains(it) }
             )
         }
     }
@@ -208,8 +210,8 @@ private fun movableElementState(
     element: MovableElementKey,
     transitionStates: List<TransitionState>,
 ): TransitionState? {
-    val content = element.contentPicker.contents
-    return elementState(transitionStates, isInContent = { content.contains(it) })
+    val contents = element.contentPicker.contents
+    return elementState(transitionStates, isInContent = { contents.contains(it) })
 }
 
 private fun movableElementContentWhenIdle(
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
index c25478b35790..471362ba19e9 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
@@ -18,6 +18,7 @@ package com.android.compose.animation.scene
 
 import androidx.compose.animation.core.LinearEasing
 import androidx.compose.animation.core.tween
+import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.size
@@ -44,9 +45,12 @@ 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.test.assertSizeIsEqualTo
+import com.android.compose.test.setContentAndCreateMainScope
 import com.android.compose.test.subjects.assertThat
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -648,4 +652,42 @@ class OverlayTest {
             }
         }
     }
+
+    @Test
+    fun overscrollingOverlay_movableElementNotInOverlay() {
+        val state =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutStateImpl(
+                    SceneA,
+                    transitions {
+                        // Make OverlayA overscrollable.
+                        overscroll(OverlayA, orientation = Orientation.Horizontal) {
+                            translate(ElementKey("elementThatDoesNotExist"), x = 10.dp)
+                        }
+                    }
+                )
+            }
+
+        val key = MovableElementKey("Foo", contents = setOf(SceneA))
+        val movableElementChildTag = "movableElementChildTag"
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    scene(SceneA) {
+                        MovableElement(key, Modifier) {
+                            content { Box(Modifier.testTag(movableElementChildTag).size(100.dp)) }
+                        }
+                    }
+                    overlay(OverlayA) { /* Does not contain the element. */ }
+                }
+            }
+
+        // Overscroll on Overlay A.
+        scope.launch { state.startTransition(transition(SceneA, OverlayA, progress = { 1.5f })) }
+        rule
+            .onNode(hasTestTag(movableElementChildTag) and inContent(SceneA))
+            .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+            .assertSizeIsEqualTo(100.dp)
+            .assertIsDisplayed()
+    }
 }
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 cd20a29a05d9..d356c25262e8 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
@@ -32,7 +32,7 @@ import com.android.compose.animation.scene.subjects.assertThat
 import com.android.compose.animation.scene.transition.link.StateLink
 import com.android.compose.animation.scene.transition.seekToScene
 import com.android.compose.test.MonotonicClockTestScope
-import com.android.compose.test.TestTransition
+import com.android.compose.test.TestSceneTransition
 import com.android.compose.test.runMonotonicClockTest
 import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
@@ -556,8 +556,8 @@ class SceneTransitionLayoutStateTest {
 
     @Test
     fun multipleTransitions() = runTest {
-        val frozenTransitions = mutableSetOf<TestTransition>()
-        fun onFreezeAndAnimate(transition: TestTransition): () -> Unit {
+        val frozenTransitions = mutableSetOf<TestSceneTransition>()
+        fun onFreezeAndAnimate(transition: TestSceneTransition): () -> Unit {
             // Instead of letting the transition finish when it is frozen, we put the transition in
             // the frozenTransitions set so that we can verify that freezeAndAnimateToCurrentState()
             // is called when expected and then we call finish() ourselves to finish the
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestOverlayTransition.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestOverlayTransition.kt
new file mode 100644
index 000000000000..646cff8b944c
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestOverlayTransition.kt
@@ -0,0 +1,112 @@
+/*
+ * 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.test
+
+import androidx.compose.foundation.gestures.Orientation
+import com.android.compose.animation.scene.ContentKey
+import com.android.compose.animation.scene.OverlayKey
+import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.SceneTransitionLayoutImpl
+import com.android.compose.animation.scene.content.state.TransitionState
+import com.android.compose.animation.scene.content.state.TransitionState.Transition
+import kotlinx.coroutines.CompletableDeferred
+
+/** A [Transition.ShowOrHideOverlay] for tests that will be finished once [finish] is called. */
+abstract class TestOverlayTransition(
+    fromScene: SceneKey,
+    overlay: OverlayKey,
+    replacedTransition: Transition?,
+) :
+    Transition.ShowOrHideOverlay(
+        overlay = overlay,
+        fromOrToScene = fromScene,
+        fromContent = fromScene,
+        toContent = overlay,
+        replacedTransition = replacedTransition,
+    ) {
+    private val finishCompletable = CompletableDeferred<Unit>()
+
+    override suspend fun run() {
+        finishCompletable.await()
+    }
+
+    /** Finish this transition. */
+    fun finish() {
+        finishCompletable.complete(Unit)
+    }
+}
+
+/** A utility to easily create a [TestOverlayTransition] in tests. */
+fun transition(
+    fromScene: SceneKey,
+    overlay: OverlayKey,
+    isEffectivelyShown: () -> Boolean = { true },
+    progress: () -> Float = { 0f },
+    progressVelocity: () -> Float = { 0f },
+    previewProgress: () -> Float = { 0f },
+    previewProgressVelocity: () -> Float = { 0f },
+    isInPreviewStage: () -> Boolean = { false },
+    interruptionProgress: () -> Float = { 0f },
+    isInitiatedByUserInput: Boolean = false,
+    isUserInputOngoing: Boolean = false,
+    isUpOrLeft: Boolean = false,
+    bouncingContent: ContentKey? = null,
+    orientation: Orientation = Orientation.Horizontal,
+    onFreezeAndAnimate: ((TestOverlayTransition) -> Unit)? = null,
+    replacedTransition: Transition? = null,
+): TestOverlayTransition {
+    return object :
+        TestOverlayTransition(fromScene, overlay, replacedTransition),
+        TransitionState.HasOverscrollProperties {
+        override val isEffectivelyShown: Boolean
+            get() = isEffectivelyShown()
+
+        override val progress: Float
+            get() = progress()
+
+        override val progressVelocity: Float
+            get() = progressVelocity()
+
+        override val previewProgress: Float
+            get() = previewProgress()
+
+        override val previewProgressVelocity: Float
+            get() = previewProgressVelocity()
+
+        override val isInPreviewStage: Boolean
+            get() = isInPreviewStage()
+
+        override val isInitiatedByUserInput: Boolean = isInitiatedByUserInput
+        override val isUserInputOngoing: Boolean = isUserInputOngoing
+        override val isUpOrLeft: Boolean = isUpOrLeft
+        override val bouncingContent: ContentKey? = bouncingContent
+        override val orientation: Orientation = orientation
+        override val absoluteDistance = 0f
+
+        override fun freezeAndAnimateToCurrentState() {
+            if (onFreezeAndAnimate != null) {
+                onFreezeAndAnimate(this)
+            } else {
+                finish()
+            }
+        }
+
+        override fun interruptionProgress(layoutImpl: SceneTransitionLayoutImpl): Float {
+            return interruptionProgress()
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestTransition.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestSceneTransition.kt
index a6a83eedb2ac..d24b895c3050 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestTransition.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestSceneTransition.kt
@@ -24,8 +24,8 @@ import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.Transition
 import kotlinx.coroutines.CompletableDeferred
 
-/** A transition for tests that will be finished once [finish] is called. */
-abstract class TestTransition(
+/** A [Transition.ChangeScene] for tests that will be finished once [finish] is called. */
+abstract class TestSceneTransition(
     fromScene: SceneKey,
     toScene: SceneKey,
     replacedTransition: Transition?,
@@ -42,7 +42,7 @@ abstract class TestTransition(
     }
 }
 
-/** A utility to easily create a [TestTransition] in tests. */
+/** A utility to easily create a [TestSceneTransition] in tests. */
 fun transition(
     from: SceneKey,
     to: SceneKey,
@@ -58,11 +58,11 @@ fun transition(
     isUpOrLeft: Boolean = false,
     bouncingContent: ContentKey? = null,
     orientation: Orientation = Orientation.Horizontal,
-    onFreezeAndAnimate: ((TestTransition) -> Unit)? = null,
+    onFreezeAndAnimate: ((TestSceneTransition) -> Unit)? = null,
     replacedTransition: Transition? = null,
-): TestTransition {
+): TestSceneTransition {
     return object :
-        TestTransition(from, to, replacedTransition), TransitionState.HasOverscrollProperties {
+        TestSceneTransition(from, to, replacedTransition), TransitionState.HasOverscrollProperties {
         override val currentScene: SceneKey
             get() = current()