Merge "Expose onTransition(Start|End) callbacks in STLState" into main

6 files changed, 134 insertions, 16 deletions

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 158256d14d1a..8153586efbca 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
@@ -232,6 +232,8 @@ fun MutableSceneTransitionLayoutState(
     canShowOverlay: (OverlayKey) -> Boolean = { true },
     canHideOverlay: (OverlayKey) -> Boolean = { true },
     canReplaceOverlay: (from: OverlayKey, to: OverlayKey) -> Boolean = { _, _ -> true },
+    onTransitionStart: (TransitionState.Transition) -> Unit = {},
+    onTransitionEnd: (TransitionState.Transition) -> Unit = {},
 ): MutableSceneTransitionLayoutState {
     return MutableSceneTransitionLayoutStateImpl(
         initialScene,
@@ -241,6 +243,8 @@ fun MutableSceneTransitionLayoutState(
         canShowOverlay,
         canHideOverlay,
         canReplaceOverlay,
+        onTransitionStart,
+        onTransitionEnd,
     )
 }
 
@@ -252,7 +256,11 @@ internal class MutableSceneTransitionLayoutStateImpl(
     internal val canChangeScene: (SceneKey) -> Boolean = { true },
     internal val canShowOverlay: (OverlayKey) -> Boolean = { true },
     internal val canHideOverlay: (OverlayKey) -> Boolean = { true },
-    internal val canReplaceOverlay: (from: OverlayKey, to: OverlayKey) -> Boolean = { _, _ -> true },
+    internal val canReplaceOverlay: (from: OverlayKey, to: OverlayKey) -> Boolean = { _, _ ->
+        true
+    },
+    private val onTransitionStart: (TransitionState.Transition) -> Unit = {},
+    private val onTransitionEnd: (TransitionState.Transition) -> Unit = {},
 ) : MutableSceneTransitionLayoutState {
     private val creationThread: Thread = Thread.currentThread()
 
@@ -367,9 +375,11 @@ internal class MutableSceneTransitionLayoutStateImpl(
             startTransitionInternal(transition, chain)
 
             // Run the transition until it is finished.
+            onTransitionStart(transition)
             transition.runInternal()
         } finally {
             finishTransition(transition)
+            onTransitionEnd(transition)
         }
     }
 
@@ -384,14 +394,10 @@ internal class MutableSceneTransitionLayoutStateImpl(
         val toContent = transition.toContent
 
         // Update the transition specs.
-        transition.transformationSpec =
-            transitions
-                .transitionSpec(fromContent, toContent, key = transition.key)
-                .transformationSpec(transition)
-        transition.previewTransformationSpec =
-            transitions
-                .transitionSpec(fromContent, toContent, key = transition.key)
-                .previewTransformationSpec(transition)
+        val spec = transitions.transitionSpec(fromContent, toContent, key = transition.key)
+        transition._cuj = spec.cuj
+        transition.transformationSpec = spec.transformationSpec(transition)
+        transition.previewTransformationSpec = spec.previewTransformationSpec(transition)
     }
 
     private fun startTransitionInternal(transition: TransitionState.Transition, chain: Boolean) {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
index 756d71c1b5cf..ff8efc28aa21 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
@@ -29,6 +29,7 @@ import com.android.compose.animation.scene.transformation.PropertyTransformation
 import com.android.compose.animation.scene.transformation.SharedElementTransformation
 import com.android.compose.animation.scene.transformation.TransformationMatcher
 import com.android.compose.animation.scene.transformation.TransformationWithRange
+import com.android.internal.jank.Cuj.CujType
 
 /** The transitions configuration of a [SceneTransitionLayout]. */
 class SceneTransitions
@@ -111,7 +112,15 @@ internal constructor(
     }
 
     private fun defaultTransition(from: ContentKey, to: ContentKey) =
-        TransitionSpecImpl(key = null, from, to, null, null, TransformationSpec.EmptyProvider)
+        TransitionSpecImpl(
+            key = null,
+            from,
+            to,
+            cuj = null,
+            previewTransformationSpec = null,
+            reversePreviewTransformationSpec = null,
+            TransformationSpec.EmptyProvider,
+        )
 
     companion object {
         internal val DefaultSwipeSpec =
@@ -147,6 +156,9 @@ internal interface TransitionSpec {
      */
     val to: ContentKey?
 
+    /** The CUJ covered by this transition. */
+    @CujType val cuj: Int?
+
     /**
      * Return a reversed version of this [TransitionSpec] for a transition going from [to] to
      * [from].
@@ -213,6 +225,7 @@ internal class TransitionSpecImpl(
     override val key: TransitionKey?,
     override val from: ContentKey?,
     override val to: ContentKey?,
+    override val cuj: Int?,
     private val previewTransformationSpec:
         ((TransitionState.Transition) -> TransformationSpecImpl)? =
         null,
@@ -226,6 +239,7 @@ internal class TransitionSpecImpl(
             key = key,
             from = to,
             to = from,
+            cuj = cuj,
             previewTransformationSpec = reversePreviewTransformationSpec,
             reversePreviewTransformationSpec = previewTransformationSpec,
             transformationSpec = { transition ->
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 fda6fab6229a..998054ef6c9e 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
@@ -25,6 +25,7 @@ import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.dp
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.transformation.Transformation
+import com.android.internal.jank.Cuj.CujType
 
 /** Define the [transitions][SceneTransitions] to be used with a [SceneTransitionLayout]. */
 fun transitions(builder: SceneTransitionsBuilder.() -> Unit): SceneTransitions {
@@ -64,6 +65,7 @@ interface SceneTransitionsBuilder {
     fun to(
         to: ContentKey,
         key: TransitionKey? = null,
+        @CujType cuj: Int? = null,
         preview: (TransitionBuilder.() -> Unit)? = null,
         reversePreview: (TransitionBuilder.() -> Unit)? = null,
         builder: TransitionBuilder.() -> Unit = {},
@@ -90,6 +92,7 @@ interface SceneTransitionsBuilder {
         from: ContentKey,
         to: ContentKey? = null,
         key: TransitionKey? = null,
+        @CujType cuj: Int? = null,
         preview: (TransitionBuilder.() -> Unit)? = null,
         reversePreview: (TransitionBuilder.() -> Unit)? = null,
         builder: TransitionBuilder.() -> Unit = {},
@@ -146,6 +149,9 @@ interface TransitionBuilder : BaseTransitionBuilder {
      */
     var swipeSpec: SpringSpec<Float>?
 
+    /** The CUJ associated to this transitions. */
+    @CujType var cuj: Int?
+
     /**
      * Define a timestamp-based range for the transformations inside [builder].
      *
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 a1649964ec13..7ca521513714 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
@@ -37,6 +37,7 @@ import com.android.compose.animation.scene.transformation.Transformation
 import com.android.compose.animation.scene.transformation.TransformationMatcher
 import com.android.compose.animation.scene.transformation.TransformationRange
 import com.android.compose.animation.scene.transformation.Translate
+import com.android.internal.jank.Cuj.CujType
 
 internal fun transitionsImpl(builder: SceneTransitionsBuilder.() -> Unit): SceneTransitions {
     val impl = SceneTransitionsBuilderImpl().apply(builder)
@@ -52,28 +53,47 @@ private class SceneTransitionsBuilderImpl : SceneTransitionsBuilder {
     override fun to(
         to: ContentKey,
         key: TransitionKey?,
+        @CujType cuj: Int?,
         preview: (TransitionBuilder.() -> Unit)?,
         reversePreview: (TransitionBuilder.() -> Unit)?,
         builder: TransitionBuilder.() -> Unit,
     ) {
-        transition(from = null, to = to, key = key, preview, reversePreview, builder)
+        transition(
+            from = null,
+            to = to,
+            key = key,
+            cuj = cuj,
+            preview = preview,
+            reversePreview = reversePreview,
+            builder = builder,
+        )
     }
 
     override fun from(
         from: ContentKey,
         to: ContentKey?,
         key: TransitionKey?,
+        @CujType cuj: Int?,
         preview: (TransitionBuilder.() -> Unit)?,
         reversePreview: (TransitionBuilder.() -> Unit)?,
         builder: TransitionBuilder.() -> Unit,
     ) {
-        transition(from = from, to = to, key = key, preview, reversePreview, builder)
+        transition(
+            from = from,
+            to = to,
+            key = key,
+            cuj = cuj,
+            preview = preview,
+            reversePreview = reversePreview,
+            builder = builder,
+        )
     }
 
     private fun transition(
         from: ContentKey?,
         to: ContentKey?,
         key: TransitionKey?,
+        @CujType cuj: Int?,
         preview: (TransitionBuilder.() -> Unit)?,
         reversePreview: (TransitionBuilder.() -> Unit)?,
         builder: TransitionBuilder.() -> Unit,
@@ -93,9 +113,10 @@ private class SceneTransitionsBuilderImpl : SceneTransitionsBuilder {
 
         val spec =
             TransitionSpecImpl(
-                key,
-                from,
-                to,
+                key = key,
+                from = from,
+                to = to,
+                cuj = cuj,
                 previewTransformationSpec = preview?.let { { t -> transformationSpec(t, it) } },
                 reversePreviewTransformationSpec =
                     reversePreview?.let { { t -> transformationSpec(t, it) } },
@@ -190,6 +211,7 @@ internal class TransitionBuilderImpl(override val transition: TransitionState.Tr
     override var spec: AnimationSpec<Float> = spring(stiffness = Spring.StiffnessLow)
     override var swipeSpec: SpringSpec<Float>? = null
     override var distance: UserActionDistance? = null
+    override var cuj: Int? = null
     private val durationMillis: Int by lazy {
         val spec = spec
         if (spec !is DurationBasedAnimationSpec) {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
index e7ca51114b93..712af56ee1bc 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
@@ -32,6 +32,7 @@ import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.TransformationSpec
 import com.android.compose.animation.scene.TransformationSpecImpl
 import com.android.compose.animation.scene.TransitionKey
+import com.android.internal.jank.Cuj.CujType
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
@@ -237,6 +238,11 @@ sealed interface TransitionState {
         /** Whether user input is currently driving the transition. */
         abstract val isUserInputOngoing: Boolean
 
+        /** The CUJ covered by this transition. */
+        @CujType
+        val cuj: Int?
+            get() = _cuj
+
         /**
          * The progress of the preview transition. This is usually in the `[0; 1]` range, but it can
          * also be less than `0` or greater than `1` when using transitions with a spring
@@ -251,13 +257,15 @@ sealed interface TransitionState {
         internal open val isInPreviewStage: Boolean = false
 
         /**
-         * The current [TransformationSpecImpl] associated to this transition.
+         * The current [TransformationSpecImpl] and other values associated to this transition from
+         * the spec.
          *
          * Important: These will be set exactly once, when this transition is
          * [started][MutableSceneTransitionLayoutStateImpl.startTransition].
          */
         internal var transformationSpec: TransformationSpecImpl = TransformationSpec.Empty
         internal var previewTransformationSpec: TransformationSpecImpl? = null
+        internal var _cuj: Int? = null
 
         /**
          * An animatable that animates from 1f to 0f. This will be used to nicely animate the sudden
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
index 5074cd5211ce..f3be5e43c294 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
@@ -491,4 +491,66 @@ class SceneTransitionLayoutStateTest {
         assertThat(state.transitionState).isIdle()
         assertThat(state.transitionState).hasCurrentScene(SceneC)
     }
+
+    @Test
+    fun trackTransitionCujs() = runTest {
+        val started = mutableSetOf<TransitionState.Transition>()
+        val finished = mutableSetOf<TransitionState.Transition>()
+        val cujWhenStarting = mutableMapOf<TransitionState.Transition, Int?>()
+        val state =
+            MutableSceneTransitionLayoutState(
+                SceneA,
+                transitions {
+                    // A <=> B.
+                    from(SceneA, to = SceneB, cuj = 1)
+
+                    // A <=> C.
+                    from(SceneA, to = SceneC, cuj = 2)
+                    from(SceneC, to = SceneA, cuj = 3)
+                },
+                onTransitionStart = { transition ->
+                    started.add(transition)
+                    cujWhenStarting[transition] = transition.cuj
+                },
+                onTransitionEnd = { finished.add(it) },
+            )
+
+        val aToB = transition(SceneA, SceneB)
+        val bToA = transition(SceneB, SceneA)
+        val aToC = transition(SceneA, SceneC)
+        val cToA = transition(SceneC, SceneA)
+
+        val animationScope = this
+        state.startTransitionImmediately(animationScope, aToB)
+        assertThat(started).containsExactly(aToB)
+        assertThat(finished).isEmpty()
+
+        state.startTransitionImmediately(animationScope, bToA)
+        assertThat(started).containsExactly(aToB, bToA)
+        assertThat(finished).isEmpty()
+
+        aToB.finish()
+        runCurrent()
+        assertThat(finished).containsExactly(aToB)
+
+        state.startTransitionImmediately(animationScope, aToC)
+        assertThat(started).containsExactly(aToB, bToA, aToC)
+        assertThat(finished).containsExactly(aToB)
+
+        state.startTransitionImmediately(animationScope, cToA)
+        assertThat(started).containsExactly(aToB, bToA, aToC, cToA)
+        assertThat(finished).containsExactly(aToB)
+
+        bToA.finish()
+        aToC.finish()
+        cToA.finish()
+        runCurrent()
+        assertThat(started).containsExactly(aToB, bToA, aToC, cToA)
+        assertThat(finished).containsExactly(aToB, bToA, aToC, cToA)
+
+        assertThat(cujWhenStarting[aToB]).isEqualTo(1)
+        assertThat(cujWhenStarting[bToA]).isEqualTo(1)
+        assertThat(cujWhenStarting[aToC]).isEqualTo(2)
+        assertThat(cujWhenStarting[cToA]).isEqualTo(3)
+    }
 }