Let IME target window per display.

Since currently in WMS only recorded one IME target window that
the IME window on top of.

For multi-sessions IME on multi-displays scenerio, we need
move WMS.mInputMethodTarget to DisplayContent.

Added unit test testInputMethodTargetUpdateWhenSwitchingOnDisplays
for verifying per IME target window update when switching on displays.

Bug: 117962777

Test: atest ActivityManagerMultiDisplayTests
Test: atest FrameworksServicesTests:ZOrderingTests
Test: atest FrameworksServicesTests:WindowContainerTraversalTests
Test: atest FrameworksServicesTests:DisplayContentTests
Change-Id: I9aa934961fb3975bd2af18599b5a2884387b5007
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 9baafcb..7600399 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -2010,9 +2010,7 @@
         clearThumbnail();
         setClientHidden(isHidden() && hiddenRequested);
 
-        if (mService.mInputMethodTarget != null && mService.mInputMethodTarget.mAppToken == this) {
-            getDisplayContent().computeImeTarget(true /* updateImeTarget */);
-        }
+        getDisplayContent().computeImeTargetIfNeeded(this);
 
         if (DEBUG_ANIM) Slog.v(TAG, "Animation done in " + this
                 + ": reportedVisible=" + reportedVisible
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 348b2af..2816855 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -157,8 +157,8 @@
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Transaction;
 import android.view.SurfaceSession;
-import android.view.WindowManagerPolicyConstants.PointerEventListener;
 import android.view.WindowManager;
+import android.view.WindowManagerPolicyConstants.PointerEventListener;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ToBooleanFunction;
@@ -496,6 +496,15 @@
      */
     WindowState mInputMethodWindow;
 
+    /**
+     * This just indicates the window the input method is on top of, not
+     * necessarily the window its input is going to.
+     */
+    WindowState mInputMethodTarget;
+
+    /** If true hold off on modifying the animation layer of mInputMethodTarget */
+    boolean mInputMethodTargetWaitingAnim;
+
     private final PointerEventDispatcher mPointerEventDispatcher;
 
     private final Consumer<WindowState> mUpdateWindowsForAnimator = w -> {
@@ -699,7 +708,7 @@
 
     private final Consumer<WindowState> mApplyPostLayoutPolicy =
             w -> mService.mPolicy.applyPostLayoutPolicyLw(w, w.mAttrs, w.getParentWindow(),
-                    mService.mInputMethodTarget);
+                    mInputMethodTarget);
 
     private final Consumer<WindowState> mApplySurfaceChangesTransaction = w -> {
         final WindowSurfacePlacer surfacePlacer = mService.mWindowPlacerLocked;
@@ -1917,7 +1926,7 @@
      * rather than directly above their target.
      */
     private boolean skipTraverseChild(WindowContainer child) {
-        if (child == mImeWindowsContainers && mService.mInputMethodTarget != null
+        if (child == mImeWindowsContainers && mInputMethodTarget != null
                 && !hasSplitScreenPrimaryStack()) {
             return true;
         }
@@ -2793,7 +2802,7 @@
         if (!focusFound) {
             final WindowState imWindow = mInputMethodWindow;
             if (imWindow != null) {
-                final WindowState prevTarget = mService.mInputMethodTarget;
+                final WindowState prevTarget = mInputMethodTarget;
 
                 final WindowState newTarget = computeImeTarget(true /* updateImeTarget*/);
                 imWindowChanged = prevTarget != newTarget;
@@ -2985,13 +2994,13 @@
             // There isn't an IME so there shouldn't be a target...That was easy!
             if (updateImeTarget) {
                 if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "Moving IM target from "
-                        + mService.mInputMethodTarget + " to null since mInputMethodWindow is null");
-                setInputMethodTarget(null, mService.mInputMethodTargetWaitingAnim);
+                        + mInputMethodTarget + " to null since mInputMethodWindow is null");
+                setInputMethodTarget(null, mInputMethodTargetWaitingAnim);
             }
             return null;
         }
 
-        final WindowState curTarget = mService.mInputMethodTarget;
+        final WindowState curTarget = mInputMethodTarget;
         if (!canUpdateImeTarget()) {
             if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "Defer updating IME target");
             return curTarget;
@@ -3018,7 +3027,7 @@
         }
 
         if (DEBUG_INPUT_METHOD && updateImeTarget) Slog.v(TAG_WM,
-                "Proposed new IME target: " + target);
+                "Proposed new IME target: " + target + " for display: " + getDisplayId());
 
         // Now, a special case -- if the last target's window is in the process of exiting, but
         // not removed, and the new target is home, keep on the last target to avoid flicker.
@@ -3039,7 +3048,7 @@
                 if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "Moving IM target from " + curTarget
                         + " to null." + (SHOW_STACK_CRAWLS ? " Callers="
                         + Debug.getCallers(4) : ""));
-                setInputMethodTarget(null, mService.mInputMethodTargetWaitingAnim);
+                setInputMethodTarget(null, mInputMethodTargetWaitingAnim);
             }
 
             return null;
@@ -3078,14 +3087,23 @@
         return target;
     }
 
+    /**
+     * Calling {@link #computeImeTarget(boolean)} to update the input method target window in
+     * the candidate app window token if needed.
+     */
+    void computeImeTargetIfNeeded(AppWindowToken candidate) {
+        if (mInputMethodTarget != null && mInputMethodTarget.mAppToken == candidate) {
+            computeImeTarget(true /* updateImeTarget */);
+        }
+    }
+
     private void setInputMethodTarget(WindowState target, boolean targetWaitingAnim) {
-        if (target == mService.mInputMethodTarget
-                && mService.mInputMethodTargetWaitingAnim == targetWaitingAnim) {
+        if (target == mInputMethodTarget && mInputMethodTargetWaitingAnim == targetWaitingAnim) {
             return;
         }
 
-        mService.mInputMethodTarget = target;
-        mService.mInputMethodTargetWaitingAnim = targetWaitingAnim;
+        mInputMethodTarget = target;
+        mInputMethodTargetWaitingAnim = targetWaitingAnim;
         assignWindowLayers(false /* setLayoutNeeded */);
     }
 
@@ -4474,7 +4492,7 @@
         mTaskStackContainers.assignLayer(t, 1);
         mAboveAppWindowsContainers.assignLayer(t, 2);
 
-        WindowState imeTarget = mService.mInputMethodTarget;
+        final WindowState imeTarget = mInputMethodTarget;
         boolean needAssignIme = true;
 
         // In the case where we have an IME target that is not in split-screen
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index a641f75..628d98f 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -609,13 +609,6 @@
      */
     final Handler mAnimationHandler = new Handler(AnimationThread.getHandler().getLooper());
 
-    /** This just indicates the window the input method is on top of, not
-     * necessarily the window its input is going to. */
-    WindowState mInputMethodTarget = null;
-
-    /** If true hold off on modifying the animation layer of mInputMethodTarget */
-    boolean mInputMethodTargetWaitingAnim;
-
     boolean mHardKeyboardAvailable;
     WindowManagerInternal.OnHardKeyboardStatusChangeListener mHardKeyboardStatusChangeListener;
     SettingsObserver mSettingsObserver;
@@ -5920,9 +5913,13 @@
         pw.print("  mGlobalConfiguration="); pw.println(mRoot.getConfiguration());
         pw.print("  mHasPermanentDpad="); pw.println(mHasPermanentDpad);
         mRoot.dumpTopFocusedDisplayId(pw);
-        if (mInputMethodTarget != null) {
-            pw.print("  mInputMethodTarget="); pw.println(mInputMethodTarget);
-        }
+        mRoot.forAllDisplays(dc -> {
+            final WindowState inputMethodTarget = dc.mInputMethodTarget;
+            if (inputMethodTarget != null) {
+                pw.print("  mInputMethodTarget in display# "); pw.print(dc.getDisplayId());
+                pw.print(' '); pw.println(inputMethodTarget);
+            }
+        });
         pw.print("  mInTouchMode="); pw.println(mInTouchMode);
         pw.print("  mLastDisplayFreezeDuration=");
                 TimeUtils.formatDuration(mLastDisplayFreezeDuration, pw);
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 99f65c3..aed8fa3 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -4476,8 +4476,9 @@
 
     @Override
     boolean needsZBoost() {
-        if (mIsImWindow && mService.mInputMethodTarget != null) {
-            final AppWindowToken appToken = mService.mInputMethodTarget.mAppToken;
+        final WindowState inputMethodTarget = getDisplayContent().mInputMethodTarget;
+        if (mIsImWindow && inputMethodTarget != null) {
+            final AppWindowToken appToken = inputMethodTarget.mAppToken;
             if (appToken != null) {
                 return appToken.needsZBoost();
             }
@@ -4607,7 +4608,7 @@
             // Likewise if we share a token with the Input method target and are ordered
             // above it but not necessarily a child (e.g. a Dialog) then we also need
             // this promotion.
-            final WindowState imeTarget = mService.mInputMethodTarget;
+            final WindowState imeTarget = getDisplayContent().mInputMethodTarget;
             boolean inTokenWithAndAboveImeTarget = imeTarget != null && imeTarget != this
                     && imeTarget.mToken == mToken && imeTarget.compareTo(this) <= 0;
             return inTokenWithAndAboveImeTarget;
@@ -4684,7 +4685,7 @@
 
     @Override
     public boolean isInputMethodTarget() {
-        return mService.mInputMethodTarget == this;
+        return getDisplayContent().mInputMethodTarget == this;
     }
 
     long getFrameNumber() {
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 dd374e9..ea93326 100644
--- a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
@@ -40,6 +40,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.doNothing;
@@ -108,7 +109,7 @@
         final WindowState imeAppTarget =
                 createWindow(null, TYPE_BASE_APPLICATION, mDisplayContent, "imeAppTarget");
 
-        mWm.mInputMethodTarget = imeAppTarget;
+        mDisplayContent.mInputMethodTarget = imeAppTarget;
 
         assertForAllWindowsOrder(Arrays.asList(
                 mWallpaperWindow,
@@ -124,8 +125,8 @@
     }
 
     @Test
-    public void testForAllWindows_WithChildWindowImeTarget() {
-        mWm.mInputMethodTarget = mChildAppWindowAbove;
+    public void testForAllWindows_WithChildWindowImeTarget() throws Exception {
+        mDisplayContent.mInputMethodTarget = mChildAppWindowAbove;
 
         assertForAllWindowsOrder(Arrays.asList(
                 mWallpaperWindow,
@@ -140,8 +141,8 @@
     }
 
     @Test
-    public void testForAllWindows_WithStatusBarImeTarget() {
-        mWm.mInputMethodTarget = mStatusBarWindow;
+    public void testForAllWindows_WithStatusBarImeTarget() throws Exception {
+        mDisplayContent.mInputMethodTarget = mStatusBarWindow;
 
         assertForAllWindowsOrder(Arrays.asList(
                 mWallpaperWindow,
@@ -568,6 +569,32 @@
         assertFalse(isOptionsPanelAtRight(landscapeDisplay.getDisplayId()));
     }
 
+    @Test
+    public void testInputMethodTargetUpdateWhenSwitchingOnDisplays() {
+        final DisplayContent newDisplay = createNewDisplay();
+
+        final WindowState appWin = createWindow(null, TYPE_APPLICATION, mDisplayContent, "appWin");
+        final WindowState appWin1 = createWindow(null, TYPE_APPLICATION, newDisplay, "appWin1");
+        appWin.setHasSurface(true);
+        appWin1.setHasSurface(true);
+
+        // Set current input method window on default display, make sure the input method target
+        // is appWin & null on the other display.
+        mDisplayContent.setInputMethodWindowLocked(mImeWindow);
+        newDisplay.setInputMethodWindowLocked(null);
+        assertTrue("appWin should be IME target window",
+                appWin.equals(mDisplayContent.mInputMethodTarget));
+        assertNull("newDisplay Ime target: ", newDisplay.mInputMethodTarget);
+
+        // Switch input method window on new display & make sure the input method target also
+        // switched as expected.
+        newDisplay.setInputMethodWindowLocked(mImeWindow);
+        mDisplayContent.setInputMethodWindowLocked(null);
+        assertTrue("appWin1 should be IME target window",
+                appWin1.equals(newDisplay.mInputMethodTarget));
+        assertNull("default display Ime target: ", mDisplayContent.mInputMethodTarget);
+    }
+
     private boolean isOptionsPanelAtRight(int displayId) {
         return (mWm.getPreferredOptionsPanelGravity(displayId) & Gravity.RIGHT) == Gravity.RIGHT;
     }
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowContainerTraversalTests.java b/services/tests/servicestests/src/com/android/server/wm/WindowContainerTraversalTests.java
index 2b8b934..fcde08e 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowContainerTraversalTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowContainerTraversalTests.java
@@ -52,7 +52,7 @@
                 WINDOWING_MODE_SPLIT_SCREEN_SECONDARY, ACTIVITY_TYPE_STANDARD,
                 TYPE_BASE_APPLICATION, mDisplayContent, "splitScreenSecondaryWindow");
 
-        mWm.mInputMethodTarget = splitScreenWindow;
+        mDisplayContent.mInputMethodTarget = splitScreenWindow;
 
         Consumer<WindowState> c = mock(Consumer.class);
         mDisplayContent.forAllWindows(c, false);
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
index 945cbb9..990ef54 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -192,7 +192,7 @@
                 mWm.getDefaultDisplayContentLocked().mAppTransition
                         .removeAppTransitionTimeoutCallbacks();
                 mWm.mH.removeMessages(WindowManagerService.H.WINDOW_FREEZE_TIMEOUT);
-                mWm.mInputMethodTarget = null;
+                mDisplayContent.mInputMethodTarget = null;
             }
 
             // Wait until everything is really cleaned up.
diff --git a/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java b/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java
index 32e4e02..0ff2a0d 100644
--- a/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/ZOrderingTests.java
@@ -211,7 +211,7 @@
 
     @Test
     public void testAssignWindowLayers_ForImeWithNoTarget() {
-        mWm.mInputMethodTarget = null;
+        mDisplayContent.mInputMethodTarget = null;
         mDisplayContent.assignChildLayers(mTransaction);
 
         // The Ime has an higher base layer than app windows and lower base layer than system
@@ -229,7 +229,7 @@
     @Test
     public void testAssignWindowLayers_ForImeWithAppTarget() {
         final WindowState imeAppTarget = createWindow("imeAppTarget");
-        mWm.mInputMethodTarget = imeAppTarget;
+        mDisplayContent.mInputMethodTarget = imeAppTarget;
 
         mDisplayContent.assignChildLayers(mTransaction);
 
@@ -255,7 +255,7 @@
                 TYPE_APPLICATION_MEDIA_OVERLAY, imeAppTarget.mToken,
                 "imeAppTargetChildBelowWindow");
 
-        mWm.mInputMethodTarget = imeAppTarget;
+        mDisplayContent.mInputMethodTarget = imeAppTarget;
         mDisplayContent.assignChildLayers(mTransaction);
 
         // Ime should be above all app windows except for child windows that are z-ordered above it
@@ -277,7 +277,7 @@
         final WindowState imeAppTarget = createWindow("imeAppTarget");
         final WindowState appAboveImeTarget = createWindow("appAboveImeTarget");
 
-        mWm.mInputMethodTarget = imeAppTarget;
+        mDisplayContent.mInputMethodTarget = imeAppTarget;
         mDisplayContent.assignChildLayers(mTransaction);
 
         // Ime should be above all app windows except for non-fullscreen app window above it and
@@ -300,7 +300,7 @@
                 mDisplayContent, "imeSystemOverlayTarget",
                 true /* ownerCanAddInternalSystemWindow */);
 
-        mWm.mInputMethodTarget = imeSystemOverlayTarget;
+        mDisplayContent.mInputMethodTarget = imeSystemOverlayTarget;
         mDisplayContent.assignChildLayers(mTransaction);
 
         // The IME target base layer is higher than all window except for the nav bar window, so the
@@ -323,7 +323,7 @@
 
     @Test
     public void testAssignWindowLayers_ForStatusBarImeTarget() {
-        mWm.mInputMethodTarget = mStatusBarWindow;
+        mDisplayContent.mInputMethodTarget = mStatusBarWindow;
         mDisplayContent.assignChildLayers(mTransaction);
 
         assertWindowHigher(mImeWindow, mChildAppWindowAbove);