Add utility to get all root tasks on a display

This is intended to replace getAllStackInfos for
task organization purposes.

Bug: 33381284
Test: updated TaskOrganizerTests
Change-Id: I1ff10de188218f1cd0b5638c297dffbf2a9a04a2

4 files changed, 65 insertions, 4 deletions

diff --git a/core/java/android/app/ITaskOrganizerController.aidl b/core/java/android/app/ITaskOrganizerController.aidl
index 5d5956e4dca4..9d6c3d64aaf2 100644
--- a/core/java/android/app/ITaskOrganizerController.aidl
+++ b/core/java/android/app/ITaskOrganizerController.aidl
@@ -52,7 +52,11 @@ interface ITaskOrganizerController {
     boolean deleteRootTask(IWindowContainer task);
 
     /** Gets direct child tasks (ordered from top-to-bottom) */
-    List<ActivityManager.RunningTaskInfo> getChildTasks(in IWindowContainer parent);
+    List<ActivityManager.RunningTaskInfo> getChildTasks(in IWindowContainer parent,
+            in int[] activityTypes);
+
+    /** Gets all root tasks on a display (ordered from top-to-bottom) */
+    List<ActivityManager.RunningTaskInfo> getRootTasks(int displayId, in int[] activityTypes);
 
     /** Get the root task which contains the current ime target */
     IWindowContainer getImeTarget(int display);
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 5f3e3a39490c..d3129f021bab 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -2795,6 +2795,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
         }
     }
 
+    // TODO(148895075): deprecate and replace with task equivalents
     @Override
     public List<ActivityManager.StackInfo> getAllStackInfos() {
         enforceCallerIsRecentsOrHasPermission(MANAGE_ACTIVITY_STACKS, "getAllStackInfos()");
@@ -2821,6 +2822,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
         }
     }
 
+    // TODO(148895075): deprecate and replace with task equivalents
     @Override
     public List<ActivityManager.StackInfo> getAllStackInfosOnDisplay(int displayId) {
         enforceCallerIsRecentsOrHasPermission(MANAGE_ACTIVITY_STACKS, "getAllStackInfos()");
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 0a0530c92a16..4b13a0c1f75d 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -43,6 +43,7 @@ import android.view.IWindowContainer;
 import android.view.SurfaceControl;
 import android.view.WindowContainerTransaction;
 
+import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.function.pooled.PooledConsumer;
 import com.android.internal.util.function.pooled.PooledLambda;
 
@@ -379,7 +380,8 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub
     }
 
     @Override
-    public List<RunningTaskInfo> getChildTasks(IWindowContainer parent) {
+    public List<RunningTaskInfo> getChildTasks(IWindowContainer parent,
+            @Nullable int[] activityTypes) {
         enforceStackPermission("getChildTasks()");
         final long ident = Binder.clearCallingIdentity();
         try {
@@ -405,6 +407,10 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub
                 for (int i = dc.getStackCount() - 1; i >= 0; --i) {
                     final ActivityStack as = dc.getStackAt(i);
                     if (as.getTile() == container) {
+                        if (activityTypes != null
+                                && !ArrayUtils.contains(activityTypes, as.getActivityType())) {
+                            continue;
+                        }
                         final RunningTaskInfo info = new RunningTaskInfo();
                         as.fillTaskInfo(info);
                         out.add(info);
@@ -417,6 +423,40 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub
         }
     }
 
+    @Override
+    public List<RunningTaskInfo> getRootTasks(int displayId, @Nullable int[] activityTypes) {
+        enforceStackPermission("getRootTasks()");
+        final long ident = Binder.clearCallingIdentity();
+        try {
+            synchronized (mGlobalLock) {
+                final DisplayContent dc =
+                        mService.mRootWindowContainer.getDisplayContent(displayId);
+                if (dc == null) {
+                    throw new IllegalArgumentException("Display " + displayId + " doesn't exist");
+                }
+                ArrayList<RunningTaskInfo> out = new ArrayList<>();
+                for (int i = dc.getStackCount() - 1; i >= 0; --i) {
+                    final ActivityStack task = dc.getStackAt(i);
+                    if (task.getTile() != null) {
+                        // a tile is supposed to look like a parent, so don't include their
+                        // "children" here. They can be accessed via getChildTasks()
+                        continue;
+                    }
+                    if (activityTypes != null
+                            && !ArrayUtils.contains(activityTypes, task.getActivityType())) {
+                        continue;
+                    }
+                    final RunningTaskInfo info = new RunningTaskInfo();
+                    task.fillTaskInfo(info);
+                    out.add(info);
+                }
+                return out;
+            }
+        } finally {
+            Binder.restoreCallingIdentity(ident);
+        }
+    }
+
     private int sanitizeAndApplyChange(WindowContainer container,
             WindowContainerTransaction.Change change) {
         if (!(container instanceof Task)) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/TaskOrganizerTests.java
index f7aa3cc9e52c..b5c25ffdb1e0 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskOrganizerTests.java
@@ -384,11 +384,19 @@ public class TaskOrganizerTests extends WindowTestsBase {
         RunningTaskInfo info2 = mWm.mAtmService.mTaskOrganizerController.createRootTask(
                 mDisplayContent.mDisplayId, WINDOWING_MODE_SPLIT_SCREEN_SECONDARY);
 
+        final int initialRootTaskCount = mWm.mAtmService.mTaskOrganizerController.getRootTasks(
+                mDisplayContent.mDisplayId, null /* activityTypes */).size();
+
         final ActivityStack stack = createTaskStackOnDisplay(
                 WINDOWING_MODE_UNDEFINED, ACTIVITY_TYPE_STANDARD, mDisplayContent);
         final ActivityStack stack2 = createTaskStackOnDisplay(
                 WINDOWING_MODE_UNDEFINED, ACTIVITY_TYPE_HOME, mDisplayContent);
 
+        // Check getRootTasks works
+        List<RunningTaskInfo> roots = mWm.mAtmService.mTaskOrganizerController.getRootTasks(
+                mDisplayContent.mDisplayId, null /* activityTypes */);
+        assertEquals(initialRootTaskCount + 2, roots.size());
+
         lastReportedTiles.clear();
         WindowContainerTransaction wct = new WindowContainerTransaction();
         wct.reparent(stack.mRemoteToken, info1.token, true /* onTop */);
@@ -415,11 +423,18 @@ public class TaskOrganizerTests extends WindowTestsBase {
 
         // Check the getChildren call
         List<RunningTaskInfo> children =
-                mWm.mAtmService.mTaskOrganizerController.getChildTasks(info1.token);
+                mWm.mAtmService.mTaskOrganizerController.getChildTasks(info1.token,
+                        null /* activityTypes */);
         assertEquals(2, children.size());
-        children = mWm.mAtmService.mTaskOrganizerController.getChildTasks(info2.token);
+        children = mWm.mAtmService.mTaskOrganizerController.getChildTasks(info2.token,
+                null /* activityTypes */);
         assertEquals(0, children.size());
 
+        // Check that getRootTasks doesn't include children of tiles
+        roots = mWm.mAtmService.mTaskOrganizerController.getRootTasks(mDisplayContent.mDisplayId,
+                null /* activityTypes */);
+        assertEquals(initialRootTaskCount, roots.size());
+
         lastReportedTiles.clear();
         wct = new WindowContainerTransaction();
         wct.reorder(stack2.mRemoteToken, true /* onTop */);