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
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 0a0530c9..4b13a0c 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.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 @@
     }
 
     @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 @@
                 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 @@
         }
     }
 
+    @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)) {