Introduce StatusBarInitializerStore

Flag: com.android.systemui.status_bar_connected_displays
Test: Unit tests in this CL
Test: Manually build and run and verify everything works as before
Bug: 367592591
Change-Id: Ia97ed6b380ba5f2a6b907bd29f14841cd3521fe8

10 files changed, 345 insertions, 6 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt
index 0ce7d70ddefa..f64387c95e3d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt
@@ -60,6 +60,7 @@ class StatusBarInitializerTest : SysuiTestCase() {
 
     val underTest =
         StatusBarInitializerImpl(
+            displayId = context.displayId,
             statusBarWindowControllerStore = windowControllerStore,
             collapsedStatusBarFragmentProvider = { mock(CollapsedStatusBarFragment::class.java) },
             creationListeners = setOf(),
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/qualifiers/Default.java b/packages/SystemUI/src/com/android/systemui/dagger/qualifiers/Default.java
new file mode 100644
index 000000000000..1950d6bac251
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/dagger/qualifiers/Default.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.dagger.qualifiers;
+
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.Retention;
+
+import javax.inject.Qualifier;
+
+@Qualifier
+@Documented
+@Retention(RUNTIME)
+public @interface Default {
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/core/StatusBarInitializer.kt b/packages/SystemUI/src/com/android/systemui/statusbar/core/StatusBarInitializer.kt
index a0629f97f675..7eff8124368b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/core/StatusBarInitializer.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/core/StatusBarInitializer.kt
@@ -18,7 +18,6 @@ package com.android.systemui.statusbar.core
 import android.app.Fragment
 import androidx.annotation.VisibleForTesting
 import com.android.systemui.CoreStartable
-import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.fragments.FragmentHostManager
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.core.StatusBarInitializer.OnStatusBarViewInitializedListener
@@ -28,8 +27,10 @@ import com.android.systemui.statusbar.phone.PhoneStatusBarViewController
 import com.android.systemui.statusbar.phone.fragment.CollapsedStatusBarFragment
 import com.android.systemui.statusbar.phone.fragment.dagger.StatusBarFragmentComponent
 import com.android.systemui.statusbar.window.StatusBarWindowControllerStore
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import java.lang.IllegalStateException
-import javax.inject.Inject
 import javax.inject.Provider
 
 /**
@@ -65,12 +66,16 @@ interface StatusBarInitializer {
             statusBarTransitions: PhoneStatusBarTransitions,
         )
     }
+
+    interface Factory {
+        fun create(displayId: Int): StatusBarInitializer
+    }
 }
 
-@SysUISingleton
 class StatusBarInitializerImpl
-@Inject
+@AssistedInject
 constructor(
+    @Assisted private val displayId: Int,
     private val statusBarWindowControllerStore: StatusBarWindowControllerStore,
     private val collapsedStatusBarFragmentProvider: Provider<CollapsedStatusBarFragment>,
     private val creationListeners: Set<@JvmSuppressWildcards OnStatusBarViewInitializedListener>,
@@ -137,4 +142,9 @@ constructor(
             )
             .commit()
     }
+
+    @AssistedFactory
+    interface Factory : StatusBarInitializer.Factory {
+        override fun create(displayId: Int): StatusBarInitializerImpl
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/core/StatusBarInitializerStore.kt b/packages/SystemUI/src/com/android/systemui/statusbar/core/StatusBarInitializerStore.kt
new file mode 100644
index 000000000000..8d044bb9ce87
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/core/StatusBarInitializerStore.kt
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.core
+
+import android.view.Display
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.display.data.repository.DisplayRepository
+import java.util.concurrent.ConcurrentHashMap
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
+
+/** Provides per display instances of [StatusBarInitializer]. */
+interface StatusBarInitializerStore {
+    /**
+     * The instance for the default/main display of the device. For example, on a phone or a tablet,
+     * the default display is the internal/built-in display of the device.
+     *
+     * Note that the id of the default display is [Display.DEFAULT_DISPLAY].
+     */
+    val defaultDisplay: StatusBarInitializer
+
+    /**
+     * Returns an instance for a specific display id.
+     *
+     * @throws IllegalArgumentException if [displayId] doesn't match the id of any existing
+     *   displays.
+     */
+    fun forDisplay(displayId: Int): StatusBarInitializer
+}
+
+@SysUISingleton
+class MultiDisplayStatusBarInitializerStore
+@Inject
+constructor(
+    @Background private val backgroundApplicationScope: CoroutineScope,
+    private val factory: StatusBarInitializer.Factory,
+    private val displayRepository: DisplayRepository,
+) : StatusBarInitializerStore, CoreStartable {
+
+    init {
+        StatusBarConnectedDisplays.assertInNewMode()
+    }
+
+    private val perDisplayInitializers = ConcurrentHashMap<Int, StatusBarInitializer>()
+
+    override val defaultDisplay: StatusBarInitializer
+        get() = forDisplay(Display.DEFAULT_DISPLAY)
+
+    override fun forDisplay(displayId: Int): StatusBarInitializer {
+        if (displayRepository.getDisplay(displayId) == null) {
+            throw IllegalArgumentException("Display with id $displayId doesn't exist.")
+        }
+        return perDisplayInitializers.computeIfAbsent(displayId) { factory.create(displayId) }
+    }
+
+    override fun start() {
+        backgroundApplicationScope.launch(
+            CoroutineName("MultiDisplayStatusBarInitializerStore#start")
+        ) {
+            displayRepository.displayRemovalEvent.collect { removedDisplayId ->
+                perDisplayInitializers.remove(removedDisplayId)
+            }
+        }
+    }
+}
+
+@SysUISingleton
+class SingleDisplayStatusBarInitializerStore
+@Inject
+constructor(factory: StatusBarInitializerImpl.Factory) : StatusBarInitializerStore {
+
+    init {
+        StatusBarConnectedDisplays.assertInLegacyMode()
+    }
+
+    private val defaultInstance = factory.create(Display.DEFAULT_DISPLAY)
+
+    override val defaultDisplay: StatusBarInitializer = defaultInstance
+
+    override fun forDisplay(displayId: Int): StatusBarInitializer = defaultDisplay
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/dagger/StatusBarPhoneModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/dagger/StatusBarPhoneModule.kt
index 5b0319883b5f..5f864e5dc53a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/dagger/StatusBarPhoneModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/dagger/StatusBarPhoneModule.kt
@@ -15,12 +15,18 @@
  */
 package com.android.systemui.statusbar.phone.dagger
 
+import android.view.Display
 import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Default
 import com.android.systemui.statusbar.CommandQueue
 import com.android.systemui.statusbar.core.CommandQueueInitializer
+import com.android.systemui.statusbar.core.MultiDisplayStatusBarInitializerStore
+import com.android.systemui.statusbar.core.SingleDisplayStatusBarInitializerStore
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
 import com.android.systemui.statusbar.core.StatusBarInitializer
 import com.android.systemui.statusbar.core.StatusBarInitializerImpl
+import com.android.systemui.statusbar.core.StatusBarInitializerStore
 import com.android.systemui.statusbar.core.StatusBarOrchestrator
 import com.android.systemui.statusbar.core.StatusBarSimpleFragment
 import com.android.systemui.statusbar.phone.CentralSurfacesCommandQueueCallbacks
@@ -47,15 +53,31 @@ interface StatusBarPhoneModule {
         impl: CentralSurfacesCommandQueueCallbacks
     ): CommandQueue.Callbacks
 
+    @Binds
+    fun initializerFactory(
+        implFactory: StatusBarInitializerImpl.Factory
+    ): StatusBarInitializer.Factory
+
     /** Binds {@link StatusBarInitializer} as a {@link CoreStartable}. */
     @Binds
     @IntoMap
     @ClassKey(StatusBarInitializerImpl::class)
-    fun bindStatusBarInitializer(impl: StatusBarInitializerImpl): CoreStartable
+    fun bindStatusBarInitializer(@Default impl: StatusBarInitializerImpl): CoreStartable
 
-    @Binds fun statusBarInitializer(impl: StatusBarInitializerImpl): StatusBarInitializer
+    @Binds fun statusBarInitializer(@Default impl: StatusBarInitializerImpl): StatusBarInitializer
 
     companion object {
+        // Dagger doesn't support providing AssistedInject types, without a qualifier. Using the
+        // Default qualifier for this reason.
+        @Default
+        @Provides
+        @SysUISingleton
+        fun statusBarInitializerImpl(
+            implFactory: StatusBarInitializerImpl.Factory
+        ): StatusBarInitializerImpl {
+            return implFactory.create(displayId = Display.DEFAULT_DISPLAY)
+        }
+
         @Provides
         @SysUISingleton
         @IntoMap
@@ -83,5 +105,32 @@ interface StatusBarPhoneModule {
                 CoreStartable.NOP
             }
         }
+
+        @Provides
+        @SysUISingleton
+        @IntoMap
+        @ClassKey(StatusBarInitializerStore::class)
+        fun initializerStoreAsCoreStartable(
+            multiDisplayStoreLazy: Lazy<MultiDisplayStatusBarInitializerStore>
+        ): CoreStartable {
+            return if (StatusBarConnectedDisplays.isEnabled) {
+                multiDisplayStoreLazy.get()
+            } else {
+                CoreStartable.NOP
+            }
+        }
+
+        @Provides
+        @SysUISingleton
+        fun initializerStore(
+            singleDisplayStoreLazy: Lazy<SingleDisplayStatusBarInitializerStore>,
+            multiDisplayStoreLazy: Lazy<MultiDisplayStatusBarInitializerStore>,
+        ): StatusBarInitializerStore {
+            return if (StatusBarConnectedDisplays.isEnabled) {
+                multiDisplayStoreLazy.get()
+            } else {
+                singleDisplayStoreLazy.get()
+            }
+        }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/core/MultiDisplayStatusBarInitializerStoreTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/core/MultiDisplayStatusBarInitializerStoreTest.kt
new file mode 100644
index 000000000000..0d1d37af7e5b
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/core/MultiDisplayStatusBarInitializerStoreTest.kt
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.core
+
+import android.platform.test.annotations.EnableFlags
+import android.view.Display
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.display.data.repository.displayRepository
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.unconfinedTestDispatcher
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+@EnableFlags(StatusBarConnectedDisplays.FLAG_NAME)
+class MultiDisplayStatusBarInitializerStoreTest : SysuiTestCase() {
+
+    private val kosmos =
+        testKosmos().also {
+            // Using unconfinedTestDispatcher to avoid having to call `runCurrent` in the tests.
+            it.testDispatcher = it.unconfinedTestDispatcher
+        }
+    private val testScope = kosmos.testScope
+    private val fakeDisplayRepository = kosmos.displayRepository
+    private val store = kosmos.multiDisplayStatusBarInitializerStore
+
+    @Before
+    fun start() {
+        store.start()
+    }
+
+    @Before
+    fun addDisplays() = runBlocking {
+        fakeDisplayRepository.addDisplay(DEFAULT_DISPLAY_ID)
+        fakeDisplayRepository.addDisplay(NON_DEFAULT_DISPLAY_ID)
+    }
+
+    @Test
+    fun forDisplay_defaultDisplay_multipleCalls_returnsSameInstance() =
+        testScope.runTest {
+            val controller = store.defaultDisplay
+
+            assertThat(store.defaultDisplay).isSameInstanceAs(controller)
+        }
+
+    @Test
+    fun forDisplay_nonDefaultDisplay_multipleCalls_returnsSameInstance() =
+        testScope.runTest {
+            val controller = store.forDisplay(NON_DEFAULT_DISPLAY_ID)
+
+            assertThat(store.forDisplay(NON_DEFAULT_DISPLAY_ID)).isSameInstanceAs(controller)
+        }
+
+    @Test
+    fun forDisplay_nonDefaultDisplay_afterDisplayRemoved_returnsNewInstance() =
+        testScope.runTest {
+            val controller = store.forDisplay(NON_DEFAULT_DISPLAY_ID)
+
+            fakeDisplayRepository.removeDisplay(NON_DEFAULT_DISPLAY_ID)
+            fakeDisplayRepository.addDisplay(NON_DEFAULT_DISPLAY_ID)
+
+            assertThat(store.forDisplay(NON_DEFAULT_DISPLAY_ID)).isNotSameInstanceAs(controller)
+        }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun forDisplay_nonExistingDisplayId_throws() =
+        testScope.runTest { store.forDisplay(NON_EXISTING_DISPLAY_ID) }
+
+    companion object {
+        private const val DEFAULT_DISPLAY_ID = Display.DEFAULT_DISPLAY
+        private const val NON_DEFAULT_DISPLAY_ID = Display.DEFAULT_DISPLAY + 1
+        private const val NON_EXISTING_DISPLAY_ID = Display.DEFAULT_DISPLAY + 2
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
index be97d9551412..40b8cdafa813 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
@@ -513,6 +513,7 @@ public class CentralSurfacesImplTest extends SysuiTestCase {
                 mLightBarController,
                 mAutoHideController,
                 new StatusBarInitializerImpl(
+                        mContext.getDisplayId(),
                         mStatusBarWindowControllerStore,
                         mCollapsedStatusBarFragmentProvider,
                         emptySet()),
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
index fcc83b3e579f..78ea70086605 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
@@ -53,6 +53,10 @@ class FakeDisplayRepository @Inject constructor() : DisplayRepository {
     private val displayAdditionEventFlow = MutableSharedFlow<Display?>(replay = 0)
     private val displayRemovalEventFlow = MutableSharedFlow<Int>(replay = 0)
 
+    suspend fun addDisplay(displayId: Int, type: Int = Display.TYPE_EXTERNAL) {
+        addDisplay(display(type, id = displayId))
+    }
+
     suspend fun addDisplay(display: Display) {
         flow.value += display
         displayAdditionEventFlow.emit(display)
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/core/FakeStatusBarInitializerFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/core/FakeStatusBarInitializerFactory.kt
new file mode 100644
index 000000000000..73ed228f5aaa
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/core/FakeStatusBarInitializerFactory.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.core
+
+import com.android.systemui.statusbar.phone.PhoneStatusBarTransitions
+import com.android.systemui.statusbar.phone.PhoneStatusBarViewController
+
+class FakeStatusBarInitializerFactory(
+    private val statusBarViewController: PhoneStatusBarViewController,
+    private val statusBarTransitions: PhoneStatusBarTransitions,
+) : StatusBarInitializer.Factory {
+
+    override fun create(displayId: Int): StatusBarInitializer =
+        FakeStatusBarInitializer(statusBarViewController, statusBarTransitions)
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/core/StatusBarInitializerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/core/StatusBarInitializerKosmos.kt
index d10320004454..7ad715ba5a89 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/core/StatusBarInitializerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/core/StatusBarInitializerKosmos.kt
@@ -16,7 +16,9 @@
 
 package com.android.systemui.statusbar.core
 
+import com.android.systemui.display.data.repository.displayRepository
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.statusbar.phone.phoneStatusBarTransitions
 import com.android.systemui.statusbar.phone.phoneStatusBarViewController
 
@@ -26,3 +28,20 @@ val Kosmos.fakeStatusBarInitializer by
     }
 
 var Kosmos.statusBarInitializer by Kosmos.Fixture { fakeStatusBarInitializer }
+
+val Kosmos.fakeStatusBarInitializerFactory by
+    Kosmos.Fixture {
+        FakeStatusBarInitializerFactory(phoneStatusBarViewController, phoneStatusBarTransitions)
+    }
+
+var Kosmos.statusBarInitializerFactory: StatusBarInitializer.Factory by
+    Kosmos.Fixture { fakeStatusBarInitializerFactory }
+
+val Kosmos.multiDisplayStatusBarInitializerStore by
+    Kosmos.Fixture {
+        MultiDisplayStatusBarInitializerStore(
+            applicationCoroutineScope,
+            fakeStatusBarInitializerFactory,
+            displayRepository,
+        )
+    }