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

13 files changed, 202 insertions, 89 deletions

diff --git a/config/boot-image-profile.txt b/config/boot-image-profile.txt
index 04b33b18dce4..731440b2fdda 100644
--- a/config/boot-image-profile.txt
+++ b/config/boot-image-profile.txt
@@ -24288,7 +24288,6 @@ HSPLandroid/inputmethodservice/AbstractInputMethodService;->onCreateInputMethodS
 HSPLandroid/inputmethodservice/IInputMethodSessionWrapper;->getInternalInputMethodSession()Landroid/view/inputmethod/InputMethodSession;
 HSPLandroid/inputmethodservice/IInputMethodWrapper$InputMethodSessionCallbackWrapper;->sessionCreated(Landroid/view/inputmethod/InputMethodSession;)V
 HSPLandroid/inputmethodservice/IInputMethodWrapper;-><init>(Landroid/inputmethodservice/AbstractInputMethodService;Landroid/view/inputmethod/InputMethod;)V
-HSPLandroid/inputmethodservice/IInputMethodWrapper;->attachToken(Landroid/os/IBinder;)V
 HSPLandroid/inputmethodservice/IInputMethodWrapper;->bindInput(Landroid/view/inputmethod/InputBinding;)V
 HSPLandroid/inputmethodservice/IInputMethodWrapper;->createSession(Landroid/view/InputChannel;Lcom/android/internal/view/IInputSessionCallback;)V
 HSPLandroid/inputmethodservice/IInputMethodWrapper;->executeMessage(Landroid/os/Message;)V
@@ -42280,11 +42279,9 @@ HSPLcom/android/internal/view/IInputContextCallback;->setSelectedText(Ljava/lang
 HSPLcom/android/internal/view/IInputContextCallback;->setTextAfterCursor(Ljava/lang/CharSequence;I)V
 HSPLcom/android/internal/view/IInputContextCallback;->setTextBeforeCursor(Ljava/lang/CharSequence;I)V
 HSPLcom/android/internal/view/IInputMethod$Stub$Proxy;->asBinder()Landroid/os/IBinder;
-HSPLcom/android/internal/view/IInputMethod$Stub$Proxy;->attachToken(Landroid/os/IBinder;)V
 HSPLcom/android/internal/view/IInputMethod$Stub;-><init>()V
 HSPLcom/android/internal/view/IInputMethod$Stub;->asInterface(Landroid/os/IBinder;)Lcom/android/internal/view/IInputMethod;
 HSPLcom/android/internal/view/IInputMethod$Stub;->onTransact(ILandroid/os/Parcel;Landroid/os/Parcel;I)Z
-HSPLcom/android/internal/view/IInputMethod;->attachToken(Landroid/os/IBinder;)V
 HSPLcom/android/internal/view/IInputMethod;->bindInput(Landroid/view/inputmethod/InputBinding;)V
 HSPLcom/android/internal/view/IInputMethod;->changeInputMethodSubtype(Landroid/view/inputmethod/InputMethodSubtype;)V
 HSPLcom/android/internal/view/IInputMethod;->createSession(Landroid/view/InputChannel;Lcom/android/internal/view/IInputSessionCallback;)V
diff --git a/core/java/android/inputmethodservice/IInputMethodWrapper.java b/core/java/android/inputmethodservice/IInputMethodWrapper.java
index 4080ee625a0f..103069474445 100644
--- a/core/java/android/inputmethodservice/IInputMethodWrapper.java
+++ b/core/java/android/inputmethodservice/IInputMethodWrapper.java
@@ -160,11 +160,10 @@ class IInputMethodWrapper extends IInputMethod.Stub
                 args.recycle();
                 return;
             }
-
             case DO_INITIALIZE_INTERNAL: {
                 SomeArgs args = (SomeArgs) msg.obj;
                 try {
-                    inputMethod.initializeInternal((IBinder) args.arg1,
+                    inputMethod.initializeInternal((IBinder) args.arg1, msg.arg1,
                             (IInputMethodPrivilegedOperations) args.arg2);
                 } finally {
                     args.recycle();
@@ -253,9 +252,10 @@ class IInputMethodWrapper extends IInputMethod.Stub
 
     @BinderThread
     @Override
-    public void initializeInternal(IBinder token, IInputMethodPrivilegedOperations privOps) {
+    public void initializeInternal(IBinder token, int displayId,
+            IInputMethodPrivilegedOperations privOps) {
         mCaller.executeOrSendMessage(
-                mCaller.obtainMessageOO(DO_INITIALIZE_INTERNAL, token, privOps));
+                mCaller.obtainMessageIOO(DO_INITIALIZE_INTERNAL, displayId, token, privOps));
     }
 
     @BinderThread
diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index 2d12b867949a..34fa5b6b9842 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -16,6 +16,7 @@
 
 package android.inputmethodservice;
 
+import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
 import static android.view.ViewGroup.LayoutParams.WRAP_CONTENT;
 import static android.view.WindowManager.LayoutParams.FLAG_DRAWS_SYSTEM_BAR_BACKGROUNDS;
@@ -461,10 +462,11 @@ public class InputMethodService extends AbstractInputMethodService {
          */
         @MainThread
         @Override
-        public final void initializeInternal(IBinder token,
+        public final void initializeInternal(IBinder token, int displayId,
                 IInputMethodPrivilegedOperations privilegedOperations) {
             mPrivOps.set(privilegedOperations);
             mImm.registerInputMethodPrivOps(token, mPrivOps);
+            updateInputMethodDisplay(displayId);
             attachToken(token);
         }
 
@@ -484,6 +486,22 @@ public class InputMethodService extends AbstractInputMethodService {
 
         /**
          * {@inheritDoc}
+         * @hide
+         */
+        @MainThread
+        @Override
+        public void updateInputMethodDisplay(int displayId) {
+            // Update display for adding IME window to the right display.
+            if (displayId != DEFAULT_DISPLAY) {
+                // TODO(b/111364446) Need to address context lifecycle issue if need to re-create
+                // for update resources & configuration correctly when show soft input
+                // in non-default display.
+                updateDisplay(displayId);
+            }
+        }
+
+        /**
+         * {@inheritDoc}
          *
          * <p>Calls {@link InputMethodService#onBindInput()} when done.</p>
          */
@@ -930,6 +948,9 @@ public class InputMethodService extends AbstractInputMethodService {
         // If the previous IME has occupied non-empty inset in the screen, we need to decide whether
         // we continue to use the same size of the inset or update it
         mShouldClearInsetOfPreviousIme = (mImm.getInputMethodWindowVisibleHeight() > 0);
+        // TODO(b/111364446) Need to address context lifecycle issue if need to re-create
+        // for update resources & configuration correctly when show soft input
+        // in non-default display.
         mInflater = (LayoutInflater)getSystemService(
                 Context.LAYOUT_INFLATER_SERVICE);
         mWindow = new SoftInputWindow(this, "InputMethod", mTheme, null, null, mDispatcherState,
diff --git a/core/java/android/view/inputmethod/InputMethod.java b/core/java/android/view/inputmethod/InputMethod.java
index f45507c2b8b3..e1600c4a92c8 100644
--- a/core/java/android/view/inputmethod/InputMethod.java
+++ b/core/java/android/view/inputmethod/InputMethod.java
@@ -89,14 +89,17 @@ public interface InputMethod {
      *
      * @param token special token for the system to identify
      *              {@link InputMethodService}
+     * @param displayId The id of the display that current IME shown.
+     *                  Used for {{@link #updateInputMethodDisplay(int)}}
      * @param privilegedOperations IPC endpoint to do some privileged
      *                             operations that are allowed only to the
      *                             current IME.
      * @hide
      */
     @MainThread
-    default void initializeInternal(IBinder token,
+    default void initializeInternal(IBinder token, int displayId,
             IInputMethodPrivilegedOperations privilegedOperations) {
+        updateInputMethodDisplay(displayId);
         attachToken(token);
     }
 
@@ -115,6 +118,15 @@ public interface InputMethod {
     public void attachToken(IBinder token);
 
     /**
+     * Update context display according to given displayId.
+     *
+     * @param displayId The id of the display that need to update for context.
+     * @hide
+     */
+    @MainThread
+    public void updateInputMethodDisplay(int displayId);
+
+    /**
      * Bind a new application environment in to the input method, so that it
      * can later start and stop input processing.
      * Typically this method is called when this input method is enabled in an
diff --git a/core/java/com/android/internal/view/IInputMethod.aidl b/core/java/com/android/internal/view/IInputMethod.aidl
index b6a654aac563..97d5a657e5c4 100644
--- a/core/java/com/android/internal/view/IInputMethod.aidl
+++ b/core/java/com/android/internal/view/IInputMethod.aidl
@@ -33,7 +33,7 @@ import com.android.internal.view.IInputSessionCallback;
  * {@hide}
  */
 oneway interface IInputMethod {
-    void initializeInternal(IBinder token, IInputMethodPrivilegedOperations privOps);
+    void initializeInternal(IBinder token, int displayId, IInputMethodPrivilegedOperations privOps);
 
     void bindInput(in InputBinding binding);
 
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 49f33e09a724..44fb9ea91dee 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -16,6 +16,7 @@
 package com.android.server.inputmethod;
 
 import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_SHOW_FOR_ALL_USERS;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG;
@@ -572,6 +573,11 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
     IBinder mCurToken;
 
     /**
+     * The displayId of current active input method.
+     */
+    int mCurTokenDisplayId = INVALID_DISPLAY;
+
+    /**
      * If non-null, this is the input method service we are currently connected
      * to.
      */
@@ -1866,7 +1872,9 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         mCurAttribute = attribute;
 
         // Check if the input method is changing.
-        if (mCurId != null && mCurId.equals(mCurMethodId)) {
+        final int displayId = mWindowManagerInternal.getDisplayIdForWindow(
+                mCurFocusedWindow);
+        if (mCurId != null && mCurId.equals(mCurMethodId) && displayId == mCurTokenDisplayId) {
             if (cs.curSession != null) {
                 // Fast case: if we are already connected to the input method,
                 // then just return it.
@@ -1930,14 +1938,20 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                 com.android.internal.R.string.input_method_binding_label);
         mCurIntent.putExtra(Intent.EXTRA_CLIENT_INTENT, PendingIntent.getActivity(
                 mContext, 0, new Intent(Settings.ACTION_INPUT_METHOD_SETTINGS), 0));
+        final int displayId = mWindowManagerInternal.getDisplayIdForWindow(mCurFocusedWindow);
+        mCurTokenDisplayId = (displayId != INVALID_DISPLAY) ? displayId : DEFAULT_DISPLAY;
+
         if (bindCurrentInputMethodServiceLocked(mCurIntent, this, IME_CONNECTION_BIND_FLAGS)) {
             mLastBindTime = SystemClock.uptimeMillis();
             mHaveConnection = true;
             mCurId = info.getId();
             mCurToken = new Binder();
             try {
-                if (DEBUG) Slog.v(TAG, "Adding window token: " + mCurToken);
-                mIWindowManager.addWindowToken(mCurToken, TYPE_INPUT_METHOD, DEFAULT_DISPLAY);
+                if (DEBUG) {
+                    Slog.v(TAG, "Adding window token: " + mCurToken + " for display: "
+                            + mCurTokenDisplayId);
+                }
+                mIWindowManager.addWindowToken(mCurToken, TYPE_INPUT_METHOD, mCurTokenDisplayId);
             } catch (RemoteException e) {
             }
             return new InputBindResult(
@@ -1964,8 +1978,9 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                     return;
                 }
                 if (DEBUG) Slog.v(TAG, "Initiating attach with token: " + mCurToken);
-                executeOrSendMessage(mCurMethod, mCaller.obtainMessageOO(
-                        MSG_INITIALIZE_IME, mCurMethod, mCurToken));
+                // Dispatch display id for InputMethodService to update context display.
+                executeOrSendMessage(mCurMethod, mCaller.obtainMessageIOO(
+                        MSG_INITIALIZE_IME, mCurTokenDisplayId, mCurMethod, mCurToken));
                 if (mCurClient != null) {
                     clearClientSessionLocked(mCurClient);
                     requestClientSessionLocked(mCurClient);
@@ -2011,15 +2026,19 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
 
         if (mCurToken != null) {
             try {
-                if (DEBUG) Slog.v(TAG, "Removing window token: " + mCurToken);
+                if (DEBUG) {
+                    Slog.v(TAG, "Removing window token: " + mCurToken + " for display: "
+                            + mCurTokenDisplayId);
+                }
                 if ((mImeWindowVis & InputMethodService.IME_ACTIVE) != 0 && savePosition) {
                     // The current IME is shown. Hence an IME switch (transition) is happening.
                     mWindowManagerInternal.saveLastInputMethodWindowForTransition();
                 }
-                mIWindowManager.removeWindowToken(mCurToken, DEFAULT_DISPLAY);
+                mIWindowManager.removeWindowToken(mCurToken, mCurTokenDisplayId);
             } catch (RemoteException e) {
             }
             mCurToken = null;
+            mCurTokenDisplayId = INVALID_DISPLAY;
         }
 
         mCurId = null;
@@ -2785,6 +2804,15 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                                 // soft input window if it is shown.
                                 if (DEBUG) Slog.v(TAG, "Unspecified window will hide input");
                                 hideCurrentInputLocked(InputMethodManager.HIDE_NOT_ALWAYS, null);
+
+                                // If focused display changed, we should unbind current method
+                                // to make app window in previous display relayout after Ime
+                                // window token removed.
+                                final int newFocusDisplayId =
+                                        mWindowManagerInternal.getDisplayIdForWindow(windowToken);
+                                if (newFocusDisplayId != mCurTokenDisplayId) {
+                                    unbindCurrentMethodLocked(false);
+                                }
                             }
                         } else if (isTextEditor && doAutoShow && (softInputMode &
                                 WindowManager.LayoutParams.SOFT_INPUT_IS_FORWARD_NAVIGATION) != 0) {
@@ -3151,7 +3179,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
      */
     @Override
     public int getInputMethodWindowVisibleHeight() {
-        return mWindowManagerInternal.getInputMethodWindowVisibleHeight();
+        return mWindowManagerInternal.getInputMethodWindowVisibleHeight(mCurTokenDisplayId);
     }
 
     @BinderThread
@@ -3352,9 +3380,12 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
             case MSG_INITIALIZE_IME:
                 args = (SomeArgs)msg.obj;
                 try {
-                    if (DEBUG) Slog.v(TAG, "Sending attach of token: " + args.arg2);
+                    if (DEBUG) {
+                        Slog.v(TAG, "Sending attach of token: " + args.arg2 + " for display: "
+                                + msg.arg1);
+                    }
                     final IBinder token = (IBinder) args.arg2;
-                    ((IInputMethod) args.arg1).initializeInternal(token,
+                    ((IInputMethod) args.arg1).initializeInternal(token, msg.arg1,
                             new InputMethodPrivilegedOperationsImpl(this, token));
                 } catch (RemoteException e) {
                 }
@@ -4516,6 +4547,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
             p.println("  mCurId=" + mCurId + " mHaveConnection=" + mHaveConnection
                     + " mBoundToMethod=" + mBoundToMethod + " mVisibleBound=" + mVisibleBound);
             p.println("  mCurToken=" + mCurToken);
+            p.println("  mCurTokenDisplayId=" + mCurTokenDisplayId);
             p.println("  mCurIntent=" + mCurIntent);
             method = mCurMethod;
             p.println("  mCurMethod=" + mCurMethod);
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 3e0429fd443e..350d6b66c1fb 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -4732,8 +4732,8 @@ public class PhoneWindowManager implements WindowManagerPolicy {
         }
         final WindowManager.LayoutParams attrs = win.getAttrs();
         final boolean isDefaultDisplay = win.isDefaultDisplay();
-        final boolean needsToOffsetInputMethodTarget = isDefaultDisplay &&
-                (win == mLastInputMethodTargetWindow && mLastInputMethodWindow != null);
+        final boolean needsToOffsetInputMethodTarget =
+                (win == mLastInputMethodTargetWindow) && (mLastInputMethodWindow != null);
         if (needsToOffsetInputMethodTarget) {
             if (DEBUG_LAYOUT) Slog.i(TAG, "Offset ime target window by the last ime window state");
             offsetInputMethodWindowLw(mLastInputMethodWindow, displayFrames);
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index d3e534ce4522..5d0101f88c69 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 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
 
     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 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
                 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 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
     }
 
     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 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
     }
 
     /**
+     * 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 86b14337139e..d92818ac611c 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 @@ class RootWindowContainer extends WindowContainer<DisplayContent> {
                 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 @@ class RootWindowContainer extends WindowContainer<DisplayContent> {
             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 ac496a84f178..793ce6047323 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 class WindowManagerInternal {
     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 @@ public abstract class WindowManagerInternal {
      * 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 679e0d870d74..2ed09ae2daf3 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 @@ public class WindowManagerService extends IWindowManager.Stub
     /** 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 @@ public class WindowManagerService extends IWindowManager.Stub
             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 @@ public class WindowManagerService extends IWindowManager.Stub
         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 @@ public class WindowManagerService extends IWindowManager.Stub
         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 @@ public class WindowManagerService extends IWindowManager.Stub
                 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 @@ public class WindowManagerService extends IWindowManager.Stub
             // 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 @@ public class WindowManagerService extends IWindowManager.Stub
 
     @Override
     public WindowManagerPolicy.WindowState getInputMethodWindowLw() {
-        return mInputMethodWindow;
+        return mRoot.getCurrentInputMethodWindow();
     }
 
     @Override
@@ -5575,10 +5570,10 @@ public class WindowManagerService extends IWindowManager.Stub
             // 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 @@ public class WindowManagerService extends IWindowManager.Stub
 
                 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 @@ public class WindowManagerService extends IWindowManager.Stub
 
             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 @@ public class WindowManagerService extends IWindowManager.Stub
         }
         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 @@ public class WindowManagerService extends IWindowManager.Stub
         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 @@ public class WindowManagerService extends IWindowManager.Stub
                 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 @@ public class WindowManagerService extends IWindowManager.Stub
         }
 
         @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 @@ public class WindowManagerService extends IWindowManager.Stub
         @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 @@ public class WindowManagerService extends IWindowManager.Stub
             }
         }
 
-        @Override
-        public boolean inputMethodClientHasFocus(IInputMethodClient client) {
-            synchronized (mWindowMap) {
-                // TODO: multi-display
-                if (getDefaultDisplayContentLocked().inputMethodClientHasFocus(client)) {
+    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;
                 }
+            }
 
-                // 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;
+            // 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;
             }
-            return false;
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 637c0eabdac4..a31f5b3ed31c 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 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
                 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
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 5db0867298b0..70e4ce419e0a 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -125,7 +125,7 @@ class WindowTestsBase {
             synchronized (sWm.mWindowMap) {
                 mWallpaperWindow = createCommonWindow(null, TYPE_WALLPAPER, "wallpaperWindow");
                 mImeWindow = createCommonWindow(null, TYPE_INPUT_METHOD, "mImeWindow");
-                sWm.mInputMethodWindow = mImeWindow;
+                mDisplayContent.mInputMethodWindow = mImeWindow;
                 mImeDialogWindow = createCommonWindow(null, TYPE_INPUT_METHOD_DIALOG,
                         "mImeDialogWindow");
                 mStatusBarWindow = createCommonWindow(null, TYPE_STATUS_BAR, "mStatusBarWindow");