Merge "Update `getFreeformTasksInZOrder` for DesktopModeTaskRepository to be multi-display aware." into main

6 files changed, 56 insertions, 54 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
index 7e0234ef8546..afc32b5c30e3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -29,6 +29,7 @@ import androidx.core.util.valueIterator
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.util.KtProtoLog
 import java.io.PrintWriter
+import java.util.ArrayList;
 import java.util.concurrent.Executor
 import java.util.function.Consumer
 
@@ -48,12 +49,12 @@ class DesktopModeTaskRepository {
         val activeTasks: ArraySet<Int> = ArraySet(),
         val visibleTasks: ArraySet<Int> = ArraySet(),
         val minimizedTasks: ArraySet<Int> = ArraySet(),
+        // Tasks currently in freeform mode, ordered from top to bottom (top is at index 0).
+        val freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
     )
 
     // Token of the current wallpaper activity, used to remove it when the last task is removed
     var wallpaperActivityToken: WindowContainerToken? = null
-    // Tasks currently in freeform mode, ordered from top to bottom (top is at index 0).
-    private val freeformTasksInZOrder = mutableListOf<Int>()
     private val activeTasksListeners = ArraySet<ActiveTasksListener>()
     // Track visible tasks separately because a task may be part of the desktop but not visible.
     private val visibleTasksListeners = ArrayMap<VisibleTasksListener, Executor>()
@@ -235,7 +236,7 @@ class DesktopModeTaskRepository {
      */
     fun getActiveNonMinimizedTasksOrderedFrontToBack(displayId: Int): List<Int> {
         val activeTasks = getActiveTasks(displayId)
-        val allTasksInZOrder = getFreeformTasksInZOrder()
+        val allTasksInZOrder = getFreeformTasksInZOrder(displayId)
         return activeTasks
                 // Don't show already minimized Tasks
                 .filter { taskId -> !isMinimizedTask(taskId) }
@@ -245,10 +246,8 @@ class DesktopModeTaskRepository {
     /**
      * Get a list of freeform tasks, ordered from top-bottom (top at index 0).
      */
-     // TODO(b/278084491): pass in display id
-    fun getFreeformTasksInZOrder(): List<Int> {
-        return freeformTasksInZOrder
-    }
+    fun getFreeformTasksInZOrder(displayId: Int): ArrayList<Int> =
+        ArrayList(displayData[displayId]?.freeformTasksInZOrder ?: emptyList())
 
     /**
      * Updates whether a freeform task with this id is visible or not and notifies listeners.
@@ -325,16 +324,16 @@ class DesktopModeTaskRepository {
     /**
      * Add (or move if it already exists) the task to the top of the ordered list.
      */
-    fun addOrMoveFreeformTaskToTop(taskId: Int) {
+     // TODO(b/342417921): Identify if there is additional checks needed to move tasks for
+     // multi-display scenarios.
+    fun addOrMoveFreeformTaskToTop(displayId: Int, taskId: Int) {
         KtProtoLog.d(
             WM_SHELL_DESKTOP_MODE,
-            "DesktopTaskRepo: add or move task to top taskId=%d",
-            taskId
+            "DesktopTaskRepo: add or move task to top: display=%d, taskId=%d",
+            displayId, taskId
         )
-        if (freeformTasksInZOrder.contains(taskId)) {
-            freeformTasksInZOrder.remove(taskId)
-        }
-        freeformTasksInZOrder.add(0, taskId)
+        displayData[displayId]?.freeformTasksInZOrder?.remove(taskId)
+        displayData.getOrCreate(displayId).freeformTasksInZOrder.add(0, taskId)
     }
 
     /** Mark a Task as minimized. */
@@ -358,17 +357,18 @@ class DesktopModeTaskRepository {
     /**
      * Remove the task from the ordered list.
      */
-    fun removeFreeformTask(taskId: Int) {
+    fun removeFreeformTask(displayId: Int, taskId: Int) {
         KtProtoLog.d(
             WM_SHELL_DESKTOP_MODE,
-            "DesktopTaskRepo: remove freeform task from ordered list taskId=%d",
-            taskId
+            "DesktopTaskRepo: remove freeform task from ordered list: display=%d, taskId=%d",
+            displayId, taskId
         )
-        freeformTasksInZOrder.remove(taskId)
+        displayData[displayId]?.freeformTasksInZOrder?.remove(taskId)
         boundsBeforeMaximizeByTaskId.remove(taskId)
         KtProtoLog.d(
             WM_SHELL_DESKTOP_MODE,
-            "DesktopTaskRepo: remaining freeform tasks: %s", freeformTasksInZOrder.toDumpString(),
+            "DesktopTaskRepo: remaining freeform tasks: %s",
+            displayData[displayId]?.freeformTasksInZOrder?.toDumpString() ?: ""
         )
     }
 
@@ -414,7 +414,6 @@ class DesktopModeTaskRepository {
         val innerPrefix = "$prefix  "
         pw.println("${prefix}DesktopModeTaskRepository")
         dumpDisplayData(pw, innerPrefix)
-        pw.println("${innerPrefix}freeformTasksInZOrder=${freeformTasksInZOrder.toDumpString()}")
         pw.println("${innerPrefix}activeTasksListeners=${activeTasksListeners.size}")
         pw.println("${innerPrefix}visibleTasksListeners=${visibleTasksListeners.size}")
     }
@@ -425,6 +424,7 @@ class DesktopModeTaskRepository {
             pw.println("${prefix}Display $displayId:")
             pw.println("${innerPrefix}activeTasks=${data.activeTasks.toDumpString()}")
             pw.println("${innerPrefix}visibleTasks=${data.visibleTasks.toDumpString()}")
+            pw.println("${innerPrefix}freeformTasksInZOrder=${data.freeformTasksInZOrder.toDumpString()}")
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index e0e2e706d649..7d2aa275a684 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -98,7 +98,7 @@ public class FreeformTaskListener implements ShellTaskOrganizer.TaskListener,
 
         if (DesktopModeStatus.canEnterDesktopMode(mContext)) {
             mDesktopModeTaskRepository.ifPresent(repository -> {
-                repository.addOrMoveFreeformTaskToTop(taskInfo.taskId);
+                repository.addOrMoveFreeformTaskToTop(taskInfo.displayId, taskInfo.taskId);
                 repository.unminimizeTask(taskInfo.displayId, taskInfo.taskId);
                 if (taskInfo.isVisible) {
                     if (repository.addActiveTask(taskInfo.displayId, taskInfo.taskId)) {
@@ -120,7 +120,7 @@ public class FreeformTaskListener implements ShellTaskOrganizer.TaskListener,
 
         if (DesktopModeStatus.canEnterDesktopMode(mContext)) {
             mDesktopModeTaskRepository.ifPresent(repository -> {
-                repository.removeFreeformTask(taskInfo.taskId);
+                repository.removeFreeformTask(taskInfo.displayId, taskInfo.taskId);
                 repository.unminimizeTask(taskInfo.displayId, taskInfo.taskId);
                 if (repository.removeActiveTask(taskInfo.taskId)) {
                     ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
@@ -167,7 +167,7 @@ public class FreeformTaskListener implements ShellTaskOrganizer.TaskListener,
                 taskInfo.taskId, taskInfo.isFocused);
         if (DesktopModeStatus.canEnterDesktopMode(mContext) && taskInfo.isFocused) {
             mDesktopModeTaskRepository.ifPresent(repository -> {
-                repository.addOrMoveFreeformTaskToTop(taskInfo.taskId);
+                repository.addOrMoveFreeformTaskToTop(taskInfo.displayId, taskInfo.taskId);
                 repository.unminimizeTask(taskInfo.displayId, taskInfo.taskId);
             });
         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
index 8f59f30da697..310ccc252469 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
@@ -363,11 +363,11 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
 
     @Test
     fun addOrMoveFreeformTaskToTop_didNotExist_addsToTop() {
-        repo.addOrMoveFreeformTaskToTop(5)
-        repo.addOrMoveFreeformTaskToTop(6)
-        repo.addOrMoveFreeformTaskToTop(7)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 5)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 6)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 7)
 
-        val tasks = repo.getFreeformTasksInZOrder()
+        val tasks = repo.getFreeformTasksInZOrder(DEFAULT_DISPLAY)
         assertThat(tasks.size).isEqualTo(3)
         assertThat(tasks[0]).isEqualTo(7)
         assertThat(tasks[1]).isEqualTo(6)
@@ -376,13 +376,13 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
 
     @Test
     fun addOrMoveFreeformTaskToTop_alreadyExists_movesToTop() {
-        repo.addOrMoveFreeformTaskToTop(5)
-        repo.addOrMoveFreeformTaskToTop(6)
-        repo.addOrMoveFreeformTaskToTop(7)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 5)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 6)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 7)
 
-        repo.addOrMoveFreeformTaskToTop(6)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 6)
 
-        val tasks = repo.getFreeformTasksInZOrder()
+        val tasks = repo.getFreeformTasksInZOrder(DEFAULT_DISPLAY)
         assertThat(tasks.size).isEqualTo(3)
         assertThat(tasks.first()).isEqualTo(6)
     }
@@ -391,7 +391,7 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
     fun removeFreeformTask_removesTaskBoundsBeforeMaximize() {
         val taskId = 1
         repo.saveBoundsBeforeMaximize(taskId, Rect(0, 0, 200, 200))
-        repo.removeFreeformTask(taskId)
+        repo.removeFreeformTask(THIRD_DISPLAY, taskId)
         assertThat(repo.removeBoundsBeforeMaximize(taskId)).isNull()
     }
 
@@ -480,31 +480,31 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
 
     @Test
     fun getActiveNonMinimizedTasksOrderedFrontToBack_returnsFreeformTasksInCorrectOrder() {
-        repo.addActiveTask(displayId = 0, taskId = 1)
-        repo.addActiveTask(displayId = 0, taskId = 2)
-        repo.addActiveTask(displayId = 0, taskId = 3)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 1)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 2)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 3)
         // The front-most task will be the one added last through addOrMoveFreeformTaskToTop
-        repo.addOrMoveFreeformTaskToTop(taskId = 3)
-        repo.addOrMoveFreeformTaskToTop(taskId = 2)
-        repo.addOrMoveFreeformTaskToTop(taskId = 1)
+        repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 3)
+        repo.addOrMoveFreeformTaskToTop(displayId = 0, taskId = 2)
+        repo.addOrMoveFreeformTaskToTop(displayId = 0, taskId = 1)
 
-        assertThat(repo.getActiveNonMinimizedTasksOrderedFrontToBack(displayId = 0)).isEqualTo(
-                listOf(1, 2, 3))
+        assertThat(repo.getActiveNonMinimizedTasksOrderedFrontToBack(displayId = 0))
+            .containsExactly(1, 2, 3).inOrder()
     }
 
     @Test
     fun getActiveNonMinimizedTasksOrderedFrontToBack_minimizedTaskNotIncluded() {
-        repo.addActiveTask(displayId = 0, taskId = 1)
-        repo.addActiveTask(displayId = 0, taskId = 2)
-        repo.addActiveTask(displayId = 0, taskId = 3)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 1)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 2)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 3)
         // The front-most task will be the one added last through addOrMoveFreeformTaskToTop
-        repo.addOrMoveFreeformTaskToTop(taskId = 3)
-        repo.addOrMoveFreeformTaskToTop(taskId = 2)
-        repo.addOrMoveFreeformTaskToTop(taskId = 1)
-        repo.minimizeTask(displayId = 0, taskId = 2)
+        repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 3)
+        repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 2)
+        repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 1)
+        repo.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = 2)
 
-        assertThat(repo.getActiveNonMinimizedTasksOrderedFrontToBack(displayId = 0)).isEqualTo(
-                listOf(1, 3))
+        assertThat(repo.getActiveNonMinimizedTasksOrderedFrontToBack(
+            displayId = DEFAULT_DISPLAY)).containsExactly(1, 3).inOrder()
     }
 
 
@@ -544,5 +544,6 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
 
     companion object {
         const val SECOND_DISPLAY = 1
+        const val THIRD_DISPLAY = 345
     }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index d8d534bec6ea..ac67bd1fedd8 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -1614,7 +1614,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task = createFreeformTask(displayId, bounds)
         whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
         desktopModeTaskRepository.addActiveTask(displayId, task.taskId)
-        desktopModeTaskRepository.addOrMoveFreeformTaskToTop(task.taskId)
+        desktopModeTaskRepository.addOrMoveFreeformTaskToTop(displayId, task.taskId)
         runningTasks.add(task)
         return task
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
index 3c488cac6edd..77f917cc28d8 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
@@ -298,7 +298,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
         val task = createFreeformTask(displayId)
         `when`(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
         desktopTaskRepo.addActiveTask(displayId, task.taskId)
-        desktopTaskRepo.addOrMoveFreeformTaskToTop(task.taskId)
+        desktopTaskRepo.addOrMoveFreeformTaskToTop(displayId, task.taskId)
         return task
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskListenerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskListenerTests.java
index cd68c6996578..3f3cafcf6375 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskListenerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskListenerTests.java
@@ -91,7 +91,8 @@ public final class FreeformTaskListenerTests extends ShellTestCase {
 
         mFreeformTaskListener.onFocusTaskChanged(task);
 
-        verify(mDesktopModeTaskRepository).addOrMoveFreeformTaskToTop(task.taskId);
+        verify(mDesktopModeTaskRepository)
+            .addOrMoveFreeformTaskToTop(task.displayId, task.taskId);
     }
 
     @Test
@@ -103,7 +104,7 @@ public final class FreeformTaskListenerTests extends ShellTestCase {
         mFreeformTaskListener.onFocusTaskChanged(fullscreenTask);
 
         verify(mDesktopModeTaskRepository, never())
-                .addOrMoveFreeformTaskToTop(fullscreenTask.taskId);
+                .addOrMoveFreeformTaskToTop(fullscreenTask.displayId, fullscreenTask.taskId);
     }
 
     @After