diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 7f944454ecdf9..f35a7e6236658 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -70,7 +70,6 @@ import static android.view.WindowManager.LayoutParams.TYPE_ACCESSIBILITY_MAGNIFI
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_STARTING;
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_BOOT_PROGRESS;
-import static android.view.WindowManager.LayoutParams.TYPE_DOCK_DIVIDER;
 import static android.view.WindowManager.LayoutParams.TYPE_DREAM;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG;
@@ -304,8 +303,8 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
     // on the IME target. We mainly have this container grouping so we can keep track of all the IME
     // window containers together and move them in-sync if/when needed. We use a subclass of
     // WindowContainer which is omitted from screen magnification, as the IME is never magnified.
-    private final NonAppWindowContainers mImeWindowsContainers =
-            new NonAppWindowContainers("mImeWindowsContainers", mWmService);
+    // TODO(display-area): is "no magnification" in the comment still true?
+    private final ImeContainer mImeWindowsContainers = new ImeContainer(mWmService);
 
     private WindowState mTmpWindow;
     private WindowState mTmpWindow2;
@@ -1995,6 +1994,10 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
         return mTaskStackContainers.getVisibleTasks();
     }
 
+    SurfaceControl getSplitScreenDividerAnchor() {
+        return mTaskStackContainers.getSplitScreenDividerAnchor();
+    }
+
     void onStackWindowingModeChanged(ActivityStack stack) {
         mTaskStackContainers.onStackWindowingModeChanged(stack);
     }
@@ -2081,7 +2084,7 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
     }
 
     boolean forAllImeWindows(ToBooleanFunction<WindowState> callback, boolean traverseTopToBottom) {
-        return mImeWindowsContainers.forAllWindows(callback, traverseTopToBottom);
+        return mImeWindowsContainers.forAllWindowForce(callback, traverseTopToBottom);
     }
 
     /**
@@ -2339,11 +2342,15 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
         throw new UnsupportedOperationException("See DisplayChildWindowContainer");
     }
 
+    void positionDisplayAt(int position, boolean includingParents) {
+        getParent().positionChildAt(position, this, includingParents);
+    }
+
     @Override
     void positionChildAt(int position, DisplayChildWindowContainer child, boolean includingParents) {
         // Children of the display are statically ordered, so the real intention here is to perform
         // the operation on the display and not the static direct children.
-        getParent().positionChildAt(position, this, includingParents);
+        positionDisplayAt(position, includingParents);
     }
 
     void positionStackAt(int position, ActivityStack child, boolean includingParents) {
@@ -4289,8 +4296,9 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
 
         @Override
         void positionChildAt(int position, ActivityStack child, boolean includingParents) {
-            if (child.getWindowConfiguration().isAlwaysOnTop()
-                    && position != POSITION_TOP && position != mChildren.size()) {
+            final boolean moveToTop = (position == POSITION_TOP || position == getChildCount());
+            final boolean moveToBottom = (position == POSITION_BOTTOM || position == 0);
+            if (child.getWindowConfiguration().isAlwaysOnTop() && !moveToTop) {
                 // This stack is always-on-top, override the default behavior.
                 Slog.w(TAG_WM, "Ignoring move of always-on-top stack=" + this + " to bottom");
 
@@ -4306,18 +4314,14 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
                 includingParents = false;
             }
             final int targetPosition = findPositionForStack(position, child, false /* adding */);
-            super.positionChildAt(targetPosition, child, includingParents);
+            super.positionChildAt(targetPosition, child, false /* includingParents */);
 
-            if (includingParents) {
-                // We still want to move the display of this stack container to top because even the
-                // target position is adjusted to non-top, the intention of the condition is to have
-                // higher z-order to gain focus (e.g. moving a task of a fullscreen stack to front
-                // in a non-top display which is using picture-in-picture mode).
-                final int topChildPosition = getChildCount() - 1;
-                if (targetPosition < topChildPosition && position >= topChildPosition) {
-                    getParent().positionChildAt(POSITION_TOP, this /* child */,
-                            true /* includingParents */);
-                }
+            if (includingParents && (moveToTop || moveToBottom)) {
+                // The DisplayContent children do not re-order, but we still want to move the
+                // display of this stack container because the intention of positioning is to have
+                // higher z-order to gain focus.
+                positionDisplayAt(moveToTop ? POSITION_TOP : POSITION_BOTTOM,
+                        true /* includingParents */);
             }
 
             setLayoutNeeded();
@@ -4636,39 +4640,12 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
             super(name, service);
         }
 
-        @Override
-        SurfaceControl.Builder makeChildSurface(WindowContainer child) {
-            final SurfaceControl.Builder builder = super.makeChildSurface(child);
-            if (child instanceof WindowToken && ((WindowToken) child).mRoundedCornerOverlay) {
-                // To draw above the ColorFade layer during the screen off transition, the
-                // rounded corner overlays need to be at the root of the surface hierarchy.
-                // TODO: move the ColorLayer into the display overlay layer such that this is not
-                // necessary anymore.
-                builder.setParent(null);
-            }
-            return builder;
-        }
-
         @Override
         void assignChildLayers(SurfaceControl.Transaction t) {
-            assignChildLayers(t, null /* imeContainer */);
-        }
-
-        void assignChildLayers(SurfaceControl.Transaction t, WindowContainer imeContainer) {
-            boolean needAssignIme = imeContainer != null
-                    && imeContainer.getSurfaceControl() != null;
+            boolean needAssignIme = mImeWindowsContainers.getSurfaceControl() != null;
             for (int j = 0; j < mChildren.size(); ++j) {
                 final WindowToken wt = mChildren.get(j);
 
-                // See {@link mSplitScreenDividerAnchor}
-                if (wt.windowType == TYPE_DOCK_DIVIDER) {
-                    wt.assignRelativeLayer(t, mTaskStackContainers.getSplitScreenDividerAnchor(), 1);
-                    continue;
-                }
-                if (wt.mRoundedCornerOverlay) {
-                    wt.assignLayer(t, WindowManagerPolicy.COLOR_FADE_LAYER + 1);
-                    continue;
-                }
                 wt.assignLayer(t, j);
                 wt.assignChildLayers(t);
 
@@ -4677,13 +4654,10 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
 
                 if (needAssignIme && layer >= mWmService.mPolicy.getWindowLayerFromTypeLw(
                         TYPE_INPUT_METHOD_DIALOG, true)) {
-                    imeContainer.assignRelativeLayer(t, wt.getSurfaceControl(), -1);
+                    mImeWindowsContainers.assignRelativeLayer(t, wt.getSurfaceControl(), -1);
                     needAssignIme = false;
                 }
             }
-            if (needAssignIme) {
-                imeContainer.assignRelativeLayer(t, getSurfaceControl(), Integer.MAX_VALUE);
-            }
         }
     }
 
@@ -4697,6 +4671,7 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
 
         @Override
         void assignChildLayers(SurfaceControl.Transaction t) {
+            mImeWindowsContainers.setNeedsLayer();
             mBelowAppWindowsContainers.assignLayer(t, 0);
             mTaskStackContainers.assignLayer(t, 1);
             mAboveAppWindowsContainers.assignLayer(t, 2);
@@ -4732,15 +4707,14 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
                         // TODO: We need to use an extra level on the app surface to ensure
                         // this is always above SurfaceView but always below attached window.
                         1);
-                needAssignIme = false;
             }
 
             // Above we have assigned layers to our children, now we ask them to assign
             // layers to their children.
             mBelowAppWindowsContainers.assignChildLayers(t);
             mTaskStackContainers.assignChildLayers(t);
-            mAboveAppWindowsContainers.assignChildLayers(t,
-                    needAssignIme ? mImeWindowsContainers : null);
+            mAboveAppWindowsContainers.assignChildLayers(t);
+            mImeWindowsContainers.assignRelativeLayer(t, getSurfaceControl(), Integer.MAX_VALUE);
             mImeWindowsContainers.assignChildLayers(t);
         }
 
@@ -4756,47 +4730,6 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
             addChild(mImeWindowsContainers, null);
         }
 
-        /**
-         * In split-screen mode we process the IME containers above the docked divider
-         * rather than directly above their target.
-         */
-        private boolean skipTraverseChild(WindowContainer child) {
-            return child == mImeWindowsContainers && mInputMethodTarget != null
-                    && !hasSplitScreenPrimaryStack();
-        }
-
-        @Override
-        boolean forAllWindows(ToBooleanFunction<WindowState> callback,
-                boolean traverseTopToBottom) {
-            // Special handling so we can process IME windows with #forAllImeWindows above their IME
-            // target, or here in order if there isn't an IME target.
-            if (traverseTopToBottom) {
-                for (int i = mChildren.size() - 1; i >= 0; --i) {
-                    final WindowContainer child = mChildren.get(i);
-                    if (skipTraverseChild(child)) {
-                        continue;
-                    }
-
-                    if (child.forAllWindows(callback, traverseTopToBottom)) {
-                        return true;
-                    }
-                }
-            } else {
-                final int count = mChildren.size();
-                for (int i = 0; i < count; i++) {
-                    final WindowContainer child = mChildren.get(i);
-                    if (skipTraverseChild(child)) {
-                        continue;
-                    }
-
-                    if (child.forAllWindows(callback, traverseTopToBottom)) {
-                        return true;
-                    }
-                }
-            }
-            return false;
-        }
-
         @Override
         void positionChildAt(int position, WindowContainer child, boolean includingParents) {
             // Children of the WindowContainers are statically ordered, so the real intention here
@@ -4910,6 +4843,68 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
         }
     }
 
+    /**
+     * Container for IME windows.
+     *
+     * This has some special behaviors:
+     * - layers assignment is ignored except if setNeedsLayer() has been called before (and no
+     *   layer has been assigned since), to facilitate assigning the layer from the IME target, or
+     *   fall back if there is no target.
+     * - the container doesn't always participate in window traversal, according to
+     *   {@link #skipImeWindowsDuringTraversal()}
+     */
+    private class ImeContainer extends NonAppWindowContainers {
+        boolean mNeedsLayer = false;
+
+        ImeContainer(WindowManagerService wms) {
+            super("ImeContainer", wms);
+        }
+
+        public void setNeedsLayer() {
+            mNeedsLayer = true;
+        }
+
+        @Override
+        boolean forAllWindows(ToBooleanFunction<WindowState> callback,
+                boolean traverseTopToBottom) {
+            final DisplayContent dc = mDisplayContent;
+            if (skipImeWindowsDuringTraversal(dc)) {
+                return false;
+            }
+            return super.forAllWindows(callback, traverseTopToBottom);
+        }
+
+        private boolean skipImeWindowsDuringTraversal(DisplayContent dc) {
+            // We skip IME windows so they're processed just above their target, except
+            // in split-screen mode where we process the IME containers above the docked divider.
+            return dc.mInputMethodTarget != null && !dc.hasSplitScreenPrimaryStack();
+        }
+
+        /** Like {@link #forAllWindows}, but ignores {@link #skipImeWindowsDuringTraversal} */
+        boolean forAllWindowForce(ToBooleanFunction<WindowState> callback,
+                boolean traverseTopToBottom) {
+            return super.forAllWindows(callback, traverseTopToBottom);
+        }
+
+        @Override
+        void assignLayer(Transaction t, int layer) {
+            if (!mNeedsLayer) {
+                return;
+            }
+            super.assignLayer(t, layer);
+            mNeedsLayer = false;
+        }
+
+        @Override
+        void assignRelativeLayer(Transaction t, SurfaceControl relativeTo, int layer) {
+            if (!mNeedsLayer) {
+                return;
+            }
+            super.assignRelativeLayer(t, relativeTo, layer);
+            mNeedsLayer = false;
+        }
+    }
+
     @Override
     SurfaceSession getSession() {
         return mSession;
@@ -5013,6 +5008,7 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
      * with {@link WindowState#assignLayer}
      */
     void assignRelativeLayerForImeTargetChild(SurfaceControl.Transaction t, WindowContainer child) {
+        mImeWindowsContainers.setNeedsLayer();
         child.assignRelativeLayer(t, mImeWindowsContainers.getSurfaceControl(), 1);
     }
 
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 223e9b9e3df29..c78de42b84ae9 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -2764,7 +2764,7 @@ public class WindowManagerService extends IWindowManager.Stub
         synchronized (mGlobalLock) {
             final DisplayContent displayContent = mRoot.getDisplayContent(displayId);
             if (displayContent != null && mRoot.getTopChild() != displayContent) {
-                mRoot.positionChildAt(WindowContainer.POSITION_TOP, displayContent,
+                displayContent.positionDisplayAt(WindowContainer.POSITION_TOP,
                         true /* includingParents */);
             }
         }
@@ -7647,7 +7647,7 @@ public class WindowManagerService extends IWindowManager.Stub
 
         final DisplayContent displayContent = touchedWindow.getDisplayContent();
         if (!displayContent.isOnTop()) {
-            displayContent.getParent().positionChildAt(WindowContainer.POSITION_TOP, displayContent,
+            displayContent.positionDisplayAt(WindowContainer.POSITION_TOP,
                     true /* includingParents */);
         }
         handleTaskFocusChange(touchedWindow.getTask());
diff --git a/services/core/java/com/android/server/wm/WindowToken.java b/services/core/java/com/android/server/wm/WindowToken.java
index 2a1e980dfb3db..b9bd4e003993f 100644
--- a/services/core/java/com/android/server/wm/WindowToken.java
+++ b/services/core/java/com/android/server/wm/WindowToken.java
@@ -17,6 +17,7 @@
 package com.android.server.wm;
 
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
+import static android.view.WindowManager.LayoutParams.TYPE_DOCK_DIVIDER;
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
 
 import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_ADD_REMOVE;
@@ -38,7 +39,9 @@ import android.annotation.CallSuper;
 import android.os.Debug;
 import android.os.IBinder;
 import android.util.proto.ProtoOutputStream;
+import android.view.SurfaceControl;
 
+import com.android.server.policy.WindowManagerPolicy;
 import com.android.server.protolog.common.ProtoLog;
 
 import java.io.PrintWriter;
@@ -256,6 +259,27 @@ class WindowToken extends WindowContainer<WindowState> {
         super.onDisplayChanged(dc);
     }
 
+    @Override
+    void assignLayer(SurfaceControl.Transaction t, int layer) {
+        if (windowType == TYPE_DOCK_DIVIDER) {
+            // See {@link DisplayContent#mSplitScreenDividerAnchor}
+            super.assignRelativeLayer(t, mDisplayContent.getSplitScreenDividerAnchor(), 1);
+        } else if (mRoundedCornerOverlay) {
+            super.assignLayer(t, WindowManagerPolicy.COLOR_FADE_LAYER + 1);
+        } else {
+            super.assignLayer(t, layer);
+        }
+    }
+
+    @Override
+    SurfaceControl.Builder makeSurface() {
+        final SurfaceControl.Builder builder = super.makeSurface();
+        if (mRoundedCornerOverlay) {
+            builder.setParent(null);
+        }
+        return builder;
+    }
+
     @CallSuper
     @Override
     public void dumpDebug(ProtoOutputStream proto, long fieldId,
@@ -315,4 +339,8 @@ class WindowToken extends WindowContainer<WindowState> {
                 mOwnerCanManageAppTokens);
         return mOwnerCanManageAppTokens && (layer > navLayer);
     }
+
+    int getWindowLayerFromType() {
+        return mWmService.mPolicy.getWindowLayerFromTypeLw(windowType, mOwnerCanManageAppTokens);
+    }
 }