Make it possible to disable shared animations (1/2)

This CL makes it possible to disable shared element animations, that are
enabled by default when an element is shared between two scenes. This
can for instance be used to disable the sharing of notifications when
going from the Shade to Lockscreen, or disabling the sharing of media
player when going from lockscreen to shade on the phone while keeping it
on tablets without having to use different keys.

See b/300867076#comment3 for videos.

Bug: 300867076
Test: atest SharedElementTest
Change-Id: I30384e064eb66f81f39ba0185f35c375f2741590

12 files changed, 366 insertions, 42 deletions

diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Element.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Element.kt
index 0cc259ab7015..a62c9840add1 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Element.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Element.kt
@@ -151,7 +151,7 @@ internal fun Modifier.element(
                 element.lastAlpha = alpha
             }
         }
-        .testTag(key.name)
+        .testTag(key.testTag)
 }
 
 private fun shouldDrawElement(
@@ -167,7 +167,8 @@ private fun shouldDrawElement(
             state.fromScene == state.toScene ||
             !layoutImpl.isTransitionReady(state) ||
             state.fromScene !in element.sceneValues ||
-            state.toScene !in element.sceneValues
+            state.toScene !in element.sceneValues ||
+            !isSharedElementEnabled(layoutImpl, state, element.key)
     ) {
         return true
     }
@@ -191,6 +192,26 @@ private fun shouldDrawElement(
     }
 }
 
+private fun isSharedElementEnabled(
+    layoutImpl: SceneTransitionLayoutImpl,
+    transition: TransitionState.Transition,
+    element: ElementKey,
+): Boolean {
+    val spec = layoutImpl.transitions.transitionSpec(transition.fromScene, transition.toScene)
+    val sharedInFromScene = spec.transformations(element, transition.fromScene).shared
+    val sharedInToScene = spec.transformations(element, transition.toScene).shared
+
+    // The sharedElement() transformation must either be null or be the same in both scenes.
+    if (sharedInFromScene != sharedInToScene) {
+        error(
+            "Different sharedElement() transformations matched $element (from=$sharedInFromScene " +
+                "to=$sharedInToScene)"
+        )
+    }
+
+    return sharedInFromScene?.enabled ?: true
+}
+
 /**
  * Chain the [com.android.compose.animation.scene.transformation.ModifierTransformation] applied
  * throughout the current transition, if any.
@@ -213,7 +234,7 @@ private fun Modifier.modifierTransformations(
 
             return layoutImpl.transitions
                 .transitionSpec(fromScene, state.toScene)
-                .transformations(element.key)
+                .transformations(element.key, scene.key)
                 .modifier
                 .fold(this) { modifier, transformation ->
                     with(transformation) {
@@ -407,17 +428,20 @@ private inline fun <T> computeValue(
     // The element is shared: interpolate between the value in fromScene and the value in toScene.
     // TODO(b/290184746): Support non linear shared paths as well as a way to make sure that shared
     // elements follow the finger direction.
-    if (fromValues != null && toValues != null) {
+    val isSharedElement = fromValues != null && toValues != null
+    if (isSharedElement && isSharedElementEnabled(layoutImpl, state, element.key)) {
         return lerp(
-            sceneValue(fromValues),
-            sceneValue(toValues),
+            sceneValue(fromValues!!),
+            sceneValue(toValues!!),
             transitionProgress,
         )
     }
 
     val transformation =
         transformation(
-            layoutImpl.transitions.transitionSpec(fromScene, toScene).transformations(element.key)
+            layoutImpl.transitions
+                .transitionSpec(fromScene, toScene)
+                .transformations(element.key, scene.key)
         )
         // If there is no transformation explicitly associated to this element value, let's use
         // the value given by the system (like the current position and size given by the layout
@@ -426,12 +450,21 @@ private inline fun <T> computeValue(
 
     // Get the transformed value, i.e. the target value at the beginning (for entering elements) or
     // end (for leaving elements) of the transition.
+    val sceneValues =
+        checkNotNull(
+            when {
+                isSharedElement && scene.key == fromScene -> fromValues
+                isSharedElement -> toValues
+                else -> fromValues ?: toValues
+            }
+        )
+
     val targetValue =
         transformation.transform(
             layoutImpl,
             scene,
             element,
-            fromValues ?: toValues!!,
+            sceneValues,
             state,
             idleValue,
         )
@@ -440,7 +473,7 @@ private inline fun <T> computeValue(
     val rangeProgress = transformation.range?.progress(transitionProgress) ?: transitionProgress
 
     // Interpolate between the value at rest and the value before entering/after leaving.
-    val isEntering = fromValues == null
+    val isEntering = scene.key == toScene
     return if (isEntering) {
         lerp(targetValue, idleValue, rangeProgress)
     } else {
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/ElementMatcher.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/ElementMatcher.kt
new file mode 100644
index 000000000000..98dbb67d7c66
--- /dev/null
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/ElementMatcher.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2023 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
+
+/** An interface to match one or more elements. */
+interface ElementMatcher {
+    /** Whether the element with key [key] in scene [scene] matches this matcher. */
+    fun matches(key: ElementKey, scene: SceneKey): Boolean
+}
+
+/**
+ * Returns an [ElementMatcher] that matches elements in [scene] also matching [this]
+ * [ElementMatcher].
+ */
+fun ElementMatcher.inScene(scene: SceneKey): ElementMatcher {
+    val delegate = this
+    val matcherScene = scene
+    return object : ElementMatcher {
+        override fun matches(key: ElementKey, scene: SceneKey): Boolean {
+            return scene == matcherScene && delegate.matches(key, scene)
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Key.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Key.kt
index f7ebe2fc6d34..b7acc48e2865 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Key.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Key.kt
@@ -16,6 +16,8 @@
 
 package com.android.compose.animation.scene
 
+import androidx.annotation.VisibleForTesting
+
 /**
  * A base class to create unique keys, associated to an [identity] that is used to check the
  * equality of two key instances.
@@ -41,6 +43,7 @@ class SceneKey(
     name: String,
     identity: Any = Object(),
 ) : Key(name, identity) {
+    @VisibleForTesting val testTag: String = "scene:$name"
 
     /** The unique [ElementKey] identifying this scene's root element. */
     val rootElementKey = ElementKey(name, identity)
@@ -61,7 +64,9 @@ class ElementKey(
      */
     val isBackground: Boolean = false,
 ) : Key(name, identity), ElementMatcher {
-    override fun matches(key: ElementKey): Boolean {
+    @VisibleForTesting val testTag: String = "element:$name"
+
+    override fun matches(key: ElementKey, scene: SceneKey): Boolean {
         return key == this
     }
 
@@ -73,7 +78,9 @@ class ElementKey(
         /** Matches any element whose [key identity][ElementKey.identity] matches [predicate]. */
         fun withIdentity(predicate: (Any) -> Boolean): ElementMatcher {
             return object : ElementMatcher {
-                override fun matches(key: ElementKey): Boolean = predicate(key.identity)
+                override fun matches(key: ElementKey, scene: SceneKey): Boolean {
+                    return predicate(key.identity)
+                }
             }
         }
     }
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Scene.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Scene.kt
index b44c8efc7ee2..3985233bd197 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Scene.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/Scene.kt
@@ -25,6 +25,7 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.layout.onPlaced
+import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.zIndex
 
@@ -45,7 +46,9 @@ internal class Scene(
 
     @Composable
     fun Content(modifier: Modifier = Modifier) {
-        Box(modifier.zIndex(zIndex).onPlaced { size = it.size }) { scope.content() }
+        Box(modifier.zIndex(zIndex).onPlaced { size = it.size }.testTag(key.testTag)) {
+            scope.content()
+        }
     }
 
     override fun toString(): String {
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/SceneTransitions.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/SceneTransitions.kt
index f567e5c3be07..75dcb2e44c13 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/SceneTransitions.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/SceneTransitions.kt
@@ -29,6 +29,7 @@ import com.android.compose.animation.scene.transformation.ModifierTransformation
 import com.android.compose.animation.scene.transformation.PropertyTransformation
 import com.android.compose.animation.scene.transformation.RangedPropertyTransformation
 import com.android.compose.animation.scene.transformation.ScaleSize
+import com.android.compose.animation.scene.transformation.SharedElementTransformation
 import com.android.compose.animation.scene.transformation.Transformation
 import com.android.compose.animation.scene.transformation.Translate
 import com.android.compose.ui.util.fastForEach
@@ -99,7 +100,8 @@ data class TransitionSpec(
     val transformations: List<Transformation>,
     val spec: AnimationSpec<Float>,
 ) {
-    private val cache = mutableMapOf<ElementKey, ElementTransformations>()
+    // TODO(b/302300957): Make sure this cache does not infinitely grow.
+    private val cache = mutableMapOf<ElementKey, MutableMap<SceneKey, ElementTransformations>>()
 
     internal fun reverse(): TransitionSpec {
         return copy(
@@ -109,12 +111,18 @@ data class TransitionSpec(
         )
     }
 
-    internal fun transformations(element: ElementKey): ElementTransformations {
-        return cache.getOrPut(element) { computeTransformations(element) }
+    internal fun transformations(element: ElementKey, scene: SceneKey): ElementTransformations {
+        return cache
+            .getOrPut(element) { mutableMapOf() }
+            .getOrPut(scene) { computeTransformations(element, scene) }
     }
 
     /** Filter [transformations] to compute the [ElementTransformations] of [element]. */
-    private fun computeTransformations(element: ElementKey): ElementTransformations {
+    private fun computeTransformations(
+        element: ElementKey,
+        scene: SceneKey,
+    ): ElementTransformations {
+        var shared: SharedElementTransformation? = null
         val modifier = mutableListOf<ModifierTransformation>()
         var offset: PropertyTransformation<Offset>? = null
         var size: PropertyTransformation<IntSize>? = null
@@ -128,16 +136,16 @@ data class TransitionSpec(
                 is Translate,
                 is EdgeTranslate,
                 is AnchoredTranslate -> {
-                    throwIfNotNull(offset, element, property = "offset")
+                    throwIfNotNull(offset, element, name = "offset")
                     offset = root as PropertyTransformation<Offset>
                 }
                 is ScaleSize,
                 is AnchoredSize -> {
-                    throwIfNotNull(size, element, property = "size")
+                    throwIfNotNull(size, element, name = "size")
                     size = root as PropertyTransformation<IntSize>
                 }
                 is Fade -> {
-                    throwIfNotNull(alpha, element, property = "alpha")
+                    throwIfNotNull(alpha, element, name = "alpha")
                     alpha = root as PropertyTransformation<Float>
                 }
                 is RangedPropertyTransformation -> onPropertyTransformation(root, current.delegate)
@@ -145,32 +153,37 @@ data class TransitionSpec(
         }
 
         transformations.fastForEach { transformation ->
-            if (!transformation.matcher.matches(element)) {
+            if (!transformation.matcher.matches(element, scene)) {
                 return@fastForEach
             }
 
             when (transformation) {
+                is SharedElementTransformation -> {
+                    throwIfNotNull(shared, element, name = "shared")
+                    shared = transformation
+                }
                 is ModifierTransformation -> modifier.add(transformation)
                 is PropertyTransformation<*> -> onPropertyTransformation(transformation)
             }
         }
 
-        return ElementTransformations(modifier, offset, size, alpha)
+        return ElementTransformations(shared, modifier, offset, size, alpha)
     }
 
     private fun throwIfNotNull(
-        previous: PropertyTransformation<*>?,
+        previous: Transformation?,
         element: ElementKey,
-        property: String,
+        name: String,
     ) {
         if (previous != null) {
-            error("$element has multiple transformations for its $property property")
+            error("$element has multiple $name transformations")
         }
     }
 }
 
 /** The transformations of an element during a transition. */
 internal class ElementTransformations(
+    val shared: SharedElementTransformation?,
     val modifier: List<ModifierTransformation>,
     val offset: PropertyTransformation<Offset>?,
     val size: PropertyTransformation<IntSize>?,
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDsl.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDsl.kt
index b465c57949ce..49669775fedd 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDsl.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDsl.kt
@@ -116,6 +116,14 @@ interface TransitionBuilder : PropertyTransformationBuilder {
     )
 
     /**
+     * Configure the shared transition when [matcher] is shared between two scenes.
+     *
+     * @param enabled whether the matched element(s) should actually be shared in this transition.
+     *   Defaults to true.
+     */
+    fun sharedElement(matcher: ElementMatcher, enabled: Boolean = true)
+
+    /**
      * Punch a hole in the element(s) matching [matcher] that has the same bounds as [bounds] and
      * using the given [shape].
      *
@@ -186,12 +194,6 @@ interface PropertyTransformationBuilder {
     fun anchoredSize(matcher: ElementMatcher, anchor: ElementKey)
 }
 
-/** An interface to match one or more elements. */
-interface ElementMatcher {
-    /** Whether the element with key [key] matches this matcher. */
-    fun matches(key: ElementKey): Boolean
-}
-
 /** The edge of a [SceneTransitionLayout]. */
 enum class Edge {
     Left,
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDslImpl.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDslImpl.kt
index 9f08cbaf873b..f1c27178391c 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDslImpl.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/TransitionDslImpl.kt
@@ -31,6 +31,7 @@ import com.android.compose.animation.scene.transformation.PropertyTransformation
 import com.android.compose.animation.scene.transformation.PunchHole
 import com.android.compose.animation.scene.transformation.RangedPropertyTransformation
 import com.android.compose.animation.scene.transformation.ScaleSize
+import com.android.compose.animation.scene.transformation.SharedElementTransformation
 import com.android.compose.animation.scene.transformation.Transformation
 import com.android.compose.animation.scene.transformation.TransformationRange
 import com.android.compose.animation.scene.transformation.Translate
@@ -110,6 +111,10 @@ internal class TransitionBuilderImpl : TransitionBuilder {
         range = null
     }
 
+    override fun sharedElement(matcher: ElementMatcher, enabled: Boolean) {
+        transformations.add(SharedElementTransformation(matcher, enabled))
+    }
+
     override fun timestampRange(
         startMillis: Int?,
         endMillis: Int?,
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/transformation/Transformation.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/transformation/Transformation.kt
index ffb9ee96898b..a65025423aee 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/transformation/Transformation.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/scene/transformation/Transformation.kt
@@ -45,6 +45,11 @@ sealed interface Transformation {
     fun reverse(): Transformation = this
 }
 
+internal class SharedElementTransformation(
+    override val matcher: ElementMatcher,
+    internal val enabled: Boolean,
+) : Transformation
+
 /** A transformation that is applied on the element during the whole transition. */
 internal interface ModifierTransformation : Transformation {
     /** Apply the transformation to [element]. */
diff --git a/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt b/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
index 8bd654585f29..328866ea76ca 100644
--- a/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
+++ b/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
@@ -224,7 +224,7 @@ class SceneTransitionLayoutTest {
 
         // In scene A, the shared element SharedFoo() is at the top end of the layout and has a size
         // of 50.dp.
-        var sharedFoo = rule.onNodeWithTag(TestElements.Foo.name, useUnmergedTree = true)
+        var sharedFoo = rule.onNodeWithTag(TestElements.Foo.testTag, useUnmergedTree = true)
         sharedFoo.assertWidthIsEqualTo(50.dp)
         sharedFoo.assertHeightIsEqualTo(50.dp)
         sharedFoo.assertPositionInRootIsEqualTo(
@@ -250,7 +250,7 @@ class SceneTransitionLayoutTest {
 
         // We need to use onAllNodesWithTag().onFirst() here given that shared elements are
         // composed and laid out in both scenes (but drawn only in one).
-        sharedFoo = rule.onAllNodesWithTag(TestElements.Foo.name).onFirst()
+        sharedFoo = rule.onAllNodesWithTag(TestElements.Foo.testTag).onFirst()
 
         // In scene B, foo is at the top start (x = 0, y = 0) of the layout and has a size of
         // 100.dp. We pause at the middle of the transition, so it should now be 75.dp given that we
@@ -284,7 +284,7 @@ class SceneTransitionLayoutTest {
         val expectedLeft = 0.dp
         val expectedSize = 100.dp + (150.dp - 100.dp) * interpolatedProgress
 
-        sharedFoo = rule.onAllNodesWithTag(TestElements.Foo.name).onFirst()
+        sharedFoo = rule.onAllNodesWithTag(TestElements.Foo.testTag).onFirst()
         assertThat((layoutState.transitionState as TransitionState.Transition).progress)
             .isEqualTo(interpolatedProgress)
         sharedFoo.assertWidthIsEqualTo(expectedSize)
diff --git a/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/TestTransition.kt b/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/TestTransition.kt
index 275149a05abf..268057fd2f2c 100644
--- a/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/TestTransition.kt
+++ b/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/TestTransition.kt
@@ -23,7 +23,11 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.test.SemanticsNodeInteraction
+import androidx.compose.ui.test.SemanticsNodeInteractionCollection
+import androidx.compose.ui.test.hasParent
+import androidx.compose.ui.test.hasTestTag
 import androidx.compose.ui.test.junit4.ComposeContentTestRule
+import androidx.compose.ui.test.onAllNodesWithTag
 import androidx.compose.ui.test.onNodeWithTag
 
 @DslMarker annotation class TransitionTestDsl
@@ -59,8 +63,21 @@ interface TransitionTestBuilder {
 
 @TransitionTestDsl
 interface TransitionTestAssertionScope {
-    /** Assert on [element]. */
-    fun onElement(element: ElementKey): SemanticsNodeInteraction
+    /**
+     * Assert on [element].
+     *
+     * Note that presence/value assertions on the returned [SemanticsNodeInteraction] will fail if 0
+     * or more than 1 elements matched [element]. If you need to assert on a shared element that
+     * will be present multiple times in the layout during transitions, either specify the [scene]
+     * in which you are matching or use [onSharedElement] instead.
+     */
+    fun onElement(element: ElementKey, scene: SceneKey? = null): SemanticsNodeInteraction
+
+    /**
+     * Assert on a shared [element]. This will throw if [element] is not shared and present only in
+     * one scene during a transition.
+     */
+    fun onSharedElement(element: ElementKey): SemanticsNodeInteractionCollection
 }
 
 /**
@@ -73,20 +90,22 @@ fun ComposeContentTestRule.testTransition(
     toSceneContent: @Composable SceneScope.() -> Unit,
     transition: TransitionBuilder.() -> Unit,
     layoutModifier: Modifier = Modifier,
+    fromScene: SceneKey = TestScenes.SceneA,
+    toScene: SceneKey = TestScenes.SceneB,
     builder: TransitionTestBuilder.() -> Unit,
 ) {
     testTransition(
-        from = TestScenes.SceneA,
-        to = TestScenes.SceneB,
+        from = fromScene,
+        to = toScene,
         transitionLayout = { currentScene, onChangeScene ->
             SceneTransitionLayout(
                 currentScene,
                 onChangeScene,
-                transitions { from(TestScenes.SceneA, to = TestScenes.SceneB, transition) },
+                transitions { from(fromScene, to = toScene, transition) },
                 layoutModifier.fillMaxSize(),
             ) {
-                scene(TestScenes.SceneA, content = fromSceneContent)
-                scene(TestScenes.SceneB, content = toSceneContent)
+                scene(fromScene, content = fromSceneContent)
+                scene(toScene, content = toSceneContent)
             }
         },
         builder,
@@ -111,8 +130,24 @@ fun ComposeContentTestRule.testTransition(
     val test = transitionTest(builder)
     val assertionScope =
         object : TransitionTestAssertionScope {
-            override fun onElement(element: ElementKey): SemanticsNodeInteraction {
-                return this@testTransition.onNodeWithTag(element.name)
+            override fun onElement(
+                element: ElementKey,
+                scene: SceneKey?
+            ): SemanticsNodeInteraction {
+                return if (scene == null) {
+                    onNodeWithTag(element.testTag)
+                } else {
+                    onNode(hasTestTag(element.testTag) and hasParent(hasTestTag(scene.testTag)))
+                }
+            }
+
+            override fun onSharedElement(element: ElementKey): SemanticsNodeInteractionCollection {
+                val interaction = onAllNodesWithTag(element.testTag)
+                val matches = interaction.fetchSemanticsNodes(atLeastOneRootRequired = false).size
+                if (matches < 2) {
+                    error("Element $element is not shared ($matches matches)")
+                }
+                return interaction
             }
         }
 
diff --git a/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/transformation/SharedElementTest.kt b/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/transformation/SharedElementTest.kt
new file mode 100644
index 000000000000..2af363860272
--- /dev/null
+++ b/packages/SystemUI/compose/core/tests/src/com/android/compose/animation/scene/transformation/SharedElementTest.kt
@@ -0,0 +1,157 @@
+/*
+ * Copyright (C) 2023 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.transformation
+
+import androidx.compose.animation.core.LinearEasing
+import androidx.compose.animation.core.tween
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.offset
+import androidx.compose.foundation.layout.size
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.test.assertPositionInRootIsEqualTo
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.unit.dp
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.compose.animation.scene.Edge
+import com.android.compose.animation.scene.TestElements
+import com.android.compose.animation.scene.TestScenes
+import com.android.compose.animation.scene.inScene
+import com.android.compose.animation.scene.testTransition
+import com.android.compose.modifiers.size
+import com.android.compose.test.assertSizeIsEqualTo
+import com.android.compose.test.onEach
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class SharedElementTest {
+    @get:Rule val rule = createComposeRule()
+
+    @Test
+    fun testSharedElement() {
+        rule.testTransition(
+            fromSceneContent = {
+                // Foo is at (10, 50) with a size of (20, 80).
+                Box(Modifier.offset(10.dp, 50.dp).element(TestElements.Foo).size(20.dp, 80.dp))
+            },
+            toSceneContent = {
+                // Foo is at (50, 70) with a size of (10, 40).
+                Box(Modifier.offset(50.dp, 70.dp).element(TestElements.Foo).size(10.dp, 40.dp))
+            },
+            transition = {
+                spec = tween(16 * 4, easing = LinearEasing)
+                // Elements should be shared by default.
+            }
+        ) {
+            before {
+                onElement(TestElements.Foo).assertPositionInRootIsEqualTo(10.dp, 50.dp)
+                onElement(TestElements.Foo).assertSizeIsEqualTo(20.dp, 80.dp)
+            }
+            at(0) {
+                onSharedElement(TestElements.Foo).onEach {
+                    assertPositionInRootIsEqualTo(10.dp, 50.dp)
+                    assertSizeIsEqualTo(20.dp, 80.dp)
+                }
+            }
+            at(16) {
+                onSharedElement(TestElements.Foo).onEach {
+                    assertPositionInRootIsEqualTo(20.dp, 55.dp)
+                    assertSizeIsEqualTo(17.5.dp, 70.dp)
+                }
+            }
+            at(32) {
+                onSharedElement(TestElements.Foo).onEach {
+                    assertPositionInRootIsEqualTo(30.dp, 60.dp)
+                    assertSizeIsEqualTo(15.dp, 60.dp)
+                }
+            }
+            at(48) {
+                onSharedElement(TestElements.Foo).onEach {
+                    assertPositionInRootIsEqualTo(40.dp, 65.dp)
+                    assertSizeIsEqualTo(12.5.dp, 50.dp)
+                }
+            }
+            after {
+                onElement(TestElements.Foo).assertPositionInRootIsEqualTo(50.dp, 70.dp)
+                onElement(TestElements.Foo).assertSizeIsEqualTo(10.dp, 40.dp)
+            }
+        }
+    }
+
+    @Test
+    fun testSharedElementDisabled() {
+        rule.testTransition(
+            fromScene = TestScenes.SceneA,
+            toScene = TestScenes.SceneB,
+            // The full layout is 100x100.
+            layoutModifier = Modifier.size(100.dp),
+            fromSceneContent = {
+                Box(Modifier.fillMaxSize()) {
+                    // Foo is at (10, 50).
+                    Box(Modifier.offset(10.dp, 50.dp).element(TestElements.Foo))
+                }
+            },
+            toSceneContent = {
+                Box(Modifier.fillMaxSize()) {
+                    // Foo is at (50, 60).
+                    Box(Modifier.offset(50.dp, 60.dp).element(TestElements.Foo))
+                }
+            },
+            transition = {
+                spec = tween(16 * 4, easing = LinearEasing)
+
+                // Disable the shared element animation.
+                sharedElement(TestElements.Foo, enabled = false)
+
+                // In SceneA, Foo leaves to the left edge.
+                translate(TestElements.Foo.inScene(TestScenes.SceneA), Edge.Left)
+
+                // In SceneB, Foo comes from the bottom edge.
+                translate(TestElements.Foo.inScene(TestScenes.SceneB), Edge.Bottom)
+            },
+        ) {
+            before { onElement(TestElements.Foo).assertPositionInRootIsEqualTo(10.dp, 50.dp) }
+            at(0) {
+                onElement(TestElements.Foo, scene = TestScenes.SceneA)
+                    .assertPositionInRootIsEqualTo(10.dp, 50.dp)
+                onElement(TestElements.Foo, scene = TestScenes.SceneB)
+                    .assertPositionInRootIsEqualTo(50.dp, 100.dp)
+            }
+            at(16) {
+                onElement(TestElements.Foo, scene = TestScenes.SceneA)
+                    .assertPositionInRootIsEqualTo(7.5.dp, 50.dp)
+                onElement(TestElements.Foo, scene = TestScenes.SceneB)
+                    .assertPositionInRootIsEqualTo(50.dp, 90.dp)
+            }
+            at(32) {
+                onElement(TestElements.Foo, scene = TestScenes.SceneA)
+                    .assertPositionInRootIsEqualTo(5.dp, 50.dp)
+                onElement(TestElements.Foo, scene = TestScenes.SceneB)
+                    .assertPositionInRootIsEqualTo(50.dp, 80.dp)
+            }
+            at(48) {
+                onElement(TestElements.Foo, scene = TestScenes.SceneA)
+                    .assertPositionInRootIsEqualTo(2.5.dp, 50.dp)
+                onElement(TestElements.Foo, scene = TestScenes.SceneB)
+                    .assertPositionInRootIsEqualTo(50.dp, 70.dp)
+            }
+            after { onElement(TestElements.Foo).assertPositionInRootIsEqualTo(50.dp, 60.dp) }
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/core/tests/src/com/android/compose/test/Selectors.kt b/packages/SystemUI/compose/core/tests/src/com/android/compose/test/Selectors.kt
new file mode 100644
index 000000000000..d6f64bfe4974
--- /dev/null
+++ b/packages/SystemUI/compose/core/tests/src/com/android/compose/test/Selectors.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2023 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.ui.test.SemanticsNodeInteraction
+import androidx.compose.ui.test.SemanticsNodeInteractionCollection
+
+/** Assert [assert] on each element of [this] [SemanticsNodeInteractionCollection]. */
+fun SemanticsNodeInteractionCollection.onEach(assert: SemanticsNodeInteraction.() -> Unit) {
+    for (i in 0 until this.fetchSemanticsNodes().size) {
+        get(i).assert()
+    }
+}