Merge "Allow TaskDisplayArea to have TaskDisplayArea children (1/n)"

15 files changed, 396 insertions, 324 deletions

diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 88914e31d5ac..e60f2e70f25c 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -2260,9 +2260,12 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
                         taskDisplayArea.getRootTask(rootTaskId));
     }
 
-    protected int getRootTaskCount() {
-        return reduceOnAllTaskDisplayAreas((taskDisplayArea, count) ->
-                count + taskDisplayArea.getRootTaskCount(), 0 /* initValue */);
+    int getRootTaskCount() {
+        final int[] count = new int[1];
+        forAllRootTasks(task -> {
+            count[0]++;
+        });
+        return count[0];
     }
 
     @VisibleForTesting
@@ -5353,19 +5356,12 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         }
 
         // Check if all task display areas have only the empty home stacks left.
-        boolean hasNonEmptyHomeStack = forAllTaskDisplayAreas(taskDisplayArea -> {
-            if (taskDisplayArea.getRootTaskCount() != 1) {
-                return true;
-            }
-            final Task stack = taskDisplayArea.getRootTaskAt(0);
-            return !stack.isActivityTypeHome() || stack.hasChild();
-        });
-        if (!hasNonEmptyHomeStack) {
+        boolean hasNonEmptyHomeStack = forAllRootTasks(stack ->
+                !stack.isActivityTypeHome() || stack.hasChild());
+        if (!hasNonEmptyHomeStack && getRootTaskCount() > 0) {
             // Release this display if only empty home stack(s) are left. This display will be
             // released along with the stack(s) removal.
-            forAllTaskDisplayAreas(taskDisplayArea -> {
-                taskDisplayArea.getRootTaskAt(0).removeIfPossible();
-            });
+            forAllRootTasks(Task::removeIfPossible);
         } else if (getTopRootTask() == null) {
             removeIfPossible();
             mRootWindowContainer.mTaskSupervisor
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index e45310a99fbd..57f263885ef2 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -594,16 +594,8 @@ class KeyguardController {
          */
         @Nullable
         private Task getRootTaskForControllingOccluding(DisplayContent display) {
-            return display.getItemFromTaskDisplayAreas(taskDisplayArea -> {
-                for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                    final Task task = taskDisplayArea.getRootTaskAt(sNdx);
-                    if (task != null && task.isFocusableAndVisible()
-                            && !task.inPinnedWindowingMode()) {
-                        return task;
-                    }
-                }
-                return null;
-            });
+            return display.getRootTask(task ->
+                    task != null && task.isFocusableAndVisible() && !task.inPinnedWindowingMode());
         }
 
         void dumpStatus(PrintWriter pw, String prefix) {
diff --git a/services/core/java/com/android/server/wm/RecentsAnimation.java b/services/core/java/com/android/server/wm/RecentsAnimation.java
index 067c772dad93..05dcd366dbac 100644
--- a/services/core/java/com/android/server/wm/RecentsAnimation.java
+++ b/services/core/java/com/android/server/wm/RecentsAnimation.java
@@ -468,14 +468,8 @@ class RecentsAnimation implements RecentsAnimationCallbacks, OnRootTaskOrderChan
      * @return The top stack that is not always-on-top.
      */
     private Task getTopNonAlwaysOnTopStack() {
-        for (int i = mDefaultTaskDisplayArea.getRootTaskCount() - 1; i >= 0; i--) {
-            final Task s = mDefaultTaskDisplayArea.getRootTaskAt(i);
-            if (s.getWindowConfiguration().isAlwaysOnTop()) {
-                continue;
-            }
-            return s;
-        }
-        return null;
+        return mDefaultTaskDisplayArea.getRootTask(task ->
+                !task.getWindowConfiguration().isAlwaysOnTop());
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 169d5f2155ed..9ec9103821d0 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -1835,19 +1835,15 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         final ArrayList<IBinder> topActivityTokens = new ArrayList<>();
         final Task topFocusedStack = getTopDisplayFocusedRootTask();
         // Traverse all displays.
-        forAllTaskDisplayAreas(taskDisplayArea -> {
-            // Traverse all stacks on a display area.
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                // Get top activity from a visible stack and add it to the list.
-                if (stack.shouldBeVisible(null /* starting */)) {
-                    final ActivityRecord top = stack.getTopNonFinishingActivity();
-                    if (top != null) {
-                        if (stack == topFocusedStack) {
-                            topActivityTokens.add(0, top.appToken);
-                        } else {
-                            topActivityTokens.add(top.appToken);
-                        }
+        forAllRootTasks(stack -> {
+            // Get top activity from a visible stack and add it to the list.
+            if (stack.shouldBeVisible(null /* starting */)) {
+                final ActivityRecord top = stack.getTopNonFinishingActivity();
+                if (top != null) {
+                    if (stack == topFocusedStack) {
+                        topActivityTokens.add(0, top.appToken);
+                    } else {
+                        topActivityTokens.add(top.appToken);
                     }
                 }
             }
@@ -1892,21 +1888,17 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         // First, found out what is currently the foreground app, so that we don't blow away the
         // previous app if this activity is being hosted by the process that is actually still the
         // foreground.
-        WindowProcessController fgApp = reduceOnAllTaskDisplayAreas((taskDisplayArea, app) -> {
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                if (isTopDisplayFocusedRootTask(stack)) {
-                    final ActivityRecord resumedActivity = stack.getResumedActivity();
-                    if (resumedActivity != null) {
-                        app = resumedActivity.app;
-                    } else if (stack.mPausingActivity != null) {
-                        app = stack.mPausingActivity.app;
-                    }
-                    break;
+        WindowProcessController fgApp = getItemFromRootTasks(stack -> {
+            if (isTopDisplayFocusedRootTask(stack)) {
+                final ActivityRecord resumedActivity = stack.getResumedActivity();
+                if (resumedActivity != null) {
+                    return resumedActivity.app;
+                } else if (stack.mPausingActivity != null) {
+                    return stack.mPausingActivity.app;
                 }
             }
-            return app;
-        }, null /* initValue */);
+            return null;
+        });
 
         // Now set this one as the previous process, only if that really makes sense to.
         if (r.hasProcess() && fgApp != null && r.app != fgApp
@@ -1923,27 +1915,21 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
             mTmpRemoteException = null;
             mTmpBoolean = false; // Set to true if an activity was started.
             final DisplayContent display = getChildAt(displayNdx);
-            display.forAllTaskDisplayAreas(displayArea -> {
+            display.forAllRootTasks(rootTask -> {
                 if (mTmpRemoteException != null) {
                     return;
                 }
 
-                for (int taskNdx = displayArea.getRootTaskCount() - 1; taskNdx >= 0; --taskNdx) {
-                    final Task rootTask = displayArea.getRootTaskAt(taskNdx);
-                    if (rootTask.getVisibility(null /*starting*/) == TASK_VISIBILITY_INVISIBLE) {
-                        break;
-                    }
-
-                    final PooledFunction c = PooledLambda.obtainFunction(
-                            RootWindowContainer::startActivityForAttachedApplicationIfNeeded, this,
-                            PooledLambda.__(ActivityRecord.class), app,
-                            rootTask.topRunningActivity());
-                    rootTask.forAllActivities(c);
-                    c.recycle();
-                    if (mTmpRemoteException != null) {
-                        return;
-                    }
+                if (rootTask.getVisibility(null /*starting*/) == TASK_VISIBILITY_INVISIBLE) {
+                    return;
                 }
+
+                final PooledFunction c = PooledLambda.obtainFunction(
+                        RootWindowContainer::startActivityForAttachedApplicationIfNeeded, this,
+                        PooledLambda.__(ActivityRecord.class), app,
+                        rootTask.topRunningActivity());
+                rootTask.forAllActivities(c);
+                c.recycle();
             });
             if (mTmpRemoteException != null) {
                 throw mTmpRemoteException;
@@ -2027,11 +2013,8 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         mCurrentUser = userId;
 
         mTaskSupervisor.mStartingUsers.add(uss);
-        forAllTaskDisplayAreas(taskDisplayArea -> {
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                stack.switchUser(userId);
-            }
+        forAllRootTasks(stack -> {
+            stack.switchUser(userId);
         });
 
         final int restoreStackId = mUserRootTaskInFront.get(userId);
@@ -2294,20 +2277,14 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
      */
     int finishTopCrashedActivities(WindowProcessController app, String reason) {
         Task focusedStack = getTopDisplayFocusedRootTask();
-        Task finishedTask = reduceOnAllTaskDisplayAreas((taskDisplayArea, task) -> {
-            // It is possible that request to finish activity might also remove its task and
-            // stack, so we need to be careful with indexes in the loop and check child count
-            // every time.
-            for (int stackNdx = 0; stackNdx < taskDisplayArea.getRootTaskCount(); ++stackNdx) {
-                final Task stack = taskDisplayArea.getRootTaskAt(stackNdx);
-                final Task t = stack.finishTopCrashedActivityLocked(app, reason);
-                if (stack == focusedStack || task == null) {
-                    task = t;
-                }
+        final Task[] finishedTask = new Task[1];
+        forAllTasks(stack -> {
+            final Task t = stack.finishTopCrashedActivityLocked(app, reason);
+            if (stack == focusedStack || finishedTask[0] == null) {
+                finishedTask[0] = t;
             }
-            return task;
-        }, null /* initValue */);
-        return finishedTask != null ? finishedTask.mTaskId : INVALID_TASK_ID;
+        });
+        return finishedTask[0] != null ? finishedTask[0].mTaskId : INVALID_TASK_ID;
     }
 
     boolean resumeFocusedTasksTopActivities() {
@@ -2330,36 +2307,32 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         for (int displayNdx = getChildCount() - 1; displayNdx >= 0; --displayNdx) {
             final DisplayContent display = getChildAt(displayNdx);
             final boolean curResult = result;
-            boolean resumedOnDisplay = display.reduceOnAllTaskDisplayAreas(
-                    (taskDisplayArea, resumed) -> {
-                        for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                            final Task rootTask = taskDisplayArea.getRootTaskAt(sNdx);
-                            final ActivityRecord topRunningActivity = rootTask.topRunningActivity();
-                            if (!rootTask.isFocusableAndVisible() || topRunningActivity == null) {
-                                continue;
-                            }
-                            if (rootTask == targetRootTask) {
-                                // Simply update the result for targetStack because the targetStack
-                                // had already resumed in above. We don't want to resume it again,
-                                // especially in some cases, it would cause a second launch failure
-                                // if app process was dead.
-                                resumed |= curResult;
-                                continue;
-                            }
-                            if (taskDisplayArea.isTopRootTask(rootTask)
-                                    && topRunningActivity.isState(RESUMED)) {
-                                // Kick off any lingering app transitions form the MoveTaskToFront
-                                // operation, but only consider the top task and stack on that
-                                // display.
-                                rootTask.executeAppTransition(targetOptions);
-                            } else {
-                                resumed |= topRunningActivity.makeActiveIfNeeded(target);
-                            }
-                        }
-                        return resumed;
-                    }, false /* initValue */);
-            result |= resumedOnDisplay;
-            if (!resumedOnDisplay) {
+            boolean[] resumedOnDisplay = new boolean[1];
+            display.forAllRootTasks(rootTask -> {
+                final ActivityRecord topRunningActivity = rootTask.topRunningActivity();
+                if (!rootTask.isFocusableAndVisible() || topRunningActivity == null) {
+                    return;
+                }
+                if (rootTask == targetRootTask) {
+                    // Simply update the result for targetStack because the targetStack
+                    // had already resumed in above. We don't want to resume it again,
+                    // especially in some cases, it would cause a second launch failure
+                    // if app process was dead.
+                    resumedOnDisplay[0] |= curResult;
+                    return;
+                }
+                if (rootTask.getDisplayArea().isTopRootTask(rootTask)
+                        && topRunningActivity.isState(RESUMED)) {
+                    // Kick off any lingering app transitions form the MoveTaskToFront
+                    // operation, but only consider the top task and stack on that
+                    // display.
+                    rootTask.executeAppTransition(targetOptions);
+                } else {
+                    resumedOnDisplay[0] |= topRunningActivity.makeActiveIfNeeded(target);
+                }
+            });
+            result |= resumedOnDisplay[0];
+            if (!resumedOnDisplay[0]) {
                 // In cases when there are no valid activities (e.g. device just booted or launcher
                 // crashed) it's possible that nothing was resumed on a display. Requesting resume
                 // of top activity in focused stack explicitly will make sure that at least home
@@ -2392,30 +2365,27 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
             }
 
             // Set the sleeping state of the stacks on the display.
-            display.forAllTaskDisplayAreas(taskDisplayArea -> {
-                for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                    final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                    if (displayShouldSleep) {
-                        stack.goToSleepIfPossible(false /* shuttingDown */);
-                    } else {
-                        stack.awakeFromSleepingLocked();
-                        if (stack.isFocusedStackOnDisplay()
-                                && !mTaskSupervisor.getKeyguardController()
-                                .isKeyguardOrAodShowing(display.mDisplayId)) {
-                            // If the keyguard is unlocked - resume immediately.
-                            // It is possible that the display will not be awake at the time we
-                            // process the keyguard going away, which can happen before the sleep
-                            // token is released. As a result, it is important we resume the
-                            // activity here.
-                            stack.resumeTopActivityUncheckedLocked(null, null);
-                        }
-                        // The visibility update must not be called before resuming the top, so the
-                        // display orientation can be updated first if needed. Otherwise there may
-                        // have redundant configuration changes due to apply outdated display
-                        // orientation (from keyguard) to activity.
-                        stack.ensureActivitiesVisible(null /* starting */, 0 /* configChanges */,
-                                false /* preserveWindows */);
+            display.forAllRootTasks(stack -> {
+                if (displayShouldSleep) {
+                    stack.goToSleepIfPossible(false /* shuttingDown */);
+                } else {
+                    stack.awakeFromSleepingLocked();
+                    if (stack.isFocusedStackOnDisplay()
+                            && !mTaskSupervisor.getKeyguardController()
+                            .isKeyguardOrAodShowing(display.mDisplayId)) {
+                        // If the keyguard is unlocked - resume immediately.
+                        // It is possible that the display will not be awake at the time we
+                        // process the keyguard going away, which can happen before the sleep
+                        // token is released. As a result, it is important we resume the
+                        // activity here.
+                        stack.resumeTopActivityUncheckedLocked(null, null);
                     }
+                    // The visibility update must not be called before resuming the top, so the
+                    // display orientation can be updated first if needed. Otherwise there may
+                    // have redundant configuration changes due to apply outdated display
+                    // orientation (from keyguard) to activity.
+                    stack.ensureActivitiesVisible(null /* starting */, 0 /* configChanges */,
+                            false /* preserveWindows */);
                 }
             });
         }
@@ -2512,13 +2482,10 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
 
     /** If displayId == INVALID_DISPLAY, this will get root task infos on all displays */
     ArrayList<RootTaskInfo> getAllRootTaskInfos(int displayId) {
-        ArrayList<RootTaskInfo> list = new ArrayList<>();
+        final ArrayList<RootTaskInfo> list = new ArrayList<>();
         if (displayId == INVALID_DISPLAY) {
-            forAllTaskDisplayAreas(taskDisplayArea -> {
-                for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                    final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                    list.add(getRootTaskInfo(stack));
-                }
+            forAllRootTasks(stack -> {
+                list.add(getRootTaskInfo(stack));
             });
             return list;
         }
@@ -2526,11 +2493,8 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         if (display == null) {
             return list;
         }
-        display.forAllTaskDisplayAreas(taskDisplayArea -> {
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                list.add(getRootTaskInfo(stack));
-            }
+        display.forAllRootTasks(stack -> {
+            list.add(getRootTaskInfo(stack));
         });
         return list;
     }
@@ -2603,10 +2567,17 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     Task findRootTaskBehind(Task rootTask) {
         final TaskDisplayArea taskDisplayArea = rootTask.getDisplayArea();
         if (taskDisplayArea != null) {
-            for (int i = taskDisplayArea.getRootTaskCount() - 1; i >= 0; i--) {
-                if (taskDisplayArea.getRootTaskAt(i) == rootTask && i > 0) {
-                    return taskDisplayArea.getRootTaskAt(i - 1);
+            final boolean[] hasFound = new boolean[1];
+            // TODO(b/175136051): should this be only the direct child root task?
+            final Task rootTaskBehind = taskDisplayArea.getRootTask(task -> {
+                if (hasFound[0]) {
+                    return true;
                 }
+                hasFound[0] = task == rootTask;
+                return false;
+            });
+            if (rootTaskBehind != null) {
+                return rootTaskBehind;
             }
         }
         throw new IllegalStateException("Failed to find a root task behind root task =" + rootTask
@@ -2748,24 +2719,16 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     // Tries to put all activity tasks to sleep. Returns true if all tasks were
     // successfully put to sleep.
     boolean putTasksToSleep(boolean allowDelay, boolean shuttingDown) {
-        return reduceOnAllTaskDisplayAreas((taskDisplayArea, result) -> {
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                // Stacks and activities could be removed while putting activities to sleep if
-                // the app process was gone. This prevents us getting exception by accessing an
-                // invalid stack index.
-                if (sNdx >= taskDisplayArea.getRootTaskCount()) {
-                    continue;
-                }
-                final Task task = taskDisplayArea.getRootTaskAt(sNdx);
-                if (allowDelay) {
-                    result &= task.goToSleepIfPossible(shuttingDown);
-                } else {
-                    task.ensureActivitiesVisible(null /* starting */, 0 /* configChanges */,
-                            !PRESERVE_WINDOWS);
-                }
+        final boolean[] result = {true};
+        forAllRootTasks(task -> {
+            if (allowDelay) {
+                result[0] &= task.goToSleepIfPossible(shuttingDown);
+            } else {
+                task.ensureActivitiesVisible(null /* starting */, 0 /* configChanges */,
+                        !PRESERVE_WINDOWS);
             }
-            return result;
-        }, true /* initValue */);
+        });
+        return result[0];
     }
 
     void handleAppCrash(WindowProcessController app) {
@@ -3021,11 +2984,11 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
                 r.getActivityType());
 
         // Return the topmost valid stack on the display.
-        for (int i = taskDisplayArea.getRootTaskCount() - 1; i >= 0; --i) {
-            final Task stack = taskDisplayArea.getRootTaskAt(i);
-            if (isValidLaunchRootTask(stack, r, windowingMode)) {
-                return stack;
-            }
+        final int targetWindowingMode = windowingMode;
+        final Task topmostValidStack = taskDisplayArea.getRootTask(stack ->
+                isValidLaunchRootTask(stack, r, targetWindowingMode));
+        if (topmostValidStack != null) {
+            return topmostValidStack;
         }
 
         // If there is no valid stack on the secondary display area - check if new dynamic stack
@@ -3260,12 +3223,8 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     }
 
     void finishVoiceTask(IVoiceInteractionSession session) {
-        forAllTaskDisplayAreas(taskDisplayArea -> {
-            final int numStacks = taskDisplayArea.getRootTaskCount();
-            for (int stackNdx = 0; stackNdx < numStacks; ++stackNdx) {
-                final Task stack = taskDisplayArea.getRootTaskAt(stackNdx);
-                stack.finishVoiceTask(session);
-            }
+        forAllRootTasks(stack -> {
+            stack.finishVoiceTask(session);
         });
     }
 
@@ -3324,20 +3283,16 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
 
     boolean allResumedActivitiesVisible() {
         boolean[] foundResumed = {false};
-        final boolean foundInvisibleResumedActivity = forAllTaskDisplayAreas(
-                taskDisplayArea -> {
-                    for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                        final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                        final ActivityRecord r = stack.getResumedActivity();
-                        if (r != null) {
-                            if (!r.nowVisible) {
-                                return true;
-                            }
-                            foundResumed[0] = true;
-                        }
-                    }
-                    return false;
-                });
+        final boolean foundInvisibleResumedActivity = forAllRootTasks(stack -> {
+            final ActivityRecord r = stack.getResumedActivity();
+            if (r != null) {
+                if (!r.nowVisible) {
+                    return true;
+                }
+                foundResumed[0] = true;
+            }
+            return false;
+        });
         if (foundInvisibleResumedActivity) {
             return false;
         }
@@ -3346,23 +3301,19 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
 
     boolean allPausedActivitiesComplete() {
         boolean[] pausing = {true};
-        final boolean hasActivityNotCompleted = forAllTaskDisplayAreas(
-                taskDisplayArea -> {
-                    for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                        final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                        final ActivityRecord r = stack.mPausingActivity;
-                        if (r != null && !r.isState(PAUSED, STOPPED, STOPPING, FINISHING)) {
-                            ProtoLog.d(WM_DEBUG_STATES, "allPausedActivitiesComplete: "
-                                    + "r=%s state=%s", r, r.getState());
-                            if (WM_DEBUG_STATES.isEnabled()) {
-                                pausing[0] = false;
-                            } else {
-                                return true;
-                            }
-                        }
-                    }
-                    return false;
-                });
+        final boolean hasActivityNotCompleted = forAllRootTasks(stack -> {
+            final ActivityRecord r = stack.mPausingActivity;
+            if (r != null && !r.isState(PAUSED, STOPPED, STOPPING, FINISHING)) {
+                ProtoLog.d(WM_DEBUG_STATES, "allPausedActivitiesComplete: "
+                        + "r=%s state=%s", r, r.getState());
+                if (WM_DEBUG_STATES.isEnabled()) {
+                    pausing[0] = false;
+                } else {
+                    return true;
+                }
+            }
+            return false;
+        });
         if (hasActivityNotCompleted) {
             return false;
         }
@@ -3413,13 +3364,10 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     }
 
     void cancelInitializingActivities() {
-        forAllTaskDisplayAreas(taskDisplayArea -> {
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                // We don't want to clear starting window for activities that aren't occluded
-                // as we need to display their starting window until they are done initializing.
-                taskDisplayArea.getRootTaskAt(sNdx).forAllOccludedActivities(
-                        ActivityRecord::cancelInitializing);
-            }
+        forAllRootTasks(task -> {
+            // We don't want to clear starting window for activities that aren't occluded
+            // as we need to display their starting window until they are done initializing.
+            task.forAllOccludedActivities(ActivityRecord::cancelInitializing);
         });
     }
 
@@ -3579,13 +3527,10 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
                 return new ArrayList<>();
             }
         } else {
-            ArrayList<ActivityRecord> activities = new ArrayList<>();
-            forAllTaskDisplayAreas(taskDisplayArea -> {
-                for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                    final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                    if (!dumpVisibleStacksOnly || stack.shouldBeVisible(null)) {
-                        activities.addAll(stack.getDumpActivitiesLocked(name));
-                    }
+            final ArrayList<ActivityRecord> activities = new ArrayList<>();
+            forAllRootTasks(stack -> {
+                if (!dumpVisibleStacksOnly || stack.shouldBeVisible(null)) {
+                    activities.addAll(stack.getDumpActivitiesLocked(name));
                 }
             });
             return activities;
@@ -3635,15 +3580,12 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
             pw.print("Display #");
             pw.print(displayContent.mDisplayId);
             pw.println(" (activities from top to bottom):");
-            displayContent.forAllTaskDisplayAreas(taskDisplayArea -> {
-                for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                    final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                    if (needSep[0]) {
-                        pw.println();
-                    }
-                    needSep[0] = stack.dump(fd, pw, dumpAll, dumpClient, dumpPackage, false);
-                    printed[0] |= needSep[0];
+            displayContent.forAllRootTasks(stack -> {
+                if (needSep[0]) {
+                    pw.println();
                 }
+                needSep[0] = stack.dump(fd, pw, dumpAll, dumpClient, dumpPackage, false);
+                printed[0] |= needSep[0];
             });
             displayContent.forAllTaskDisplayAreas(taskDisplayArea -> {
                 printed[0] |= printThisActivity(pw, taskDisplayArea.getFocusedActivity(),
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index ff7e68eb383e..ca34e9da6b50 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -3878,6 +3878,13 @@ class Task extends WindowContainer<WindowContainer> {
     }
 
     @Override
+    void forAllRootTasks(Consumer<Task> callback, boolean traverseTopToBottom) {
+        if (isRootTask()) {
+            callback.accept(this);
+        }
+    }
+
+    @Override
     boolean forAllTasks(Function<Task, Boolean> callback) {
         if (super.forAllTasks(callback)) return true;
         return callback.apply(this);
@@ -3902,12 +3909,29 @@ class Task extends WindowContainer<WindowContainer> {
     }
 
     @Override
+    boolean forAllRootTasks(Function<Task, Boolean> callback, boolean traverseTopToBottom) {
+        return isRootTask() ? callback.apply(this) : false;
+    }
+
+    @Override
     Task getTask(Predicate<Task> callback, boolean traverseTopToBottom) {
         final Task t = super.getTask(callback, traverseTopToBottom);
         if (t != null) return t;
         return callback.test(this) ? this : null;
     }
 
+    @Nullable
+    @Override
+    Task getRootTask(Predicate<Task> callback, boolean traverseTopToBottom) {
+        return isRootTask() && callback.test(this) ? this : null;
+    }
+
+    @Nullable
+    @Override
+    <R> R getItemFromRootTasks(Function<Task, R> callback, boolean traverseTopToBottom) {
+        return isRootTask() ? callback.apply(this) : null;
+    }
+
     /**
      * @param canAffectSystemUiFlags If false, all windows in this task can not affect SystemUI
      *                               flags. See {@link WindowState#canAffectSystemUiFlags()}.
diff --git a/services/core/java/com/android/server/wm/TaskDisplayArea.java b/services/core/java/com/android/server/wm/TaskDisplayArea.java
index 9425602763c5..1c589b3d94f4 100644
--- a/services/core/java/com/android/server/wm/TaskDisplayArea.java
+++ b/services/core/java/com/android/server/wm/TaskDisplayArea.java
@@ -474,7 +474,7 @@ final class TaskDisplayArea extends DisplayArea<Task> {
     private int findMinPositionForRootTask(Task rootTask) {
         int minPosition = POSITION_BOTTOM;
         for (int i = 0; i < mChildren.size(); ++i) {
-            if (getPriority(getRootTaskAt(i)) < getPriority(rootTask)) {
+            if (getPriority(mChildren.get(i)) < getPriority(rootTask)) {
                 minPosition = i;
             } else {
                 break;
@@ -495,7 +495,7 @@ final class TaskDisplayArea extends DisplayArea<Task> {
 
     private int findMaxPositionForRootTask(Task rootTask) {
         for (int i = mChildren.size() - 1; i >= 0; --i) {
-            final Task curr = getRootTaskAt(i);
+            final Task curr = mChildren.get(i);
             // Since a stack could be repositioned while still being one of the children, we check
             // if 'curr' is the same stack and skip it if so
             final boolean sameRootTask = curr == rootTask;
@@ -894,14 +894,9 @@ final class TaskDisplayArea extends DisplayArea<Task> {
         }
     }
 
+    @Nullable
     Task getRootTask(int rootTaskId) {
-        for (int i = getRootTaskCount() - 1; i >= 0; --i) {
-            final Task stack = getRootTaskAt(i);
-            if (stack.getRootTaskId() == rootTaskId) {
-                return stack;
-            }
-        }
-        return null;
+        return getRootTask(stack -> stack.getRootTaskId() == rootTaskId);
     }
 
     /**
@@ -1065,15 +1060,15 @@ final class TaskDisplayArea extends DisplayArea<Task> {
             // Only split-screen windowing modes can do this currently...
             return null;
         }
-        for (int i = getRootTaskCount() - 1; i >= 0; --i) {
-            final Task t = getRootTaskAt(i);
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final Task t = mChildren.get(i);
             if (!t.mCreatedByOrganizer || t.getRequestedOverrideWindowingMode() != windowingMode) {
                 continue;
             }
             // If not already set, pick a launch root which is not the one we are launching into.
             if (mLaunchRootTask == null) {
-                for (int j = 0, n = getRootTaskCount(); j < n; ++j) {
-                    final Task tt = getRootTaskAt(j);
+                for (int j = 0, n = mChildren.size(); j < n; ++j) {
+                    final Task tt = mChildren.get(j);
                     if (tt.mCreatedByOrganizer && tt != t) {
                         mLaunchRootTask = tt;
                         break;
@@ -1128,8 +1123,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
             return mPreferredTopFocusableRootTask;
         }
 
-        for (int i = getRootTaskCount() - 1; i >= 0; --i) {
-            final Task stack = getRootTaskAt(i);
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final Task stack = mChildren.get(i);
             if (stack.isFocusableAndVisible()) {
                 return stack;
             }
@@ -1143,8 +1138,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
                 ? currentFocus.getWindowingMode() : WINDOWING_MODE_UNDEFINED;
 
         Task candidate = null;
-        for (int i = getRootTaskCount() - 1; i >= 0; --i) {
-            final Task rootTask = getRootTaskAt(i);
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final Task rootTask = mChildren.get(i);
             if (ignoreCurrent && rootTask == currentFocus) {
                 continue;
             }
@@ -1223,8 +1218,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
     }
 
     boolean allResumedActivitiesComplete() {
-        for (int stackNdx = getRootTaskCount() - 1; stackNdx >= 0; --stackNdx) {
-            final ActivityRecord r = getRootTaskAt(stackNdx).getResumedActivity();
+        for (int stackNdx = mChildren.size() - 1; stackNdx >= 0; --stackNdx) {
+            final ActivityRecord r = mChildren.get(stackNdx).getResumedActivity();
             if (r != null && !r.isState(RESUMED)) {
                 return false;
             }
@@ -1251,8 +1246,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
      */
     boolean pauseBackTasks(boolean userLeaving, ActivityRecord resuming) {
         boolean someActivityPaused = false;
-        for (int stackNdx = getRootTaskCount() - 1; stackNdx >= 0; --stackNdx) {
-            final Task stack = getRootTaskAt(stackNdx);
+        for (int stackNdx = mChildren.size() - 1; stackNdx >= 0; --stackNdx) {
+            final Task stack = mChildren.get(stackNdx);
             final ActivityRecord resumedActivity = stack.getResumedActivity();
             if (resumedActivity != null
                     && (stack.getVisibility(resuming) != TASK_VISIBILITY_VISIBLE
@@ -1272,8 +1267,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
     void findTaskLocked(final ActivityRecord r, final boolean isPreferredDisplayArea,
             RootWindowContainer.FindTaskResult result) {
         mTmpFindTaskResult.clear();
-        for (int stackNdx = getRootTaskCount() - 1; stackNdx >= 0; --stackNdx) {
-            final Task stack = getRootTaskAt(stackNdx);
+        for (int stackNdx = mChildren.size() - 1; stackNdx >= 0; --stackNdx) {
+            final Task stack = mChildren.get(stackNdx);
             if (!r.hasCompatibleActivityType(stack) && stack.isLeafTask()) {
                 ProtoLog.d(WM_DEBUG_TASKS, "Skipping stack: (mismatch activity/stack) "
                         + "%s", stack);
@@ -1582,8 +1577,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
 
         // Look in other focusable stacks.
         if (topRunning == null) {
-            for (int i = getRootTaskCount() - 1; i >= 0; --i) {
-                final Task stack = getRootTaskAt(i);
+            for (int i = mChildren.size() - 1; i >= 0; --i) {
+                final Task stack = mChildren.get(i);
                 // Only consider focusable stacks other than the current focused one.
                 if (stack == focusedStack || !stack.isTopActivityFocusable()) {
                     continue;
@@ -1608,11 +1603,11 @@ final class TaskDisplayArea extends DisplayArea<Task> {
     }
 
     protected int getRootTaskCount() {
-        return mChildren.size();
-    }
-
-    protected Task getRootTaskAt(int index) {
-        return mChildren.get(index);
+        final int[] count = new int[1];
+        forAllRootTasks(task -> {
+            count[0]++;
+        });
+        return count[0];
     }
 
     @Nullable
@@ -1711,9 +1706,9 @@ final class TaskDisplayArea extends DisplayArea<Task> {
         // Find the next position where the stack should be placed
         final boolean isRootTask = rootTask.isRootTask();
         final int numRootTasks =
-                isRootTask ? getRootTaskCount() : rootTask.getParent().getChildCount();
+                isRootTask ? mChildren.size() : rootTask.getParent().getChildCount();
         for (int rootTaskNdx = 0; rootTaskNdx < numRootTasks; rootTaskNdx++) {
-            final Task s = isRootTask ? getRootTaskAt(rootTaskNdx)
+            final Task s = isRootTask ? mChildren.get(rootTaskNdx)
                     : (Task) rootTask.getParent().getChildAt(rootTaskNdx);
             if (s == rootTask) {
                 continue;
@@ -1836,8 +1831,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
             boolean preserveWindows, boolean notifyClients, boolean userLeaving) {
         mAtmService.mTaskSupervisor.beginActivityVisibilityUpdate();
         try {
-            for (int stackNdx = getRootTaskCount() - 1; stackNdx >= 0; --stackNdx) {
-                final Task stack = getRootTaskAt(stackNdx);
+            for (int stackNdx = mChildren.size() - 1; stackNdx >= 0; --stackNdx) {
+                final Task stack = mChildren.get(stackNdx);
                 stack.ensureActivitiesVisible(starting, configChanges, preserveWindows,
                         notifyClients, userLeaving);
             }
@@ -1871,13 +1866,13 @@ final class TaskDisplayArea extends DisplayArea<Task> {
         // related WindowContainer will also be removed. So, we set display area as removed after
         // reparenting stack finished.
         // Keep the order from bottom to top.
-        int numStacks = getRootTaskCount();
+        int numStacks = mChildren.size();
 
         final boolean splitScreenActivated = toDisplayArea.isSplitScreenModeActivated();
         final Task rootStack = splitScreenActivated ? toDisplayArea
                 .getTopRootTaskInWindowingMode(WINDOWING_MODE_SPLIT_SCREEN_SECONDARY) : null;
         for (int stackNdx = 0; stackNdx < numStacks; stackNdx++) {
-            final Task stack = getRootTaskAt(stackNdx);
+            final Task stack = mChildren.get(stackNdx);
             // Always finish non-standard type stacks and stacks created by a organizer.
             // TODO: For stacks created by organizer, consider reparenting children tasks if the use
             //       case arises in the future.
@@ -1897,8 +1892,8 @@ final class TaskDisplayArea extends DisplayArea<Task> {
             }
             // Stacks may be removed from this display. Ensure each stack will be processed
             // and the loop will end.
-            stackNdx -= numStacks - getRootTaskCount();
-            numStacks = getRootTaskCount();
+            stackNdx -= numStacks - mChildren.size();
+            numStacks = mChildren.size();
         }
         if (lastReparentedStack != null && splitScreenActivated) {
             if (!lastReparentedStack.supportsSplitScreenWindowingMode()) {
diff --git a/services/core/java/com/android/server/wm/TaskLaunchParamsModifier.java b/services/core/java/com/android/server/wm/TaskLaunchParamsModifier.java
index 9d36b84431f3..a7306c97bd2a 100644
--- a/services/core/java/com/android/server/wm/TaskLaunchParamsModifier.java
+++ b/services/core/java/com/android/server/wm/TaskLaunchParamsModifier.java
@@ -795,17 +795,13 @@ class TaskLaunchParamsModifier implements LaunchParamsModifier {
     private void adjustBoundsToAvoidConflictInDisplay(@NonNull DisplayContent display,
             @NonNull Rect inOutBounds) {
         final List<Rect> taskBoundsToCheck = new ArrayList<>();
-        display.forAllTaskDisplayAreas(taskDisplayArea -> {
-            int numStacks = taskDisplayArea.getRootTaskCount();
-            for (int sNdx = 0; sNdx < numStacks; ++sNdx) {
-                final Task task = taskDisplayArea.getRootTaskAt(sNdx);
-                if (!task.inFreeformWindowingMode()) {
-                    continue;
-                }
+        display.forAllRootTasks(task -> {
+            if (!task.inFreeformWindowingMode()) {
+                return;
+            }
 
-                for (int j = 0; j < task.getChildCount(); ++j) {
-                    taskBoundsToCheck.add(task.getChildAt(j).getBounds());
-                }
+            for (int j = 0; j < task.getChildCount(); ++j) {
+                taskBoundsToCheck.add(task.getChildAt(j).getBounds());
             }
         }, false /* traverseTopToBottom */);
         adjustBoundsToAvoidConflict(display.getBounds(), taskBoundsToCheck, inOutBounds);
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 2095d1925e4b..ca44f0d60226 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -687,16 +687,13 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
                 if (dc == null) {
                     throw new IllegalArgumentException("Display " + displayId + " doesn't exist");
                 }
-                ArrayList<RunningTaskInfo> out = new ArrayList<>();
-                dc.forAllTaskDisplayAreas(taskDisplayArea -> {
-                    for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                        final Task task = taskDisplayArea.getRootTaskAt(sNdx);
-                        if (activityTypes != null
-                                && !ArrayUtils.contains(activityTypes, task.getActivityType())) {
-                            continue;
-                        }
-                        out.add(task.getTaskInfo());
+                final ArrayList<RunningTaskInfo> out = new ArrayList<>();
+                dc.forAllRootTasks(task -> {
+                    if (activityTypes != null
+                            && !ArrayUtils.contains(activityTypes, task.getActivityType())) {
+                        return;
                     }
+                    out.add(task.getTaskInfo());
                 });
                 return out;
             }
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index adc7a227861a..6e0fec1c80fa 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -1664,6 +1664,39 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
     }
 
     /**
+     * For all root tasks at or below this container call the callback.
+     *
+     * @param callback Calls the {@link ToBooleanFunction#apply} method for each root task found and
+     *                 stops the search if {@link ToBooleanFunction#apply} returns {@code true}.
+     */
+    boolean forAllRootTasks(Function<Task, Boolean> callback) {
+        return forAllRootTasks(callback, true /* traverseTopToBottom */);
+    }
+
+    boolean forAllRootTasks(Function<Task, Boolean> callback, boolean traverseTopToBottom) {
+        int count = mChildren.size();
+        if (traverseTopToBottom) {
+            for (int i = count - 1; i >= 0; --i) {
+                if (mChildren.get(i).forAllRootTasks(callback)) {
+                    return true;
+                }
+            }
+        } else {
+            for (int i = 0; i < count; i++) {
+                if (mChildren.get(i).forAllRootTasks(callback)) {
+                    return true;
+                }
+                // Root tasks may be removed from this display. Ensure each task will be processed
+                // and the loop will end.
+                int newCount = mChildren.size();
+                i -= count - newCount;
+                count = newCount;
+            }
+        }
+        return false;
+    }
+
+    /**
      * For all tasks at or below this container call the callback.
      *
      * @param callback Callback to be called for every task.
@@ -1698,6 +1731,33 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
         }
     }
 
+    /**
+     * For all root tasks at or below this container call the callback.
+     *
+     * @param callback Callback to be called for every root task.
+     */
+    void forAllRootTasks(Consumer<Task> callback) {
+        forAllRootTasks(callback, true /* traverseTopToBottom */);
+    }
+
+    void forAllRootTasks(Consumer<Task> callback, boolean traverseTopToBottom) {
+        int count = mChildren.size();
+        if (traverseTopToBottom) {
+            for (int i = count - 1; i >= 0; --i) {
+                mChildren.get(i).forAllRootTasks(callback, traverseTopToBottom);
+            }
+        } else {
+            for (int i = 0; i < count; i++) {
+                mChildren.get(i).forAllRootTasks(callback, traverseTopToBottom);
+                // Root tasks may be removed from this display. Ensure each task will be processed
+                // and the loop will end.
+                int newCount = mChildren.size();
+                i -= count - newCount;
+                count = newCount;
+            }
+        }
+    }
+
     Task getTaskAbove(Task t) {
         return getTask(
                 (above) -> true, t, false /*includeBoundary*/, false /*traverseTopToBottom*/);
@@ -1780,6 +1840,44 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
         return null;
     }
 
+    /**
+     * Gets a root task in a branch of the tree.
+     *
+     * @param callback called to test if this is the task that should be returned.
+     * @return The root task if found or null.
+     */
+    @Nullable
+    Task getRootTask(Predicate<Task> callback) {
+        return getRootTask(callback, true /*traverseTopToBottom*/);
+    }
+
+    @Nullable
+    Task getRootTask(Predicate<Task> callback, boolean traverseTopToBottom) {
+        int count = mChildren.size();
+        if (traverseTopToBottom) {
+            for (int i = count - 1; i >= 0; --i) {
+                final Task t = mChildren.get(i).getRootTask(callback, traverseTopToBottom);
+                if (t != null) {
+                    return t;
+                }
+            }
+        } else {
+            for (int i = 0; i < count; i++) {
+                final Task t = mChildren.get(i).getRootTask(callback, traverseTopToBottom);
+                if (t != null) {
+                    return t;
+                }
+                // Root tasks may be removed from this display. Ensure each task will be processed
+                // and the loop will end.
+                int newCount = mChildren.size();
+                i -= count - newCount;
+                count = newCount;
+            }
+        }
+
+        return null;
+    }
+
     private Task processGetTaskWithBoundary(Predicate<Task> callback,
             WindowContainer boundary, boolean includeBoundary, boolean traverseTopToBottom,
             boolean[] boundaryFound, WindowContainer wc) {
@@ -1973,6 +2071,47 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
     }
 
     /**
+     * Finds the first non {@code null} return value from calling the callback on all root
+     * {@link Task} at or below this container.
+     * @param callback Applies on each root {@link Task} found and stops the search if it
+     *                 returns non {@code null}.
+     * @param traverseTopToBottom If {@code true}, traverses the hierarchy from top-to-bottom in
+     *                            terms of z-order, else from bottom-to-top.
+     * @return the first returned object that is not {@code null}. Returns {@code null} if not
+     *         found.
+     */
+    @Nullable
+    <R> R getItemFromRootTasks(Function<Task, R> callback, boolean traverseTopToBottom) {
+        int count = mChildren.size();
+        if (traverseTopToBottom) {
+            for (int i = count - 1; i >= 0; --i) {
+                R result = (R) mChildren.get(i).getItemFromRootTasks(callback, traverseTopToBottom);
+                if (result != null) {
+                    return result;
+                }
+            }
+        } else {
+            for (int i = 0; i < count; i++) {
+                R result = (R) mChildren.get(i).getItemFromRootTasks(callback, traverseTopToBottom);
+                if (result != null) {
+                    return result;
+                }
+                // Root tasks may be removed from this display. Ensure each task will be processed
+                // and the loop will end.
+                int newCount = mChildren.size();
+                i -= count - newCount;
+                count = newCount;
+            }
+        }
+        return null;
+    }
+
+    @Nullable
+    <R> R getItemFromRootTasks(Function<Task, R> callback) {
+        return getItemFromRootTasks(callback, true /* traverseTopToBottom */);
+    }
+
+    /**
      * Returns 1, 0, or -1 depending on if this container is greater than, equal to, or lesser than
      * the input container in terms of z-order.
      */
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java
index 99bd0d7198c0..466b86117fdc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityDisplayTests.java
@@ -248,7 +248,7 @@ public class ActivityDisplayTests extends WindowTestsBase {
         int topPosition = taskDisplayArea.getRootTaskCount() - 1;
         // Ensure the new alwaysOnTop stack is put below the pinned stack, but on top of the
         // existing alwaysOnTop stack.
-        assertEquals(anotherAlwaysOnTopStack, taskDisplayArea.getRootTaskAt(topPosition - 1));
+        assertEquals(topPosition - 1, taskDisplayArea.getIndexOf(anotherAlwaysOnTopStack));
 
         final Task nonAlwaysOnTopStack = taskDisplayArea.createRootTask(
                 WINDOWING_MODE_FREEFORM, ACTIVITY_TYPE_STANDARD, true /* onTop */);
@@ -256,7 +256,7 @@ public class ActivityDisplayTests extends WindowTestsBase {
         topPosition = taskDisplayArea.getRootTaskCount() - 1;
         // Ensure the non-alwaysOnTop stack is put below the three alwaysOnTop stacks, but above the
         // existing other non-alwaysOnTop stacks.
-        assertEquals(nonAlwaysOnTopStack, taskDisplayArea.getRootTaskAt(topPosition - 3));
+        assertEquals(topPosition - 3, taskDisplayArea.getIndexOf(nonAlwaysOnTopStack));
 
         anotherAlwaysOnTopStack.setAlwaysOnTop(false);
         taskDisplayArea.positionChildAt(POSITION_TOP, anotherAlwaysOnTopStack,
@@ -264,16 +264,16 @@ public class ActivityDisplayTests extends WindowTestsBase {
         assertFalse(anotherAlwaysOnTopStack.isAlwaysOnTop());
         // Ensure, when always on top is turned off for a stack, the stack is put just below all
         // other always on top stacks.
-        assertEquals(anotherAlwaysOnTopStack, taskDisplayArea.getRootTaskAt(topPosition - 2));
+        assertEquals(topPosition - 2, taskDisplayArea.getIndexOf(anotherAlwaysOnTopStack));
         anotherAlwaysOnTopStack.setAlwaysOnTop(true);
 
         // Ensure always on top state changes properly when windowing mode changes.
         anotherAlwaysOnTopStack.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
         assertFalse(anotherAlwaysOnTopStack.isAlwaysOnTop());
-        assertEquals(anotherAlwaysOnTopStack, taskDisplayArea.getRootTaskAt(topPosition - 2));
+        assertEquals(topPosition - 2, taskDisplayArea.getIndexOf(anotherAlwaysOnTopStack));
         anotherAlwaysOnTopStack.setWindowingMode(WINDOWING_MODE_FREEFORM);
         assertTrue(anotherAlwaysOnTopStack.isAlwaysOnTop());
-        assertEquals(anotherAlwaysOnTopStack, taskDisplayArea.getRootTaskAt(topPosition - 1));
+        assertEquals(topPosition - 1, taskDisplayArea.getIndexOf(anotherAlwaysOnTopStack));
 
         final Task dreamStack = taskDisplayArea.createRootTask(
                 WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_DREAM, true /* onTop */);
@@ -282,7 +282,7 @@ public class ActivityDisplayTests extends WindowTestsBase {
         topPosition = taskDisplayArea.getRootTaskCount() - 1;
         // Ensure dream shows above all activities, including PiP
         assertEquals(dreamStack, taskDisplayArea.getTopRootTask());
-        assertEquals(pinnedStack, taskDisplayArea.getRootTaskAt(topPosition - 1));
+        assertEquals(topPosition - 1, taskDisplayArea.getIndexOf(pinnedStack));
 
         final Task assistStack = taskDisplayArea.createRootTask(
                 WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_ASSISTANT, true /* onTop */);
@@ -294,8 +294,8 @@ public class ActivityDisplayTests extends WindowTestsBase {
         // is false and on top of everything when true.
         final boolean isAssistantOnTop = mContext.getResources()
                 .getBoolean(com.android.internal.R.bool.config_assistantOnTopOfDream);
-        assertEquals(assistStack, taskDisplayArea.getRootTaskAt(
-                    isAssistantOnTop ? topPosition : topPosition - 4));
+        assertEquals(isAssistantOnTop ? topPosition : topPosition - 4,
+                taskDisplayArea.getIndexOf(assistStack));
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
index ef2e88913914..5695412394ce 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStarterTests.java
@@ -512,11 +512,8 @@ public class ActivityStarterTests extends WindowTestsBase {
     }
 
     private void assertNoTasks(DisplayContent display) {
-        display.forAllTaskDisplayAreas(taskDisplayArea -> {
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                assertFalse(stack.hasChild());
-            }
+        display.forAllRootTasks(stack -> {
+            assertFalse(stack.hasChild());
         });
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
index 9af2b96f3f78..814f23049e3b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
@@ -92,7 +92,7 @@ public class RootWindowContainerTests extends WindowTestsBase {
     public void testAllPausedActivitiesComplete() {
         DisplayContent displayContent = mWm.mRoot.getDisplayContent(DEFAULT_DISPLAY);
         TaskDisplayArea taskDisplayArea = displayContent.getDefaultTaskDisplayArea();
-        Task stack = taskDisplayArea.getRootTaskAt(0);
+        Task stack = taskDisplayArea.getRootTask(task -> true, false /* traverseTopToBottom */);
         ActivityRecord activity = createActivityRecord(displayContent);
         stack.mPausingActivity = activity;
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/RunningTasksTest.java b/services/tests/wmtests/src/com/android/server/wm/RunningTasksTest.java
index 6821d47d5b1a..e2be39b6eb15 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RunningTasksTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RunningTasksTest.java
@@ -34,6 +34,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Build/Install/Run:
@@ -70,9 +71,12 @@ public class RunningTasksTest extends WindowTestsBase {
 
         final int numTasks = 10;
         int activeTime = 0;
+        final List<Task> rootTasks = new ArrayList<>();
+        display.getDefaultTaskDisplayArea().forAllRootTasks(task -> {
+            rootTasks.add(task);
+        }, false /* traverseTopToBottom */);
         for (int i = 0; i < numTasks; i++) {
-            createTask(display.getDefaultTaskDisplayArea().getRootTaskAt(i % numStacks),
-                    ".Task" + i, i, activeTime++, null);
+            createTask(rootTasks.get(i % numStacks), ".Task" + i, i, activeTime++, null);
         }
 
         // Ensure that the latest tasks were returned in order of decreasing last active time,
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index f123bc16e52c..3c8268533b6e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -741,11 +741,10 @@ public class WindowOrganizerTests extends WindowTestsBase {
     }
 
     private List<Task> getTasksCreatedByOrganizer(DisplayContent dc) {
-        ArrayList<Task> out = new ArrayList<>();
-        dc.forAllTaskDisplayAreas(taskDisplayArea -> {
-            for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                final Task t = taskDisplayArea.getRootTaskAt(sNdx);
-                if (t.mCreatedByOrganizer) out.add(t);
+        final ArrayList<Task> out = new ArrayList<>();
+        dc.forAllRootTasks(task -> {
+            if (task.mCreatedByOrganizer) {
+                out.add(task);
             }
         });
         return out;
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index d64b46d925ff..e13a5952ccac 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -1121,12 +1121,9 @@ class WindowTestsBase extends SystemServiceTestsBase {
             mService.mTaskOrganizerController.setLaunchRoot(mDisplayId,
                     mSecondary.mRemoteToken.toWindowContainerToken());
             DisplayContent dc = mService.mRootWindowContainer.getDisplayContent(mDisplayId);
-            dc.forAllTaskDisplayAreas(taskDisplayArea -> {
-                for (int sNdx = taskDisplayArea.getRootTaskCount() - 1; sNdx >= 0; --sNdx) {
-                    final Task stack = taskDisplayArea.getRootTaskAt(sNdx);
-                    if (!WindowConfiguration.isSplitScreenWindowingMode(stack.getWindowingMode())) {
-                        stack.reparent(mSecondary, POSITION_BOTTOM);
-                    }
+            dc.forAllRootTasks(stack -> {
+                if (!WindowConfiguration.isSplitScreenWindowingMode(stack.getWindowingMode())) {
+                    stack.reparent(mSecondary, POSITION_BOTTOM);
                 }
             });
         }