Introduce sharedElement.elevateInContent (1/2)

This CL introduces a new elevateInContent parameter to draw an element
above all other composables in its content (scene/overlay). It
allows to prevent elements from being clipped by a parent (like a
scrollable list) while still being drawn in the same content, therefore
keeping a relatively similar zIndex in the whole SceneTransitionLayout
compared to other scenes and overlays.

The first version of this CL had `elevateInContent` be a simple
`Boolean`. However, doing so would require us to instrument all scenes
and overlays to always use Modifier.container(), and all elements to use
Modifier.drawInContainer(), which are most of the time not necessary.
Making elevateInContent by a `ContentKey?` allows to only compose these
modifiers when necessary.

Note that using this parameter can currently lead to some strange issues
where text is not drawn (see b/374257277). I expect this to be fixed in
the Compose libraries directly in the future.

Bug: 373799480
Test: atest ElevateInContentScreenshotTest
Flag: com.android.systemui.scene_container
Change-Id: Ifa9e65ade0bc7bab01c80c0eb77c5424db13047f
diff --git a/packages/SystemUI/compose/scene/Android.bp b/packages/SystemUI/compose/scene/Android.bp
index af1172b..682c49cfd 100644
--- a/packages/SystemUI/compose/scene/Android.bp
+++ b/packages/SystemUI/compose/scene/Android.bp
@@ -40,6 +40,8 @@
     static_libs: [
         "androidx.compose.runtime_runtime",
         "androidx.compose.material3_material3",
+
+        "PlatformComposeCore",
     ],
 
     kotlincflags: ["-Xjvm-default=all"],
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 ebe1df4..f14622f 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
@@ -48,10 +48,13 @@
 import androidx.compose.ui.util.fastCoerceIn
 import androidx.compose.ui.util.fastForEachReversed
 import androidx.compose.ui.util.lerp
+import com.android.compose.animation.scene.Element.State
 import com.android.compose.animation.scene.content.Content
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.transformation.PropertyTransformation
 import com.android.compose.animation.scene.transformation.SharedElementTransformation
+import com.android.compose.modifiers.thenIf
+import com.android.compose.ui.graphics.drawInContainer
 import com.android.compose.ui.util.lerp
 import kotlin.math.roundToInt
 import kotlinx.coroutines.launch
@@ -146,10 +149,58 @@
     // TODO(b/341072461): Revert this and read the current transitions in ElementNode directly once
     // we can ensure that SceneTransitionLayoutImpl will compose new contents first.
     val currentTransitionStates = layoutImpl.state.transitionStates
-    return then(ElementModifier(layoutImpl, currentTransitionStates, content, key))
+    return thenIf(layoutImpl.state.isElevationPossible(content.key, key)) {
+            Modifier.maybeElevateInContent(layoutImpl, content, key, currentTransitionStates)
+        }
+        .then(ElementModifier(layoutImpl, currentTransitionStates, content, key))
         .testTag(key.testTag)
 }
 
+private fun Modifier.maybeElevateInContent(
+    layoutImpl: SceneTransitionLayoutImpl,
+    content: Content,
+    key: ElementKey,
+    transitionStates: List<TransitionState>,
+): Modifier {
+    fun isSharedElement(
+        stateByContent: Map<ContentKey, State>,
+        transition: TransitionState.Transition,
+    ): Boolean {
+        fun inFromContent() = transition.fromContent in stateByContent
+        fun inToContent() = transition.toContent in stateByContent
+        fun inCurrentScene() = transition.currentScene in stateByContent
+
+        return if (transition is TransitionState.Transition.ReplaceOverlay) {
+            (inFromContent() && (inToContent() || inCurrentScene())) ||
+                (inToContent() && inCurrentScene())
+        } else {
+            inFromContent() && inToContent()
+        }
+    }
+
+    return drawInContainer(
+        content.containerState,
+        enabled = {
+            val stateByContent = layoutImpl.elements.getValue(key).stateByContent
+            val state = elementState(transitionStates, isInContent = { it in stateByContent })
+
+            state is TransitionState.Transition &&
+                state.transformationSpec
+                    .transformations(key, content.key)
+                    .shared
+                    ?.elevateInContent == content.key &&
+                isSharedElement(stateByContent, state) &&
+                isSharedElementEnabled(key, state) &&
+                shouldPlaceElement(
+                    layoutImpl,
+                    content.key,
+                    layoutImpl.elements.getValue(key),
+                    state,
+                )
+        },
+    )
+}
+
 /**
  * An element associated to [ElementNode]. Note that this element does not support updates as its
  * arguments should always be the same.
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
index a9a8668..e1e2411 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
@@ -19,13 +19,16 @@
 import android.util.Log
 import androidx.annotation.VisibleForTesting
 import androidx.compose.runtime.Stable
+import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.util.fastAll
+import androidx.compose.ui.util.fastAny
 import androidx.compose.ui.util.fastFilter
 import androidx.compose.ui.util.fastForEach
 import com.android.compose.animation.scene.content.state.TransitionState
+import com.android.compose.animation.scene.transformation.SharedElementTransformation
 import com.android.compose.animation.scene.transition.link.LinkedTransition
 import com.android.compose.animation.scene.transition.link.StateLink
 import kotlin.math.absoluteValue
@@ -271,6 +274,14 @@
         mutableStateOf(listOf(TransitionState.Idle(initialScene, initialOverlays)))
         private set
 
+    /**
+     * The flattened list of [SharedElementTransformation] within all the transitions in
+     * [transitionStates].
+     */
+    private val transformationsWithElevation: List<SharedElementTransformation> by derivedStateOf {
+        transformationsWithElevation(transitionStates)
+    }
+
     override val currentScene: SceneKey
         get() = transitionState.currentScene
 
@@ -743,6 +754,42 @@
 
         animate()
     }
+
+    private fun transformationsWithElevation(
+        transitionStates: List<TransitionState>
+    ): List<SharedElementTransformation> {
+        return buildList {
+            transitionStates.fastForEach { state ->
+                if (state !is TransitionState.Transition) {
+                    return@fastForEach
+                }
+
+                state.transformationSpec.transformations.fastForEach { transformation ->
+                    if (
+                        transformation is SharedElementTransformation &&
+                            transformation.elevateInContent != null
+                    ) {
+                        add(transformation)
+                    }
+                }
+            }
+        }
+    }
+
+    /**
+     * Return whether we might need to elevate [element] (or any element if [element] is `null`) in
+     * [content].
+     *
+     * This is used to compose `Modifier.container()` and `Modifier.drawInContainer()` only when
+     * necessary, for performance.
+     */
+    internal fun isElevationPossible(content: ContentKey, element: ElementKey?): Boolean {
+        if (transformationsWithElevation.isEmpty()) return false
+        return transformationsWithElevation.fastAny { transformation ->
+            transformation.elevateInContent == content &&
+                (element == null || transformation.matcher.matches(element, content))
+        }
+    }
 }
 
 private const val TAG = "SceneTransitionLayoutState"
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 e825c6e..dc26b6b 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
@@ -204,8 +204,17 @@
      *
      * @param enabled whether the matched element(s) should actually be shared in this transition.
      *   Defaults to true.
+     * @param elevateInContent the content in which we should elevate the element when it is shared,
+     *   drawing above all other composables of that content. If `null` (the default), we will
+     *   simply draw this element in its original location. If not `null`, it has to be either the
+     *   [fromContent][TransitionState.Transition.fromContent] or
+     *   [toContent][TransitionState.Transition.toContent] of the transition.
      */
-    fun sharedElement(matcher: ElementMatcher, enabled: Boolean = true)
+    fun sharedElement(
+        matcher: ElementMatcher,
+        enabled: Boolean = true,
+        elevateInContent: ContentKey? = null,
+    )
 
     /**
      * Adds the transformations in [builder] but in reversed order. This allows you to partially
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 a5ad999..269d91b0 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
@@ -249,8 +249,22 @@
         reversed = false
     }
 
-    override fun sharedElement(matcher: ElementMatcher, enabled: Boolean) {
-        transformations.add(SharedElementTransformation(matcher, enabled))
+    override fun sharedElement(
+        matcher: ElementMatcher,
+        enabled: Boolean,
+        elevateInContent: ContentKey?,
+    ) {
+        check(
+            elevateInContent == null ||
+                elevateInContent == transition.fromContent ||
+                elevateInContent == transition.toContent
+        ) {
+            "elevateInContent (${elevateInContent?.debugName}) should be either fromContent " +
+                "(${transition.fromContent.debugName}) or toContent " +
+                "(${transition.toContent.debugName})"
+        }
+
+        transformations.add(SharedElementTransformation(matcher, enabled, elevateInContent))
     }
 
     override fun timestampRange(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Content.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Content.kt
index c8407b1..8187e39 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Content.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Content.kt
@@ -51,6 +51,9 @@
 import com.android.compose.animation.scene.element
 import com.android.compose.animation.scene.modifiers.noResizeDuringTransitions
 import com.android.compose.animation.scene.nestedScrollToScene
+import com.android.compose.modifiers.thenIf
+import com.android.compose.ui.graphics.ContainerState
+import com.android.compose.ui.graphics.container
 
 /** A content defined in a [SceneTransitionLayout], i.e. a scene or an overlay. */
 @Stable
@@ -62,6 +65,7 @@
     zIndex: Float,
 ) {
     internal val scope = ContentScopeImpl(layoutImpl, content = this)
+    val containerState = ContainerState()
 
     var content by mutableStateOf(content)
     var zIndex by mutableFloatStateOf(zIndex)
@@ -82,6 +86,9 @@
                     val placeable = measurable.measure(constraints)
                     layout(placeable.width, placeable.height) { placeable.place(0, 0) }
                 }
+                .thenIf(layoutImpl.state.isElevationPossible(content = key, element = null)) {
+                    Modifier.container(containerState)
+                }
                 .testTag(key.testTag)
         ) {
             scope.content()
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt
index 9bb3023..de7f418 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Transformation.kt
@@ -52,6 +52,7 @@
 internal class SharedElementTransformation(
     override val matcher: ElementMatcher,
     internal val enabled: Boolean,
+    internal val elevateInContent: ContentKey?,
 ) : Transformation
 
 /** A transformation that changes the value of an element property, like its size or offset. */