Merge "[flexiglass] Make sure scene and userActions are in-sync" into main
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 22566e7..9c9e6c6 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
@@ -34,6 +34,8 @@
 import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.SceneTransitionLayout
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
 import com.android.compose.animation.scene.observableTransitionState
 import com.android.systemui.ribbon.ui.composable.BottomRightCornerRibbon
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
@@ -56,7 +58,6 @@
  *   must have entries in this map.
  * @param modifier A modifier.
  */
-@OptIn(ExperimentalComposeUiApi::class)
 @Composable
 fun SceneContainer(
     viewModel: SceneContainerViewModel,
@@ -66,8 +67,6 @@
 ) {
     val coroutineScope = rememberCoroutineScope()
     val currentSceneKey: SceneKey by viewModel.currentScene.collectAsStateWithLifecycle()
-    val currentDestinations by
-        viewModel.currentDestinationScenes(coroutineScope).collectAsStateWithLifecycle()
     val state: MutableSceneTransitionLayoutState = remember {
         MutableSceneTransitionLayoutState(
             initialScene = currentSceneKey,
@@ -88,20 +87,19 @@
         onDispose { viewModel.setTransitionState(null) }
     }
 
+    val userActionsBySceneKey: Map<SceneKey, Map<UserAction, UserActionResult>> =
+        sceneByKey.values.associate { scene ->
+            val userActions by scene.destinationScenes.collectAsStateWithLifecycle(emptyMap())
+            val resolvedUserActions = viewModel.resolveSceneFamilies(userActions)
+            scene.key to resolvedUserActions
+        }
+
     Box(
         modifier = Modifier.fillMaxSize(),
     ) {
         SceneTransitionLayout(state = state, modifier = modifier.fillMaxSize()) {
             sceneByKey.forEach { (sceneKey, composableScene) ->
-                scene(
-                    key = sceneKey,
-                    userActions =
-                        if (sceneKey == currentSceneKey) {
-                            currentDestinations
-                        } else {
-                            viewModel.resolveSceneFamilies(composableScene.destinationScenes.value)
-                        },
-                ) {
+                scene(key = sceneKey, userActions = checkNotNull(userActionsBySceneKey[sceneKey])) {
                     with(composableScene) {
                         this@scene.Content(
                             modifier = Modifier.element(sceneKey.rootElementKey).fillMaxSize(),
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
index cb4d96f..39b3662 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -133,7 +133,6 @@
                 sceneInteractor = sceneInteractor,
                 falsingInteractor = kosmos.falsingInteractor,
                 powerInteractor = kosmos.powerInteractor,
-                scenes = kosmos.scenes,
             )
             .apply { setTransitionState(transitionState) }
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
index 5c30379..ea95aab 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
@@ -28,10 +28,8 @@
 import com.android.systemui.power.data.repository.fakePowerRepository
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.domain.interactor.sceneInteractor
-import com.android.systemui.scene.fakeScenes
 import com.android.systemui.scene.sceneContainerConfig
 import com.android.systemui.scene.sceneKeys
-import com.android.systemui.scene.scenes
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.fakeSceneDataSource
 import com.android.systemui.testKosmos
@@ -66,7 +64,6 @@
                 sceneInteractor = sceneInteractor,
                 falsingInteractor = kosmos.falsingInteractor,
                 powerInteractor = kosmos.powerInteractor,
-                scenes = kosmos.scenes,
             )
     }
 
@@ -217,23 +214,4 @@
 
             assertThat(isVisible).isFalse()
         }
-
-    @Test
-    fun currentDestinationScenes_onlyTheCurrentSceneIsCollected() =
-        testScope.runTest {
-            val unused by collectLastValue(underTest.currentDestinationScenes(backgroundScope))
-            val currentScene by collectLastValue(sceneInteractor.currentScene)
-            kosmos.fakeScenes.forEach { scene ->
-                fakeSceneDataSource.changeScene(toScene = scene.key)
-                runCurrent()
-                assertThat(currentScene).isEqualTo(scene.key)
-
-                assertThat(scene.isDestinationScenesBeingCollected).isTrue()
-                kosmos.fakeScenes
-                    .filter { it.key != scene.key }
-                    .forEach { otherScene ->
-                        assertThat(otherScene.isDestinationScenesBeingCollected).isFalse()
-                    }
-            }
-        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
index d380251..a28222e 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
@@ -26,17 +26,12 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.scene.domain.interactor.SceneInteractor
-import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.utils.coroutines.flow.flatMapLatestConflated
 import javax.inject.Inject
-import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.stateIn
 
 /** Models UI state for the scene container. */
 @SysUISingleton
@@ -46,7 +41,6 @@
     private val sceneInteractor: SceneInteractor,
     private val falsingInteractor: FalsingInteractor,
     private val powerInteractor: PowerInteractor,
-    scenes: Set<@JvmSuppressWildcards Scene>,
 ) {
     /**
      * Keys of all scenes in the container.
@@ -62,25 +56,6 @@
     /** Whether the container is visible. */
     val isVisible: StateFlow<Boolean> = sceneInteractor.isVisible
 
-    private val destinationScenesBySceneKey =
-        scenes.associate { scene ->
-            scene.key to scene.destinationScenes.flatMapLatestConflated { replaceSceneFamilies(it) }
-        }
-
-    fun currentDestinationScenes(
-        scope: CoroutineScope,
-    ): StateFlow<Map<UserAction, UserActionResult>> {
-        return currentScene
-            .flatMapLatestConflated { currentSceneKey ->
-                checkNotNull(destinationScenesBySceneKey[currentSceneKey])
-            }
-            .stateIn(
-                scope = scope,
-                started = SharingStarted.WhileSubscribed(),
-                initialValue = emptyMap(),
-            )
-    }
-
     /**
      * Binds the given flow so the system remembers it.
      *