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() {