Merge "Make PropertyTransformation public" into main
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 f14622f..63c5d7a 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,7 +48,6 @@
 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
@@ -163,7 +162,7 @@
     transitionStates: List<TransitionState>,
 ): Modifier {
     fun isSharedElement(
-        stateByContent: Map<ContentKey, State>,
+        stateByContent: Map<ContentKey, Element.State>,
         transition: TransitionState.Transition,
     ): Boolean {
         fun inFromContent() = transition.fromContent in stateByContent
@@ -1281,14 +1280,14 @@
                 checkNotNull(if (currentContent == toContent) toState else fromState)
             val idleValue = contentValue(overscrollState)
             val targetValue =
-                propertySpec.transform(
-                    layoutImpl,
-                    currentContent,
-                    element,
-                    overscrollState,
-                    transition,
-                    idleValue,
-                )
+                with(propertySpec) {
+                    layoutImpl.propertyTransformationScope.transform(
+                        currentContent,
+                        element.key,
+                        transition,
+                        idleValue,
+                    )
+                }
 
             // Make sure we don't read progress if values are the same and we don't need to
             // interpolate, so we don't invalidate the phase where this is read.
@@ -1376,24 +1375,26 @@
         val idleValue = contentValue(contentState)
         val isEntering = content == toContent
         val previewTargetValue =
-            previewTransformation.transform(
-                layoutImpl,
-                content,
-                element,
-                contentState,
-                transition,
-                idleValue,
-            )
+            with(previewTransformation) {
+                layoutImpl.propertyTransformationScope.transform(
+                    content,
+                    element.key,
+                    transition,
+                    idleValue,
+                )
+            }
 
         val targetValueOrNull =
-            transformation?.transform(
-                layoutImpl,
-                content,
-                element,
-                contentState,
-                transition,
-                idleValue,
-            )
+            transformation?.let { transformation ->
+                with(transformation) {
+                    layoutImpl.propertyTransformationScope.transform(
+                        content,
+                        element.key,
+                        transition,
+                        idleValue,
+                    )
+                }
+            }
 
         // Make sure we don't read progress if values are the same and we don't need to interpolate,
         // so we don't invalidate the phase where this is read.
@@ -1460,7 +1461,14 @@
 
     val idleValue = contentValue(contentState)
     val targetValue =
-        transformation.transform(layoutImpl, content, element, contentState, transition, idleValue)
+        with(transformation) {
+            layoutImpl.propertyTransformationScope.transform(
+                content,
+                element.key,
+                transition,
+                idleValue,
+            )
+        }
 
     // Make sure we don't read progress if values are the same and we don't need to interpolate, so
     // we don't invalidate the phase where this is read.
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
index d58d3f24..e93cf8f7 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
@@ -129,6 +129,7 @@
     private val verticalDraggableHandler: DraggableHandlerImpl
 
     internal val elementStateScope = ElementStateScopeImpl(this)
+    internal val propertyTransformationScope = PropertyTransformationScopeImpl(this)
     private var _userActionDistanceScope: UserActionDistanceScope? = null
     internal val userActionDistanceScope: UserActionDistanceScope
         get() =
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
index 690c809..8457481 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/UserActionDistanceScopeImpl.kt
@@ -18,6 +18,8 @@
 
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.unit.IntSize
+import androidx.compose.ui.unit.LayoutDirection
+import com.android.compose.animation.scene.transformation.PropertyTransformationScope
 
 internal class ElementStateScopeImpl(private val layoutImpl: SceneTransitionLayoutImpl) :
     ElementStateScope {
@@ -46,3 +48,15 @@
     override val fontScale: Float
         get() = layoutImpl.density.fontScale
 }
+
+internal class PropertyTransformationScopeImpl(private val layoutImpl: SceneTransitionLayoutImpl) :
+    PropertyTransformationScope, ElementStateScope by layoutImpl.elementStateScope {
+    override val density: Float
+        get() = layoutImpl.density.density
+
+    override val fontScale: Float
+        get() = layoutImpl.density.fontScale
+
+    override val layoutDirection: LayoutDirection
+        get() = layoutImpl.layoutDirection
+}
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 c5a3067c..5936d25 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
@@ -18,10 +18,8 @@
 
 import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.ContentKey
-import com.android.compose.animation.scene.Element
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
 
 /** Anchor the size of an element to the size of another element. */
@@ -31,19 +29,15 @@
     private val anchorWidth: Boolean,
     private val anchorHeight: Boolean,
 ) : PropertyTransformation<IntSize> {
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: IntSize,
     ): IntSize {
         fun anchorSizeIn(content: ContentKey): IntSize {
             val size =
-                layoutImpl.elements[anchor]?.stateByContent?.get(content)?.targetSize?.takeIf {
-                    it != Element.SizeUnspecified
-                }
+                anchor.targetSize(content)
                     ?: throwMissingAnchorException(
                         transformation = "AnchoredSize",
                         anchor = anchor,
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
index 86e06ab..0a59dfe 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/AnchoredTranslate.kt
@@ -17,12 +17,9 @@
 package com.android.compose.animation.scene.transformation
 
 import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.geometry.isSpecified
 import com.android.compose.animation.scene.ContentKey
-import com.android.compose.animation.scene.Element
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
 
 /** Anchor the translation of an element to another element. */
@@ -30,11 +27,9 @@
     override val matcher: ElementMatcher,
     private val anchor: ElementKey,
 ) : PropertyTransformation<Offset> {
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: Offset,
     ): Offset {
@@ -46,18 +41,13 @@
             )
         }
 
-        val anchor = layoutImpl.elements[anchor] ?: throwException(content = null)
-        fun anchorOffsetIn(content: ContentKey): Offset? {
-            return anchor.stateByContent[content]?.targetOffset?.takeIf { it.isSpecified }
-        }
-
         // [element] will move the same amount as [anchor] does.
         // TODO(b/290184746): Also support anchors that are not shared but translated because of
         // other transformations, like an edge translation.
         val anchorFromOffset =
-            anchorOffsetIn(transition.fromContent) ?: throwException(transition.fromContent)
+            anchor.targetOffset(transition.fromContent) ?: throwException(transition.fromContent)
         val anchorToOffset =
-            anchorOffsetIn(transition.toContent) ?: throwException(transition.toContent)
+            anchor.targetOffset(transition.toContent) ?: throwException(transition.toContent)
         val offset = anchorToOffset - anchorFromOffset
 
         return if (content == transition.toContent) {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt
index 7f86479..7223dad 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/DrawScale.kt
@@ -18,10 +18,9 @@
 
 import androidx.compose.ui.geometry.Offset
 import com.android.compose.animation.scene.ContentKey
-import com.android.compose.animation.scene.Element
+import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
 import com.android.compose.animation.scene.Scale
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
 
 /**
@@ -35,11 +34,9 @@
     private val pivot: Offset = Offset.Unspecified,
 ) : PropertyTransformation<Scale> {
 
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: Scale,
     ): Scale {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt
index 031f50e..4ae07c5 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/EdgeTranslate.kt
@@ -19,9 +19,8 @@
 import androidx.compose.ui.geometry.Offset
 import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.Edge
-import com.android.compose.animation.scene.Element
+import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
 
 /** Translate an element from an edge of the layout. */
@@ -30,21 +29,18 @@
     private val edge: Edge,
     private val startsOutsideLayoutBounds: Boolean = true,
 ) : PropertyTransformation<Offset> {
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: Offset,
     ): Offset {
-        val sceneSize = layoutImpl.content(content).targetSize
-        val elementSize = stateInContent.targetSize
-        if (elementSize == Element.SizeUnspecified) {
-            return value
-        }
+        val sceneSize =
+            content.targetSize()
+                ?: error("Content ${content.debugName} does not have a target size")
+        val elementSize = element.targetSize(content) ?: return value
 
-        return when (edge.resolve(layoutImpl.layoutDirection)) {
+        return when (edge.resolve(layoutDirection)) {
             Edge.Resolved.Top ->
                 if (startsOutsideLayoutBounds) {
                     Offset(value.x, -elementSize.height.toFloat())
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt
index 078aa0f..c11ec97 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Fade.kt
@@ -17,18 +17,15 @@
 package com.android.compose.animation.scene.transformation
 
 import com.android.compose.animation.scene.ContentKey
-import com.android.compose.animation.scene.Element
+import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
 
 /** Fade an element in or out. */
 internal class Fade(override val matcher: ElementMatcher) : PropertyTransformation<Float> {
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: Float,
     ): Float {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt
index 5f3fdaf..a159a5b 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/ScaleSize.kt
@@ -18,9 +18,8 @@
 
 import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.ContentKey
-import com.android.compose.animation.scene.Element
+import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
 import kotlin.math.roundToInt
 
@@ -33,11 +32,9 @@
     private val width: Float = 1f,
     private val height: Float = 1f,
 ) : PropertyTransformation<IntSize> {
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: IntSize,
     ): IntSize {
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 de7f418..d38067d 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
@@ -18,13 +18,15 @@
 
 import androidx.compose.animation.core.Easing
 import androidx.compose.animation.core.LinearEasing
+import androidx.compose.ui.unit.Density
+import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.util.fastCoerceAtLeast
 import androidx.compose.ui.util.fastCoerceAtMost
 import androidx.compose.ui.util.fastCoerceIn
 import com.android.compose.animation.scene.ContentKey
-import com.android.compose.animation.scene.Element
+import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
+import com.android.compose.animation.scene.ElementStateScope
 import com.android.compose.animation.scene.content.state.TransitionState
 
 /** A transformation applied to one or more elements during a transition. */
@@ -56,22 +58,30 @@
 ) : Transformation
 
 /** A transformation that changes the value of an element property, like its size or offset. */
-internal sealed interface PropertyTransformation<T> : Transformation {
+interface PropertyTransformation<T> : Transformation {
     /**
-     * Transform [value], i.e. the value of the transformed property without this transformation.
+     * Return the transformed value for the given property, i.e.:
+     * - the value at progress = 0% for elements that are entering the layout (i.e. elements in the
+     *   content we are transitioning to).
+     * - the value at progress = 100% for elements that are leaving the layout (i.e. elements in the
+     *   content we are transitioning from).
+     *
+     * The returned value will be interpolated using the [transition] progress and [value], the
+     * value of the property when we are idle.
      */
-    // TODO(b/290184746): Figure out a public API for custom transformations that don't have access
-    // to these internal classes.
-    fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: T,
     ): T
 }
 
+interface PropertyTransformationScope : Density, ElementStateScope {
+    /** The current [direction][LayoutDirection] of the layout. */
+    val layoutDirection: LayoutDirection
+}
+
 /**
  * A [PropertyTransformation] associated to a range. This is a helper class so that normal
  * implementations of [PropertyTransformation] don't have to take care of reversing their range when
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
index 7014271..af0a6ed 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
@@ -21,10 +21,9 @@
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.dp
 import com.android.compose.animation.scene.ContentKey
-import com.android.compose.animation.scene.Element
+import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementMatcher
 import com.android.compose.animation.scene.OverscrollScope
-import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
 
 internal class Translate(
@@ -32,15 +31,13 @@
     private val x: Dp = 0.dp,
     private val y: Dp = 0.dp,
 ) : PropertyTransformation<Offset> {
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: Offset,
     ): Offset {
-        return with(layoutImpl.density) { Offset(value.x + x.toPx(), value.y + y.toPx()) }
+        return Offset(value.x + x.toPx(), value.y + y.toPx())
     }
 }
 
@@ -51,11 +48,9 @@
 ) : PropertyTransformation<Offset> {
     private val cachedOverscrollScope = CachedOverscrollScope()
 
-    override fun transform(
-        layoutImpl: SceneTransitionLayoutImpl,
+    override fun PropertyTransformationScope.transform(
         content: ContentKey,
-        element: Element,
-        stateInContent: Element.State,
+        element: ElementKey,
         transition: TransitionState.Transition,
         value: Offset,
     ): Offset {
@@ -64,7 +59,7 @@
         // that this method was invoked after performing this check.
         val overscrollProperties = transition as TransitionState.HasOverscrollProperties
         val overscrollScope =
-            cachedOverscrollScope.getFromCacheOrCompute(layoutImpl.density, overscrollProperties)
+            cachedOverscrollScope.getFromCacheOrCompute(density = this, overscrollProperties)
 
         return Offset(x = value.x + overscrollScope.x(), y = value.y + overscrollScope.y())
     }
@@ -75,7 +70,7 @@
  * [TransitionState.HasOverscrollProperties]. This helps avoid recreating a scope every frame
  * whenever an overscroll transition is computed.
  */
-private class CachedOverscrollScope() {
+private class CachedOverscrollScope {
     private var previousScope: OverscrollScope? = null
     private var previousDensity: Density? = null
     private var previousOverscrollProperties: TransitionState.HasOverscrollProperties? = null