Merge "Add cache validation to ensure the validity"
diff --git a/core/java/android/widget/Magnifier.java b/core/java/android/widget/Magnifier.java
index 3280d47..8e2786d 100644
--- a/core/java/android/widget/Magnifier.java
+++ b/core/java/android/widget/Magnifier.java
@@ -215,11 +215,11 @@
             final Point windowCoords = getCurrentClampedWindowCoordinates();
             final InternalPopupWindow currentWindowInstance = mWindow;
             sPixelCopyHandlerThread.getThreadHandler().post(() -> {
-                if (mWindow != currentWindowInstance) {
-                    // The magnifier was dismissed (and maybe shown again) in the meantime.
-                    return;
-                }
                 synchronized (mLock) {
+                    if (mWindow != currentWindowInstance) {
+                        // The magnifier was dismissed (and maybe shown again) in the meantime.
+                        return;
+                    }
                     mWindow.setContentPositionForNextDraw(windowCoords.x, windowCoords.y);
                 }
             });
diff --git a/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java b/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java
index f1b7eec..ca1b489 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java
@@ -261,6 +261,12 @@
         return mPages.get(0).mColumns;
     }
 
+    public int getNumVisibleTiles() {
+        if (mPages.size() == 0) return 0;
+        TilePage currentPage = mPages.get(getCurrentItem());
+        return currentPage.mRecords.size();
+    }
+
     public void startTileReveal(Set<String> tileSpecs, final Runnable postAnimation) {
         if (tileSpecs.isEmpty() || mPages.size() < 2 || getScrollX() != 0 || !beginFakeDrag()) {
             // Do not start the reveal animation unless there are tiles to animate, multiple
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java b/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java
index 2a4bb60..3744d7d 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java
@@ -155,6 +155,7 @@
         TouchAnimator.Builder translationYBuilder = new Builder();
 
         if (mQsPanel.getHost() == null) return;
+        if (mQuickQsPanel.getTileLayout().getNumVisibleTiles() < 1) return;
         Collection<QSTile> tiles = mQsPanel.getHost().getTiles();
         int count = 0;
         int[] loc1 = new int[2];
@@ -169,6 +170,7 @@
         QSTileLayout tileLayout = mQsPanel.getTileLayout();
         mAllViews.add((View) tileLayout);
         int height = mQs.getView() != null ? mQs.getView().getMeasuredHeight() : 0;
+        int width = mQs.getView() != null ? mQs.getView().getMeasuredWidth() : 0;
         int heightDiff = height - mQs.getHeader().getBottom()
                 + mQs.getHeader().getPaddingBottom();
         firstPageBuilder.addFloat(tileLayout, "translationY", heightDiff, 0);
@@ -181,7 +183,9 @@
             }
             final View tileIcon = tileView.getIcon().getIconView();
             View view = mQs.getView();
-            if (count < mNumQuickTiles && mAllowFancy) {
+
+            // This case: less tiles to animate in small displays.
+            if (count < mQuickQsPanel.getTileLayout().getNumVisibleTiles() && mAllowFancy) {
                 // Quick tiles.
                 QSTileView quickTileView = mQuickQsPanel.getTileView(tile);
                 if (quickTileView == null) continue;
@@ -192,18 +196,26 @@
                 final int xDiff = loc2[0] - loc1[0];
                 final int yDiff = loc2[1] - loc1[1];
                 lastXDiff = loc1[0] - lastX;
-                // Move the quick tile right from its location to the new one.
-                translationXBuilder.addFloat(quickTileView, "translationX", 0, xDiff);
-                translationYBuilder.addFloat(quickTileView, "translationY", 0, yDiff);
 
-                // Counteract the parent translation on the tile. So we have a static base to
-                // animate the label position off from.
-                //firstPageBuilder.addFloat(tileView, "translationY", mQsPanel.getHeight(), 0);
+                if (count < tileLayout.getNumVisibleTiles()) {
+                    // Move the quick tile right from its location to the new one.
+                    translationXBuilder.addFloat(quickTileView, "translationX", 0, xDiff);
+                    translationYBuilder.addFloat(quickTileView, "translationY", 0, yDiff);
 
-                // Move the real tile from the quick tile position to its final
-                // location.
-                translationXBuilder.addFloat(tileView, "translationX", -xDiff, 0);
-                translationYBuilder.addFloat(tileView, "translationY", -yDiff, 0);
+                    // Counteract the parent translation on the tile. So we have a static base to
+                    // animate the label position off from.
+                    //firstPageBuilder.addFloat(tileView, "translationY", mQsPanel.getHeight(), 0);
+
+                    // Move the real tile from the quick tile position to its final
+                    // location.
+                    translationXBuilder.addFloat(tileView, "translationX", -xDiff, 0);
+                    translationYBuilder.addFloat(tileView, "translationY", -yDiff, 0);
+
+                } else { // These tiles disappear when expanding
+                    firstPageBuilder.addFloat(quickTileView, "alpha", 1, 0);
+                    translationYBuilder.addFloat(quickTileView, "translationY", 0, yDiff);
+                    translationXBuilder.addFloat(quickTileView, "translationX", 0, xDiff + width);
+                }
 
                 mQuickQsViews.add(tileView.getIconWithBackground());
                 mAllViews.add(tileView.getIcon());
@@ -218,10 +230,9 @@
                 final int xDiff = loc2[0] - loc1[0];
                 final int yDiff = loc2[1] - loc1[1];
 
-                firstPageBuilder.addFloat(tileView, "translationY", heightDiff, 0);
-                translationXBuilder.addFloat(tileView, "translationX", -xDiff, 0);
+                firstPageBuilder.addFloat(tileView, "translationY", -heightDiff, 0);
                 translationYBuilder.addFloat(tileView, "translationY", -yDiff, 0);
-                translationYBuilder.addFloat(tileIcon, "translationY", -yDiff, 0);
+                translationXBuilder.addFloat(tileView, "translationX", -xDiff, 0);
 
                 mAllViews.add(tileIcon);
             } else {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
index 3fc258b..6c330b0 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
@@ -663,5 +663,7 @@
         void setListening(boolean listening);
 
         default void setExpansion(float expansion) {}
+
+        int getNumVisibleTiles();
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java b/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java
index 1c50f79..556786a 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java
@@ -18,18 +18,17 @@
 
 import android.content.Context;
 import android.content.res.Configuration;
+import android.graphics.Rect;
 import android.util.AttributeSet;
 import android.view.Gravity;
 import android.view.View;
 import android.widget.LinearLayout;
-import android.widget.Space;
 
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
 import com.android.systemui.plugins.qs.QSTile;
 import com.android.systemui.plugins.qs.QSTile.SignalState;
 import com.android.systemui.plugins.qs.QSTile.State;
-import com.android.systemui.plugins.qs.QSTileView;
 import com.android.systemui.qs.customize.QSCustomizer;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.tuner.TunerService.Tunable;
@@ -43,6 +42,7 @@
 public class QuickQSPanel extends QSPanel {
 
     public static final String NUM_QUICK_TILES = "sysui_qqs_count";
+    private static final String TAG = "QuickQSPanel";
 
     private boolean mDisabledByPolicy;
     private static int mDefaultMaxTiles;
@@ -178,121 +178,95 @@
         super.setVisibility(visibility);
     }
 
-    private static class HeaderTileLayout extends LinearLayout implements QSTileLayout {
+    private static class HeaderTileLayout extends TileLayout {
 
-        protected final ArrayList<TileRecord> mRecords = new ArrayList<>();
         private boolean mListening;
-        /** Size of the QS tile (width & height). */
-        private int mTileDimensionSize;
 
         public HeaderTileLayout(Context context) {
             super(context);
             setClipChildren(false);
             setClipToPadding(false);
-
-            mTileDimensionSize = mContext.getResources().getDimensionPixelSize(
-                    R.dimen.qs_quick_tile_size);
-            updateLayoutParams();
         }
 
         @Override
         protected void onConfigurationChanged(Configuration newConfig) {
             super.onConfigurationChanged(newConfig);
-            updateLayoutParams();
+            updateResources();
+        }
+
+        @Override
+        public void onFinishInflate(){
+            updateResources();
         }
 
         private void updateLayoutParams() {
-            setGravity(Gravity.CENTER);
             int width = getResources().getDimensionPixelSize(R.dimen.qs_quick_layout_width);
-            LayoutParams lp = new LayoutParams(width, LayoutParams.MATCH_PARENT);
+            LinearLayout.LayoutParams lp = new LinearLayout.LayoutParams(width, LayoutParams.MATCH_PARENT);
             lp.gravity = Gravity.CENTER_HORIZONTAL;
             setLayoutParams(lp);
         }
 
-        /**
-         * Returns {@link LayoutParams} based on the given {@code spaceWidth}. If the width is 0,
-         * then we're going to have the space expand to take up as much space as possible. If the
-         * width is non-zero, we want the inter-tile spacers to be fixed.
-         */
-        private LayoutParams generateSpaceLayoutParams() {
-            LayoutParams lp = new LayoutParams(0, mTileDimensionSize);
-            lp.weight = 1;
-            lp.gravity = Gravity.CENTER;
-            return lp;
-        }
-
-        @Override
-        public void setListening(boolean listening) {
-            if (mListening == listening) return;
-            mListening = listening;
-            for (TileRecord record : mRecords) {
-                record.tile.setListening(this, mListening);
-            }
-        }
-
-        @Override
-        public void addTile(TileRecord tile) {
-            if (getChildCount() != 0) {
-                addView(new Space(mContext), getChildCount(), generateSpaceLayoutParams());
-            }
-
-            addView(tile.tileView, getChildCount(), generateTileLayoutParams());
-            mRecords.add(tile);
-            tile.tile.setListening(this, mListening);
-        }
-
         private LayoutParams generateTileLayoutParams() {
-            LayoutParams lp = new LayoutParams(mTileDimensionSize, mTileDimensionSize);
-            lp.gravity = Gravity.CENTER;
+            LayoutParams lp = new LayoutParams(mCellWidth, mCellHeight);
             return lp;
         }
 
         @Override
-        public void removeTile(TileRecord tile) {
-            int childIndex = getChildIndex(tile.tileView);
-            // Remove the tile.
-            removeViewAt(childIndex);
-            if (getChildCount() != 0) {
-                // Remove its spacer as well.
-                removeViewAt(childIndex);
-            }
-            mRecords.remove(tile);
-            tile.tile.setListening(this, false);
-        }
-
-        private int getChildIndex(QSTileView tileView) {
-            final int childViewCount = getChildCount();
-            for (int i = 0; i < childViewCount; i++) {
-                if (getChildAt(i) == tileView) {
-                    return i;
-                }
-            }
-            return -1;
+        protected void addTileView(TileRecord tile) {
+            addView(tile.tileView, getChildCount(), generateTileLayoutParams());
         }
 
         @Override
-        public int getOffsetTop(TileRecord tile) {
-            return 0;
+        protected void onLayout(boolean changed, int l, int t, int r, int b) {
+            // We only care about clipping on the right side
+            Rect bounds = new Rect(0, 0, r - l, 10000);
+            setClipBounds(bounds);
+
+            calculateColumns();
+
+            for (int i = 0; i < mRecords.size(); i++) {
+                mRecords.get(i).tileView.setVisibility( i < mColumns ? View.VISIBLE : View.GONE);
+            }
+
+            setAccessibilityOrder();
+            layoutTileRecords(mColumns);
         }
 
         @Override
         public boolean updateResources() {
-            // No resources here.
+            mCellWidth = mContext.getResources().getDimensionPixelSize(R.dimen.qs_quick_tile_size);
+            mCellHeight = mCellWidth;
+
+            updateLayoutParams();
+
             return false;
         }
 
-        @Override
-        public boolean hasOverlappingRendering() {
-            return false;
-        }
+        private boolean calculateColumns() {
+            int prevNumColumns = mColumns;
+            int maxTiles = mRecords.size();
 
-        @Override
-        protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
-            super.onMeasure(widthMeasureSpec, heightMeasureSpec);
-            if (hideOverflowingChildren(widthMeasureSpec)) {
-                return; // Rely on visibility change to trigger remeasure.
+            if (maxTiles == 0){ // Early return during setup
+                mColumns = 0;
+                return true;
             }
 
+            final int availableWidth = getMeasuredWidth() - getPaddingStart() - getPaddingEnd();
+            final int leftoverWithespace = availableWidth - maxTiles * mCellWidth;
+            final int smallestHorizontalMarginNeeded = leftoverWithespace / (maxTiles - 1);
+
+            if (smallestHorizontalMarginNeeded > 0){
+                mCellMarginHorizontal = smallestHorizontalMarginNeeded;
+                mColumns = maxTiles;
+            } else{
+                mColumns = mCellWidth == 0 ? 1 :
+                        Math.min(maxTiles, availableWidth / mCellWidth );
+                mCellMarginHorizontal = (availableWidth - mColumns * mCellWidth) / (mColumns - 1);
+            }
+            return mColumns != prevNumColumns;
+        }
+
+        private void setAccessibilityOrder() {
             if (mRecords != null && mRecords.size() > 0) {
                 View previousView = this;
                 for (TileRecord record : mRecords) {
@@ -306,31 +280,28 @@
             }
         }
 
-        /**
-         * Hide child views that would otherwise be clipped.
-         * @return {@code true} if any child visibilities have changed.
-         */
-        private boolean hideOverflowingChildren(int widthMeasureSpec) {
-            if (getChildCount() == 0) {
-                return false;
+        @Override
+        protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
+            // Measure each QS tile.
+            for (TileRecord record : mRecords) {
+                if (record.tileView.getVisibility() == GONE) continue;
+                record.tileView.measure(exactly(mCellWidth), exactly(mCellHeight));
             }
-            boolean childVisibilityChanged = false;
-            int widthRemaining = MeasureSpec.getSize(widthMeasureSpec)
-                - getChildAt(0).getMeasuredWidth() - getPaddingStart() - getPaddingEnd();
-            for (int i = 2; i < getChildCount(); i += 2) {
-                View tileChild = getChildAt(i);
-                LayoutParams lp = (LayoutParams) tileChild.getLayoutParams();
-                // All Space views have 0 width; only tiles contribute to the total width.
-                widthRemaining = widthRemaining
-                    - tileChild.getMeasuredWidth() - lp.getMarginEnd() - lp.getMarginStart();
-                int newVisibility = widthRemaining < 0 ? View.GONE : View.VISIBLE;
-                if (tileChild.getVisibility() != newVisibility) {
-                    tileChild.setVisibility(newVisibility);
-                    getChildAt(i - 1).setVisibility(newVisibility); // Hide spacer as well.
-                    childVisibilityChanged = true;
-                }
-            }
-            return childVisibilityChanged;
+
+            int height = mCellHeight;
+            if (height < 0) height = 0;
+
+            setMeasuredDimension(MeasureSpec.getSize(widthMeasureSpec), height);
+        }
+
+        @Override
+        public int getNumVisibleTiles() {
+            return mColumns;
+        }
+
+        @Override
+        protected int getColumnStart(int column) {
+            return getPaddingStart() + column *  (mCellWidth + mCellMarginHorizontal);
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java b/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java
index 45d63e0..c67165e 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java
@@ -56,6 +56,10 @@
     public void addTile(TileRecord tile) {
         mRecords.add(tile);
         tile.tile.setListening(this, mListening);
+        addTileView(tile);
+    }
+
+    protected void addTileView(TileRecord tile) {
         addView(tile.tileView);
     }
 
@@ -120,19 +124,18 @@
         return false;
     }
 
-    private static int exactly(int size) {
+    protected static int exactly(int size) {
         return MeasureSpec.makeMeasureSpec(size, MeasureSpec.EXACTLY);
     }
 
-    @Override
-    protected void onLayout(boolean changed, int l, int t, int r, int b) {
-        final int w = getWidth();
+
+    protected void layoutTileRecords(int numRecords) {
         final boolean isRtl = getLayoutDirection() == LAYOUT_DIRECTION_RTL;
         int row = 0;
         int column = 0;
 
         // Layout each QS tile.
-        for (int i = 0; i < mRecords.size(); i++, column++) {
+        for (int i = 0; i < numRecords; i++, column++) {
             // If we reached the last column available to layout a tile, wrap back to the next row.
             if (column == mColumns) {
                 column = 0;
@@ -147,12 +150,22 @@
         }
     }
 
+    @Override
+    protected void onLayout(boolean changed, int l, int t, int r, int b) {
+        layoutTileRecords(mRecords.size());
+    }
+
     private int getRowTop(int row) {
         return row * (mCellHeight + mCellMarginVertical) + mCellMarginTop;
     }
 
-    private int getColumnStart(int column) {
+    protected int getColumnStart(int column) {
         return getPaddingStart() + mSidePadding + mCellMarginHorizontal / 2 +
                 column *  (mCellWidth + mCellMarginHorizontal);
     }
+
+    @Override
+    public int getNumVisibleTiles() {
+        return mRecords.size();
+    }
 }
diff --git a/services/core/java/com/android/server/am/ActivityDisplay.java b/services/core/java/com/android/server/am/ActivityDisplay.java
index 3568a47..aa5a2e0 100644
--- a/services/core/java/com/android/server/am/ActivityDisplay.java
+++ b/services/core/java/com/android/server/am/ActivityDisplay.java
@@ -138,6 +138,10 @@
         return new DisplayWindowController(mDisplay, this);
     }
 
+    DisplayWindowController getWindowContainerController() {
+        return mWindowContainerController;
+    }
+
     void updateBounds() {
         mDisplay.getSize(mTmpDisplaySize);
         setBounds(0, 0, mTmpDisplaySize.x, mTmpDisplaySize.y);
@@ -837,7 +841,7 @@
         if (mStacks.isEmpty() && mRemoved) {
             mWindowContainerController.removeContainer();
             mWindowContainerController = null;
-            mSupervisor.releaseActivityDisplayLocked(mDisplayId);
+            mSupervisor.removeChild(this);
         }
     }
 
diff --git a/services/core/java/com/android/server/am/ActivityStack.java b/services/core/java/com/android/server/am/ActivityStack.java
index 78fef65..355d890 100644
--- a/services/core/java/com/android/server/am/ActivityStack.java
+++ b/services/core/java/com/android/server/am/ActivityStack.java
@@ -775,6 +775,11 @@
                 true /* includingParents */);
     }
 
+    void positionChildWindowContainerAtBottom(TaskRecord child) {
+        mWindowContainerController.positionChildAtBottom(child.getWindowContainerController(),
+                true /* includingParents */);
+    }
+
     /**
      * Returns whether to defer the scheduling of the multi-window mode.
      */
@@ -2859,8 +2864,7 @@
         final int position = getAdjustedPositionForTask(task, mTaskHistory.size(), starting);
         mTaskHistory.add(position, task);
         updateTaskMovement(task, true);
-        mWindowContainerController.positionChildAtTop(task.getWindowContainerController(),
-                true /* includingParents */);
+        positionChildWindowContainerAtTop(task);
     }
 
     private void insertTaskAtBottom(TaskRecord task) {
@@ -2869,8 +2873,7 @@
         final int position = getAdjustedPositionForTask(task, 0, null);
         mTaskHistory.add(position, task);
         updateTaskMovement(task, true);
-        mWindowContainerController.positionChildAtBottom(task.getWindowContainerController(),
-                true /* includingParents */);
+        positionChildWindowContainerAtBottom(task);
     }
 
     void startActivityLocked(ActivityRecord r, ActivityRecord focusedTopActivity,
@@ -3141,8 +3144,7 @@
                     p.reparent(targetTask, 0 /* position - bottom */, "resetTargetTaskIfNeeded");
                 }
 
-                mWindowContainerController.positionChildAtBottom(
-                        targetTask.getWindowContainerController(), false /* includingParents */);
+                positionChildWindowContainerAtBottom(targetTask);
                 replyChainEnd = -1;
             } else if (forceReset || finishOnTaskLaunch || clearWhenTaskReset) {
                 // If the activity should just be removed -- either
@@ -3277,8 +3279,7 @@
                         if (DEBUG_TASKS) Slog.v(TAG_TASKS, "Pulling activity " + p
                                 + " from " + srcPos + " in to resetting task " + task);
                     }
-                    mWindowContainerController.positionChildAtTop(
-                            task.getWindowContainerController(), true /* includingParents */);
+                    positionChildWindowContainerAtTop(task);
 
                     // Now we've moved it in to place...  but what if this is
                     // a singleTop activity and we have put it on top of another
@@ -5239,8 +5240,7 @@
         addTask(task, toTop ? MAX_VALUE : 0, true /* schedulePictureInPictureModeChange */, reason);
         if (toTop) {
             // TODO: figure-out a way to remove this call.
-            mWindowContainerController.positionChildAtTop(task.getWindowContainerController(),
-                    true /* includingParents */);
+            positionChildWindowContainerAtTop(task);
         }
     }
 
diff --git a/services/core/java/com/android/server/am/ActivityStackSupervisor.java b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
index 4cfcbee..1ffdc67 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -176,7 +176,10 @@
 import com.android.server.am.ActivityStack.ActivityState;
 import com.android.server.wm.ActivityTaskManagerInternal.SleepToken;
 import com.android.server.wm.ConfigurationContainer;
+import com.android.server.wm.DisplayWindowController;
 import com.android.server.wm.PinnedStackWindowController;
+import com.android.server.wm.RootWindowContainerController;
+import com.android.server.wm.RootWindowContainerListener;
 import com.android.server.wm.WindowManagerService;
 
 import java.io.FileDescriptor;
@@ -190,7 +193,7 @@
 import java.util.Set;
 
 public class ActivityStackSupervisor extends ConfigurationContainer implements DisplayListener,
-        RecentTasks.Callbacks {
+        RecentTasks.Callbacks, RootWindowContainerListener {
     private static final String TAG = TAG_WITH_CLASS_NAME ? "ActivityStackSupervisor" : TAG_AM;
     private static final String TAG_FOCUS = TAG + POSTFIX_FOCUS;
     private static final String TAG_IDLE = TAG + POSTFIX_IDLE;
@@ -416,9 +419,14 @@
     /** Stack id of the front stack when user switched, indexed by userId. */
     SparseIntArray mUserStackInFront = new SparseIntArray(2);
 
-    // TODO: There should be an ActivityDisplayController coordinating am/wm interaction.
-    /** Mapping from displayId to display current state */
-    private final SparseArray<ActivityDisplay> mActivityDisplays = new SparseArray<>();
+    /** Reference to default display so we can quickly look it up. */
+    private ActivityDisplay mDefaultDisplay;
+
+    /**
+     * List of displays which contain activities, sorted by z-order.
+     * The last entry in the list is the topmost.
+     */
+    private final ArrayList<ActivityDisplay> mActivityDisplays = new ArrayList<>();
 
     private final SparseArray<IntArray> mDisplayAccessUIDs = new SparseArray<>();
 
@@ -453,7 +461,7 @@
 
     @Override
     protected ActivityDisplay getChildAt(int index) {
-        return mActivityDisplays.valueAt(index);
+        return mActivityDisplays.get(index);
     }
 
     @Override
@@ -531,13 +539,6 @@
     private final FindTaskResult mTmpFindTaskResult = new FindTaskResult();
 
     /**
-     * Temp storage for display ids sorted in focus order.
-     * Maps position to id. Using {@link SparseIntArray} instead of {@link ArrayList} because
-     * it's more efficient, as the number of displays is usually small.
-     */
-    private SparseIntArray mTmpOrderedDisplayIds = new SparseIntArray();
-
-    /**
      * Used to keep track whether app visibilities got changed since the last pause. Useful to
      * determine whether to invoke the task stack change listener after pausing.
      */
@@ -569,6 +570,8 @@
 
     private boolean mInitialized;
 
+    private RootWindowContainerController mWindowContainerController;
+
     /**
      * Description of a request to start a new activity, which has been held
      * due to app switches being disabled.
@@ -612,6 +615,11 @@
         mService = service;
     }
 
+    @VisibleForTesting
+    void setWindowContainerController(RootWindowContainerController controller) {
+        mWindowContainerController = controller;
+    }
+
     public void initialize() {
         if (mInitialized) {
             return;
@@ -664,38 +672,57 @@
     void setWindowManager(WindowManagerService wm) {
         mWindowManager = wm;
         getKeyguardController().setWindowManager(wm);
+        setWindowContainerController(new RootWindowContainerController(this));
 
-        mDisplayManager =
-                (DisplayManager) mService.mContext.getSystemService(Context.DISPLAY_SERVICE);
+        mDisplayManager = mService.mContext.getSystemService(DisplayManager.class);
         mDisplayManager.registerDisplayListener(this, null);
         mDisplayManagerInternal = LocalServices.getService(DisplayManagerInternal.class);
 
-        Display[] displays = mDisplayManager.getDisplays();
+        final Display[] displays = mDisplayManager.getDisplays();
         for (int displayNdx = displays.length - 1; displayNdx >= 0; --displayNdx) {
             final Display display = displays[displayNdx];
-            ActivityDisplay activityDisplay = new ActivityDisplay(this, display);
-            mActivityDisplays.put(display.getDisplayId(), activityDisplay);
+            final ActivityDisplay activityDisplay = new ActivityDisplay(this, display);
+            if (activityDisplay.mDisplayId == DEFAULT_DISPLAY) {
+                mDefaultDisplay = activityDisplay;
+            }
+            addChild(activityDisplay, ActivityDisplay.POSITION_TOP);
             calculateDefaultMinimalSizeOfResizeableTasks(activityDisplay);
         }
 
         final ActivityDisplay defaultDisplay = getDefaultDisplay();
         mHomeStack = mLastFocusedStack = defaultDisplay.getOrCreateStack(
                 WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_HOME, ON_TOP);
+        positionChildAt(defaultDisplay, ActivityDisplay.POSITION_TOP);
+    }
+
+    /** Change the z-order of the given display. */
+    private void positionChildAt(ActivityDisplay display, int position) {
+        if (position >= mActivityDisplays.size()) {
+            position = mActivityDisplays.size() - 1;
+        } else if (position < 0) {
+            position = 0;
+        }
+
+        if (mActivityDisplays.isEmpty()) {
+            mActivityDisplays.add(display);
+        } else if (mActivityDisplays.get(position) != display) {
+            mActivityDisplays.remove(display);
+            mActivityDisplays.add(position, display);
+        }
+    }
+
+    @Override
+    public void onChildPositionChanged(DisplayWindowController childController, int position) {
+        // Assume AM lock is held from positionChildAt of controller in each hierarchy.
+        final ActivityDisplay display = getActivityDisplay(childController.getDisplayId());
+        if (display != null) {
+            positionChildAt(display, position);
+        }
     }
 
     ActivityStack getTopDisplayFocusedStack() {
-        mWindowManager.getDisplaysInFocusOrder(mTmpOrderedDisplayIds);
-
-        for (int i = mTmpOrderedDisplayIds.size() - 1; i >= 0; --i) {
-            final int displayId = mTmpOrderedDisplayIds.get(i);
-            final ActivityDisplay display = mActivityDisplays.get(displayId);
-
-            // If WindowManagerService has encountered the display before we have, ignore as there
-            // will be no stacks present and therefore no activities.
-            if (display == null) {
-                continue;
-            }
-            final ActivityStack focusedStack = display.getFocusedStack();
+        for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
+            final ActivityStack focusedStack = mActivityDisplays.get(i).getFocusedStack();
             if (focusedStack != null) {
                 return focusedStack;
             }
@@ -718,16 +745,8 @@
         }
         // The top focused stack might not have a resumed activity yet - look on all displays in
         // focus order.
-        mWindowManager.getDisplaysInFocusOrder(mTmpOrderedDisplayIds);
-        for (int i = mTmpOrderedDisplayIds.size() - 1; i >= 0; --i) {
-            final int displayId = mTmpOrderedDisplayIds.get(i);
-            final ActivityDisplay display = mActivityDisplays.get(displayId);
-
-            // If WindowManagerService has encountered the display before we have, ignore as there
-            // will be no stacks present and therefore no activities.
-            if (display == null) {
-                continue;
-            }
+        for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
+            final ActivityDisplay display = mActivityDisplays.get(i);
             final ActivityRecord resumedActivityOnDisplay = display.getResumedActivity();
             if (resumedActivityOnDisplay != null) {
                 return resumedActivityOnDisplay;
@@ -848,7 +867,7 @@
 
         int numDisplays = mActivityDisplays.size();
         for (int displayNdx = 0; displayNdx < numDisplays; ++displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 final TaskRecord task = stack.taskForIdLocked(id);
@@ -905,7 +924,7 @@
     ActivityRecord isInAnyStackLocked(IBinder token) {
         int numDisplays = mActivityDisplays.size();
         for (int displayNdx = 0; displayNdx < numDisplays; ++displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 final ActivityRecord r = stack.isInStackLocked(token);
@@ -947,7 +966,7 @@
         mWindowManager.deferSurfaceLayout();
         try {
             for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-                final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+                final ActivityDisplay display = mActivityDisplays.get(displayNdx);
                 for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                     final ActivityStack stack = display.getChildAt(stackNdx);
                     final List<TaskRecord> tasks = stack.getAllTasks();
@@ -1011,7 +1030,7 @@
         final String processName = app.processName;
         boolean didSomething = false;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 if (!isTopDisplayFocusedStack(stack)) {
@@ -1046,7 +1065,7 @@
 
     boolean allResumedActivitiesIdle() {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 if (!isTopDisplayFocusedStack(stack) || stack.numActivities() == 0) {
@@ -1067,7 +1086,7 @@
 
     boolean allResumedActivitiesComplete() {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 if (isTopDisplayFocusedStack(stack)) {
@@ -1090,7 +1109,7 @@
     private boolean allResumedActivitiesVisible() {
         boolean foundResumed = false;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 final ActivityRecord r = stack.getResumedActivity();
@@ -1116,7 +1135,7 @@
     boolean pauseBackStacks(boolean userLeaving, ActivityRecord resuming, boolean dontWait) {
         boolean someActivityPaused = false;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            someActivityPaused |= mActivityDisplays.valueAt(displayNdx)
+            someActivityPaused |= mActivityDisplays.get(displayNdx)
                     .pauseBackStacks(userLeaving, resuming, dontWait);
         }
         return someActivityPaused;
@@ -1125,7 +1144,7 @@
     boolean allPausedActivitiesComplete() {
         boolean pausing = true;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 final ActivityRecord r = stack.mPausingActivity;
@@ -1145,7 +1164,7 @@
 
     void cancelInitializingActivities() {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.cancelInitializingActivities();
@@ -1266,17 +1285,8 @@
         }
 
         // Look in other non-focused and non-home stacks.
-        mWindowManager.getDisplaysInFocusOrder(mTmpOrderedDisplayIds);
-
-        for (int i = mTmpOrderedDisplayIds.size() - 1; i >= 0; --i) {
-            final int displayId = mTmpOrderedDisplayIds.get(i);
-            final ActivityDisplay display = mActivityDisplays.get(displayId);
-
-            // If WindowManagerService has encountered the display before we have, ignore as there
-            // will be no stacks present and therefore no activities.
-            if (display == null) {
-                continue;
-            }
+        for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
+            final ActivityDisplay display = mActivityDisplays.get(i);
 
             // TODO: We probably want to consider the top fullscreen stack as we could have a pinned
             // stack on top.
@@ -1757,7 +1767,7 @@
             boolean noResumedActivities = true;
             boolean allFocusedProcessesDiffer = true;
             for (int displayNdx = 0; displayNdx < mActivityDisplays.size(); ++displayNdx) {
-                final ActivityDisplay activityDisplay = mActivityDisplays.valueAt(displayNdx);
+                final ActivityDisplay activityDisplay = mActivityDisplays.get(displayNdx);
                 final ActivityRecord resumedActivity = activityDisplay.getResumedActivity();
                 final WindowProcessController resumedActivityProcess =
                     resumedActivity == null ? null : resumedActivity.app;
@@ -1927,7 +1937,7 @@
     void updateUIDsPresentOnDisplay() {
         mDisplayAccessUIDs.clear();
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay activityDisplay = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay activityDisplay = mActivityDisplays.get(displayNdx);
             // Only bother calculating the whitelist for private displays
             if (activityDisplay.isPrivate()) {
                 mDisplayAccessUIDs.append(
@@ -2173,7 +2183,7 @@
     boolean handleAppDiedLocked(WindowProcessController app) {
         boolean hasVisibleActivities = false;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 hasVisibleActivities |= stack.handleAppDiedLocked(app);
@@ -2184,7 +2194,7 @@
 
     void closeSystemDialogsLocked() {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.closeSystemDialogsLocked();
@@ -2213,7 +2223,7 @@
             boolean doit, boolean evenPersistent, int userId) {
         boolean didSomething = false;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 if (stack.finishDisabledPackageActivitiesLocked(
@@ -2235,7 +2245,7 @@
         // hosted by the process that is actually still the foreground.
         WindowProcessController fgApp = null;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 if (isTopDisplayFocusedStack(stack)) {
@@ -2277,7 +2287,7 @@
 
         // Resume all top activities in focused stacks on all displays.
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             final ActivityStack focusedStack = display.getFocusedStack();
             if (focusedStack == null) {
                 continue;
@@ -2296,7 +2306,7 @@
 
     void updateActivityApplicationInfoLocked(ApplicationInfo aInfo) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.updateActivityApplicationInfoLocked(aInfo);
@@ -2314,7 +2324,7 @@
         TaskRecord finishedTask = null;
         ActivityStack focusedStack = getTopDisplayFocusedStack();
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             // 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 < display.getChildCount(); ++stackNdx) {
@@ -2330,7 +2340,7 @@
 
     void finishVoiceTask(IVoiceInteractionSession session) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             final int numStacks = display.getChildCount();
             for (int stackNdx = 0; stackNdx < numStacks; ++stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
@@ -2417,7 +2427,7 @@
 
     protected <T extends ActivityStack> T getStack(int stackId) {
         for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
-            final T stack = mActivityDisplays.valueAt(i).getStack(stackId);
+            final T stack = mActivityDisplays.get(i).getStack(stackId);
             if (stack != null) {
                 return stack;
             }
@@ -2428,7 +2438,7 @@
     /** @see ActivityDisplay#getStack(int, int) */
     private <T extends ActivityStack> T getStack(int windowingMode, int activityType) {
         for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
-            final T stack = mActivityDisplays.valueAt(i).getStack(windowingMode, activityType);
+            final T stack = mActivityDisplays.get(i).getStack(windowingMode, activityType);
             if (stack != null) {
                 return stack;
             }
@@ -2642,19 +2652,12 @@
         }
 
         // Now look through all displays
-        mWindowManager.getDisplaysInFocusOrder(mTmpOrderedDisplayIds);
-        for (int i = mTmpOrderedDisplayIds.size() - 1; i >= 0; --i) {
-            final int displayId = mTmpOrderedDisplayIds.get(i);
-            if (displayId == preferredDisplay.mDisplayId) {
+        for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
+            final ActivityDisplay display = mActivityDisplays.get(i);
+            if (display == preferredDisplay) {
                 // We've already checked this one
                 continue;
             }
-            // If a display is registered in WM, it must also be available in AM.
-            final ActivityDisplay display = getActivityDisplayOrCreateLocked(displayId);
-            if (display == null) {
-                // Looks like the display no longer exists in the system...
-                continue;
-            }
             final ActivityStack nextFocusableStack = display.getNextFocusableStack(currentFocus,
                     ignoreCurrent);
             if (nextFocusableStack != null) {
@@ -2676,13 +2679,12 @@
      * @return Next valid {@link ActivityStack}, null if not found.
      */
     ActivityStack getNextValidLaunchStackLocked(@NonNull ActivityRecord r, int currentFocus) {
-        mWindowManager.getDisplaysInFocusOrder(mTmpOrderedDisplayIds);
-        for (int i = mTmpOrderedDisplayIds.size() - 1; i >= 0; --i) {
-            final int displayId = mTmpOrderedDisplayIds.get(i);
-            if (displayId == currentFocus) {
+        for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
+            final ActivityDisplay display = mActivityDisplays.get(i);
+            if (display.mDisplayId == currentFocus) {
                 continue;
             }
-            final ActivityStack stack = getValidLaunchStackOnDisplay(displayId, r,
+            final ActivityStack stack = getValidLaunchStackOnDisplay(display.mDisplayId, r,
                     null /* options */);
             if (stack != null) {
                 return stack;
@@ -3081,13 +3083,13 @@
      */
     void removeStacksInWindowingModes(int... windowingModes) {
         for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
-            mActivityDisplays.valueAt(i).removeStacksInWindowingModes(windowingModes);
+            mActivityDisplays.get(i).removeStacksInWindowingModes(windowingModes);
         }
     }
 
     void removeStacksWithActivityTypes(int... activityTypes) {
         for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
-            mActivityDisplays.valueAt(i).removeStacksWithActivityTypes(activityTypes);
+            mActivityDisplays.get(i).removeStacksWithActivityTypes(activityTypes);
         }
     }
 
@@ -3463,7 +3465,7 @@
         ActivityRecord affinityMatch = null;
         if (DEBUG_TASKS) Slog.d(TAG_TASKS, "Looking for task of " + r);
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 if (!r.hasCompatibleActivityType(stack)) {
@@ -3500,7 +3502,7 @@
     ActivityRecord findActivityLocked(Intent intent, ActivityInfo info,
             boolean compareIntentFilters) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 final ActivityRecord ar = stack.findActivityLocked(
@@ -3515,7 +3517,7 @@
 
     boolean hasAwakeDisplay() {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             if (!display.shouldSleep()) {
                 return true;
             }
@@ -3543,7 +3545,7 @@
 
     void prepareForShutdownLocked() {
         for (int i = 0; i < mActivityDisplays.size(); i++) {
-            createSleepTokenLocked("shutdown", mActivityDisplays.keyAt(i));
+            createSleepTokenLocked("shutdown", mActivityDisplays.get(i).mDisplayId);
         }
     }
 
@@ -3586,7 +3588,7 @@
     void applySleepTokensLocked(boolean applyToStacks) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
             // Set the sleeping state of the display.
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             final boolean displayShouldSleep = display.shouldSleep();
             if (displayShouldSleep == display.isSleeping()) {
                 continue;
@@ -3666,7 +3668,7 @@
     private boolean putStacksToSleepLocked(boolean allowDelay, boolean shuttingDown) {
         boolean allSleep = true;
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 if (allowDelay) {
@@ -3697,7 +3699,7 @@
 
     void handleAppCrashLocked(WindowProcessController app) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.handleAppCrashLocked(app);
@@ -3746,7 +3748,7 @@
         try {
             // First the front stacks. In case any are not fullscreen and are in front of home.
             for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-                final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+                final ActivityDisplay display = mActivityDisplays.get(displayNdx);
                 for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                     final ActivityStack stack = display.getChildAt(stackNdx);
                     stack.ensureActivitiesVisibleLocked(starting, configChanges, preserveWindows,
@@ -3760,7 +3762,7 @@
 
     void addStartingWindowsForVisibleActivities(boolean taskSwitch) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.addStartingWindowsForVisibleActivities(taskSwitch);
@@ -3778,7 +3780,7 @@
         }
         mTaskLayersChanged = false;
         for (int displayNdx = 0; displayNdx < mActivityDisplays.size(); displayNdx++) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             int baseLayer = 0;
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
@@ -3789,7 +3791,7 @@
 
     void clearOtherAppTimeTrackers(AppTimeTracker except) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.clearOtherAppTimeTrackers(except);
@@ -3799,7 +3801,7 @@
 
     void scheduleDestroyAllActivities(WindowProcessController app, String reason) {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.scheduleDestroyActivities(app, reason);
@@ -3818,7 +3820,7 @@
         // let's iterate through the tasks and release the oldest one.
         final int numDisplays = mActivityDisplays.size();
         for (int displayNdx = 0; displayNdx < numDisplays; ++displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             final int stackCount = display.getChildCount();
             // Step through all stacks starting from behind, to hit the oldest things first.
             for (int stackNdx = 0; stackNdx < stackCount; stackNdx++) {
@@ -3849,7 +3851,7 @@
 
         mStartingUsers.add(uss);
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 stack.switchUserLocked(userId);
@@ -3953,7 +3955,7 @@
 
     void validateTopActivitiesLocked() {
         for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 final ActivityRecord r = stack.topRunningActivityLocked();
@@ -3984,7 +3986,7 @@
 
     public void dumpDisplays(PrintWriter pw) {
         for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(i);
+            final ActivityDisplay display = mActivityDisplays.get(i);
             pw.print("[id:" + display.mDisplayId + " stacks:");
             display.dumpStacks(pw);
             pw.print("]");
@@ -3998,7 +4000,7 @@
         pw.println("mCurTaskIdForUser=" + mCurTaskIdForUser);
         pw.print(prefix); pw.println("mUserStackInFront=" + mUserStackInFront);
         for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(i);
+            final ActivityDisplay display = mActivityDisplays.get(i);
             display.dump(pw, prefix);
         }
         if (!mWaitingForActivityVisible.isEmpty()) {
@@ -4018,7 +4020,7 @@
         final long token = proto.start(fieldId);
         super.writeToProto(proto, CONFIGURATION_CONTAINER, false /* trim */);
         for (int displayNdx = 0; displayNdx < mActivityDisplays.size(); ++displayNdx) {
-            final ActivityDisplay activityDisplay = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay activityDisplay = mActivityDisplays.get(displayNdx);
             activityDisplay.writeToProto(proto, DISPLAYS);
         }
         getKeyguardController().writeToProto(proto, KEYGUARD_CONTROLLER);
@@ -4047,7 +4049,7 @@
         pw.print(prefix); pw.println("Display override configurations:");
         final int displayCount = mActivityDisplays.size();
         for (int i = 0; i < displayCount; i++) {
-            final ActivityDisplay activityDisplay = mActivityDisplays.valueAt(i);
+            final ActivityDisplay activityDisplay = mActivityDisplays.get(i);
             pw.print(prefix); pw.print("  "); pw.print(activityDisplay.mDisplayId); pw.print(": ");
                     pw.println(activityDisplay.getOverrideConfiguration());
         }
@@ -4065,7 +4067,7 @@
             ArrayList<ActivityRecord> activities = new ArrayList<>();
             int numDisplays = mActivityDisplays.size();
             for (int displayNdx = 0; displayNdx < numDisplays; ++displayNdx) {
-                final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+                final ActivityDisplay display = mActivityDisplays.get(displayNdx);
                 for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                     final ActivityStack stack = display.getChildAt(stackNdx);
                     if (!dumpVisibleStacksOnly || stack.shouldBeVisible(null)) {
@@ -4096,11 +4098,11 @@
             boolean dumpClient, String dumpPackage) {
         boolean printed = false;
         boolean needSep = false;
-        for (int displayNdx = 0; displayNdx < mActivityDisplays.size(); ++displayNdx) {
-            ActivityDisplay activityDisplay = mActivityDisplays.valueAt(displayNdx);
+        for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
+            ActivityDisplay activityDisplay = mActivityDisplays.get(displayNdx);
             pw.print("Display #"); pw.print(activityDisplay.mDisplayId);
                     pw.println(" (activities from top to bottom):");
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 pw.println();
@@ -4299,12 +4301,18 @@
 
     // TODO: Look into consolidating with getActivityDisplayOrCreateLocked()
     ActivityDisplay getActivityDisplay(int displayId) {
-        return mActivityDisplays.get(displayId);
+        for (int i = mActivityDisplays.size() - 1; i >= 0; --i) {
+            final ActivityDisplay activityDisplay = mActivityDisplays.get(i);
+            if (activityDisplay.mDisplayId == displayId) {
+                return activityDisplay;
+            }
+        }
+        return null;
     }
 
     // TODO(multi-display): Look at all callpoints to make sure they make sense in multi-display.
     ActivityDisplay getDefaultDisplay() {
-        return mActivityDisplays.get(DEFAULT_DISPLAY);
+        return mDefaultDisplay;
     }
 
     /**
@@ -4313,7 +4321,7 @@
      */
     // TODO: Look into consolidating with getActivityDisplay()
     ActivityDisplay getActivityDisplayOrCreateLocked(int displayId) {
-        ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+        ActivityDisplay activityDisplay = getActivityDisplay(displayId);
         if (activityDisplay != null) {
             return activityDisplay;
         }
@@ -4328,15 +4336,23 @@
         }
         // The display hasn't been added to ActivityManager yet, create a new record now.
         activityDisplay = new ActivityDisplay(this, display);
-        attachDisplay(activityDisplay);
+        addChild(activityDisplay, ActivityDisplay.POSITION_BOTTOM);
         calculateDefaultMinimalSizeOfResizeableTasks(activityDisplay);
         mWindowManager.onDisplayAdded(displayId);
         return activityDisplay;
     }
 
     @VisibleForTesting
-    void attachDisplay(ActivityDisplay display) {
-        mActivityDisplays.put(display.mDisplayId, display);
+    void addChild(ActivityDisplay activityDisplay, int position) {
+        positionChildAt(activityDisplay, position);
+        mWindowContainerController.positionChildAt(
+                activityDisplay.getWindowContainerController(), position);
+    }
+
+    void removeChild(ActivityDisplay activityDisplay) {
+        // The caller must tell the controller of {@link ActivityDisplay} to release its container
+        // {@link DisplayContent}. That is done in {@link ActivityDisplay#releaseSelfIfNeeded}).
+        mActivityDisplays.remove(activityDisplay);
     }
 
     private void calculateDefaultMinimalSizeOfResizeableTasks(ActivityDisplay display) {
@@ -4351,7 +4367,7 @@
         }
 
         synchronized (mService.mGlobalLock) {
-            final ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+            final ActivityDisplay activityDisplay = getActivityDisplay(displayId);
             if (activityDisplay == null) {
                 return;
             }
@@ -4362,14 +4378,9 @@
         }
     }
 
-    void releaseActivityDisplayLocked(int displayId) {
-        mActivityDisplays.remove(displayId);
-    }
-
-
     private void handleDisplayChanged(int displayId) {
         synchronized (mService.mGlobalLock) {
-            ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+            ActivityDisplay activityDisplay = getActivityDisplay(displayId);
             // TODO: The following code block should be moved into {@link ActivityDisplay}.
             if (activityDisplay != null) {
                 // The window policy is responsible for stopping activities on the default display
@@ -4392,7 +4403,7 @@
     }
 
     SleepToken createSleepTokenLocked(String tag, int displayId) {
-        ActivityDisplay display = mActivityDisplays.get(displayId);
+        final ActivityDisplay display = getActivityDisplay(displayId);
         if (display == null) {
             throw new IllegalArgumentException("Invalid display: " + displayId);
         }
@@ -4406,7 +4417,7 @@
     private void removeSleepTokenLocked(SleepTokenImpl token) {
         mSleepTokens.remove(token);
 
-        ActivityDisplay display = mActivityDisplays.get(token.mDisplayId);
+        final ActivityDisplay display = getActivityDisplay(token.mDisplayId);
         if (display != null) {
             display.mAllSleepTokens.remove(token);
             if (display.mAllSleepTokens.isEmpty()) {
@@ -4429,7 +4440,7 @@
 
     private StackInfo getStackInfo(ActivityStack stack) {
         final int displayId = stack.mDisplayId;
-        final ActivityDisplay display = mActivityDisplays.get(displayId);
+        final ActivityDisplay display = getActivityDisplay(displayId);
         StackInfo info = new StackInfo();
         stack.getWindowContainerBounds(info.bounds);
         info.displayId = displayId;
@@ -4483,7 +4494,7 @@
     ArrayList<StackInfo> getAllStackInfosLocked() {
         ArrayList<StackInfo> list = new ArrayList<>();
         for (int displayNdx = 0; displayNdx < mActivityDisplays.size(); ++displayNdx) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = mActivityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 list.add(getStackInfo(stack));
@@ -4765,14 +4776,12 @@
     }
 
     ActivityStack findStackBehind(ActivityStack stack) {
-        // TODO(multi-display): We are only looking for stacks on the default display.
-        final ActivityDisplay display = mActivityDisplays.get(DEFAULT_DISPLAY);
-        if (display == null) {
-            return null;
-        }
-        for (int i = display.getChildCount() - 1; i >= 0; i--) {
-            if (display.getChildAt(i) == stack && i > 0) {
-                return display.getChildAt(i - 1);
+        final ActivityDisplay display = getActivityDisplay(stack.mDisplayId);
+        if (display != null) {
+            for (int i = display.getChildCount() - 1; i >= 0; i--) {
+                if (display.getChildAt(i) == stack && i > 0) {
+                    return display.getChildAt(i - 1);
+                }
             }
         }
         throw new IllegalStateException("Failed to find a stack behind stack=" + stack
@@ -4904,7 +4913,7 @@
         final ActivityStack topFocusedStack = getTopDisplayFocusedStack();
         // Traverse all displays.
         for (int i = mActivityDisplays.size() - 1; i >= 0; i--) {
-            final ActivityDisplay display = mActivityDisplays.valueAt(i);
+            final ActivityDisplay display = mActivityDisplays.get(i);
             // Traverse all stacks on a display.
             for (int j = display.getChildCount() - 1; j >= 0; --j) {
                 final ActivityStack stack = display.getChildAt(j);
diff --git a/services/core/java/com/android/server/am/RunningTasks.java b/services/core/java/com/android/server/am/RunningTasks.java
index 7008cee..d878f51 100644
--- a/services/core/java/com/android/server/am/RunningTasks.java
+++ b/services/core/java/com/android/server/am/RunningTasks.java
@@ -16,13 +16,9 @@
 
 package com.android.server.am;
 
-import static com.android.server.am.ActivityManagerDebugConfig.TAG_AM;
-import static com.android.server.am.ActivityManagerDebugConfig.TAG_WITH_CLASS_NAME;
-
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.WindowConfiguration.ActivityType;
 import android.app.WindowConfiguration.WindowingMode;
-import android.util.SparseArray;
 
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -45,7 +41,7 @@
     private final ArrayList<TaskRecord> mTmpStackTasks = new ArrayList<>();
 
     void getTasks(int maxNum, List<RunningTaskInfo> list, @ActivityType int ignoreActivityType,
-            @WindowingMode int ignoreWindowingMode, SparseArray<ActivityDisplay> activityDisplays,
+            @WindowingMode int ignoreWindowingMode, ArrayList<ActivityDisplay> activityDisplays,
             int callingUid, boolean allowed) {
         // Return early if there are no tasks to fetch
         if (maxNum <= 0) {
@@ -56,7 +52,7 @@
         mTmpSortedSet.clear();
         final int numDisplays = activityDisplays.size();
         for (int displayNdx = 0; displayNdx < numDisplays; ++displayNdx) {
-            final ActivityDisplay display = activityDisplays.valueAt(displayNdx);
+            final ActivityDisplay display = activityDisplays.get(displayNdx);
             for (int stackNdx = display.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
                 final ActivityStack stack = display.getChildAt(stackNdx);
                 mTmpStackTasks.clear();
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index ba46737..32fa9bf 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -921,6 +921,11 @@
     }
 
     @Override
+    DisplayWindowController getController() {
+        return (DisplayWindowController) super.getController();
+    }
+
+    @Override
     public Display getDisplay() {
         return mDisplay;
     }
diff --git a/services/core/java/com/android/server/wm/DisplayWindowController.java b/services/core/java/com/android/server/wm/DisplayWindowController.java
index 74a8a35..76b6dbe 100644
--- a/services/core/java/com/android/server/wm/DisplayWindowController.java
+++ b/services/core/java/com/android/server/wm/DisplayWindowController.java
@@ -73,6 +73,10 @@
         // override configuration propagation to just here.
     }
 
+    public int getDisplayId() {
+        return mDisplayId;
+    }
+
     /**
      * Positions the task stack at the given position in the task stack container.
      */
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index d8cbb26..86b14337 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -174,24 +174,6 @@
         return null;
     }
 
-    /**
-     * Get an array with display ids ordered by focus priority - last items should be given
-     * focus first. Sparse array just maps position to displayId.
-     */
-    void getDisplaysInFocusOrder(SparseIntArray displaysInFocusOrder) {
-        displaysInFocusOrder.clear();
-
-        final int size = mChildren.size();
-        for (int i = 0; i < size; ++i) {
-            final DisplayContent displayContent = mChildren.get(i);
-            if (displayContent.isRemovalDeferred()) {
-                // Don't report displays that are going to be removed soon.
-                continue;
-            }
-            displaysInFocusOrder.put(i, displayContent.getDisplayId());
-        }
-    }
-
     DisplayContent getDisplayContent(int displayId) {
         for (int i = mChildren.size() - 1; i >= 0; --i) {
             final DisplayContent current = mChildren.get(i);
@@ -1098,6 +1080,25 @@
     }
 
     @Override
+    void positionChildAt(int position, DisplayContent child, boolean includingParents) {
+        super.positionChildAt(position, child, includingParents);
+        final RootWindowContainerController controller = getController();
+        if (controller != null) {
+            controller.onChildPositionChanged(child, position);
+        }
+    }
+
+    void positionChildAt(int position, DisplayContent child) {
+        // Only called from controller so no need to notify the change to controller.
+        super.positionChildAt(position, child, false /* includingParents */);
+    }
+
+    @Override
+    RootWindowContainerController getController() {
+        return (RootWindowContainerController) super.getController();
+    }
+
+    @Override
     void scheduleAnimation() {
         mService.scheduleAnimationLocked();
     }
diff --git a/services/core/java/com/android/server/wm/RootWindowContainerController.java b/services/core/java/com/android/server/wm/RootWindowContainerController.java
new file mode 100644
index 0000000..93be6e9
--- /dev/null
+++ b/services/core/java/com/android/server/wm/RootWindowContainerController.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wm;
+
+/**
+ * Controller for the root container. This is created by activity manager to link activity
+ * stack supervisor to the root window container they use in window manager.
+ */
+public class RootWindowContainerController
+        extends WindowContainerController<RootWindowContainer, RootWindowContainerListener> {
+
+    public RootWindowContainerController(RootWindowContainerListener listener) {
+        super(listener, WindowManagerService.getInstance());
+        synchronized (mWindowMap) {
+            mRoot.setController(this);
+        }
+    }
+
+    void onChildPositionChanged(DisplayContent child, int position) {
+        // This callback invokes to AM directly so here assumes AM lock is held. If there is another
+        // path called only with WM lock, it should change to use handler to post or move outside of
+        // WM lock with adding AM lock.
+        mListener.onChildPositionChanged(child.getController(), position);
+    }
+
+    /** Move the display to the given position. */
+    public void positionChildAt(DisplayWindowController child, int position) {
+        synchronized (mWindowMap) {
+            mContainer.positionChildAt(position, child.mContainer);
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/wm/RootWindowContainerListener.java b/services/core/java/com/android/server/wm/RootWindowContainerListener.java
new file mode 100644
index 0000000..f413e3f7
--- /dev/null
+++ b/services/core/java/com/android/server/wm/RootWindowContainerListener.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wm;
+
+/**
+ * Interface used by the creator of {@link RootWindowContainerController} to notify the changes to
+ * the display container in activity manager.
+ */
+public interface RootWindowContainerListener extends WindowContainerListener {
+    /** Called when the z-order of display is changed. */
+    void onChildPositionChanged(DisplayWindowController childController, int position);
+}
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index f8d0c72..e80a47e 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -5345,17 +5345,6 @@
         mWindowPlacerLocked.performSurfacePlacement();
     }
 
-    /**
-     * Get an array with display ids ordered by focus priority - last items should be given
-     * focus first. Sparse array just maps position to displayId.
-     */
-    // TODO: Maintain display list in focus order in ActivityManager and remove this call.
-    public void getDisplaysInFocusOrder(SparseIntArray displaysInFocusOrder) {
-        synchronized(mWindowMap) {
-            mRoot.getDisplaysInFocusOrder(displaysInFocusOrder);
-        }
-    }
-
     @Override
     public void setOverscan(int displayId, int left, int top, int right, int bottom) {
         if (mContext.checkCallingOrSelfPermission(
diff --git a/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java b/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
index 9fae43a..82d7ab8 100644
--- a/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
+++ b/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
@@ -63,6 +63,7 @@
 import static org.robolectric.Shadows.shadowOf;
 import static org.robolectric.shadow.api.Shadow.extract;
 import static org.testng.Assert.fail;
+import static org.testng.Assert.expectThrows;
 
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 import static java.util.Collections.emptyList;
@@ -1941,6 +1942,29 @@
         task.markCancel();
     }
 
+    @Test
+    public void testHandleCancel_callsMarkCancelAndWaitCancel() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = spy(createKeyValueBackupTask(transportMock, PACKAGE_1));
+        doNothing().when(task).waitCancel();
+
+        task.handleCancel(true);
+
+        InOrder inOrder = inOrder(task);
+        inOrder.verify(task).markCancel();
+        inOrder.verify(task).waitCancel();
+    }
+
+    @Test
+    public void testHandleCancel_whenCancelAllFalse_throws() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        expectThrows(IllegalArgumentException.class, () -> task.handleCancel(false));
+    }
+
     private void runTask(KeyValueBackupTask task) {
         // Pretend we are not on the main-thread to prevent RemoteCall from complaining
         mShadowMainLooper.setCurrentThread(false);
diff --git a/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java b/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java
index 20df2ae..1aa80c8 100644
--- a/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java
+++ b/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java
@@ -33,9 +33,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Matchers.anyInt;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
@@ -46,7 +44,6 @@
 import android.app.WaitResult;
 import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
-import android.util.SparseIntArray;
 
 import androidx.test.filters.MediumTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -54,7 +51,6 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.invocation.InvocationOnMock;
 
 import java.util.ArrayList;
 
@@ -246,22 +242,6 @@
                 null /* target */, null /* targetOptions */);
     }
 
-    @Test
-    public void testTopRunningActivityLockedWithNonExistentDisplay() throws Exception {
-        // Create display that ActivityManagerService does not know about
-        final int unknownDisplayId = 100;
-
-        doAnswer((InvocationOnMock invocationOnMock) -> {
-            final SparseIntArray displayIds = invocationOnMock.<SparseIntArray>getArgument(0);
-            displayIds.put(0, 0);
-            displayIds.put(1, unknownDisplayId);
-            return null;
-        }).when(mSupervisor.mWindowManager).getDisplaysInFocusOrder(any());
-
-        // Supervisor should skip over the non-existent display.
-        assertEquals(null, mSupervisor.topRunningActivityLocked());
-    }
-
     /**
      * Verifies that removal of activity with task and stack is done correctly.
      */
@@ -339,12 +319,6 @@
         final ActivityRecord activity = new ActivityBuilder(mService).setCreateTask(true)
                 .setStack(stack).build();
 
-        doAnswer((InvocationOnMock invocationOnMock) -> {
-            final SparseIntArray displayIds = invocationOnMock.<SparseIntArray>getArgument(0);
-            displayIds.put(0, display.mDisplayId);
-            return null;
-        }).when(mSupervisor.mWindowManager).getDisplaysInFocusOrder(any());
-
         // Make sure the top running activity is not affected when keyguard is not locked
         assertEquals(activity, mService.mStackSupervisor.topRunningActivityLocked());
         assertEquals(activity, mService.mStackSupervisor.topRunningActivityLocked(
diff --git a/services/tests/servicestests/src/com/android/server/am/ActivityTestsBase.java b/services/tests/servicestests/src/com/android/server/am/ActivityTestsBase.java
index 9c0b525..aef5537 100644
--- a/services/tests/servicestests/src/com/android/server/am/ActivityTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/am/ActivityTestsBase.java
@@ -35,7 +35,6 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.when;
 
 import android.app.ActivityManagerInternal;
 import android.app.ActivityOptions;
@@ -45,7 +44,6 @@
 import com.android.server.wm.DisplayWindowController;
 
 import org.junit.Rule;
-import org.mockito.Mockito;
 import org.mockito.invocation.InvocationOnMock;
 
 import android.app.IApplicationThread;
@@ -63,26 +61,22 @@
 import android.os.UserHandle;
 import android.service.voice.IVoiceInteractionSession;
 import android.testing.DexmakerShareClassLoaderRule;
-import android.util.SparseIntArray;
 
 import androidx.test.InstrumentationRegistry;
 
 import com.android.internal.app.IVoiceInteractor;
 import com.android.server.AttributeCache;
 import com.android.server.wm.AppWindowContainerController;
-import com.android.server.wm.DisplayWindowController;
 import com.android.server.wm.PinnedStackWindowController;
+import com.android.server.wm.RootWindowContainerController;
 import com.android.server.wm.StackWindowController;
 import com.android.server.wm.TaskWindowContainerController;
 import com.android.server.wm.WindowManagerService;
 import com.android.server.wm.WindowTestUtils;
-import com.android.server.uri.UriGrantsManagerInternal;
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Rule;
 import org.mockito.MockitoAnnotations;
-import org.mockito.invocation.InvocationOnMock;
 
 import java.util.List;
 
@@ -500,13 +494,14 @@
                     (DisplayManager) mService.mContext.getSystemService(Context.DISPLAY_SERVICE);
             mWindowManager = prepareMockWindowManager();
             mKeyguardController = mock(KeyguardController.class);
+            setWindowContainerController(mock(RootWindowContainerController.class));
         }
 
         @Override
         public void initialize() {
             super.initialize();
             mDisplay = spy(new TestActivityDisplay(this, DEFAULT_DISPLAY));
-            attachDisplay(mDisplay);
+            addChild(mDisplay, ActivityDisplay.POSITION_TOP);
         }
 
         @Override
@@ -576,12 +571,6 @@
             return null;
         }).when(service).inSurfaceTransaction(any());
 
-        doAnswer((InvocationOnMock invocationOnMock) -> {
-            final SparseIntArray displayIds = invocationOnMock.<SparseIntArray>getArgument(0);
-            displayIds.put(0, 0);
-            return null;
-        }).when(service).getDisplaysInFocusOrder(any());
-
         return service;
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/am/RecentTasksTest.java b/services/tests/servicestests/src/com/android/server/am/RecentTasksTest.java
index ba82487..5195214 100644
--- a/services/tests/servicestests/src/com/android/server/am/RecentTasksTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/RecentTasksTest.java
@@ -872,8 +872,8 @@
             super.initialize();
             mDisplay = new TestActivityDisplay(this, DEFAULT_DISPLAY);
             mOtherDisplay = new TestActivityDisplay(this, DEFAULT_DISPLAY);
-            attachDisplay(mOtherDisplay);
-            attachDisplay(mDisplay);
+            addChild(mOtherDisplay, ActivityDisplay.POSITION_TOP);
+            addChild(mDisplay, ActivityDisplay.POSITION_TOP);
         }
 
         @Override
@@ -1045,7 +1045,7 @@
 
         @Override
         void getTasks(int maxNum, List<RunningTaskInfo> list, int ignoreActivityType,
-                int ignoreWindowingMode, SparseArray<ActivityDisplay> activityDisplays,
+                int ignoreWindowingMode, ArrayList<ActivityDisplay> activityDisplays,
                 int callingUid, boolean allowed) {
             lastAllowed = allowed;
             super.getTasks(maxNum, list, ignoreActivityType, ignoreWindowingMode, activityDisplays,
diff --git a/services/tests/servicestests/src/com/android/server/am/RunningTasksTest.java b/services/tests/servicestests/src/com/android/server/am/RunningTasksTest.java
index 283c027..d56c6a6 100644
--- a/services/tests/servicestests/src/com/android/server/am/RunningTasksTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/RunningTasksTest.java
@@ -68,9 +68,9 @@
     public void testCollectTasksByLastActiveTime() throws Exception {
         // Create a number of stacks with tasks (of incrementing active time)
         final ActivityStackSupervisor supervisor = mService.mStackSupervisor;
-        final SparseArray<ActivityDisplay> displays = new SparseArray<>();
+        final ArrayList<ActivityDisplay> displays = new ArrayList<>();
         final ActivityDisplay display = new TestActivityDisplay(supervisor, DEFAULT_DISPLAY);
-        displays.put(DEFAULT_DISPLAY, display);
+        displays.add(display);
 
         final int numStacks = 2;
         for (int stackIndex = 0; stackIndex < numStacks; stackIndex++) {
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 0d40c5e..b330304 100644
--- a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
@@ -446,24 +446,6 @@
         assertEquals(anotherAlwaysOnTopStack, mDisplayContent.getStacks().get(topPosition - 1));
     }
 
-    /**
-     * Test that WM does not report displays to AM that are pending to be removed.
-     */
-    @Test
-    public void testDontReportDeferredRemoval() {
-        // Create a display and add an animating window to it.
-        final DisplayContent dc = createNewDisplay();
-        final WindowState window = createWindow(null /* parent */, TYPE_BASE_APPLICATION, dc, "w");
-        window.mAnimatingExit = true;
-        // Request display removal, it should be deferred.
-        dc.removeIfPossible();
-        // Request ordered display ids from WM.
-        final SparseIntArray orderedDisplayIds = new SparseIntArray();
-        sWm.getDisplaysInFocusOrder(orderedDisplayIds);
-        // Make sure that display that is marked for removal is not reported.
-        assertEquals(-1, orderedDisplayIds.indexOfValue(dc.getDisplayId()));
-    }
-
     @Test
     public void testDisplayCutout_rot0() throws Exception {
         synchronized (sWm.getWindowManagerLock()) {