Merge "Move shade to the default display when the keyguard is visible" into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt
index ef1ae09..fd9f5f0 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.display.data.repository.display
 import com.android.systemui.display.data.repository.displayRepository
+import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.kosmos.useUnconfinedTestDispatcher
 import com.android.systemui.testKosmos
@@ -38,17 +39,31 @@
 class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
     private val kosmos = testKosmos().useUnconfinedTestDispatcher()
     private val testScope = kosmos.testScope
+    private val keyguardRepository = kosmos.fakeKeyguardRepository
     private val displayRepository = kosmos.displayRepository
-    val underTest = StatusBarTouchShadeDisplayPolicy(displayRepository, testScope.backgroundScope)
+
+    private fun createUnderTest(
+        shadeOnDefaultDisplayWhenLocked: Boolean = false
+    ): StatusBarTouchShadeDisplayPolicy {
+        return StatusBarTouchShadeDisplayPolicy(
+            displayRepository,
+            keyguardRepository,
+            testScope.backgroundScope,
+            shadeOnDefaultDisplayWhenLocked = shadeOnDefaultDisplayWhenLocked,
+        )
+    }
 
     @Test
     fun displayId_defaultToDefaultDisplay() {
+        val underTest = createUnderTest()
+
         assertThat(underTest.displayId.value).isEqualTo(Display.DEFAULT_DISPLAY)
     }
 
     @Test
     fun onStatusBarTouched_called_updatesDisplayId() =
         testScope.runTest {
+            val underTest = createUnderTest()
             val displayId by collectLastValue(underTest.displayId)
 
             displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
@@ -60,6 +75,7 @@
     @Test
     fun onStatusBarTouched_notExistentDisplay_displayIdNotUpdated() =
         testScope.runTest {
+            val underTest = createUnderTest()
             val displayIds by collectValues(underTest.displayId)
             assertThat(displayIds).isEqualTo(listOf(Display.DEFAULT_DISPLAY))
 
@@ -72,6 +88,7 @@
     @Test
     fun onStatusBarTouched_afterDisplayRemoved_goesBackToDefaultDisplay() =
         testScope.runTest {
+            val underTest = createUnderTest()
             val displayId by collectLastValue(underTest.displayId)
 
             displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
@@ -83,4 +100,40 @@
 
             assertThat(displayId).isEqualTo(Display.DEFAULT_DISPLAY)
         }
+
+    @Test
+    fun onStatusBarTouched_afterKeyguardVisible_goesBackToDefaultDisplay() =
+        testScope.runTest {
+            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
+            val displayId by collectLastValue(underTest.displayId)
+
+            displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
+            underTest.onStatusBarTouched(2)
+
+            assertThat(displayId).isEqualTo(2)
+
+            keyguardRepository.setKeyguardShowing(true)
+
+            assertThat(displayId).isEqualTo(Display.DEFAULT_DISPLAY)
+        }
+
+    @Test
+    fun onStatusBarTouched_afterKeyguardHides_goesBackToPreviousDisplay() =
+        testScope.runTest {
+            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
+            val displayId by collectLastValue(underTest.displayId)
+
+            displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
+            underTest.onStatusBarTouched(2)
+
+            assertThat(displayId).isEqualTo(2)
+
+            keyguardRepository.setKeyguardShowing(true)
+
+            assertThat(displayId).isEqualTo(Display.DEFAULT_DISPLAY)
+
+            keyguardRepository.setKeyguardShowing(false)
+
+            assertThat(displayId).isEqualTo(2)
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt
index a002aa5..0954e5e 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt
@@ -49,6 +49,7 @@
 import dagger.multibindings.ClassKey
 import dagger.multibindings.IntoMap
 import javax.inject.Provider
+import javax.inject.Qualifier
 
 /**
  * Module responsible for managing display-specific components and resources for the notification
@@ -237,9 +238,23 @@
             CoreStartable.NOP
         }
     }
+
+    @Provides
+    @ShadeOnDefaultDisplayWhenLocked
+    fun provideShadeOnDefaultDisplayWhenLocked(): Boolean = true
 }
 
 @Module
 internal interface OptionalShadeDisplayAwareBindings {
     @BindsOptionalOf fun bindOptionalOfWindowRootView(): WindowRootView
 }
+
+/**
+ * Annotates the boolean value that defines whether the shade window should go back to the default
+ * display when the keyguard is visible.
+ *
+ * As of today (Dec 2024), This is a configuration parameter provided in the dagger graph as the
+ * final policy around keyguard display is still under discussion, and will be evaluated based on
+ * how well this solution behaves from the performance point of view.
+ */
+@Qualifier @Retention(AnnotationRetention.RUNTIME) annotation class ShadeOnDefaultDisplayWhenLocked
diff --git a/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt b/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt
index 22e9487..30b086f 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt
@@ -22,34 +22,55 @@
 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.keyguard.data.repository.KeyguardRepository
+import com.android.systemui.shade.ShadeOnDefaultDisplayWhenLocked
 import com.android.systemui.shade.shared.flag.ShadeWindowGoesAround
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 
 /**
  * Moves the shade on the last display that received a status bar touch.
  *
- * If the display is removed, falls back to the default one.
+ * If the display is removed, falls back to the default one. When [shadeOnDefaultDisplayWhenLocked]
+ * is true, the shade falls back to the default display when the keyguard is visible.
  */
 @SysUISingleton
 class StatusBarTouchShadeDisplayPolicy
 @Inject
-constructor(displayRepository: DisplayRepository, @Background val backgroundScope: CoroutineScope) :
-    ShadeDisplayPolicy {
-    override val name: String
-        get() = "status_bar_latest_touch"
+constructor(
+    displayRepository: DisplayRepository,
+    keyguardRepository: KeyguardRepository,
+    @Background val backgroundScope: CoroutineScope,
+    @ShadeOnDefaultDisplayWhenLocked val shadeOnDefaultDisplayWhenLocked: Boolean,
+) : ShadeDisplayPolicy {
+    override val name: String = "status_bar_latest_touch"
 
     private val currentDisplayId = MutableStateFlow(Display.DEFAULT_DISPLAY)
     private val availableDisplayIds: StateFlow<Set<Int>> = displayRepository.displayIds
 
-    override val displayId: StateFlow<Int>
-        get() = currentDisplayId
+    override val displayId: StateFlow<Int> =
+        if (shadeOnDefaultDisplayWhenLocked) {
+            keyguardRepository.isKeyguardShowing
+                .combine(currentDisplayId) { isKeyguardShowing, currentDisplayId ->
+                    if (isKeyguardShowing) {
+                        Display.DEFAULT_DISPLAY
+                    } else {
+                        currentDisplayId
+                    }
+                }
+                .stateIn(backgroundScope, SharingStarted.WhileSubscribed(), currentDisplayId.value)
+        } else {
+            currentDisplayId
+        }
 
     private var removalListener: Job? = null