Fix STLImpl scene update

See b/317014852 for context. I was unfortunately not able to create a
small sample to reproduce, and therefore to come up with a test. I will
return to understanding this more once I'm done with more urgent tasks.

Bug: 317014852
Test: Manual, see Flexiglass steps from chat linked in
 b/317014852#comment1
Flag: N/A
Change-Id: I12498ca3b464bc422478bdb2a01bd1a5885e04b3

2 files changed, 10 insertions, 2 deletions

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 3608e374fdbc..63fe9e98764d 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
@@ -280,6 +280,10 @@ internal fun SceneTransitionLayoutForTesting(
             .also { onLayoutImpl?.invoke(it) }
     }
 
+    // TODO(b/317014852): Move this into the SideEffect {} again once STLImpl.scenes is not a
+    // SnapshotStateMap anymore.
+    layoutImpl.updateScenes(scenes)
+
     val targetSceneChannel = remember { Channel<SceneKey>(Channel.CONFLATED) }
     SideEffect {
         if (state != layoutImpl.state) {
@@ -293,7 +297,6 @@ internal fun SceneTransitionLayoutForTesting(
         (state as SceneTransitionLayoutStateImpl).transitions = transitions
         layoutImpl.density = density
         layoutImpl.edgeDetector = edgeDetector
-        layoutImpl.updateScenes(scenes)
 
         state.transitions = transitions
 
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 c99c3250bbb1..45e1a0fa8f77 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
@@ -46,7 +46,12 @@ internal class SceneTransitionLayoutImpl(
     builder: SceneTransitionLayoutScope.() -> Unit,
     coroutineScope: CoroutineScope,
 ) {
-    internal val scenes = mutableMapOf<SceneKey, Scene>()
+    /**
+     * The map of [Scene]s.
+     *
+     * TODO(b/317014852): Make this a normal MutableMap instead.
+     */
+    internal val scenes = SnapshotStateMap<SceneKey, Scene>()
 
     /**
      * The map of [Element]s.