Merge changes I43f0f9e5,Iba860162,Id2678cf2 into main

* changes:
  Allow anchorSize to anchor only the width or height
  Fix ElementNode update/recycling
  Do not throw when element is present in neither from- or toScene
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 2b11952..31604a6 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
@@ -213,22 +213,11 @@
     override fun onDetach() {
         super.onDetach()
         removeNodeFromSceneValues()
+        maybePruneMaps(layoutImpl, element, sceneValues)
     }
 
     private fun removeNodeFromSceneValues() {
         sceneValues.nodes.remove(this)
-
-        // If element is not composed from this scene anymore, remove the scene values. This works
-        // because [onAttach] is called before [onDetach], so if an element is moved from the UI
-        // tree we will first add the new code location then remove the old one.
-        if (sceneValues.nodes.isEmpty()) {
-            element.sceneValues.remove(sceneValues.scene)
-        }
-
-        // If the element is not composed in any scene, remove it from the elements map.
-        if (element.sceneValues.isEmpty()) {
-            layoutImpl.elements.remove(element.key)
-        }
     }
 
     fun update(
@@ -237,12 +226,16 @@
         element: Element,
         sceneValues: Element.TargetValues,
     ) {
+        check(layoutImpl == this.layoutImpl && scene == this.scene)
         removeNodeFromSceneValues()
-        this.layoutImpl = layoutImpl
-        this.scene = scene
+
+        val prevElement = this.element
+        val prevSceneValues = this.sceneValues
         this.element = element
         this.sceneValues = sceneValues
+
         addNodeToSceneValues()
+        maybePruneMaps(layoutImpl, prevElement, prevSceneValues)
     }
 
     override fun ContentDrawScope.draw() {
@@ -261,6 +254,28 @@
             }
         }
     }
+
+    companion object {
+        private fun maybePruneMaps(
+            layoutImpl: SceneTransitionLayoutImpl,
+            element: Element,
+            sceneValues: Element.TargetValues,
+        ) {
+            // If element is not composed from this scene anymore, remove the scene values. This
+            // works because [onAttach] is called before [onDetach], so if an element is moved from
+            // the UI tree we will first add the new code location then remove the old one.
+            if (
+                sceneValues.nodes.isEmpty() && element.sceneValues[sceneValues.scene] == sceneValues
+            ) {
+                element.sceneValues.remove(sceneValues.scene)
+
+                // If the element is not composed in any scene, remove it from the elements map.
+                if (element.sceneValues.isEmpty() && layoutImpl.elements[element.key] == element) {
+                    layoutImpl.elements.remove(element.key)
+                }
+            }
+        }
+    }
 }
 
 private fun shouldDrawElement(
@@ -615,7 +630,9 @@
     val toValues = element.sceneValues[toScene]
 
     if (fromValues == null && toValues == null) {
-        error("This should not happen, element $element is neither in $fromScene or $toScene")
+        // TODO(b/311600838): Throw an exception instead once layers of disposed elements are not
+        // run anymore.
+        return lastValue()
     }
 
     // The element is shared: interpolate between the value in fromScene and the value in toScene.
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 f820074..dfa2a9a 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
@@ -226,12 +226,17 @@
     )
 
     /**
-     * Scale the element(s) matching [matcher] so that it grows/shrinks to the same size as [anchor]
-     * .
+     * Scale the element(s) matching [matcher] so that it grows/shrinks to the same size as
+     * [anchor].
      *
      * Note: This currently only works if [anchor] is a shared element of this transition.
      */
-    fun anchoredSize(matcher: ElementMatcher, anchor: ElementKey)
+    fun anchoredSize(
+        matcher: ElementMatcher,
+        anchor: ElementKey,
+        anchorWidth: Boolean = true,
+        anchorHeight: Boolean = true,
+    )
 }
 
 /** The edge of a [SceneTransitionLayout]. */
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 8c0a5a3..8f4a36e 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
@@ -178,7 +178,12 @@
         transformation(DrawScale(matcher, scaleX, scaleY, pivot))
     }
 
-    override fun anchoredSize(matcher: ElementMatcher, anchor: ElementKey) {
-        transformation(AnchoredSize(matcher, anchor))
+    override fun anchoredSize(
+        matcher: ElementMatcher,
+        anchor: ElementKey,
+        anchorWidth: Boolean,
+        anchorHeight: Boolean,
+    ) {
+        transformation(AnchoredSize(matcher, anchor, anchorWidth, anchorHeight))
     }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
index 95385d5..40c814e 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredSize.kt
@@ -29,6 +29,8 @@
 internal class AnchoredSize(
     override val matcher: ElementMatcher,
     private val anchor: ElementKey,
+    private val anchorWidth: Boolean,
+    private val anchorHeight: Boolean,
 ) : PropertyTransformation<IntSize> {
     override fun transform(
         layoutImpl: SceneTransitionLayoutImpl,
@@ -41,7 +43,10 @@
         fun anchorSizeIn(scene: SceneKey): IntSize {
             val size = layoutImpl.elements[anchor]?.sceneValues?.get(scene)?.targetSize
             return if (size != null && size != Element.SizeUnspecified) {
-                size
+                IntSize(
+                    width = if (anchorWidth) size.width else value.width,
+                    height = if (anchorHeight) size.height else value.height,
+                )
             } else {
                 value
             }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
index ce3e1db..439dc00 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
@@ -17,16 +17,21 @@
 package com.android.compose.animation.scene
 
 import androidx.compose.animation.core.tween
+import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.offset
 import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.pager.HorizontalPager
+import androidx.compose.foundation.pager.PagerState
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
@@ -36,6 +41,9 @@
 import androidx.compose.ui.unit.dp
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertThrows
 import org.junit.Rule
 import org.junit.Test
@@ -430,6 +438,97 @@
     }
 
     @Test
+    @OptIn(ExperimentalFoundationApi::class)
+    fun elementModifierNodeIsRecycledInLazyLayouts() = runTest {
+        val nPages = 2
+        val pagerState = PagerState(currentPage = 0) { nPages }
+        var nullableLayoutImpl: SceneTransitionLayoutImpl? = null
+
+        // This is how we scroll a pager inside a test, as explained in b/315457147#comment2.
+        lateinit var scrollScope: CoroutineScope
+        fun scrollToPage(page: Int) {
+            var animationFinished by mutableStateOf(false)
+            rule.runOnIdle {
+                scrollScope.launch {
+                    pagerState.scrollToPage(page)
+                    animationFinished = true
+                }
+            }
+            rule.waitUntil(timeoutMillis = 10_000) { animationFinished }
+        }
+
+        rule.setContent {
+            scrollScope = rememberCoroutineScope()
+
+            SceneTransitionLayoutForTesting(
+                currentScene = TestScenes.SceneA,
+                onChangeScene = {},
+                transitions = remember { transitions {} },
+                state = remember { SceneTransitionLayoutState(TestScenes.SceneA) },
+                edgeDetector = DefaultEdgeDetector,
+                modifier = Modifier,
+                transitionInterceptionThreshold = 0f,
+                onLayoutImpl = { nullableLayoutImpl = it },
+            ) {
+                scene(TestScenes.SceneA) {
+                    // The pages are full-size and beyondBoundsPageCount is 0, so at rest only one
+                    // page should be composed.
+                    HorizontalPager(
+                        pagerState,
+                        beyondBoundsPageCount = 0,
+                    ) { page ->
+                        when (page) {
+                            0 -> Box(Modifier.element(TestElements.Foo).fillMaxSize())
+                            1 -> Box(Modifier.fillMaxSize())
+                            else -> error("page $page < nPages $nPages")
+                        }
+                    }
+                }
+            }
+        }
+
+        assertThat(nullableLayoutImpl).isNotNull()
+        val layoutImpl = nullableLayoutImpl!!
+
+        // There is only Foo in the elements map.
+        assertThat(layoutImpl.elements.keys).containsExactly(TestElements.Foo)
+        val element = layoutImpl.elements.getValue(TestElements.Foo)
+        val sceneValues = element.sceneValues
+        assertThat(sceneValues.keys).containsExactly(TestScenes.SceneA)
+
+        // Get the ElementModifier node that should be reused later on when coming back to this
+        // page.
+        val nodes = sceneValues.getValue(TestScenes.SceneA).nodes
+        assertThat(nodes).hasSize(1)
+        val node = nodes.single()
+
+        // Go to the second page.
+        scrollToPage(1)
+        rule.waitForIdle()
+
+        assertThat(nodes).isEmpty()
+        assertThat(sceneValues).isEmpty()
+        assertThat(layoutImpl.elements).isEmpty()
+
+        // Go back to the first page.
+        scrollToPage(0)
+        rule.waitForIdle()
+
+        assertThat(layoutImpl.elements.keys).containsExactly(TestElements.Foo)
+        val newElement = layoutImpl.elements.getValue(TestElements.Foo)
+        val newSceneValues = newElement.sceneValues
+        assertThat(newElement).isNotEqualTo(element)
+        assertThat(newSceneValues).isNotEqualTo(sceneValues)
+        assertThat(newSceneValues.keys).containsExactly(TestScenes.SceneA)
+
+        // The ElementModifier node should be the same as before.
+        val newNodes = newSceneValues.getValue(TestScenes.SceneA).nodes
+        assertThat(newNodes).hasSize(1)
+        val newNode = newNodes.single()
+        assertThat(newNode).isSameInstanceAs(node)
+    }
+
+    @Test
     fun existingElementsDontRecomposeWhenTransitionStateChanges() {
         var fooCompositions = 0
 
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredSizeTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredSizeTest.kt
index 8ef6757..e555a01 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredSizeTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredSizeTest.kt
@@ -85,4 +85,50 @@
             after { onElement(TestElements.Bar).assertDoesNotExist() }
         }
     }
+
+    @Test
+    fun testAnchoredWidthOnly() {
+        rule.testTransition(
+            fromSceneContent = { Box(Modifier.size(100.dp, 100.dp).element(TestElements.Foo)) },
+            toSceneContent = {
+                Box(Modifier.size(50.dp, 50.dp).element(TestElements.Foo))
+                Box(Modifier.size(200.dp, 60.dp).element(TestElements.Bar))
+            },
+            transition = {
+                spec = tween(16 * 4, easing = LinearEasing)
+                anchoredSize(TestElements.Bar, TestElements.Foo, anchorHeight = false)
+            },
+        ) {
+            before { onElement(TestElements.Bar).assertDoesNotExist() }
+            at(0) { onElement(TestElements.Bar).assertSizeIsEqualTo(100.dp, 60.dp) }
+            at(16) { onElement(TestElements.Bar).assertSizeIsEqualTo(125.dp, 60.dp) }
+            at(32) { onElement(TestElements.Bar).assertSizeIsEqualTo(150.dp, 60.dp) }
+            at(48) { onElement(TestElements.Bar).assertSizeIsEqualTo(175.dp, 60.dp) }
+            at(64) { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 60.dp) }
+            after { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 60.dp) }
+        }
+    }
+
+    @Test
+    fun testAnchoredHeightOnly() {
+        rule.testTransition(
+            fromSceneContent = { Box(Modifier.size(100.dp, 100.dp).element(TestElements.Foo)) },
+            toSceneContent = {
+                Box(Modifier.size(50.dp, 50.dp).element(TestElements.Foo))
+                Box(Modifier.size(200.dp, 60.dp).element(TestElements.Bar))
+            },
+            transition = {
+                spec = tween(16 * 4, easing = LinearEasing)
+                anchoredSize(TestElements.Bar, TestElements.Foo, anchorWidth = false)
+            },
+        ) {
+            before { onElement(TestElements.Bar).assertDoesNotExist() }
+            at(0) { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 100.dp) }
+            at(16) { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 90.dp) }
+            at(32) { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 80.dp) }
+            at(48) { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 70.dp) }
+            at(64) { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 60.dp) }
+            after { onElement(TestElements.Bar).assertSizeIsEqualTo(200.dp, 60.dp) }
+        }
+    }
 }