Cache reference to stack objects we use a lot in WM.

This helps with fast access since we are no longer using stack ids
to look-up stack objects.

Test: Existing tests pass.
Test: go/wm-smoke
Change-Id: I98cb6f4eb1e6d6acb81b2077aee03265a075a227
diff --git a/services/core/java/com/android/server/am/ActivityStackSupervisor.java b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
index bac71c7..3890ae5 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -101,7 +101,6 @@
 import android.app.Activity;
 import android.app.ActivityManager;
 import android.app.ActivityManager.RunningTaskInfo;
-import android.app.ActivityManager.StackId;
 import android.app.ActivityManager.StackInfo;
 import android.app.ActivityManagerInternal.SleepToken;
 import android.app.ActivityOptions;
@@ -3891,8 +3890,8 @@
                 final ActivityStack stack = stacks.get(stackNdx);
                 pw.println();
                 pw.println("  Stack #" + stack.mStackId
-                        + ": type=" + activityTypeToString(getActivityType())
-                        + " mode=" + windowingModeToString(getWindowingMode()));
+                        + ": type=" + activityTypeToString(stack.getActivityType())
+                        + " mode=" + windowingModeToString(stack.getWindowingMode()));
                 pw.println("  mFullscreen=" + stack.mFullscreen);
                 pw.println("  isSleeping=" + stack.shouldSleepActivities());
                 pw.println("  mBounds=" + stack.mBounds);
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 1b463c7..5d03493 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -1310,8 +1310,7 @@
 
                 // Notify the pinned stack upon all windows drawn. If there was an animation in
                 // progress then this signal will resume that animation.
-                final TaskStack pinnedStack =
-                        mDisplayContent.getStack(WINDOWING_MODE_PINNED);
+                final TaskStack pinnedStack = mDisplayContent.getPinnedStack();
                 if (pinnedStack != null) {
                     pinnedStack.onAllWindowsDrawn();
                 }
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 05afd55..03fdc96 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -16,6 +16,7 @@
 
 package com.android.server.wm;
 
+import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_UNDEFINED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
@@ -116,7 +117,6 @@
 
 import android.annotation.CallSuper;
 import android.annotation.NonNull;
-import android.app.ActivityManager.StackId;
 import android.content.pm.PackageManager;
 import android.content.res.CompatibilityInfo;
 import android.content.res.Configuration;
@@ -296,10 +296,6 @@
     /** Window tokens that are in the process of exiting, but still on screen for animations. */
     final ArrayList<WindowToken> mExitingTokens = new ArrayList<>();
 
-    /** A special TaskStack with id==HOME_STACK_ID that moves to the bottom whenever any TaskStack
-     * (except a future lockscreen TaskStack) moves to the top. */
-    private TaskStack mHomeStack = null;
-
     /** Detect user tapping outside of current focused task bounds .*/
     TaskTapPointerEventListener mTapDetector;
 
@@ -979,7 +975,7 @@
 
             // In the presence of the PINNED stack or System Alert
             // windows we unfortunately can not seamlessly rotate.
-            if (getStack(WINDOWING_MODE_PINNED) != null) {
+            if (hasPinnedStack()) {
                 mayRotateSeamlessly = false;
             }
             for (int i = 0; i < mService.mSessions.size(); i++) {
@@ -1450,20 +1446,31 @@
     }
 
     TaskStack getHomeStack() {
-        if (mHomeStack == null && mDisplayId == DEFAULT_DISPLAY) {
-            Slog.e(TAG_WM, "getHomeStack: Returning null from this=" + this);
-        }
-        return mHomeStack;
+        return mTaskStackContainers.getHomeStack();
     }
 
-    TaskStack getStackById(int stackId) {
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            final TaskStack stack = mTaskStackContainers.get(i);
-            if (stack.mStackId == stackId) {
-                return stack;
-            }
-        }
-        return null;
+    /**
+     * @return The primary split-screen stack, but only if it is visible, and {@code null} otherwise.
+     */
+    TaskStack getSplitScreenPrimaryStackStack() {
+        TaskStack stack = mTaskStackContainers.getSplitScreenPrimaryStackStack();
+        return (stack != null && stack.isVisible()) ? stack : null;
+    }
+
+    /**
+     * Like {@link #getSplitScreenPrimaryStackStack}, but also returns the stack if it's currently
+     * not visible.
+     */
+    TaskStack getSplitScreenPrimaryStackStackIgnoringVisibility() {
+        return mTaskStackContainers.getSplitScreenPrimaryStackStack();
+    }
+
+    TaskStack getPinnedStack() {
+        return mTaskStackContainers.getPinnedStack();
+    }
+
+    private boolean hasPinnedStack() {
+        return mTaskStackContainers.getPinnedStack() != null;
     }
 
     /**
@@ -1479,29 +1486,16 @@
      * activity type. Null is no compatible stack on the display.
      */
     TaskStack getStack(int windowingMode, int activityType) {
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            final TaskStack stack = mTaskStackContainers.get(i);
-            if (stack.isCompatible(windowingMode, activityType)) {
-                return stack;
-            }
-        }
-        return null;
+        return mTaskStackContainers.getStack(windowingMode, activityType);
     }
 
     @VisibleForTesting
-    int getStackCount() {
-        return mTaskStackContainers.size();
+    TaskStack getTopStack() {
+        return mTaskStackContainers.getTopStack();
     }
 
-    @VisibleForTesting
-    int getStackPosition(int windowingMode, int activityType) {
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            final TaskStack stack = mTaskStackContainers.get(i);
-            if (stack.isCompatible(windowingMode, activityType)) {
-                return i;
-            }
-        }
-        return -1;
+    void onStackWindowingModeChanged(TaskStack stack) {
+        mTaskStackContainers.onStackWindowingModeChanged(stack);
     }
 
     @Override
@@ -1522,8 +1516,8 @@
      * bounds were updated.
      */
     void updateStackBoundsAfterConfigChange(@NonNull List<Integer> changedStackList) {
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            final TaskStack stack = mTaskStackContainers.get(i);
+        for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(i);
             if (stack.updateBoundsAfterConfigChange()) {
                 changedStackList.add(stack.mStackId);
             }
@@ -1532,7 +1526,7 @@
         // If there was no pinned stack, we still need to notify the controller of the display info
         // update as a result of the config change.  We do this here to consolidate the flow between
         // changes when there is and is not a stack.
-        if (getStack(WINDOWING_MODE_PINNED) == null) {
+        if (!hasPinnedStack()) {
             mPinnedStackControllerLocked.onDisplayInfoChanged();
         }
     }
@@ -1631,8 +1625,8 @@
         mDisplay.getDisplayInfo(mDisplayInfo);
         mDisplay.getMetrics(mDisplayMetrics);
 
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            mTaskStackContainers.get(i).updateDisplayInfo(null);
+        for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+            mTaskStackContainers.getChildAt(i).updateDisplayInfo(null);
         }
     }
 
@@ -1753,24 +1747,12 @@
         out.set(mContentRect);
     }
 
-    TaskStack addStackToDisplay(int stackId, boolean onTop, StackWindowController controller) {
+    TaskStack createStack(int stackId, boolean onTop, StackWindowController controller) {
         if (DEBUG_STACK) Slog.d(TAG_WM, "Create new stackId=" + stackId + " on displayId="
                 + mDisplayId);
 
-        TaskStack stack = getStackById(stackId);
-        if (stack != null) {
-            // It's already attached to the display...clear mDeferRemoval, set controller, and move
-            // stack to appropriate z-order on display as needed.
-            stack.mDeferRemoval = false;
-            stack.setController(controller);
-            // We're not moving the display to front when we're adding stacks, only when
-            // requested to change the position of stack explicitly.
-            mTaskStackContainers.positionChildAt(onTop ? POSITION_TOP : POSITION_BOTTOM, stack,
-                    false /* includingParents */);
-        } else {
-            stack = new TaskStack(mService, stackId, controller);
-            mTaskStackContainers.addStackToDisplay(stack, onTop);
-        }
+        final TaskStack stack = new TaskStack(mService, stackId, controller);
+        mTaskStackContainers.addStackToDisplay(stack, onTop);
 
         if (stack.inSplitScreenPrimaryWindowingMode()) {
             mDividerControllerLocked.notifyDockedStackExistsChanged(true);
@@ -1789,7 +1771,7 @@
                     + " to its current displayId=" + mDisplayId);
         }
 
-        prevDc.mTaskStackContainers.removeStackFromDisplay(stack);
+        prevDc.mTaskStackContainers.removeChild(stack);
         mTaskStackContainers.addStackToDisplay(stack, onTop);
     }
 
@@ -1823,8 +1805,8 @@
     }
 
     int taskIdFromPoint(int x, int y) {
-        for (int stackNdx = mTaskStackContainers.size() - 1; stackNdx >= 0; --stackNdx) {
-            final TaskStack stack = mTaskStackContainers.get(stackNdx);
+        for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(stackNdx);
             final int taskId = stack.taskIdFromPoint(x, y);
             if (taskId != -1) {
                 return taskId;
@@ -1840,8 +1822,8 @@
     Task findTaskForResizePoint(int x, int y) {
         final int delta = dipToPixel(RESIZE_HANDLE_WIDTH_IN_DP, mDisplayMetrics);
         mTmpTaskForResizePointSearchResult.reset();
-        for (int stackNdx = mTaskStackContainers.size() - 1; stackNdx >= 0; --stackNdx) {
-            final TaskStack stack = mTaskStackContainers.get(stackNdx);
+        for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(stackNdx);
             if (!stack.getWindowConfiguration().canResizeTask()) {
                 return null;
             }
@@ -1863,8 +1845,8 @@
             mTouchExcludeRegion.set(mBaseDisplayRect);
             final int delta = dipToPixel(RESIZE_HANDLE_WIDTH_IN_DP, mDisplayMetrics);
             mTmpRect2.setEmpty();
-            for (int stackNdx = mTaskStackContainers.size() - 1; stackNdx >= 0; --stackNdx) {
-                final TaskStack stack = mTaskStackContainers.get(stackNdx);
+            for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+                final TaskStack stack = mTaskStackContainers.getChildAt(stackNdx);
                 stack.setTouchExcludeRegion(
                         focusedTask, delta, mTouchExcludeRegion, mContentRect, mTmpRect2);
             }
@@ -1895,7 +1877,7 @@
             mTouchExcludeRegion.op(mTmpRegion, Region.Op.UNION);
         }
         // TODO(multi-display): Support docked stacks on secondary displays.
-        if (mDisplayId == DEFAULT_DISPLAY && getDockedStackLocked() != null) {
+        if (mDisplayId == DEFAULT_DISPLAY && getSplitScreenPrimaryStackStack() != null) {
             mDividerControllerLocked.getTouchRegion(mTmpRect);
             mTmpRegion.set(mTmpRect);
             mTouchExcludeRegion.op(mTmpRegion, Op.UNION);
@@ -1912,8 +1894,8 @@
     }
 
     private void resetAnimationBackgroundAnimator() {
-        for (int stackNdx = mTaskStackContainers.size() - 1; stackNdx >= 0; --stackNdx) {
-            mTaskStackContainers.get(stackNdx).resetAnimationBackgroundAnimator();
+        for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            mTaskStackContainers.getChildAt(stackNdx).resetAnimationBackgroundAnimator();
         }
     }
 
@@ -1984,8 +1966,8 @@
             float dividerAnimationTarget) {
         boolean updated = false;
 
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            final TaskStack stack = mTaskStackContainers.get(i);
+        for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(i);
             if (stack == null || !stack.isAdjustedForIme()) {
                 continue;
             }
@@ -2013,8 +1995,8 @@
 
     boolean clearImeAdjustAnimation() {
         boolean changed = false;
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            final TaskStack stack = mTaskStackContainers.get(i);
+        for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(i);
             if (stack != null && stack.isAdjustedForIme()) {
                 stack.resetAdjustedForIme(true /* adjustBoundsNow */);
                 changed  = true;
@@ -2024,8 +2006,8 @@
     }
 
     void beginImeAdjustAnimation() {
-        for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-            final TaskStack stack = mTaskStackContainers.get(i);
+        for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(i);
             if (stack.isVisible() && stack.isAdjustedForIme()) {
                 stack.beginImeAdjustAnimation();
             }
@@ -2054,8 +2036,8 @@
         // - If IME is not visible, divider is not moved and is normal width.
 
         if (imeVisible && dockVisible && (imeOnTop || imeOnBottom) && !dockMinimized) {
-            for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-                final TaskStack stack = mTaskStackContainers.get(i);
+            for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+                final TaskStack stack = mTaskStackContainers.getChildAt(i);
                 final boolean isDockedOnBottom = stack.getDockSide() == DOCKED_BOTTOM;
                 if (stack.isVisible() && (imeOnBottom || isDockedOnBottom)
                         && stack.inSplitScreenWindowingMode()) {
@@ -2067,8 +2049,8 @@
             mDividerControllerLocked.setAdjustedForIme(
                     imeOnBottom /*ime*/, true /*divider*/, true /*animate*/, imeWin, imeHeight);
         } else {
-            for (int i = mTaskStackContainers.size() - 1; i >= 0; --i) {
-                final TaskStack stack = mTaskStackContainers.get(i);
+            for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+                final TaskStack stack = mTaskStackContainers.getChildAt(i);
                 stack.resetAdjustedForIme(!dockVisible);
             }
             mDividerControllerLocked.setAdjustedForIme(
@@ -2099,8 +2081,8 @@
     }
 
     void prepareFreezingTaskBounds() {
-        for (int stackNdx = mTaskStackContainers.size() - 1; stackNdx >= 0; --stackNdx) {
-            final TaskStack stack = mTaskStackContainers.get(stackNdx);
+        for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(stackNdx);
             stack.prepareFreezingTaskBounds();
         }
     }
@@ -2159,22 +2141,22 @@
         final long token = proto.start(fieldId);
         super.writeToProto(proto, WINDOW_CONTAINER);
         proto.write(ID, mDisplayId);
-        for (int stackNdx = mTaskStackContainers.size() - 1; stackNdx >= 0; --stackNdx) {
-            final TaskStack stack = mTaskStackContainers.get(stackNdx);
+        for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(stackNdx);
             stack.writeToProto(proto, STACKS);
         }
         mDividerControllerLocked.writeToProto(proto, DOCKED_STACK_DIVIDER_CONTROLLER);
         mPinnedStackControllerLocked.writeToProto(proto, PINNED_STACK_CONTROLLER);
-        for (int i = mAboveAppWindowsContainers.size() - 1; i >= 0; --i) {
-            final WindowToken windowToken = mAboveAppWindowsContainers.get(i);
+        for (int i = mAboveAppWindowsContainers.getChildCount() - 1; i >= 0; --i) {
+            final WindowToken windowToken = mAboveAppWindowsContainers.getChildAt(i);
             windowToken.writeToProto(proto, ABOVE_APP_WINDOWS);
         }
-        for (int i = mBelowAppWindowsContainers.size() - 1; i >= 0; --i) {
-            final WindowToken windowToken = mBelowAppWindowsContainers.get(i);
+        for (int i = mBelowAppWindowsContainers.getChildCount() - 1; i >= 0; --i) {
+            final WindowToken windowToken = mBelowAppWindowsContainers.getChildAt(i);
             windowToken.writeToProto(proto, BELOW_APP_WINDOWS);
         }
-        for (int i = mImeWindowsContainers.size() - 1; i >= 0; --i) {
-            final WindowToken windowToken = mImeWindowsContainers.get(i);
+        for (int i = mImeWindowsContainers.getChildCount() - 1; i >= 0; --i) {
+            final WindowToken windowToken = mImeWindowsContainers.getChildAt(i);
             windowToken.writeToProto(proto, IME_WINDOWS);
         }
         proto.write(DPI, mBaseDisplayDensity);
@@ -2220,8 +2202,8 @@
 
         pw.println();
         pw.println(prefix + "Application tokens in top down Z order:");
-        for (int stackNdx = mTaskStackContainers.size() - 1; stackNdx >= 0; --stackNdx) {
-            final TaskStack stack = mTaskStackContainers.get(stackNdx);
+        for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            final TaskStack stack = mTaskStackContainers.getChildAt(stackNdx);
             stack.dump(prefix + "  ", pw);
         }
 
@@ -2240,6 +2222,22 @@
         pw.println();
         mDimLayerController.dump(prefix, pw);
         pw.println();
+
+        // Dump stack references
+        final TaskStack homeStack = getHomeStack();
+        if (homeStack != null) {
+            pw.println(prefix + "homeStack=" + homeStack.getName());
+        }
+        final TaskStack pinnedStack = getPinnedStack();
+        if (pinnedStack != null) {
+            pw.println(prefix + "pinnedStack=" + pinnedStack.getName());
+        }
+        final TaskStack splitScreenPrimaryStack = getSplitScreenPrimaryStackStack();
+        if (splitScreenPrimaryStack != null) {
+            pw.println(prefix + "splitScreenPrimaryStack=" + splitScreenPrimaryStack.getName());
+        }
+
+        pw.println();
         mDividerControllerLocked.dump(prefix, pw);
         pw.println();
         mPinnedStackControllerLocked.dump(prefix, pw);
@@ -2265,22 +2263,6 @@
         return stack != null && stack.isVisible();
     }
 
-    /**
-     * @return The docked stack, but only if it is visible, and {@code null} otherwise.
-     */
-    TaskStack getDockedStackLocked() {
-        final TaskStack stack = getStack(WINDOWING_MODE_SPLIT_SCREEN_PRIMARY);
-        return (stack != null && stack.isVisible()) ? stack : null;
-    }
-
-    /**
-     * Like {@link #getDockedStackLocked}, but also returns the docked stack if it's currently not
-     * visible.
-     */
-    TaskStack getDockedStackIgnoringVisibility() {
-        return getStack(WINDOWING_MODE_SPLIT_SCREEN_PRIMARY);
-    }
-
     /** Find the visible, touch-deliverable window under the given point */
     WindowState getTouchableWinAtPointLocked(float xf, float yf) {
         final int x = (int) xf;
@@ -3357,14 +3339,6 @@
      */
     static class DisplayChildWindowContainer<E extends WindowContainer> extends WindowContainer<E> {
 
-        int size() {
-            return mChildren.size();
-        }
-
-        E get(int index) {
-            return mChildren.get(index);
-        }
-
         @Override
         boolean fillsParent() {
             return true;
@@ -3382,25 +3356,108 @@
      */
     private final class TaskStackContainers extends DisplayChildWindowContainer<TaskStack> {
 
+        // Cached reference to some special stacks we tend to get a lot so we don't need to loop
+        // through the list to find them.
+        private TaskStack mHomeStack = null;
+        private TaskStack mPinnedStack = null;
+        private TaskStack mSplitScreenPrimaryStack = null;
+
+        /**
+         * Returns the topmost stack on the display that is compatible with the input windowing mode
+         * and activity type. Null is no compatible stack on the display.
+         */
+        TaskStack getStack(int windowingMode, int activityType) {
+            if (activityType == ACTIVITY_TYPE_HOME) {
+                return mHomeStack;
+            }
+            if (windowingMode == WINDOWING_MODE_PINNED) {
+                return mPinnedStack;
+            } else if (windowingMode == WINDOWING_MODE_SPLIT_SCREEN_PRIMARY) {
+                return mSplitScreenPrimaryStack;
+            }
+            for (int i = mTaskStackContainers.getChildCount() - 1; i >= 0; --i) {
+                final TaskStack stack = mTaskStackContainers.getChildAt(i);
+                if (stack.isCompatible(windowingMode, activityType)) {
+                    return stack;
+                }
+            }
+            return null;
+        }
+
+        @VisibleForTesting
+        TaskStack getTopStack() {
+            return mTaskStackContainers.getChildCount() > 0
+                    ? mTaskStackContainers.getChildAt(mTaskStackContainers.getChildCount() - 1) : null;
+        }
+
+        TaskStack getHomeStack() {
+            if (mHomeStack == null && mDisplayId == DEFAULT_DISPLAY) {
+                Slog.e(TAG_WM, "getHomeStack: Returning null from this=" + this);
+            }
+            return mHomeStack;
+        }
+
+        TaskStack getPinnedStack() {
+            return mPinnedStack;
+        }
+
+        TaskStack getSplitScreenPrimaryStackStack() {
+            return mSplitScreenPrimaryStack;
+        }
+
         /**
          * Adds the stack to this container.
-         * @see WindowManagerService#addStackToDisplay(int, int, boolean)
+         * @see DisplayContent#createStack(int, boolean, StackWindowController)
          */
         void addStackToDisplay(TaskStack stack, boolean onTop) {
-            if (stack.isActivityTypeHome()) {
-                if (mHomeStack != null) {
-                    throw new IllegalArgumentException("attachStack: HOME_STACK_ID (0) not first.");
-                }
-                mHomeStack = stack;
-            }
+            addStackReferenceIfNeeded(stack);
             addChild(stack, onTop);
             stack.onDisplayChanged(DisplayContent.this);
         }
 
-        /** Removes the stack from its container and prepare for changing the parent. */
-        void removeStackFromDisplay(TaskStack stack) {
-            removeChild(stack);
-            stack.onRemovedFromDisplay();
+        void onStackWindowingModeChanged(TaskStack stack) {
+            removeStackReferenceIfNeeded(stack);
+            addStackReferenceIfNeeded(stack);
+            if (stack == mPinnedStack && getTopStack() != stack) {
+                // Looks like this stack changed windowing mode to pinned. Move it to the top.
+                positionChildAt(POSITION_TOP, stack, false /* includingParents */);
+            }
+        }
+
+        private void addStackReferenceIfNeeded(TaskStack stack) {
+            if (stack.isActivityTypeHome()) {
+                if (mHomeStack != null) {
+                    throw new IllegalArgumentException("addStackReferenceIfNeeded: home stack="
+                            + mHomeStack + " already exist on display=" + this + " stack=" + stack);
+                }
+                mHomeStack = stack;
+            }
+            final int windowingMode = stack.getWindowingMode();
+            if (windowingMode == WINDOWING_MODE_PINNED) {
+                if (mPinnedStack != null) {
+                    throw new IllegalArgumentException("addStackReferenceIfNeeded: pinned stack="
+                            + mPinnedStack + " already exist on display=" + this
+                            + " stack=" + stack);
+                }
+                mPinnedStack = stack;
+            } else if (windowingMode == WINDOWING_MODE_SPLIT_SCREEN_PRIMARY) {
+                if (mSplitScreenPrimaryStack != null) {
+                    throw new IllegalArgumentException("addStackReferenceIfNeeded:"
+                            + " split-screen-primary" + " stack=" + mSplitScreenPrimaryStack
+                            + " already exist on display=" + this + " stack=" + stack);
+                }
+                mSplitScreenPrimaryStack = stack;
+            }
+        }
+
+        private void removeStackReferenceIfNeeded(TaskStack stack) {
+            if (stack == mHomeStack) {
+                mHomeStack = null;
+            } else if (stack == mPinnedStack) {
+                mPinnedStack = null;
+            } else if (stack == mSplitScreenPrimaryStack) {
+                mSplitScreenPrimaryStack = null;
+            }
         }
 
         private void addChild(TaskStack stack, boolean toTop) {
@@ -3410,6 +3467,11 @@
             setLayoutNeeded();
         }
 
+        @Override
+        protected void removeChild(TaskStack stack) {
+            super.removeChild(stack);
+            removeStackReferenceIfNeeded(stack);
+        }
 
         @Override
         boolean isOnTop() {
@@ -3452,8 +3514,7 @@
                     : requestedPosition >= topChildPosition;
             int targetPosition = requestedPosition;
 
-            if (toTop && stack.getWindowingMode() != WINDOWING_MODE_PINNED
-                    && getStack(WINDOWING_MODE_PINNED) != null) {
+            if (toTop && stack.getWindowingMode() != WINDOWING_MODE_PINNED && hasPinnedStack()) {
                 // The pinned stack is always the top most stack (always-on-top) when it is present.
                 TaskStack topStack = mChildren.get(topChildPosition);
                 if (topStack.getWindowingMode() != WINDOWING_MODE_PINNED) {
diff --git a/services/core/java/com/android/server/wm/DockedStackDividerController.java b/services/core/java/com/android/server/wm/DockedStackDividerController.java
index 629af66..52526e2 100644
--- a/services/core/java/com/android/server/wm/DockedStackDividerController.java
+++ b/services/core/java/com/android/server/wm/DockedStackDividerController.java
@@ -16,7 +16,6 @@
 
 package com.android.server.wm;
 
-import static android.app.WindowConfiguration.WINDOWING_MODE_SPLIT_SCREEN_PRIMARY;
 import static android.app.WindowConfiguration.WINDOWING_MODE_SPLIT_SCREEN_SECONDARY;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
@@ -320,7 +319,7 @@
         if (mWindow == null) {
             return;
         }
-        TaskStack stack = mDisplayContent.getDockedStackIgnoringVisibility();
+        TaskStack stack = mDisplayContent.getSplitScreenPrimaryStackStackIgnoringVisibility();
 
         // If the stack is invisible, we policy force hide it in WindowAnimator.shouldForceHide
         final boolean visible = stack != null;
@@ -360,7 +359,7 @@
     }
 
     void positionDockedStackedDivider(Rect frame) {
-        TaskStack stack = mDisplayContent.getDockedStackLocked();
+        TaskStack stack = mDisplayContent.getSplitScreenPrimaryStackStack();
         if (stack == null) {
             // Unfortunately we might end up with still having a divider, even though the underlying
             // stack was already removed. This is because we are on AM thread and the removal of the
@@ -457,7 +456,7 @@
         long animDuration = 0;
         if (animate) {
             final TaskStack stack =
-                    mDisplayContent.getStack(WINDOWING_MODE_SPLIT_SCREEN_PRIMARY);
+                    mDisplayContent.getSplitScreenPrimaryStackStackIgnoringVisibility();
             final long transitionDuration = isAnimationMaximizing()
                     ? mService.mAppTransition.getLastClipRevealTransitionDuration()
                     : DEFAULT_APP_TRANSITION_DURATION;
@@ -511,7 +510,8 @@
     void registerDockedStackListener(IDockedStackListener listener) {
         mDockedStackListeners.register(listener);
         notifyDockedDividerVisibilityChanged(wasVisible());
-        notifyDockedStackExistsChanged(mDisplayContent.getDockedStackIgnoringVisibility() != null);
+        notifyDockedStackExistsChanged(
+                mDisplayContent.getSplitScreenPrimaryStackStackIgnoringVisibility() != null);
         notifyDockedStackMinimizedChanged(mMinimizedDock, false /* animate */,
                 isHomeStackResizable());
         notifyAdjustedForImeChanged(mAdjustedForIme, 0 /* animDuration */);
@@ -530,7 +530,7 @@
         final TaskStack stack = targetWindowingMode != WINDOWING_MODE_UNDEFINED
                 ? mDisplayContent.getStack(targetWindowingMode)
                 : null;
-        final TaskStack dockedStack = mDisplayContent.getDockedStackLocked();
+        final TaskStack dockedStack = mDisplayContent.getSplitScreenPrimaryStackStack();
         boolean visibleAndValid = visible && stack != null && dockedStack != null;
         if (visibleAndValid) {
             stack.getDimBounds(mTmpRect);
@@ -598,7 +598,7 @@
     }
 
     private void checkMinimizeChanged(boolean animate) {
-        if (mDisplayContent.getDockedStackIgnoringVisibility() == null) {
+        if (mDisplayContent.getSplitScreenPrimaryStackStackIgnoringVisibility() == null) {
             return;
         }
         final TaskStack homeStack = mDisplayContent.getHomeStack();
@@ -760,7 +760,7 @@
     }
 
     private boolean setMinimizedDockedStack(boolean minimized) {
-        final TaskStack stack = mDisplayContent.getDockedStackIgnoringVisibility();
+        final TaskStack stack = mDisplayContent.getSplitScreenPrimaryStackStackIgnoringVisibility();
         notifyDockedStackMinimizedChanged(minimized, false /* animate */, isHomeStackResizable());
         return stack != null && stack.setAdjustedForMinimizedDock(minimized ? 1f : 0f);
     }
@@ -811,8 +811,7 @@
     }
 
     private boolean animateForMinimizedDockedStack(long now) {
-        final TaskStack stack =
-                mDisplayContent.getStack(WINDOWING_MODE_SPLIT_SCREEN_PRIMARY);
+        final TaskStack stack = mDisplayContent.getSplitScreenPrimaryStackStackIgnoringVisibility();
         if (!mAnimationStarted) {
             mAnimationStarted = true;
             mAnimationStartTime = now;
diff --git a/services/core/java/com/android/server/wm/PinnedStackController.java b/services/core/java/com/android/server/wm/PinnedStackController.java
index ef31598..365366a 100644
--- a/services/core/java/com/android/server/wm/PinnedStackController.java
+++ b/services/core/java/com/android/server/wm/PinnedStackController.java
@@ -417,8 +417,7 @@
                             false /* useCurrentMinEdgeSize */);
                 }
                 final Rect animatingBounds = mTmpAnimatingBoundsRect;
-                final TaskStack pinnedStack =
-                        mDisplayContent.getStack(WINDOWING_MODE_PINNED);
+                final TaskStack pinnedStack = mDisplayContent.getPinnedStack();
                 if (pinnedStack != null) {
                     pinnedStack.getAnimationOrCurrentBounds(animatingBounds);
                 } else {
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 7832f5d..fd57470 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -411,17 +411,6 @@
         }
     }
 
-    TaskStack getStackById(int stackId) {
-        for (int i = mChildren.size() - 1; i >= 0; i--) {
-            final DisplayContent dc = mChildren.get(i);
-            final TaskStack stack = dc.getStackById(stackId);
-            if (stack != null) {
-                return stack;
-            }
-        }
-        return null;
-    }
-
     TaskStack getStack(int windowingMode, int activityType) {
         for (int i = mChildren.size() - 1; i >= 0; i--) {
             final DisplayContent dc = mChildren.get(i);
diff --git a/services/core/java/com/android/server/wm/StackWindowController.java b/services/core/java/com/android/server/wm/StackWindowController.java
index c0e5fd4..1fda832 100644
--- a/services/core/java/com/android/server/wm/StackWindowController.java
+++ b/services/core/java/com/android/server/wm/StackWindowController.java
@@ -43,7 +43,7 @@
 public class StackWindowController
         extends WindowContainerController<TaskStack, StackWindowListener> {
 
-    final int mStackId;
+    private final int mStackId;
 
     private final H mHandler;
 
@@ -72,7 +72,7 @@
                         + " to unknown displayId=" + displayId);
             }
 
-            dc.addStackToDisplay(stackId, onTop, this);
+            dc.createStack(stackId, onTop, this);
             getRawBounds(outBounds);
         }
     }
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 6c6934f..891d637a 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -419,7 +419,7 @@
         return mFillsParent
                 || !inSplitScreenSecondaryWindowingMode()
                 || displayContent == null
-                || displayContent.getDockedStackIgnoringVisibility() != null;
+                || displayContent.getSplitScreenPrimaryStackStackIgnoringVisibility() != null;
     }
 
     /** Original bounds of the task if applicable, otherwise fullscreen rect. */
diff --git a/services/core/java/com/android/server/wm/TaskStack.java b/services/core/java/com/android/server/wm/TaskStack.java
index 3780d19..d170b6f 100644
--- a/services/core/java/com/android/server/wm/TaskStack.java
+++ b/services/core/java/com/android/server/wm/TaskStack.java
@@ -294,7 +294,7 @@
         if (mFillsParent
                 || !inSplitScreenSecondaryWindowingMode()
                 || mDisplayContent == null
-                || mDisplayContent.getDockedStackLocked() != null) {
+                || mDisplayContent.getSplitScreenPrimaryStackStack() != null) {
             return true;
         }
         return false;
@@ -673,6 +673,16 @@
         }
     }
 
+    @Override
+    public void onConfigurationChanged(Configuration newParentConfig) {
+        final int prevWindowingMode = getWindowingMode();
+        super.onConfigurationChanged(newParentConfig);
+        if (mDisplayContent != null && prevWindowingMode != getWindowingMode()) {
+            mDisplayContent.onStackWindowingModeChanged(this);
+        }
+    }
+
+    @Override
     void onDisplayChanged(DisplayContent dc) {
         if (mDisplayContent != null) {
             throw new IllegalStateException("onDisplayChanged: Already attached");
@@ -683,7 +693,7 @@
                 "animation background stackId=" + mStackId);
 
         Rect bounds = null;
-        final TaskStack dockedStack = dc.getDockedStackIgnoringVisibility();
+        final TaskStack dockedStack = dc.getSplitScreenPrimaryStackStackIgnoringVisibility();
         if (inSplitScreenPrimaryWindowingMode()
                 || (dockedStack != null && inSplitScreenSecondaryWindowingMode()
                         && !dockedStack.fillsParent())) {
@@ -762,7 +772,8 @@
             return;
         }
 
-        final TaskStack dockedStack = mDisplayContent.getDockedStackIgnoringVisibility();
+        final TaskStack dockedStack =
+                mDisplayContent.getSplitScreenPrimaryStackStackIgnoringVisibility();
         if (dockedStack == null) {
             // Not sure why you are calling this method when there is no docked stack...
             throw new IllegalStateException(
@@ -889,17 +900,12 @@
     }
 
     @Override
-    void removeImmediately() {
-        super.removeImmediately();
+    void onParentSet() {
+        if (getParent() != null || mDisplayContent == null) {
+            return;
+        }
 
-        onRemovedFromDisplay();
-    }
-
-    /**
-     * Removes the stack it from its current parent, so it can be either destroyed completely or
-     * re-parented.
-     */
-    void onRemovedFromDisplay() {
+        // Looks like the stack was removed from the display. Go ahead and clean things up.
         mDisplayContent.mDimLayerController.removeDimLayerUser(this);
         EventLog.writeEvent(EventLogTags.WM_STACK_REMOVED, mStackId);
 
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 40923c8..1b0825e 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -73,12 +73,12 @@
 
 
     @Override
-    final protected int getChildCount() {
+    protected int getChildCount() {
         return mChildren.size();
     }
 
     @Override
-    final protected E getChildAt(int index) {
+    protected E getChildAt(int index) {
         return mChildren.get(index);
     }
 
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 4279d2e..fd0de86 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -3391,7 +3391,8 @@
 
             // Notify whether the docked stack exists for the current user
             final DisplayContent displayContent = getDefaultDisplayContentLocked();
-            final TaskStack stack = displayContent.getDockedStackIgnoringVisibility();
+            final TaskStack stack =
+                    displayContent.getSplitScreenPrimaryStackStackIgnoringVisibility();
             displayContent.mDividerControllerLocked.notifyDockedStackExistsChanged(
                     stack != null && stack.hasTaskForUser(newUserId));
 
@@ -7132,7 +7133,7 @@
     public int getDockedStackSide() {
         synchronized (mWindowMap) {
             final TaskStack dockedStack = getDefaultDisplayContentLocked()
-                    .getDockedStackIgnoringVisibility();
+                    .getSplitScreenPrimaryStackStackIgnoringVisibility();
             return dockedStack == null ? DOCKED_INVALID : dockedStack.getDockSide();
         }
     }
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 4cb2a9d..1d43c75 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2361,7 +2361,8 @@
                             // just in case they have the divider at an unstable position. Better
                             // also reset drag resizing state, because the owner can't do it
                             // anymore.
-                            final TaskStack stack = dc.getDockedStackIgnoringVisibility();
+                            final TaskStack stack =
+                                    dc.getSplitScreenPrimaryStackStackIgnoringVisibility();
                             if (stack != null) {
                                 stack.resetDockedStackToMiddle();
                             }
diff --git a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
index d362c3b..27c5eab 100644
--- a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
@@ -348,20 +348,17 @@
      */
     @Test
     public void testPinnedStackLocation() {
-        createStackControllerOnStackOnDisplay(
-                WINDOWING_MODE_PINNED, ACTIVITY_TYPE_STANDARD, mDisplayContent);
-        final int initialStackCount = mDisplayContent.getStackCount();
-        // Ensure that the pinned stack was placed at the end
-        assertEquals(initialStackCount - 1,
-                mDisplayContent.getStackPosition(WINDOWING_MODE_PINNED, ACTIVITY_TYPE_STANDARD));
+        final TaskStack pinnedStack = createStackControllerOnStackOnDisplay(
+                WINDOWING_MODE_PINNED, ACTIVITY_TYPE_STANDARD, mDisplayContent).mContainer;
+        // Ensure that the pinned stack is the top stack
+        assertEquals(pinnedStack, mDisplayContent.getPinnedStack());
+        assertEquals(pinnedStack, mDisplayContent.getTopStack());
         // By default, this should try to create a new stack on top
-        createTaskStackOnDisplay(mDisplayContent);
-        final int afterStackCount = mDisplayContent.getStackCount();
-        // Make sure the stack count has increased
-        assertEquals(initialStackCount + 1, afterStackCount);
+        final TaskStack otherStack = createTaskStackOnDisplay(mDisplayContent);
+        // Ensure that the other stack is on the display.
+        assertEquals(mDisplayContent, otherStack.getDisplayContent());
         // Ensure that the pinned stack is still on top
-        assertEquals(afterStackCount - 1,
-                mDisplayContent.getStackPosition(WINDOWING_MODE_PINNED, ACTIVITY_TYPE_STANDARD));
+        assertEquals(pinnedStack, mDisplayContent.getTopStack());
     }
 
     /**