Merge changes I73ee8e68,I5da93401 into main

* changes:
  Introduce StatusBarWindowControllerStore
  Create DisplayWindowPropertiesRepository
diff --git a/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt b/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
index b56ed8c..589dbf9 100644
--- a/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
@@ -24,6 +24,8 @@
 import com.android.systemui.display.data.repository.DisplayRepositoryImpl
 import com.android.systemui.display.data.repository.DisplayScopeRepository
 import com.android.systemui.display.data.repository.DisplayScopeRepositoryImpl
+import com.android.systemui.display.data.repository.DisplayWindowPropertiesRepository
+import com.android.systemui.display.data.repository.DisplayWindowPropertiesRepositoryImpl
 import com.android.systemui.display.data.repository.FocusedDisplayRepository
 import com.android.systemui.display.data.repository.FocusedDisplayRepositoryImpl
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor
@@ -58,6 +60,11 @@
 
     @Binds fun displayScopeRepository(impl: DisplayScopeRepositoryImpl): DisplayScopeRepository
 
+    @Binds
+    fun displayWindowPropertiesRepository(
+        impl: DisplayWindowPropertiesRepositoryImpl
+    ): DisplayWindowPropertiesRepository
+
     companion object {
         @Provides
         @SysUISingleton
@@ -72,5 +79,19 @@
                 CoreStartable.NOP
             }
         }
+
+        @Provides
+        @SysUISingleton
+        @IntoMap
+        @ClassKey(DisplayWindowPropertiesRepository::class)
+        fun displayWindowPropertiesRepoAsCoreStartable(
+            repoLazy: Lazy<DisplayWindowPropertiesRepositoryImpl>
+        ): CoreStartable {
+            return if (StatusBarConnectedDisplays.isEnabled) {
+                return repoLazy.get()
+            } else {
+                CoreStartable.NOP
+            }
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepository.kt
new file mode 100644
index 0000000..88d3a28
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepository.kt
@@ -0,0 +1,115 @@
+/*
+ * 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.display.data.repository
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.view.Display
+import android.view.WindowManager
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.display.shared.model.DisplayWindowProperties
+import com.android.systemui.res.R
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
+import com.google.common.collect.HashBasedTable
+import com.google.common.collect.Table
+import java.io.PrintWriter
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
+
+/** Provides per display instances of [DisplayWindowProperties]. */
+interface DisplayWindowPropertiesRepository {
+
+    /**
+     * Returns a [DisplayWindowProperties] instance for a given display id and window type.
+     *
+     * @throws IllegalArgumentException if no display with the given display id exists.
+     */
+    fun get(
+        displayId: Int,
+        @WindowManager.LayoutParams.WindowType windowType: Int,
+    ): DisplayWindowProperties
+}
+
+@SysUISingleton
+class DisplayWindowPropertiesRepositoryImpl
+@Inject
+constructor(
+    @Background private val backgroundApplicationScope: CoroutineScope,
+    private val globalContext: Context,
+    private val globalWindowManager: WindowManager,
+    private val displayRepository: DisplayRepository,
+) : DisplayWindowPropertiesRepository, CoreStartable {
+
+    init {
+        StatusBarConnectedDisplays.assertInNewMode()
+    }
+
+    private val properties: Table<Int, Int, DisplayWindowProperties> = HashBasedTable.create()
+
+    override fun get(
+        displayId: Int,
+        @WindowManager.LayoutParams.WindowType windowType: Int,
+    ): DisplayWindowProperties {
+        val display =
+            displayRepository.getDisplay(displayId)
+                ?: throw IllegalArgumentException("Display with id $displayId doesn't exist")
+        return properties.get(displayId, windowType)
+            ?: create(display, windowType).also { properties.put(displayId, windowType, it) }
+    }
+
+    override fun start() {
+        backgroundApplicationScope.launch(
+            CoroutineName("DisplayWindowPropertiesRepositoryImpl#start")
+        ) {
+            displayRepository.displayRemovalEvent.collect { removedDisplayId ->
+                properties.row(removedDisplayId).clear()
+            }
+        }
+    }
+
+    private fun create(display: Display, windowType: Int): DisplayWindowProperties {
+        val displayId = display.displayId
+        return if (displayId == Display.DEFAULT_DISPLAY) {
+            // For the default display, we can just reuse the global/application properties.
+            // Creating a window context is expensive, therefore we avoid it.
+            DisplayWindowProperties(
+                displayId = displayId,
+                windowType = windowType,
+                context = globalContext,
+                windowManager = globalWindowManager,
+            )
+        } else {
+            val context = createWindowContext(display, windowType)
+            @SuppressLint("NonInjectedService") // Need to manually get the service
+            val windowManager = context.getSystemService(WindowManager::class.java) as WindowManager
+            DisplayWindowProperties(displayId, windowType, context, windowManager)
+        }
+    }
+
+    private fun createWindowContext(display: Display, windowType: Int): Context =
+        globalContext.createWindowContext(display, windowType, /* options= */ null).also {
+            it.setTheme(R.style.Theme_SystemUI)
+        }
+
+    override fun dump(pw: PrintWriter, args: Array<out String>) {
+        pw.write("perDisplayContexts: $properties")
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/display/shared/model/DisplayWindowProperties.kt b/packages/SystemUI/src/com/android/systemui/display/shared/model/DisplayWindowProperties.kt
new file mode 100644
index 0000000..6acc296
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/display/shared/model/DisplayWindowProperties.kt
@@ -0,0 +1,43 @@
+/*
+ * 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.display.shared.model
+
+import android.content.Context
+import android.view.WindowManager
+
+/** Represents a display specific group of window related properties. */
+data class DisplayWindowProperties(
+    /** The id of the display associated with this instance. */
+    val displayId: Int,
+    /**
+     * The window type that was used to create the [Context] in this instance, using
+     * [Context.createWindowContext]. This is the window type that can be used when adding views to
+     * the [WindowManager] associated with this instance.
+     */
+    @WindowManager.LayoutParams.WindowType val windowType: Int,
+    /**
+     * The display specific [Context] created using [Context.createWindowContext] with window type
+     * associated with this instance.
+     */
+    val context: Context,
+
+    /**
+     * The display specific [WindowManager] instance to be used when adding windows of the type
+     * associated with this instance.
+     */
+    val windowManager: WindowManager,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
index cf238d5..cd1642e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
@@ -22,15 +22,20 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.log.LogBuffer
 import com.android.systemui.log.LogBufferFactory
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
 import com.android.systemui.statusbar.data.StatusBarDataLayerModule
 import com.android.systemui.statusbar.phone.LightBarController
 import com.android.systemui.statusbar.phone.StatusBarSignalPolicy
 import com.android.systemui.statusbar.phone.ongoingcall.OngoingCallController
 import com.android.systemui.statusbar.phone.ongoingcall.OngoingCallLog
 import com.android.systemui.statusbar.ui.SystemBarUtilsProxyImpl
+import com.android.systemui.statusbar.window.MultiDisplayStatusBarWindowControllerStore
+import com.android.systemui.statusbar.window.SingleDisplayStatusBarWindowControllerStore
 import com.android.systemui.statusbar.window.StatusBarWindowController
 import com.android.systemui.statusbar.window.StatusBarWindowControllerImpl
+import com.android.systemui.statusbar.window.StatusBarWindowControllerStore
 import dagger.Binds
+import dagger.Lazy
 import dagger.Module
 import dagger.Provides
 import dagger.multibindings.ClassKey
@@ -62,13 +67,19 @@
     @ClassKey(StatusBarSignalPolicy::class)
     abstract fun bindStatusBarSignalPolicy(impl: StatusBarSignalPolicy): CoreStartable
 
+    @Binds
+    @SysUISingleton
+    abstract fun statusBarWindowControllerFactory(
+        implFactory: StatusBarWindowControllerImpl.Factory
+    ): StatusBarWindowController.Factory
+
     companion object {
 
         @Provides
         @SysUISingleton
-        fun statusBarWindowController(
-            context: Context?,
-            viewCaptureAwareWindowManager: ViewCaptureAwareWindowManager?,
+        fun defaultStatusBarWindowController(
+            context: Context,
+            viewCaptureAwareWindowManager: ViewCaptureAwareWindowManager,
             factory: StatusBarWindowControllerImpl.Factory,
         ): StatusBarWindowController {
             return factory.create(context, viewCaptureAwareWindowManager)
@@ -76,6 +87,33 @@
 
         @Provides
         @SysUISingleton
+        fun windowControllerStore(
+            multiDisplayImplLazy: Lazy<MultiDisplayStatusBarWindowControllerStore>,
+            singleDisplayImplLazy: Lazy<SingleDisplayStatusBarWindowControllerStore>,
+        ): StatusBarWindowControllerStore {
+            return if (StatusBarConnectedDisplays.isEnabled) {
+                multiDisplayImplLazy.get()
+            } else {
+                singleDisplayImplLazy.get()
+            }
+        }
+
+        @Provides
+        @SysUISingleton
+        @IntoMap
+        @ClassKey(MultiDisplayStatusBarWindowControllerStore::class)
+        fun multiDisplayControllerStoreAsCoreStartable(
+            storeLazy: Lazy<MultiDisplayStatusBarWindowControllerStore>
+        ): CoreStartable {
+            return if (StatusBarConnectedDisplays.isEnabled) {
+                storeLazy.get()
+            } else {
+                CoreStartable.NOP
+            }
+        }
+
+        @Provides
+        @SysUISingleton
         @OngoingCallLog
         fun provideOngoingCallLogBuffer(factory: LogBufferFactory): LogBuffer {
             return factory.create("OngoingCall", 75)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowController.kt
index 421e5c4..e8dc934 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowController.kt
@@ -16,8 +16,10 @@
 
 package com.android.systemui.statusbar.window
 
+import android.content.Context
 import android.view.View
 import android.view.ViewGroup
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
 import com.android.systemui.animation.ActivityTransitionAnimator
 import com.android.systemui.fragments.FragmentHostManager
 import java.util.Optional
@@ -73,4 +75,11 @@
      *   this#setForceStatusBarVisible} together and use some sort of ranking system instead.
      */
     fun setOngoingProcessRequiresStatusBarVisible(visible: Boolean)
+
+    interface Factory {
+        fun create(
+            context: Context,
+            viewCaptureAwareWindowManager: ViewCaptureAwareWindowManager,
+        ): StatusBarWindowController
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowControllerImpl.java
index 1ee7cf3..d709e5a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowControllerImpl.java
@@ -354,11 +354,13 @@
     }
 
     @AssistedFactory
-    public interface Factory {
+    public interface Factory extends StatusBarWindowController.Factory {
         /** Creates a new instance. */
+        @NonNull
+        @Override
         StatusBarWindowControllerImpl create(
-                Context context,
-                ViewCaptureAwareWindowManager viewCaptureAwareWindowManager);
+                @NonNull Context context,
+                @NonNull ViewCaptureAwareWindowManager viewCaptureAwareWindowManager);
     }
 
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowControllerStore.kt b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowControllerStore.kt
new file mode 100644
index 0000000..5f30b37
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowControllerStore.kt
@@ -0,0 +1,117 @@
+/*
+ * 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.window
+
+import android.view.Display
+import android.view.WindowManager
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
+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 com.android.systemui.display.data.repository.DisplayWindowPropertiesRepository
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
+import java.util.concurrent.ConcurrentHashMap
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
+
+/** Store that allows to retrieve per display instances of [StatusBarWindowController]. */
+interface StatusBarWindowControllerStore {
+    /**
+     * 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: StatusBarWindowController
+
+    /**
+     * 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): StatusBarWindowController
+}
+
+@SysUISingleton
+class MultiDisplayStatusBarWindowControllerStore
+@Inject
+constructor(
+    @Background private val backgroundApplicationScope: CoroutineScope,
+    private val controllerFactory: StatusBarWindowController.Factory,
+    private val displayWindowPropertiesRepository: DisplayWindowPropertiesRepository,
+    private val viewCaptureAwareWindowManagerFactory: ViewCaptureAwareWindowManager.Factory,
+    private val displayRepository: DisplayRepository,
+) : StatusBarWindowControllerStore, CoreStartable {
+
+    init {
+        StatusBarConnectedDisplays.assertInNewMode()
+    }
+
+    private val perDisplayControllers = ConcurrentHashMap<Int, StatusBarWindowController>()
+
+    override fun start() {
+        backgroundApplicationScope.launch(CoroutineName("StatusBarWindowController#start")) {
+            displayRepository.displayRemovalEvent.collect { displayId ->
+                perDisplayControllers.remove(displayId)
+            }
+        }
+    }
+
+    override val defaultDisplay: StatusBarWindowController
+        get() = forDisplay(Display.DEFAULT_DISPLAY)
+
+    override fun forDisplay(displayId: Int): StatusBarWindowController {
+        if (displayRepository.getDisplay(displayId) == null) {
+            throw IllegalArgumentException("Display with id $displayId doesn't exist.")
+        }
+        return perDisplayControllers.computeIfAbsent(displayId) {
+            createControllerForDisplay(displayId)
+        }
+    }
+
+    private fun createControllerForDisplay(displayId: Int): StatusBarWindowController {
+        val statusBarDisplayContext =
+            displayWindowPropertiesRepository.get(
+                displayId = displayId,
+                windowType = WindowManager.LayoutParams.TYPE_STATUS_BAR,
+            )
+        val viewCaptureAwareWindowManager =
+            viewCaptureAwareWindowManagerFactory.create(statusBarDisplayContext.windowManager)
+        return controllerFactory.create(
+            statusBarDisplayContext.context,
+            viewCaptureAwareWindowManager,
+        )
+    }
+}
+
+@SysUISingleton
+class SingleDisplayStatusBarWindowControllerStore
+@Inject
+constructor(private val controller: StatusBarWindowController) : StatusBarWindowControllerStore {
+
+    init {
+        StatusBarConnectedDisplays.assertInLegacyMode()
+    }
+
+    override val defaultDisplay = controller
+
+    override fun forDisplay(displayId: Int) = controller
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryImplTest.kt
new file mode 100644
index 0000000..ff3186a
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryImplTest.kt
@@ -0,0 +1,158 @@
+/*
+ * 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.display.data.repository
+
+import android.content.testableContext
+import android.platform.test.annotations.EnableFlags
+import android.view.Display
+import android.view.mockWindowManager
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.display.shared.model.DisplayWindowProperties
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.unconfinedTestDispatcher
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
+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
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+
+@EnableFlags(StatusBarConnectedDisplays.FLAG_NAME)
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class DisplayWindowPropertiesRepositoryImplTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos().also { it.testDispatcher = it.unconfinedTestDispatcher }
+    private val fakeDisplayRepository = kosmos.displayRepository
+    private val testScope = kosmos.testScope
+
+    private val applicationContext = kosmos.testableContext
+    private val applicationWindowManager = kosmos.mockWindowManager
+
+    private val repo =
+        DisplayWindowPropertiesRepositoryImpl(
+            kosmos.applicationCoroutineScope,
+            applicationContext,
+            applicationWindowManager,
+            fakeDisplayRepository,
+        )
+
+    @Before
+    fun start() {
+        repo.start()
+    }
+
+    @Before
+    fun addDisplays() = runBlocking {
+        fakeDisplayRepository.addDisplay(createDisplay(DEFAULT_DISPLAY_ID))
+        fakeDisplayRepository.addDisplay(createDisplay(NON_DEFAULT_DISPLAY_ID))
+    }
+
+    @Test
+    fun get_defaultDisplayId_returnsDefaultProperties() =
+        testScope.runTest {
+            val displayContext = repo.get(DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO)
+
+            assertThat(displayContext)
+                .isEqualTo(
+                    DisplayWindowProperties(
+                        displayId = DEFAULT_DISPLAY_ID,
+                        windowType = WINDOW_TYPE_FOO,
+                        context = applicationContext,
+                        windowManager = applicationWindowManager,
+                    )
+                )
+        }
+
+    @Test
+    fun get_nonDefaultDisplayId_returnsNewStatusBarContext() =
+        testScope.runTest {
+            val displayContext = repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO)
+
+            assertThat(displayContext.context).isNotSameInstanceAs(applicationContext)
+        }
+
+    @Test
+    fun get_nonDefaultDisplayId_returnsNewWindowManager() =
+        testScope.runTest {
+            val displayContext = repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO)
+
+            assertThat(displayContext.windowManager).isNotSameInstanceAs(applicationWindowManager)
+        }
+
+    @Test
+    fun get_multipleCallsForDefaultDisplay_returnsSameInstance() =
+        testScope.runTest {
+            val displayContext = repo.get(DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO)
+
+            assertThat(repo.get(DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO))
+                .isSameInstanceAs(displayContext)
+        }
+
+    @Test
+    fun get_multipleCallsForNonDefaultDisplay_returnsSameInstance() =
+        testScope.runTest {
+            val displayContext = repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO)
+
+            assertThat(repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO))
+                .isSameInstanceAs(displayContext)
+        }
+
+    @Test
+    fun get_multipleCalls_differentType_returnsNewInstance() =
+        testScope.runTest {
+            val displayContext = repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO)
+
+            assertThat(repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_BAR))
+                .isNotSameInstanceAs(displayContext)
+        }
+
+    @Test
+    fun get_afterDisplayRemoved_returnsNewInstance() =
+        testScope.runTest {
+            val displayContext = repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO)
+
+            fakeDisplayRepository.removeDisplay(NON_DEFAULT_DISPLAY_ID)
+            fakeDisplayRepository.addDisplay(createDisplay(NON_DEFAULT_DISPLAY_ID))
+
+            assertThat(repo.get(NON_DEFAULT_DISPLAY_ID, WINDOW_TYPE_FOO))
+                .isNotSameInstanceAs(displayContext)
+        }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun get_nonExistingDisplayId_throws() =
+        testScope.runTest { repo.get(NON_EXISTING_DISPLAY_ID, WINDOW_TYPE_FOO) }
+
+    private fun createDisplay(displayId: Int) =
+        mock<Display> { on { getDisplayId() } doReturn displayId }
+
+    companion object {
+        private const val DEFAULT_DISPLAY_ID = Display.DEFAULT_DISPLAY
+        private const val NON_DEFAULT_DISPLAY_ID = DEFAULT_DISPLAY_ID + 1
+        private const val NON_EXISTING_DISPLAY_ID = DEFAULT_DISPLAY_ID + 2
+        private const val WINDOW_TYPE_FOO = 123
+        private const val WINDOW_TYPE_BAR = 321
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/MultiDisplayStatusBarWindowControllerStoreTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/MultiDisplayStatusBarWindowControllerStoreTest.kt
new file mode 100644
index 0000000..faaa4c4
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/MultiDisplayStatusBarWindowControllerStoreTest.kt
@@ -0,0 +1,120 @@
+/*
+ * 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.window
+
+import android.platform.test.annotations.EnableFlags
+import android.view.Display
+import android.view.WindowManager
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
+import com.android.app.viewcapture.mockViewCaptureAwareWindowManager
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.display.data.repository.displayRepository
+import com.android.systemui.display.data.repository.fakeDisplayWindowPropertiesRepository
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.unconfinedTestDispatcher
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
+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
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+@EnableFlags(StatusBarConnectedDisplays.FLAG_NAME)
+class MultiDisplayStatusBarWindowControllerStoreTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos().also { it.testDispatcher = it.unconfinedTestDispatcher }
+    private val testScope = kosmos.testScope
+    private val fakeDisplayRepository = kosmos.displayRepository
+
+    private val store =
+        MultiDisplayStatusBarWindowControllerStore(
+            backgroundApplicationScope = kosmos.applicationCoroutineScope,
+            controllerFactory = kosmos.fakeStatusBarWindowControllerFactory,
+            displayWindowPropertiesRepository = kosmos.fakeDisplayWindowPropertiesRepository,
+            viewCaptureAwareWindowManagerFactory =
+                object : ViewCaptureAwareWindowManager.Factory {
+                    override fun create(
+                        windowManager: WindowManager
+                    ): ViewCaptureAwareWindowManager {
+                        return kosmos.mockViewCaptureAwareWindowManager
+                    }
+                },
+            displayRepository = fakeDisplayRepository,
+        )
+
+    @Before
+    fun start() {
+        store.start()
+    }
+
+    @Before
+    fun addDisplays() = runBlocking {
+        fakeDisplayRepository.addDisplay(createDisplay(DEFAULT_DISPLAY_ID))
+        fakeDisplayRepository.addDisplay(createDisplay(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(createDisplay(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) }
+
+    private fun createDisplay(displayId: Int): Display = mock {
+        on { getDisplayId() } doReturn displayId
+    }
+
+    companion object {
+        private const val DEFAULT_DISPLAY_ID = Display.DEFAULT_DISPLAY
+        private const val NON_DEFAULT_DISPLAY_ID = DEFAULT_DISPLAY_ID + 1
+        private const val NON_EXISTING_DISPLAY_ID = DEFAULT_DISPLAY_ID + 2
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/app/viewcapture/ViewCaptureAwareWindowManagerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/app/viewcapture/ViewCaptureAwareWindowManagerKosmos.kt
new file mode 100644
index 0000000..e1c6699
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/app/viewcapture/ViewCaptureAwareWindowManagerKosmos.kt
@@ -0,0 +1,26 @@
+/*
+ * 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.app.viewcapture
+
+import com.android.systemui.kosmos.Kosmos
+import org.mockito.kotlin.mock
+
+val Kosmos.mockViewCaptureAwareWindowManager by
+    Kosmos.Fixture { mock<ViewCaptureAwareWindowManager>() }
+
+var Kosmos.viewCaptureAwareWindowManager: ViewCaptureAwareWindowManager by
+    Kosmos.Fixture { mockViewCaptureAwareWindowManager }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryKosmos.kt
new file mode 100644
index 0000000..ff4ba61
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryKosmos.kt
@@ -0,0 +1,26 @@
+/*
+ * 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.display.data.repository
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testDispatcher
+
+val Kosmos.fakeDisplayScopeRepository by
+    Kosmos.Fixture { FakeDisplayScopeRepository(testDispatcher) }
+
+var Kosmos.displayScopeRepository: DisplayScopeRepository by
+    Kosmos.Fixture { fakeDisplayScopeRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryKosmos.kt
new file mode 100644
index 0000000..65b18c1
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryKosmos.kt
@@ -0,0 +1,25 @@
+/*
+ * 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.display.data.repository
+
+import com.android.systemui.kosmos.Kosmos
+
+val Kosmos.fakeDisplayWindowPropertiesRepository by
+    Kosmos.Fixture { FakeDisplayWindowPropertiesRepository() }
+
+var Kosmos.displayWindowPropertiesRepository: DisplayWindowPropertiesRepository by
+    Kosmos.Fixture { fakeDisplayWindowPropertiesRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayScopeRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayScopeRepository.kt
new file mode 100644
index 0000000..3c25924
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayScopeRepository.kt
@@ -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.display.data.repository
+
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+
+class FakeDisplayScopeRepository(private val dispatcher: CoroutineDispatcher) :
+    DisplayScopeRepository {
+
+    private val perDisplayScopes = mutableMapOf<Int, CoroutineScope>()
+
+    override fun scopeForDisplay(displayId: Int): CoroutineScope {
+        return perDisplayScopes.computeIfAbsent(displayId) { CoroutineScope(dispatcher) }
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayWindowPropertiesRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayWindowPropertiesRepository.kt
new file mode 100644
index 0000000..9282f27
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayWindowPropertiesRepository.kt
@@ -0,0 +1,37 @@
+/*
+ * 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.display.data.repository
+
+import com.android.systemui.display.shared.model.DisplayWindowProperties
+import com.google.common.collect.HashBasedTable
+import org.mockito.kotlin.mock
+
+class FakeDisplayWindowPropertiesRepository : DisplayWindowPropertiesRepository {
+
+    private val properties = HashBasedTable.create<Int, Int, DisplayWindowProperties>()
+
+    override fun get(displayId: Int, windowType: Int): DisplayWindowProperties {
+        return properties.get(displayId, windowType)
+            ?: DisplayWindowProperties(
+                    displayId = displayId,
+                    windowType = windowType,
+                    context = mock(),
+                    windowManager = mock(),
+                )
+                .also { properties.put(displayId, windowType, it) }
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/FakeStatusBarWindowControllerFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/FakeStatusBarWindowControllerFactory.kt
new file mode 100644
index 0000000..10f328b
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/FakeStatusBarWindowControllerFactory.kt
@@ -0,0 +1,27 @@
+/*
+ * 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.window
+
+import android.content.Context
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
+
+class FakeStatusBarWindowControllerFactory : StatusBarWindowController.Factory {
+    override fun create(
+        context: Context,
+        viewCaptureAwareWindowManager: ViewCaptureAwareWindowManager,
+    ) = FakeStatusBarWindowController()
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/FakeStatusBarWindowControllerStore.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/FakeStatusBarWindowControllerStore.kt
new file mode 100644
index 0000000..d19e322
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/FakeStatusBarWindowControllerStore.kt
@@ -0,0 +1,31 @@
+/*
+ * 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.window
+
+import android.view.Display
+
+class FakeStatusBarWindowControllerStore : StatusBarWindowControllerStore {
+
+    private val perDisplayControllers = mutableMapOf<Int, FakeStatusBarWindowController>()
+
+    override val defaultDisplay
+        get() = forDisplay(Display.DEFAULT_DISPLAY)
+
+    override fun forDisplay(displayId: Int): StatusBarWindowController {
+        return perDisplayControllers.computeIfAbsent(displayId) { FakeStatusBarWindowController() }
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/StatusBarWindowControllerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/StatusBarWindowControllerKosmos.kt
index c198b35..6c6f243 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/StatusBarWindowControllerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/window/StatusBarWindowControllerKosmos.kt
@@ -21,3 +21,15 @@
 val Kosmos.fakeStatusBarWindowController by Kosmos.Fixture { FakeStatusBarWindowController() }
 
 var Kosmos.statusBarWindowController by Kosmos.Fixture { fakeStatusBarWindowController }
+
+val Kosmos.fakeStatusBarWindowControllerStore by
+    Kosmos.Fixture { FakeStatusBarWindowControllerStore() }
+
+var Kosmos.statusBarWindowControllerStore: StatusBarWindowControllerStore by
+    Kosmos.Fixture { fakeStatusBarWindowControllerStore }
+
+val Kosmos.fakeStatusBarWindowControllerFactory by
+    Kosmos.Fixture { FakeStatusBarWindowControllerFactory() }
+
+var Kosmos.statusBarWindowControllerFactory: StatusBarWindowController.Factory by
+    Kosmos.Fixture { fakeStatusBarWindowControllerFactory }