WindowManager: Emit windows from forAllWindows closer to Z-order

Otherwise IME windows will be reported as under the docked divider.

Bug: 77526241
Test: Manual. WindowContainerTraversalTests.
Change-Id: I008e253464b36c61973f49c7992352fb96649dc1

diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 9621edd..c0dc750 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java

@@ -1510,6 +1510,10 @@
         return (stack != null && stack.isVisible()) ? stack : null;
     }
 
+    boolean hasSplitScreenPrimaryStack() {
+        return getSplitScreenPrimaryStack() != null;
+    }
+
     /**
      * Like {@link #getSplitScreenPrimaryStack}, but also returns the stack if it's currently
      * not visible.
@@ -1613,6 +1617,18 @@
         mService.mWindowsChanged = true;
     }
 
+    /**
+     * In split-screen mode we process the IME containers above the docked divider
+     * rather than directly above their target.
+     */
+    private boolean skipTraverseChild(WindowContainer child) {
+        if (child == mImeWindowsContainers && mService.mInputMethodTarget != null
+                && !hasSplitScreenPrimaryStack()) {
+            return true;
+        }
+        return false;
+    }
+
     @Override
     boolean forAllWindows(ToBooleanFunction<WindowState> callback, boolean traverseTopToBottom) {
         // Special handling so we can process IME windows with #forAllImeWindows above their IME
@@ -1620,11 +1636,10 @@
         if (traverseTopToBottom) {
             for (int i = mChildren.size() - 1; i >= 0; --i) {
                 final DisplayChildWindowContainer child = mChildren.get(i);
-                if (child == mImeWindowsContainers && mService.mInputMethodTarget != null) {
-                    // In this case the Ime windows will be processed above their target so we skip
-                    // here.
+                if (skipTraverseChild(child)) {
                     continue;
                 }
+
                 if (child.forAllWindows(callback, traverseTopToBottom)) {
                     return true;
                 }
@@ -1633,11 +1648,10 @@
             final int count = mChildren.size();
             for (int i = 0; i < count; i++) {
                 final DisplayChildWindowContainer child = mChildren.get(i);
-                if (child == mImeWindowsContainers && mService.mInputMethodTarget != null) {
-                    // In this case the Ime windows will be processed above their target so we skip
-                    // here.
+                if (skipTraverseChild(child)) {
                     continue;
                 }
+
                 if (child.forAllWindows(callback, traverseTopToBottom)) {
                     return true;
                 }

diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 0de3c92..c710c97 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java

@@ -3997,32 +3997,33 @@
         return false;
     }
 
+    private boolean applyImeWindowsIfNeeded(ToBooleanFunction<WindowState> callback,
+            boolean traverseTopToBottom) {
+        // If this window is the current IME target, so we need to process the IME windows
+        // directly above it. The exception is if we are in split screen
+        // in which case we process the IME at the DisplayContent level to
+        // ensure it is above the docked divider.
+        if (isInputMethodTarget() && !inSplitScreenWindowingMode()) {
+            if (getDisplayContent().forAllImeWindows(callback, traverseTopToBottom)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private boolean applyInOrderWithImeWindows(ToBooleanFunction<WindowState> callback,
             boolean traverseTopToBottom) {
         if (traverseTopToBottom) {
-            if (isInputMethodTarget()) {
-                // This window is the current IME target, so we need to process the IME windows
-                // directly above it.
-                if (getDisplayContent().forAllImeWindows(callback, traverseTopToBottom)) {
-                    return true;
-                }
-            }
-            if (callback.apply(this)) {
+            if (applyImeWindowsIfNeeded(callback, traverseTopToBottom)
+                    || callback.apply(this)) {
                 return true;
             }
         } else {
-            if (callback.apply(this)) {
+            if (callback.apply(this)
+                    || applyImeWindowsIfNeeded(callback, traverseTopToBottom)) {
                 return true;
             }
-            if (isInputMethodTarget()) {
-                // This window is the current IME target, so we need to process the IME windows
-                // directly above it.
-                if (getDisplayContent().forAllImeWindows(callback, traverseTopToBottom)) {
-                    return true;
-                }
-            }
         }
-
         return false;
     }