[flexiglass] Fixes compose-enabled sysuig build.

Bug: 279501596
Test: built and deployed sysuig with compose build time flag enabled,
disabled, and Compose gallery when the flag is disabled.

Change-Id: I8c48055d90d9f9fa6a994ed035f71d472b1639fd
diff --git a/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/scene/shared/page/SceneModule.kt b/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/scene/ui/composable/SceneModule.kt
similarity index 94%
rename from packages/SystemUI/compose/facade/disabled/src/com/android/systemui/scene/shared/page/SceneModule.kt
rename to packages/SystemUI/compose/facade/disabled/src/com/android/systemui/scene/ui/composable/SceneModule.kt
index 18c9513..24064b1 100644
--- a/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/scene/shared/page/SceneModule.kt
+++ b/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/scene/ui/composable/SceneModule.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.page
+package com.android.systemui.scene.ui.composable
 
 import com.android.systemui.scene.shared.model.Scene
 import dagger.Module
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ui/composable/SceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ui/composable/SceneModule.kt
index 530706e..ee53ece 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ui/composable/SceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ui/composable/SceneModule.kt
@@ -17,7 +17,7 @@
 package com.android.systemui.scene.ui.composable
 
 import com.android.systemui.bouncer.ui.composable.BouncerScene
-import com.android.systemui.keyguard.ui.composable.LockScreenScene
+import com.android.systemui.keyguard.ui.composable.LockscreenScene
 import com.android.systemui.qs.ui.composable.QuickSettingsScene
 import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.shade.ui.composable.ShadeScene
@@ -30,7 +30,7 @@
     fun scenes(
         bouncer: BouncerScene,
         gone: GoneScene,
-        lockScreen: LockScreenScene,
+        lockScreen: LockscreenScene,
         qs: QuickSettingsScene,
         shade: ShadeScene,
     ): Set<Scene> {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
index 6f6d0f9..f48bab9 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
@@ -55,11 +55,13 @@
 class BouncerScene
 @Inject
 constructor(
-    private val viewModel: BouncerViewModel,
+    private val viewModelFactory: BouncerViewModel.Factory,
 ) : ComposableScene {
     override val key = SceneKey.Bouncer
 
-    override fun destinationScenes(): StateFlow<Map<UserAction, SceneModel>> =
+    override fun destinationScenes(
+        containerName: String,
+    ): StateFlow<Map<UserAction, SceneModel>> =
         MutableStateFlow<Map<UserAction, SceneModel>>(
                 mapOf(
                     UserAction.Back to SceneModel(SceneKey.Lockscreen),
@@ -67,7 +69,11 @@
             )
             .asStateFlow()
 
-    @Composable override fun Content(modifier: Modifier) = BouncerScene(viewModel, modifier)
+    @Composable
+    override fun Content(
+        containerName: String,
+        modifier: Modifier,
+    ) = BouncerScene(viewModelFactory.create(containerName), modifier)
 }
 
 @Composable
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
index ab7bc26..7c07a8b 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
@@ -52,25 +52,32 @@
 @Inject
 constructor(
     @Application private val applicationScope: CoroutineScope,
-    private val viewModel: LockscreenSceneViewModel,
+    private val viewModelFactory: LockscreenSceneViewModel.Factory,
 ) : ComposableScene {
     override val key = SceneKey.Lockscreen
 
-    override fun destinationScenes(): StateFlow<Map<UserAction, SceneModel>> =
-        viewModel.upDestinationSceneKey
-            .map { pageKey -> destinationScenes(up = pageKey) }
-            .stateIn(
-                scope = applicationScope,
-                started = SharingStarted.Eagerly,
-                initialValue = destinationScenes(up = viewModel.upDestinationSceneKey.value)
-            )
+    private var unsafeViewModel: LockscreenSceneViewModel? = null
+
+    override fun destinationScenes(
+        containerName: String,
+    ): StateFlow<Map<UserAction, SceneModel>> =
+        getOrCreateViewModelSingleton(containerName).let { viewModel ->
+            viewModel.upDestinationSceneKey
+                .map { pageKey -> destinationScenes(up = pageKey) }
+                .stateIn(
+                    scope = applicationScope,
+                    started = SharingStarted.Eagerly,
+                    initialValue = destinationScenes(up = viewModel.upDestinationSceneKey.value)
+                )
+        }
 
     @Composable
     override fun Content(
+        containerName: String,
         modifier: Modifier,
     ) {
         LockscreenScene(
-            viewModel = viewModel,
+            viewModel = getOrCreateViewModelSingleton(containerName),
             modifier = modifier,
         )
     }
@@ -83,6 +90,13 @@
             UserAction.Swipe(Direction.DOWN) to SceneModel(SceneKey.Shade)
         )
     }
+
+    private fun getOrCreateViewModelSingleton(
+        containerName: String,
+    ): LockscreenSceneViewModel {
+        return unsafeViewModel
+            ?: viewModelFactory.create(containerName).also { unsafeViewModel = it }
+    }
 }
 
 @Composable
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
index 130395a..58db37e 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
@@ -44,11 +44,13 @@
 class QuickSettingsScene
 @Inject
 constructor(
-    private val viewModel: QuickSettingsSceneViewModel,
+    private val viewModelFactory: QuickSettingsSceneViewModel.Factory,
 ) : ComposableScene {
     override val key = SceneKey.QuickSettings
 
-    override fun destinationScenes(): StateFlow<Map<UserAction, SceneModel>> =
+    override fun destinationScenes(
+        containerName: String,
+    ): StateFlow<Map<UserAction, SceneModel>> =
         MutableStateFlow<Map<UserAction, SceneModel>>(
                 mapOf(
                     UserAction.Swipe(Direction.UP) to SceneModel(SceneKey.Shade),
@@ -58,10 +60,11 @@
 
     @Composable
     override fun Content(
+        containerName: String,
         modifier: Modifier,
     ) {
         QuickSettingsScene(
-            viewModel = viewModel,
+            viewModel = viewModelFactory.create(containerName),
             modifier = modifier,
         )
     }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
index a213666..6f3363e 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
@@ -22,5 +22,5 @@
 
 /** Compose-capable extension of [Scene]. */
 interface ComposableScene : Scene {
-    @Composable fun Content(modifier: Modifier)
+    @Composable fun Content(containerName: String, modifier: Modifier)
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
index 0070552..b387463 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
@@ -41,7 +41,9 @@
 class GoneScene @Inject constructor() : ComposableScene {
     override val key = SceneKey.Gone
 
-    override fun destinationScenes(): StateFlow<Map<UserAction, SceneModel>> =
+    override fun destinationScenes(
+        containerName: String,
+    ): StateFlow<Map<UserAction, SceneModel>> =
         MutableStateFlow<Map<UserAction, SceneModel>>(
                 mapOf(
                     UserAction.Swipe(Direction.DOWN) to SceneModel(SceneKey.Shade),
@@ -51,6 +53,7 @@
 
     @Composable
     override fun Content(
+        containerName: String,
         modifier: Modifier,
     ) {
         /*
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 f8a73d5..5e07610 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
@@ -75,6 +75,7 @@
             if (key == currentSceneKey) {
                 Scene(
                     scene = composableScene,
+                    containerName = viewModel.containerName,
                     onSceneChanged = viewModel::setCurrentScene,
                     modifier = Modifier.fillMaxSize(),
                 )
@@ -87,6 +88,7 @@
 @Composable
 private fun Scene(
     scene: ComposableScene,
+    containerName: String,
     onSceneChanged: (SceneModel) -> Unit,
     modifier: Modifier = Modifier,
 ) {
@@ -97,11 +99,12 @@
             modifier = Modifier.align(Alignment.Center),
         ) {
             scene.Content(
+                containerName = containerName,
                 modifier = Modifier,
             )
 
             val destinationScenes: Map<UserAction, SceneModel> by
-                scene.destinationScenes().collectAsState()
+                scene.destinationScenes(containerName).collectAsState()
             val swipeLeftDestinationScene = destinationScenes[UserAction.Swipe(Direction.LEFT)]
             val swipeUpDestinationScene = destinationScenes[UserAction.Swipe(Direction.UP)]
             val swipeRightDestinationScene = destinationScenes[UserAction.Swipe(Direction.RIGHT)]
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
index 5a09204..e4513d0 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
@@ -48,25 +48,32 @@
 @Inject
 constructor(
     @Application private val applicationScope: CoroutineScope,
-    private val viewModel: ShadeSceneViewModel,
+    private val viewModelFactory: ShadeSceneViewModel.Factory,
 ) : ComposableScene {
     override val key = SceneKey.Shade
 
-    override fun destinationScenes(): StateFlow<Map<UserAction, SceneModel>> =
-        viewModel.upDestinationSceneKey
-            .map { sceneKey -> destinationScenes(up = sceneKey) }
-            .stateIn(
-                scope = applicationScope,
-                started = SharingStarted.Eagerly,
-                initialValue = destinationScenes(up = viewModel.upDestinationSceneKey.value),
-            )
+    private var unsafeViewModel: ShadeSceneViewModel? = null
+
+    override fun destinationScenes(
+        containerName: String,
+    ): StateFlow<Map<UserAction, SceneModel>> =
+        getOrCreateViewModelSingleton(containerName).let { viewModel ->
+            viewModel.upDestinationSceneKey
+                .map { sceneKey -> destinationScenes(up = sceneKey) }
+                .stateIn(
+                    scope = applicationScope,
+                    started = SharingStarted.Eagerly,
+                    initialValue = destinationScenes(up = viewModel.upDestinationSceneKey.value),
+                )
+        }
 
     @Composable
     override fun Content(
+        containerName: String,
         modifier: Modifier,
     ) {
         ShadeScene(
-            viewModel = viewModel,
+            viewModel = getOrCreateViewModelSingleton(containerName),
             modifier = modifier,
         )
     }
@@ -79,6 +86,13 @@
             UserAction.Swipe(Direction.DOWN) to SceneModel(SceneKey.QuickSettings),
         )
     }
+
+    private fun getOrCreateViewModelSingleton(
+        containerName: String,
+    ): ShadeSceneViewModel {
+        return unsafeViewModel
+            ?: viewModelFactory.create(containerName).also { unsafeViewModel = it }
+    }
 }
 
 @Composable
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
index eaa8ed5..c6528d0 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
@@ -20,6 +20,8 @@
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.bouncer.domain.interactor.BouncerInteractor
 import com.android.systemui.dagger.qualifiers.Application
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
@@ -34,7 +36,7 @@
     @Application private val applicationContext: Context,
     @Application private val applicationScope: CoroutineScope,
     interactorFactory: BouncerInteractor.Factory,
-    containerName: String,
+    @Assisted containerName: String,
 ) {
     private val interactor: BouncerInteractor = interactorFactory.create(containerName)
 
@@ -94,4 +96,11 @@
             else -> null
         }
     }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(
+            containerName: String,
+        ): BouncerViewModel
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt b/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
index 0ed8b21..6d7455f 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.scene
 
-import com.android.systemui.scene.shared.page.SceneModule
+import com.android.systemui.scene.ui.composable.SceneModule
 import dagger.Module
 
 @Module(
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scene.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scene.kt
index 435ff4b..354de8a 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scene.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scene.kt
@@ -59,7 +59,7 @@
      * The API is designed such that it's possible to emit ever-changing values for each
      * [UserAction] to enable, disable, or change the destination scene of a given user action.
      */
-    fun destinationScenes(): StateFlow<Map<UserAction, SceneModel>> =
+    fun destinationScenes(containerName: String): StateFlow<Map<UserAction, SceneModel>> =
         MutableStateFlow(emptyMap<UserAction, SceneModel>()).asStateFlow()
 }
 
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 afc0531..a4daafc 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
@@ -29,7 +29,7 @@
 @AssistedInject
 constructor(
     private val interactor: SceneInteractor,
-    @Assisted private val containerName: String,
+    @Assisted val containerName: String,
 ) {
     /**
      * Keys of all scenes in the container.