Merge "Fix race in DisplayRepository tests" into main
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
index 1f5878b..a327e4a 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
@@ -28,7 +28,6 @@
 import android.view.Display
 import com.android.app.tracing.FlowTracing.traceEach
 import com.android.app.tracing.traceSection
-import com.android.systemui.Flags
 import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
@@ -51,7 +50,6 @@
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.scan
-import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 
 /** Provides a [Flow] of [Display] as returned by [DisplayManager]. */
@@ -102,7 +100,7 @@
     private val displayManager: DisplayManager,
     @Background backgroundHandler: Handler,
     @Background bgApplicationScope: CoroutineScope,
-    @Background backgroundCoroutineDispatcher: CoroutineDispatcher
+    @Background backgroundCoroutineDispatcher: CoroutineDispatcher,
 ) : DisplayRepository {
     private val allDisplayEvents: Flow<DisplayEvent> =
         conflatedCallbackFlow {
@@ -139,70 +137,55 @@
     override val displayAdditionEvent: Flow<Display?> =
         allDisplayEvents.filterIsInstance<DisplayEvent.Added>().map { getDisplay(it.displayId) }
 
-    // TODO: b/345472038 - Delete after the flag is ramped up.
-    private val oldEnabledDisplays: Flow<Set<Display>> =
-        allDisplayEvents
-            .map { getDisplays() }
-            .shareIn(bgApplicationScope, started = SharingStarted.WhileSubscribed(), replay = 1)
+    // This is necessary because there might be multiple displays, and we could
+    // have missed events for those added before this process or flow started.
+    // Note it causes a binder call from the main thread (it's traced).
+    private val initialDisplays: Set<Display> =
+        traceSection("$TAG#initialDisplays") { displayManager.displays?.toSet() ?: emptySet() }
+    private val initialDisplayIds = initialDisplays.map { display -> display.displayId }.toSet()
 
     /** Propagate to the listeners only enabled displays */
     private val enabledDisplayIds: Flow<Set<Int>> =
-        if (Flags.enableEfficientDisplayRepository()) {
-                allDisplayEvents
-                    .scan(initial = emptySet()) { previousIds: Set<Int>, event: DisplayEvent ->
-                        val id = event.displayId
-                        when (event) {
-                            is DisplayEvent.Removed -> previousIds - id
-                            is DisplayEvent.Added,
-                            is DisplayEvent.Changed -> previousIds + id
-                        }
-                    }
-                    .distinctUntilChanged()
-                    .stateIn(
-                        bgApplicationScope,
-                        SharingStarted.WhileSubscribed(),
-                        // This is necessary because there might be multiple displays, and we could
-                        // have missed events for those added before this process or flow started.
-                        // Note it causes a binder call from the main thread (it's traced).
-                        getDisplays().map { display -> display.displayId }.toSet(),
-                    )
-            } else {
-                oldEnabledDisplays.map { enabledDisplaysSet ->
-                    enabledDisplaysSet.map { it.displayId }.toSet()
+        allDisplayEvents
+            .scan(initial = initialDisplayIds) { previousIds: Set<Int>, event: DisplayEvent ->
+                val id = event.displayId
+                when (event) {
+                    is DisplayEvent.Removed -> previousIds - id
+                    is DisplayEvent.Added,
+                    is DisplayEvent.Changed -> previousIds + id
                 }
             }
+            .distinctUntilChanged()
+            .stateIn(bgApplicationScope, SharingStarted.WhileSubscribed(), initialDisplayIds)
             .debugLog("enabledDisplayIds")
 
     private val defaultDisplay by lazy {
         getDisplay(Display.DEFAULT_DISPLAY) ?: error("Unable to get default display.")
     }
+
     /**
      * Represents displays that went though the [DisplayListener.onDisplayAdded] callback.
      *
      * Those are commonly the ones provided by [DisplayManager.getDisplays] by default.
      */
     private val enabledDisplays: Flow<Set<Display>> =
-        if (Flags.enableEfficientDisplayRepository()) {
-            enabledDisplayIds
-                .mapElementsLazily { displayId -> getDisplay(displayId) }
-                .onEach {
-                    if (it.isEmpty()) Log.wtf(TAG, "No enabled displays. This should never happen.")
-                }
-                .flowOn(backgroundCoroutineDispatcher)
-                .debugLog("enabledDisplays")
-                .stateIn(
-                    bgApplicationScope,
-                    started = SharingStarted.WhileSubscribed(),
-                    // This triggers a single binder call on the UI thread per process. The
-                    // alternative would be to use sharedFlows, but they are prohibited due to
-                    // performance concerns.
-                    // Ultimately, this is a trade-off between a one-time UI thread binder call and
-                    // the constant overhead of sharedFlows.
-                    initialValue = getDisplays()
-                )
-        } else {
-            oldEnabledDisplays
-        }
+        enabledDisplayIds
+            .mapElementsLazily { displayId -> getDisplay(displayId) }
+            .onEach {
+                if (it.isEmpty()) Log.wtf(TAG, "No enabled displays. This should never happen.")
+            }
+            .flowOn(backgroundCoroutineDispatcher)
+            .debugLog("enabledDisplays")
+            .stateIn(
+                bgApplicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                // This triggers a single binder call on the UI thread per process. The
+                // alternative would be to use sharedFlows, but they are prohibited due to
+                // performance concerns.
+                // Ultimately, this is a trade-off between a one-time UI thread binder call and
+                // the constant overhead of sharedFlows.
+                initialValue = initialDisplays,
+            )
 
     /**
      * Represents displays that went though the [DisplayListener.onDisplayAdded] callback.
@@ -211,10 +194,7 @@
      */
     override val displays: Flow<Set<Display>> = enabledDisplays
 
-    private fun getDisplays(): Set<Display> =
-        traceSection("$TAG#getDisplays()") { displayManager.displays?.toSet() ?: emptySet() }
-
-    private val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
+    val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
     private val ignoredDisplayIds: Flow<Set<Int>> = _ignoredDisplayIds.debugLog("ignoredDisplayIds")
 
     private fun getInitialConnectedDisplays(): Set<Int> =
@@ -271,7 +251,7 @@
                 // the flow starts being collected. This is to ensure the call to get displays (an
                 // IPC) happens in the background instead of when this object
                 // is instantiated.
-                initialValue = emptySet()
+                initialValue = emptySet(),
             )
 
     private val connectedExternalDisplayIds: Flow<Set<Int>> =
@@ -308,7 +288,7 @@
                         TAG,
                         "combining enabled=$enabledDisplaysIds, " +
                             "connectedExternalDisplayIds=$connectedExternalDisplayIds, " +
-                            "ignored=$ignoredDisplayIds"
+                            "ignored=$ignoredDisplayIds",
                     )
                 }
                 connectedExternalDisplayIds - enabledDisplaysIds - ignoredDisplayIds
@@ -382,7 +362,7 @@
             val previousSet: Set<T>,
             // Caches T values from the previousSet that were already converted to V
             val valueMap: Map<T, V>,
-            val resultSet: Set<V>
+            val resultSet: Set<V>,
         )
 
         val emptyInitialState = State(emptySet<T>(), emptyMap(), emptySet<V>())
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
index 76539d7..633efd8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
@@ -63,21 +63,27 @@
     private val defaultDisplay =
         display(type = TYPE_INTERNAL, id = DEFAULT_DISPLAY, state = Display.STATE_ON)
 
-    private lateinit var displayRepository: DisplayRepositoryImpl
+    // This is Lazy as displays could be set before the instance is created, and we want to verify
+    // that the initial state (soon after construction) contains the expected ones set in every
+    // test.
+    private val displayRepository: DisplayRepositoryImpl by lazy {
+        DisplayRepositoryImpl(
+                displayManager,
+                testHandler,
+                TestScope(UnconfinedTestDispatcher()),
+                UnconfinedTestDispatcher(),
+            )
+            .also {
+                verify(displayManager, never()).registerDisplayListener(any(), any())
+                // It needs to be called, just once, for the initial value.
+                verify(displayManager).getDisplays()
+            }
+    }
 
     @Before
     fun setup() {
         setDisplays(listOf(defaultDisplay))
         setAllDisplaysIncludingDisabled(DEFAULT_DISPLAY)
-        displayRepository =
-            DisplayRepositoryImpl(
-                displayManager,
-                testHandler,
-                TestScope(UnconfinedTestDispatcher()),
-                UnconfinedTestDispatcher()
-            )
-        verify(displayManager, never()).registerDisplayListener(any(), any())
-        verify(displayManager, never()).getDisplays(any())
     }
 
     @Test
@@ -502,7 +508,7 @@
             .registerDisplayListener(
                 connectedDisplayListener.capture(),
                 eq(testHandler),
-                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED)
+                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED),
             )
         return flowValue
     }
@@ -522,7 +528,7 @@
                     DisplayManager.EVENT_FLAG_DISPLAY_ADDED or
                         DisplayManager.EVENT_FLAG_DISPLAY_CHANGED or
                         DisplayManager.EVENT_FLAG_DISPLAY_REMOVED
-                )
+                ),
             )
     }