Fix SysUIStateInteractor for lazy states of external displays

Iterating the instances in the repository was not always correct: as instances were created lazily, it might have happened that the SysUI State for an external display was not there.

Now we're iterating all displays instead.

+ a performance optimization: changes from the interactor are committed at the end, after all the StateChanges are applied to each SysUI state. This makes sure to minimize the state propagated to launcher.

Bug: 362719719
Bug: 398011576
Test: SysUIStatePerDisplayInteractorTest
Flag: com.android.systemui.shade_window_goes_around
Change-Id: Ib0178a332a1e1f05d70ba3da52334497ba766557

6 files changed, 62 insertions, 27 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/common/domain/interactor/SysUIStatePerDisplayInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/common/domain/interactor/SysUIStatePerDisplayInteractorTest.kt
index ed9cd98a825a..f64f13d4a9b5 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/common/domain/interactor/SysUIStatePerDisplayInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/common/domain/interactor/SysUIStatePerDisplayInteractorTest.kt
@@ -19,6 +19,7 @@ package com.android.systemui.common.domain.interactor
 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.model.StateChange
 import com.android.systemui.model.fakeSysUIStatePerDisplayRepository
 import com.android.systemui.model.sysUiStateFactory
@@ -26,6 +27,7 @@ import com.android.systemui.model.sysuiStateInteractor
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
+import kotlinx.coroutines.runBlocking
 import org.junit.Before
 import org.junit.runner.RunWith
 
@@ -49,6 +51,13 @@ class SysUIStatePerDisplayInteractorTest : SysuiTestCase() {
             add(1, state1)
             add(2, state2)
         }
+        runBlocking {
+            kosmos.displayRepository.apply {
+                addDisplay(0)
+                addDisplay(1)
+                addDisplay(2)
+            }
+        }
     }
 
     @Test
diff --git a/packages/SystemUI/src/com/android/systemui/common/domain/interactor/SysUIStateDisplaysInteractor.kt b/packages/SystemUI/src/com/android/systemui/common/domain/interactor/SysUIStateDisplaysInteractor.kt
index 097d50bb8f9d..9db7b50905f8 100644
--- a/packages/SystemUI/src/com/android/systemui/common/domain/interactor/SysUIStateDisplaysInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/common/domain/interactor/SysUIStateDisplaysInteractor.kt
@@ -16,7 +16,9 @@
 
 package com.android.systemui.common.domain.interactor
 
+import android.util.Log
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.display.data.repository.DisplayRepository
 import com.android.systemui.display.data.repository.PerDisplayRepository
 import com.android.systemui.model.StateChange
 import com.android.systemui.model.SysUiState
@@ -26,20 +28,36 @@ import javax.inject.Inject
 @SysUISingleton
 class SysUIStateDisplaysInteractor
 @Inject
-constructor(private val sysUIStateRepository: PerDisplayRepository<SysUiState>) {
+constructor(
+    private val sysUIStateRepository: PerDisplayRepository<SysUiState>,
+    private val displayRepository: DisplayRepository,
+) {
 
     /**
      * Sets the flags on the given [targetDisplayId] based on the [stateChanges], while making sure
      * that those flags are not set in any other display.
      */
     fun setFlagsExclusivelyToDisplay(targetDisplayId: Int, stateChanges: StateChange) {
-        sysUIStateRepository.forEachInstance { displayId, instance ->
-            if (displayId == targetDisplayId) {
-                stateChanges.applyTo(instance)
-            } else {
-                stateChanges.clearAllChangedFlagsIn(instance)
-            }
+        if (SysUiState.DEBUG) {
+            Log.d(TAG, "Setting flags $stateChanges only for display $targetDisplayId")
         }
+        displayRepository.displays.value
+            .mapNotNull { sysUIStateRepository[it.displayId] }
+            .apply {
+                // Let's first modify all states, without committing changes ...
+                forEach { displaySysUIState ->
+                    if (displaySysUIState.displayId == targetDisplayId) {
+                        stateChanges.applyTo(displaySysUIState)
+                    } else {
+                        stateChanges.clearFrom(displaySysUIState)
+                    }
+                }
+                // ... And commit changes at the end
+                forEach { sysuiState -> sysuiState.commitUpdate() }
+            }
     }
-}
 
+    private companion object {
+        const val TAG = "SysUIStateInteractor"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/PerDisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/PerDisplayRepository.kt
index 04f245e91914..36d3eb51283a 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/PerDisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/PerDisplayRepository.kt
@@ -91,18 +91,6 @@ interface PerDisplayRepository<T> {
 
     /** Debug name for this repository, mainly for tracing and logging. */
     val debugName: String
-
-    /**
-     * Invokes the specified action on each instance held by this repository.
-     *
-     * The action will receive the displayId and the instance associated with that display.
-     * If there is no instance for the display, the action is not called.
-     */
-    fun forEachInstance(action: (Int, T) -> Unit) {
-        displayIds.forEach { displayId ->
-            get(displayId)?.let { instance -> action(displayId, instance) }
-        }
-    }
 }
 
 /**
diff --git a/packages/SystemUI/src/com/android/systemui/model/SysUIStateChange.kt b/packages/SystemUI/src/com/android/systemui/model/SysUIStateChange.kt
index f29b15730986..aaed606f8fb2 100644
--- a/packages/SystemUI/src/com/android/systemui/model/SysUIStateChange.kt
+++ b/packages/SystemUI/src/com/android/systemui/model/SysUIStateChange.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.model
 
+import com.android.systemui.shared.system.QuickStepContract.getSystemUiStateString
+
 /**
  * Represents a set of state changes. A bit can either be set to `true` or `false`.
  *
@@ -43,13 +45,16 @@ class StateChange {
 
     fun hasChanges() = flagsToSet != 0L || flagsToClear != 0L
 
-    /** Applies all changed flags to [sysUiState]. */
+    /**
+     * Applies all changed flags to [sysUiState].
+     *
+     * Note this doesn't call [SysUiState.commitUpdate].
+     */
     fun applyTo(sysUiState: SysUiState) {
         iterateBits(flagsToSet or flagsToClear) { bit ->
             val isBitSetInNewState = flagsToSet and bit != 0L
             sysUiState.setFlag(bit, isBitSetInNewState)
         }
-        sysUiState.commitUpdate()
     }
 
     fun applyTo(sysUiState: Long): Long {
@@ -69,14 +74,25 @@ class StateChange {
         }
     }
 
-    /** Clears all the flags changed in a [sysUiState] */
-    fun clearAllChangedFlagsIn(sysUiState: SysUiState) {
+    /**
+     * Clears all the flags changed in a [sysUiState].
+     *
+     * Note this doesn't call [SysUiState.commitUpdate].
+     */
+    fun clearFrom(sysUiState: SysUiState) {
         iterateBits(flagsToSet or flagsToClear) { bit -> sysUiState.setFlag(bit, false) }
-        sysUiState.commitUpdate()
     }
 
     fun clear() {
         flagsToSet = 0
         flagsToClear = 0
     }
+
+    override fun toString(): String {
+        return """StateChange(flagsToSet=${getSystemUiStateString(flagsToSet)}, flagsToClear=${
+            getSystemUiStateString(
+                flagsToClear
+            )
+        })"""
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt b/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt
index 53105b2c0f6a..663355941613 100644
--- a/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt
+++ b/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt
@@ -74,6 +74,9 @@ interface SysUiState : Dumpable {
      */
     fun destroy()
 
+    /** The display ID this instances is associated with */
+    val displayId: Int
+
     companion object {
         const val DEBUG: Boolean = false
     }
@@ -84,7 +87,7 @@ private const val TAG = "SysUIState"
 class SysUiStateImpl
 @AssistedInject
 constructor(
-    @Assisted private val displayId: Int,
+    @Assisted override val displayId: Int,
     private val sceneContainerPlugin: SceneContainerPlugin?,
     private val dumpManager: DumpManager,
     private val stateDispatcher: SysUIStateDispatcher,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt
index 00deaafd7009..dbd1c9ce56fe 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt
@@ -19,6 +19,7 @@ package com.android.systemui.model
 import android.view.Display
 import com.android.systemui.common.domain.interactor.SysUIStateDisplaysInteractor
 import com.android.systemui.display.data.repository.FakePerDisplayRepository
+import com.android.systemui.display.data.repository.displayRepository
 import com.android.systemui.dump.dumpManager
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
@@ -45,5 +46,5 @@ val Kosmos.sysUiStateFactory by Fixture {
 val Kosmos.fakeSysUIStatePerDisplayRepository by Fixture { FakePerDisplayRepository<SysUiState>() }
 
 val Kosmos.sysuiStateInteractor by Fixture {
-    SysUIStateDisplaysInteractor(fakeSysUIStatePerDisplayRepository)
+    SysUIStateDisplaysInteractor(fakeSysUIStatePerDisplayRepository, displayRepository)
 }