Update peristent repository on visiblity change

Persistent repo should be updated on visiblity change too as there might
be cases where only the visiblity of tasks is impacted. i.e Use minimize
button and unminimize task.

Also handle the case where tasks are marked as not visible eventhough
they should be persisted as visible. E.g leaving desktop and rebooting.

Bug: 376036170
Test: atest DesktopRepositoryTest
Flag: com.android.window.flags.enable_desktop_windowing_persistence
Change-Id: I08f74c944220b6f51d1edb905210665b79eace4f

3 files changed, 53 insertions, 4 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
index fdb01ab0765d..5648feb96b06 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
@@ -294,6 +294,9 @@ class DesktopRepository (
                 taskId, isVisible, displayId)
             logD("VisibleTaskCount has changed from %d to %d", prevCount, newCount)
             notifyVisibleTaskListeners(displayId, newCount)
+            if (Flags.enableDesktopWindowingPersistence()) {
+                updatePersistentRepository(displayId)
+            }
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
index 2d11e02bd3c6..9e646f430c98 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
@@ -50,7 +50,9 @@ class DesktopPersistentRepository(
         DataStoreFactory.create(
             serializer = DesktopPersistentRepositoriesSerializer,
             produceFile = { context.dataStoreFile(DESKTOP_REPOSITORIES_DATASTORE_FILE) },
-            scope = bgCoroutineScope))
+            scope = bgCoroutineScope,
+        ),
+    )
 
     /** Provides `dataStore.data` flow and handles exceptions thrown during collection */
     private val dataStoreFlow: Flow<DesktopPersistentRepositories> =
@@ -116,7 +118,11 @@ class DesktopPersistentRepository(
                 val desktop =
                     getDesktop(currentRepository, desktopId)
                         .toBuilder()
-                        .updateTaskStates(visibleTasks, minimizedTasks)
+                        .updateTaskStates(
+                            visibleTasks,
+                            minimizedTasks,
+                            freeformTasksInZOrder,
+                        )
                         .updateZOrder(freeformTasksInZOrder)
 
                 desktopPersistentRepositories
@@ -169,9 +175,21 @@ class DesktopPersistentRepository(
 
         private fun Desktop.Builder.updateTaskStates(
             visibleTasks: ArraySet<Int>,
-            minimizedTasks: ArraySet<Int>
+            minimizedTasks: ArraySet<Int>,
+            freeformTasksInZOrder: ArrayList<Int>,
         ): Desktop.Builder {
             clearTasksByTaskId()
+
+            // Handle the case where tasks are not marked as visible but are meant to be visible
+            // after reboot. E.g. User moves out of desktop when there are multiple tasks are
+            // visible, they will be marked as not visible afterwards. This ensures that they are
+            // still persisted as visible.
+            // TODO - b/350476823: Remove this logic once repository holds expanded tasks
+            if (freeformTasksInZOrder.size > visibleTasks.size + minimizedTasks.size &&
+                visibleTasks.isEmpty()
+            ) {
+                visibleTasks.addAll(freeformTasksInZOrder.filterNot { it in minimizedTasks })
+            }
             putAllTasksByTaskId(
                 visibleTasks.associateWith {
                     createDesktopTask(it, state = DesktopTaskState.VISIBLE)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
index 38d9ab96e7fd..414c1a658b95 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
@@ -53,6 +53,7 @@ import org.mockito.Mockito.inOrder
 import org.mockito.Mockito.spy
 import org.mockito.kotlin.any
 import org.mockito.kotlin.never
+import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 
@@ -227,6 +228,33 @@ class DesktopRepositoryTest : ShellTestCase() {
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+    fun updateTaskVisibility_multipleTasks_persistsVisibleTasks() =
+        runTest(StandardTestDispatcher()) {
+            repo.updateTask(DEFAULT_DISPLAY, taskId = 1, isVisible = true)
+            repo.updateTask(DEFAULT_DISPLAY, taskId = 2, isVisible = true)
+
+            inOrder(persistentRepository).run {
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(arrayOf(1)),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf()
+                    )
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(arrayOf(1, 2)),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf()
+                    )
+            }
+        }
+
+    @Test
     fun isOnlyVisibleNonClosingTask_singleVisibleClosingTask() {
         repo.updateTask(DEFAULT_DISPLAY, taskId = 1, isVisible = true)
         repo.addClosingTask(DEFAULT_DISPLAY, 1)
@@ -614,7 +642,7 @@ class DesktopRepositoryTest : ShellTestCase() {
                         minimizedTasks = ArraySet(),
                         freeformTasksInZOrder = arrayListOf(7, 6, 5)
                     )
-                verify(persistentRepository)
+                verify(persistentRepository, times(2))
                     .addOrUpdateDesktop(
                         DEFAULT_USER_ID,
                         DEFAULT_DESKTOP_ID,