Allow anchorSize to anchor only the width or height

Bug: 308961608
Test: AnchorSizeTest
Flag: N/A
Change-Id: I43f0f9e5344edbbf8d1730e215a074d4fc0b4a03

4 files changed, 67 insertions, 6 deletions

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 f820074ec3d1..dfa2a9a18e91 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 @@ interface PropertyTransformationBuilder {
     )
 
     /**
-     * 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 8c0a5a394331..8f4a36e47212 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 @@ internal class TransitionBuilderImpl : TransitionBuilder {
         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 95385d51cb25..40c814e0f25c 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 @@ import com.android.compose.animation.scene.TransitionState
 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 @@ internal class AnchoredSize(
         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/transformation/AnchoredSizeTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredSizeTest.kt
index 8ef6757d33bd..e555a01d42fd 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 @@ class AnchoredSizeTest {
             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) }
+        }
+    }
 }