Move onChangeScene and transitions to STLState (1/2)

This CL is a refactoring that moves onChangeScene and transitions
to SceneTransitionLayoutState. It also removes the
SceneTransitionLayoutState() factory to create a STLState, and replaces
it with the updateSceneTransitionLayoutState() Composable to create a
STLState from a hoisted source of truth. In ag/25812691, I'm adding a
MutableSceneTransitionLayoutState that can replace current usages of the
SceneTransitionLayoutState() factory outside composition (currently used
only in tests. This means that the tests don't compile with only this
CL and this CL will be submitted together with ag/25812691.

Note that the name updateSceneTransitionLayoutState() was chosen to be
consistent with Compose updateTransition().

Test: PlatformComposeSceneTransitionLayoutTests
Bug: 318794193
Flag: N/A
Change-Id: Ie7519ec4a3c92afb31f6c5fc58074e9d2131d2f5

6 files changed, 107 insertions, 51 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
index 249b3e14ec72..d47527a0a191 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
@@ -33,11 +33,11 @@ import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.SceneTransitionLayout
-import com.android.compose.animation.scene.SceneTransitionLayoutState
 import com.android.compose.animation.scene.Swipe
 import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.observableTransitionState
 import com.android.compose.animation.scene.transitions
+import com.android.compose.animation.scene.updateSceneTransitionLayoutState
 import com.android.systemui.communal.shared.model.CommunalSceneKey
 import com.android.systemui.communal.shared.model.ObservableCommunalTransitionState
 import com.android.systemui.communal.ui.viewmodel.BaseCommunalViewModel
@@ -76,7 +76,13 @@ fun CommunalContainer(
         viewModel.currentScene
             .transform { value -> emit(value.toTransitionSceneKey()) }
             .collectAsState(TransitionSceneKey.Blank)
-    val sceneTransitionLayoutState = remember { SceneTransitionLayoutState(currentScene) }
+    val sceneTransitionLayoutState =
+        updateSceneTransitionLayoutState(
+            currentScene,
+            onChangeScene = { viewModel.onSceneChanged(it.toCommunalSceneKey()) },
+            transitions = sceneTransitions,
+        )
+
     // Don't show hub mode UI if keyguard is present. This is important since we're in the shade,
     // which can be opened from many locations.
     val isKeyguardShowing by viewModel.isKeyguardVisible.collectAsState(initial = false)
@@ -98,12 +104,9 @@ fun CommunalContainer(
 
     Box(modifier = modifier.fillMaxSize()) {
         SceneTransitionLayout(
-            modifier = Modifier.fillMaxSize(),
-            currentScene = currentScene,
-            onChangeScene = { sceneKey -> viewModel.onSceneChanged(sceneKey.toCommunalSceneKey()) },
-            transitions = sceneTransitions,
             state = sceneTransitionLayoutState,
-            edgeDetector = FixedSizeEdgeDetector(ContainerDimensions.EdgeSwipeSize)
+            modifier = Modifier.fillMaxSize(),
+            edgeDetector = FixedSizeEdgeDetector(ContainerDimensions.EdgeSwipeSize),
         ) {
             scene(
                 TransitionSceneKey.Blank,
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
index 4eb9089dc589..c35202cd830a 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
@@ -25,7 +25,6 @@ import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.remember
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
@@ -38,11 +37,11 @@ import com.android.compose.animation.scene.Edge as SceneTransitionEdge
 import com.android.compose.animation.scene.ObservableTransitionState as SceneTransitionObservableTransitionState
 import com.android.compose.animation.scene.SceneKey as SceneTransitionSceneKey
 import com.android.compose.animation.scene.SceneTransitionLayout
-import com.android.compose.animation.scene.SceneTransitionLayoutState
 import com.android.compose.animation.scene.Swipe
 import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.UserAction as SceneTransitionUserAction
 import com.android.compose.animation.scene.observableTransitionState
+import com.android.compose.animation.scene.updateSceneTransitionLayoutState
 import com.android.systemui.ribbon.ui.composable.BottomRightCornerRibbon
 import com.android.systemui.scene.shared.model.Direction
 import com.android.systemui.scene.shared.model.Edge
@@ -82,7 +81,12 @@ fun SceneContainer(
     val currentScene = checkNotNull(sceneByKey[currentSceneKey])
     val currentDestinations: Map<UserAction, SceneModel> by
         currentScene.destinationScenes.collectAsState()
-    val state = remember { SceneTransitionLayoutState(currentSceneKey.toTransitionSceneKey()) }
+    val state =
+        updateSceneTransitionLayoutState(
+            currentSceneKey.toTransitionSceneKey(),
+            onChangeScene = viewModel::onSceneChanged,
+            transitions = SceneContainerTransitions,
+        )
 
     DisposableEffect(viewModel, state) {
         viewModel.setTransitionState(state.observableTransitionState().map { it.toModel() })
@@ -93,9 +97,6 @@ fun SceneContainer(
         modifier = Modifier.fillMaxSize(),
     ) {
         SceneTransitionLayout(
-            currentScene = currentSceneKey.toTransitionSceneKey(),
-            onChangeScene = viewModel::onSceneChanged,
-            transitions = SceneContainerTransitions,
             state = state,
             modifier =
                 modifier
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
index 338557d0942e..05efec5353b3 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
@@ -312,7 +312,7 @@ internal class SceneGestureHandler(
             // immediately go back B => A.
             if (targetScene != swipeTransition._currentScene) {
                 swipeTransition._currentScene = targetScene
-                layoutImpl.onChangeScene(targetScene.key)
+                layoutImpl.state.onChangeScene(targetScene.key)
             }
 
             animateOffset(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
index 84fade8937ff..a9d828a23f89 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -19,7 +19,6 @@ package com.android.compose.animation.scene
 import androidx.annotation.FloatRange
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.remember
@@ -29,7 +28,40 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Shape
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
 import androidx.compose.ui.platform.LocalDensity
-import kotlinx.coroutines.channels.Channel
+
+/**
+ * [SceneTransitionLayout] is a container that automatically animates its content whenever its state
+ * changes.
+ *
+ * Note: You should use [androidx.compose.animation.AnimatedContent] instead of
+ * [SceneTransitionLayout] if it fits your need. Use [SceneTransitionLayout] over AnimatedContent if
+ * you need support for swipe gestures, shared elements or transitions defined declaratively outside
+ * UI code.
+ *
+ * @param state the state of this layout.
+ * @param edgeDetector the edge detector used to detect which edge a swipe is started from, if any.
+ * @param transitionInterceptionThreshold used during a scene transition. For the scene to be
+ *   intercepted, the progress value must be above the threshold, and below (1 - threshold).
+ * @param scenes the configuration of the different scenes of this layout.
+ * @see updateSceneTransitionLayoutState
+ */
+@Composable
+fun SceneTransitionLayout(
+    state: SceneTransitionLayoutState,
+    modifier: Modifier = Modifier,
+    edgeDetector: EdgeDetector = DefaultEdgeDetector,
+    @FloatRange(from = 0.0, to = 0.5) transitionInterceptionThreshold: Float = 0f,
+    scenes: SceneTransitionLayoutScope.() -> Unit,
+) {
+    SceneTransitionLayoutForTesting(
+        state,
+        modifier,
+        edgeDetector,
+        transitionInterceptionThreshold,
+        onLayoutImpl = null,
+        scenes,
+    )
+}
 
 /**
  * [SceneTransitionLayout] is a container that automatically animates its content whenever
@@ -45,7 +77,6 @@ import kotlinx.coroutines.channels.Channel
  *   This is called when the user commits a transition to a new scene because of a [UserAction], for
  *   instance by triggering back navigation or by swiping to a new scene.
  * @param transitions the definition of the transitions used to animate a change of scene.
- * @param state the observable state of this layout.
  * @param edgeDetector the edge detector used to detect which edge a swipe is started from, if any.
  * @param transitionInterceptionThreshold used during a scene transition. For the scene to be
  *   intercepted, the progress value must be above the threshold, and below (1 - threshold).
@@ -57,20 +88,16 @@ fun SceneTransitionLayout(
     onChangeScene: (SceneKey) -> Unit,
     transitions: SceneTransitions,
     modifier: Modifier = Modifier,
-    state: SceneTransitionLayoutState = remember { SceneTransitionLayoutState(currentScene) },
     edgeDetector: EdgeDetector = DefaultEdgeDetector,
     @FloatRange(from = 0.0, to = 0.5) transitionInterceptionThreshold: Float = 0f,
     scenes: SceneTransitionLayoutScope.() -> Unit,
 ) {
-    SceneTransitionLayoutForTesting(
-        currentScene,
-        onChangeScene,
-        modifier,
-        transitions,
+    val state = updateSceneTransitionLayoutState(currentScene, onChangeScene, transitions)
+    SceneTransitionLayout(
         state,
+        modifier,
         edgeDetector,
         transitionInterceptionThreshold,
-        onLayoutImpl = null,
         scenes,
     )
 }
@@ -346,11 +373,8 @@ enum class SwipeDirection(val orientation: Orientation) {
  */
 @Composable
 internal fun SceneTransitionLayoutForTesting(
-    currentScene: SceneKey,
-    onChangeScene: (SceneKey) -> Unit,
+    state: SceneTransitionLayoutState,
     modifier: Modifier = Modifier,
-    transitions: SceneTransitions = transitions {},
-    state: SceneTransitionLayoutState = remember { SceneTransitionLayoutState(currentScene) },
     edgeDetector: EdgeDetector = DefaultEdgeDetector,
     transitionInterceptionThreshold: Float = 0f,
     onLayoutImpl: ((SceneTransitionLayoutImpl) -> Unit)? = null,
@@ -361,7 +385,6 @@ internal fun SceneTransitionLayoutForTesting(
     val layoutImpl = remember {
         SceneTransitionLayoutImpl(
                 state = state as SceneTransitionLayoutStateImpl,
-                onChangeScene = onChangeScene,
                 density = density,
                 edgeDetector = edgeDetector,
                 transitionInterceptionThreshold = transitionInterceptionThreshold,
@@ -375,7 +398,6 @@ internal fun SceneTransitionLayoutForTesting(
     // SnapshotStateMap anymore.
     layoutImpl.updateScenes(scenes)
 
-    val targetSceneChannel = remember { Channel<SceneKey>(Channel.CONFLATED) }
     SideEffect {
         if (state != layoutImpl.state) {
             error(
@@ -384,23 +406,8 @@ internal fun SceneTransitionLayoutForTesting(
             )
         }
 
-        layoutImpl.onChangeScene = onChangeScene
-        (state as SceneTransitionLayoutStateImpl).transitions = transitions
         layoutImpl.density = density
         layoutImpl.edgeDetector = edgeDetector
-
-        state.transitions = transitions
-
-        targetSceneChannel.trySend(currentScene)
-    }
-
-    LaunchedEffect(targetSceneChannel) {
-        for (newKey in targetSceneChannel) {
-            // Inspired by AnimateAsState.kt: let's poll the last value to avoid being one frame
-            // late.
-            val newKey = targetSceneChannel.tryReceive().getOrNull() ?: newKey
-            animateToScene(layoutImpl.state, newKey)
-        }
     }
 
     layoutImpl.Content(modifier)
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 0227aba94b53..92117285b81a 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
@@ -49,7 +49,6 @@ internal typealias MovableElementContent =
 @Stable
 internal class SceneTransitionLayoutImpl(
     internal val state: SceneTransitionLayoutStateImpl,
-    internal var onChangeScene: (SceneKey) -> Unit,
     internal var density: Density,
     internal var edgeDetector: EdgeDetector,
     internal var transitionInterceptionThreshold: Float,
@@ -244,7 +243,7 @@ internal class SceneTransitionLayoutImpl(
                 // TODO(b/290184746): Make sure that this works with SystemUI once we use
                 // SceneTransitionLayout in Flexiglass.
                 scene(state.transitionState.currentScene).userActions[Back]?.let { backScene ->
-                    BackHandler { onChangeScene(backScene) }
+                    BackHandler { state.onChangeScene(backScene) }
                 }
 
                 Box {
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 0607aa148157..4f51ccb8be7d 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
@@ -16,12 +16,21 @@
 
 package com.android.compose.animation.scene
 
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
+import kotlinx.coroutines.channels.Channel
 
-/** The state of a [SceneTransitionLayout]. */
+/**
+ * The state of a [SceneTransitionLayout].
+ *
+ * @see updateSceneTransitionLayoutState
+ */
 @Stable
 sealed interface SceneTransitionLayoutState {
     /**
@@ -46,9 +55,45 @@ sealed interface SceneTransitionLayoutState {
     fun isTransitioningBetween(scene: SceneKey, other: SceneKey): Boolean
 }
 
-/** Create a new [SceneTransitionLayoutState] that is currently idle at scene [currentScene]. */
-fun SceneTransitionLayoutState(currentScene: SceneKey): SceneTransitionLayoutState {
-    return SceneTransitionLayoutStateImpl(currentScene, SceneTransitions.Empty)
+/**
+ * Sets up a [SceneTransitionLayoutState] and keeps it synced with [currentScene], [onChangeScene]
+ * and [transitions]. New transitions will automatically be started whenever [currentScene] is
+ * changed.
+ *
+ * @param currentScene the current scene
+ * @param onChangeScene a mutator that should set [currentScene] to the given scene when called.
+ *   This is called when the user commits a transition to a new scene because of a [UserAction], for
+ *   instance by triggering back navigation or by swiping to a new scene.
+ * @param transitions the definition of the transitions used to animate a change of scene.
+ */
+@Composable
+fun updateSceneTransitionLayoutState(
+    currentScene: SceneKey,
+    onChangeScene: (SceneKey) -> Unit,
+    transitions: SceneTransitions = transitions {},
+): SceneTransitionLayoutState {
+    val state = remember {
+        SceneTransitionLayoutStateImpl(currentScene, transitions, onChangeScene)
+    }
+
+    val targetSceneChannel = remember { Channel<SceneKey>(Channel.CONFLATED) }
+    SideEffect {
+        state.onChangeScene = onChangeScene
+        state.transitions = transitions
+
+        targetSceneChannel.trySend(currentScene)
+    }
+
+    LaunchedEffect(targetSceneChannel) {
+        for (newKey in targetSceneChannel) {
+            // Inspired by AnimateAsState.kt: let's poll the last value to avoid being one frame
+            // late.
+            val newKey = targetSceneChannel.tryReceive().getOrNull() ?: newKey
+            animateToScene(state, newKey)
+        }
+    }
+
+    return state
 }
 
 @Stable
@@ -112,6 +157,7 @@ sealed interface TransitionState {
 internal class SceneTransitionLayoutStateImpl(
     initialScene: SceneKey,
     internal var transitions: SceneTransitions,
+    internal var onChangeScene: (SceneKey) -> Unit,
 ) : SceneTransitionLayoutState {
     override var transitionState: TransitionState by
         mutableStateOf(TransitionState.Idle(initialScene))