Support IME Window to show in secondary display.

1) Moving WMS.setInputMethodWindowLocked to DisplayContent,
   each display can have its own IME window.
2) Add getDisplayIdFromWindow in WindowManagerInternal,
   used for InputMethodManagerService to know which display
   for given IME window token.
3) Support add / remove IME window according displayId.
4) Modify WMS.inputMethodClientHasFocus to traverse all active display
   if inputMethodClient focused.
5) Add displayId parameter for IInputMethod.initializeInternal to
   update context display then client can addView to right display.

Note: 1) There should be zero behavior difference as long as the target
         app is running on the default display.
      2) The current implementation is not final and there are still
         chances that the current IME may not work well or even crash
	 depending on how the IME is implemented.

Bug: 111364446
Test: manual, use ActivityView & launch Messages in VirtualDisplay,
      tap search icon to see if soft input keyboard shown &
      app window size is adjusted by soft input.
Change-Id: I8da315936caebdc8b2c16cff4e24192c06743251
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index d3e534c..5d0101f 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -409,6 +409,11 @@
 
     private InputMonitor mInputMonitor;
 
+    /**
+     * The input method window for this display.
+     */
+    WindowState mInputMethodWindow;
+
     private final Consumer<WindowState> mUpdateWindowsForAnimator = w -> {
         WindowStateAnimator winAnimator = w.mWinAnimator;
         final AppWindowToken atoken = w.mAppToken;
@@ -2107,18 +2112,16 @@
                 mTouchExcludeRegion.op(mTmpRect2, Region.Op.UNION);
             }
         }
-        final WindowState inputMethod = mService.mInputMethodWindow;
-        if (inputMethod != null && inputMethod.isVisibleLw()) {
+        if (mInputMethodWindow != null && mInputMethodWindow.isVisibleLw()) {
             // If the input method is visible and the user is typing, we don't want these touch
             // events to be intercepted and used to change focus. This would likely cause a
             // disappearance of the input method.
-            inputMethod.getTouchableRegion(mTmpRegion);
-            if (inputMethod.getDisplayId() == mDisplayId) {
+            mInputMethodWindow.getTouchableRegion(mTmpRegion);
+            if (mInputMethodWindow.getDisplayId() == mDisplayId) {
                 mTouchExcludeRegion.op(mTmpRegion, Op.UNION);
             } else {
                 // IME is on a different display, so we need to update its tap detector.
-                // TODO(multidisplay): Remove when IME will always appear on same display.
-                inputMethod.getDisplayContent().setTouchExcludeRegion(null /* focusedTask */);
+                setTouchExcludeRegion(null /* focusedTask */);
             }
         }
         for (int i = mTapExcludedWindows.size() - 1; i >= 0; i--) {
@@ -2257,7 +2260,7 @@
     }
 
     void adjustForImeIfNeeded() {
-        final WindowState imeWin = mService.mInputMethodWindow;
+        final WindowState imeWin = mInputMethodWindow;
         final boolean imeVisible = imeWin != null && imeWin.isVisibleLw() && imeWin.isDisplayedLw()
                 && !mDividerControllerLocked.isImeHideRequested();
         final boolean dockVisible = isStackVisible(WINDOWING_MODE_SPLIT_SCREEN_PRIMARY);
@@ -2640,12 +2643,21 @@
     }
 
     /**
+     * Set input method window for the display.
+     * @param win Set when window added or Null when destroyed.
+     */
+    void setInputMethodWindowLocked(WindowState win) {
+        mInputMethodWindow = win;
+        computeImeTarget(true /* updateImeTarget */);
+    }
+
+    /**
      * Determine and return the window that should be the IME target.
      * @param updateImeTarget If true the system IME target will be updated to match what we found.
      * @return The window that should be used as the IME target or null if there isn't any.
      */
     WindowState computeImeTarget(boolean updateImeTarget) {
-        if (mService.mInputMethodWindow == null) {
+        if (mInputMethodWindow == null) {
             // 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 "
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 86b14337..d92818a 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -681,10 +681,11 @@
                 i--;
                 WindowState win = mService.mDestroySurface.get(i);
                 win.mDestroying = false;
-                if (mService.mInputMethodWindow == win) {
-                    mService.setInputMethodWindowLocked(null);
+                final DisplayContent displayContent = win.getDisplayContent();
+                if (displayContent.mInputMethodWindow == win) {
+                    displayContent.setInputMethodWindowLocked(null);
                 }
-                if (win.getDisplayContent().mWallpaperController.isWallpaperTarget(win)) {
+                if (displayContent.mWallpaperController.isWallpaperTarget(win)) {
                     wallpaperDestroyed = true;
                 }
                 win.destroySurfaceUnchecked();
@@ -1113,4 +1114,18 @@
             callback.accept(mChildren.get(i));
         }
     }
+
+    /**
+     * Get current topmost focused IME window in system.
+     * Will look on all displays in current Z-order.
+     */
+    WindowState getCurrentInputMethodWindow() {
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final DisplayContent displayContent = mChildren.get(i);
+            if (displayContent.mInputMethodWindow != null) {
+                return displayContent.mInputMethodWindow;
+            }
+        }
+        return null;
+    }
 }
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index ac496a8..793ce60 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -335,9 +335,9 @@
     public abstract void registerAppTransitionListener(AppTransitionListener listener);
 
     /**
-     * Retrieves a height of input method window.
+     * Retrieves a height of input method window for given display.
      */
-    public abstract int getInputMethodWindowVisibleHeight();
+    public abstract int getInputMethodWindowVisibleHeight(int displayId);
 
     /**
       * Saves last input method window for transition.
@@ -447,4 +447,9 @@
      * Returns {@code true} if a process that is identified by {@code client} has IME focus.
      */
     public abstract boolean inputMethodClientHasFocus(IInputMethodClient client);
+
+    /**
+     * Return the display Id for given window.
+     */
+    public abstract int getDisplayIdForWindow(IBinder windowToken);
 }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 679e0d8..2ed09ae 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -650,8 +650,6 @@
     /** If true hold off on modifying the animation layer of mInputMethodTarget */
     boolean mInputMethodTargetWaitingAnim;
 
-    WindowState mInputMethodWindow = null;
-
     boolean mHardKeyboardAvailable;
     WindowManagerInternal.OnHardKeyboardStatusChangeListener mHardKeyboardStatusChangeListener;
     SettingsObserver mSettingsObserver;
@@ -1414,7 +1412,7 @@
             win.mToken.addWindow(win);
             if (type == TYPE_INPUT_METHOD) {
                 win.mGivenInsetsPending = true;
-                setInputMethodWindowLocked(win);
+                displayContent.setInputMethodWindowLocked(win);
                 imMayMove = false;
             } else if (type == TYPE_INPUT_METHOD_DIALOG) {
                 displayContent.computeImeTarget(true /* updateImeTarget */);
@@ -1687,8 +1685,9 @@
         mWindowsChanged = true;
         if (DEBUG_WINDOW_MOVEMENT) Slog.v(TAG_WM, "Final remove of window: " + win);
 
-        if (mInputMethodWindow == win) {
-            setInputMethodWindowLocked(null);
+        final DisplayContent displayContent = win.getDisplayContent();
+        if (displayContent.mInputMethodWindow == win) {
+            displayContent.setInputMethodWindowLocked(null);
         }
 
         final WindowToken token = win.mToken;
@@ -1733,13 +1732,6 @@
         dc.getInputMonitor().updateInputWindowsLw(true /*force*/);
     }
 
-    void setInputMethodWindowLocked(WindowState win) {
-        mInputMethodWindow = win;
-        final DisplayContent dc = win != null
-                ? win.getDisplayContent() : getDefaultDisplayContentLocked();
-        dc.computeImeTarget(true /* updateImeTarget */);
-    }
-
     private void updateHiddenWhileSuspendedState(ArraySet<String> packages, boolean suspended) {
         synchronized (mWindowMap) {
             mRoot.updateHiddenWhileSuspendedState(packages, suspended);
@@ -2058,8 +2050,10 @@
                 if ((result & WindowManagerGlobal.RELAYOUT_RES_FIRST_TIME) != 0) {
                     focusMayChange = isDefaultDisplay;
                 }
-                if (win.mAttrs.type == TYPE_INPUT_METHOD && mInputMethodWindow == null) {
-                    setInputMethodWindowLocked(win);
+                final DisplayContent displayContent = win.getDisplayContent();
+                if (win.mAttrs.type == TYPE_INPUT_METHOD
+                        && displayContent.mInputMethodWindow == null) {
+                    displayContent.setInputMethodWindowLocked(win);
                     imMayMove = true;
                 }
                 win.adjustStartingWindowFlags();
@@ -2222,8 +2216,9 @@
             // of a transaction to avoid artifacts.
             win.mAnimatingExit = true;
         } else {
-            if (mInputMethodWindow == win) {
-                setInputMethodWindowLocked(null);
+            final DisplayContent displayContent = win.getDisplayContent();
+            if (displayContent.mInputMethodWindow == win) {
+                displayContent.setInputMethodWindowLocked(null);
             }
             boolean stopped = win.mAppToken != null ? win.mAppToken.mAppStopped : true;
             // We set mDestroying=true so AppWindowToken#notifyAppStopped in-to destroy surfaces
@@ -2828,7 +2823,7 @@
 
     @Override
     public WindowManagerPolicy.WindowState getInputMethodWindowLw() {
-        return mInputMethodWindow;
+        return mRoot.getCurrentInputMethodWindow();
     }
 
     @Override
@@ -5575,10 +5570,10 @@
             // change message pending.
             mH.removeMessages(H.REPORT_FOCUS_CHANGE);
             mH.sendEmptyMessage(H.REPORT_FOCUS_CHANGE);
-            // TODO(multidisplay): Focused windows on default display only.
-            final DisplayContent displayContent = getDefaultDisplayContentLocked();
+            final DisplayContent displayContent = (newFocus != null) ? newFocus.getDisplayContent()
+                    : getDefaultDisplayContentLocked();
             boolean imWindowChanged = false;
-            if (mInputMethodWindow != null) {
+            if (displayContent.mInputMethodWindow != null) {
                 final WindowState prevTarget = mInputMethodTarget;
 
                 final WindowState newTarget =
@@ -5587,10 +5582,11 @@
 
                 if (mode != UPDATE_FOCUS_WILL_ASSIGN_LAYERS
                         && mode != UPDATE_FOCUS_WILL_PLACE_SURFACES) {
-                    final int prevImeAnimLayer = mInputMethodWindow.mWinAnimator.mAnimLayer;
+                    final int prevImeAnimLayer =
+                            displayContent.mInputMethodWindow.mWinAnimator.mAnimLayer;
                     displayContent.assignWindowLayers(false /* setLayoutNeeded */);
-                    imWindowChanged |=
-                            prevImeAnimLayer != mInputMethodWindow.mWinAnimator.mAnimLayer;
+                    imWindowChanged |= prevImeAnimLayer
+                            != displayContent.mInputMethodWindow.mWinAnimator.mAnimLayer;
                 }
             }
 
@@ -5613,7 +5609,7 @@
 
             int focusChanged = mPolicy.focusChangedLw(oldFocus, newFocus);
 
-            if (imWindowChanged && oldFocus != mInputMethodWindow) {
+            if (imWindowChanged && oldFocus != displayContent.mInputMethodWindow) {
                 // Focus of the input method window changed. Perform layout if needed.
                 if (mode == UPDATE_FOCUS_PLACING_SURFACES) {
                     displayContent.performLayout(true /*initial*/,  updateInputWindows);
@@ -6033,8 +6029,15 @@
         }
         synchronized (mWindowMap) {
             final Region r = new Region();
-            if (mInputMethodWindow != null) {
-                mInputMethodWindow.getTouchableRegion(r);
+            // TODO(b/111080190): this method is only return the recent focused IME touch region,
+            // For Multi-Session IME, will need to add API for given display Id to
+            // get the right IME touch region.
+            for (int i = mRoot.mChildren.size() - 1; i >= 0; --i) {
+                final DisplayContent displayContent = mRoot.mChildren.get(i);
+                if (displayContent.mInputMethodWindow != null) {
+                    displayContent.mInputMethodWindow.getTouchableRegion(r);
+                    return r;
+                }
             }
             return r;
         }
@@ -6217,8 +6220,9 @@
         if (mFocusedApp != null) {
             mFocusedApp.writeNameToProto(proto, FOCUSED_APP);
         }
-        if (mInputMethodWindow != null) {
-            mInputMethodWindow.writeIdentifierToProto(proto, INPUT_METHOD_WINDOW);
+        final WindowState imeWindow = mRoot.getCurrentInputMethodWindow();
+        if (imeWindow != null) {
+            imeWindow.writeIdentifierToProto(proto, INPUT_METHOD_WINDOW);
         }
         proto.write(DISPLAY_FROZEN, mDisplayFrozen);
         final DisplayContent defaultDisplayContent = getDefaultDisplayContentLocked();
@@ -6388,8 +6392,9 @@
                 pw.print("  mLastStatusBarVisibility=0x");
                         pw.println(Integer.toHexString(mLastStatusBarVisibility));
             }
-            if (mInputMethodWindow != null) {
-                pw.print("  mInputMethodWindow="); pw.println(mInputMethodWindow);
+            final WindowState imeWindow = mRoot.getCurrentInputMethodWindow();
+            if (imeWindow != null) {
+                pw.print("  mInputMethodWindow="); pw.println(imeWindow);
             }
             mWindowPlacerLocked.dump(pw, "  ");
             mRoot.mWallpaperController.dump(pw, "  ");
@@ -7287,10 +7292,9 @@
         }
 
         @Override
-        public int getInputMethodWindowVisibleHeight() {
+        public int getInputMethodWindowVisibleHeight(int displayId) {
             synchronized (mWindowMap) {
-                // TODO(multi-display): Have caller pass in the display they are interested in.
-                final DisplayContent dc = getDefaultDisplayContentLocked();
+                final DisplayContent dc = mRoot.getDisplayContent(displayId);
                 return dc.mDisplayFrames.getInputMethodWindowVisibleHeight();
             }
         }
@@ -7298,8 +7302,9 @@
         @Override
         public void saveLastInputMethodWindowForTransition() {
             synchronized (mWindowMap) {
-                if (mInputMethodWindow != null) {
-                    mPolicy.setLastInputMethodWindowLw(mInputMethodWindow, mInputMethodTarget);
+                final WindowState imeWindow = mRoot.getCurrentInputMethodWindow();
+                if (imeWindow != null) {
+                    mPolicy.setLastInputMethodWindowLw(imeWindow, mInputMethodTarget);
                 }
             }
         }
@@ -7406,28 +7411,42 @@
             }
         }
 
-        @Override
-        public boolean inputMethodClientHasFocus(IInputMethodClient client) {
-            synchronized (mWindowMap) {
-                // TODO: multi-display
-                if (getDefaultDisplayContentLocked().inputMethodClientHasFocus(client)) {
-                    return true;
-                }
-
-                // Okay, how about this...  what is the current focus?
-                // It seems in some cases we may not have moved the IM
-                // target window, such as when it was in a pop-up window,
-                // so let's also look at the current focus.  (An example:
-                // go to Gmail, start searching so the keyboard goes up,
-                // press home.  Sometimes the IME won't go down.)
-                // Would be nice to fix this more correctly, but it's
-                // way at the end of a release, and this should be good enough.
-                if (mCurrentFocus != null && mCurrentFocus.mSession.mClient != null
-                        && mCurrentFocus.mSession.mClient.asBinder() == client.asBinder()) {
+    public boolean inputMethodClientHasFocus(IInputMethodClient client) {
+        boolean hasFocus;
+        synchronized (mWindowMap) {
+            // Check all displays if any input method window has focus.
+            for (int i = mRoot.mChildren.size() - 1; i >= 0; --i) {
+                final DisplayContent displayContent = mRoot.mChildren.get(i);
+                if (displayContent.inputMethodClientHasFocus(client)) {
                     return true;
                 }
             }
-            return false;
+
+            // Okay, how about this...  what is the current focus?
+            // It seems in some cases we may not have moved the IM
+            // target window, such as when it was in a pop-up window,
+            // so let's also look at the current focus.  (An example:
+            // go to Gmail, start searching so the keyboard goes up,
+            // press home.  Sometimes the IME won't go down.)
+            // Would be nice to fix this more correctly, but it's
+            // way at the end of a release, and this should be good enough.
+            if (mCurrentFocus != null && mCurrentFocus.mSession.mClient != null
+                    && mCurrentFocus.mSession.mClient.asBinder() == client.asBinder()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+        @Override
+        public int getDisplayIdForWindow(IBinder windowToken) {
+            synchronized (mWindowMap) {
+                final WindowState window = mWindowMap.get(windowToken);
+                if (window != null) {
+                    return window.getDisplayContent().getDisplayId();
+                }
+                return Display.INVALID_DISPLAY;
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 637c0ea..a31f5b3 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -845,7 +845,7 @@
                 mWindowFrames.mContainingFrame.bottom =
                         mWindowFrames.mContainingFrame.top + frozen.height();
             }
-            final WindowState imeWin = mService.mInputMethodWindow;
+            final WindowState imeWin = mService.mRoot.getCurrentInputMethodWindow();
             // IME is up and obscuring this window. Adjust the window position so it is visible.
             if (imeWin != null && imeWin.isVisibleNow() && isInputMethodTarget()) {
                 if (inFreeformWindowingMode() && mWindowFrames.mContainingFrame.bottom