2 files changed, 19 insertions, 16 deletions

diff --git a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
index e862c7e96200..ad9995097639 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
@@ -39,6 +39,7 @@ import android.os.UserHandle;
 import android.provider.Settings;
 import android.util.EventLog;
 import android.util.Slog;
+import android.view.Display;
 import android.view.WindowManager;
 import android.view.inputmethod.InputMethod;
 import android.view.inputmethod.InputMethodInfo;
@@ -78,6 +79,7 @@ final class InputMethodBindingController {
     @GuardedBy("ImfLock.class") @Nullable private IInputMethodInvoker mCurMethod;
     @GuardedBy("ImfLock.class") private int mCurMethodUid = Process.INVALID_UID;
     @GuardedBy("ImfLock.class") @Nullable private IBinder mCurToken;
+    @GuardedBy("ImfLock.class") private int mCurTokenDisplayId = Display.INVALID_DISPLAY;
     @GuardedBy("ImfLock.class") private int mCurSeq;
     @GuardedBy("ImfLock.class") private boolean mVisibleBound;
     @GuardedBy("ImfLock.class") private boolean mSupportsStylusHw;
@@ -193,6 +195,17 @@ final class InputMethodBindingController {
     }
 
     /**
+     * Returns the displayId associated with {@link #getCurToken()}.
+     *
+     * @return the displayId associated with {@link #getCurToken()}. {@link Display#INVALID_DISPLAY}
+     *         while {@link #getCurToken()} returns {@code null}
+     */
+    @GuardedBy("ImfLock.class")
+    int getCurTokenDisplayId() {
+        return mCurTokenDisplayId;
+    }
+
+    /**
      * The Intent used to connect to the current input method.
      */
     @GuardedBy("ImfLock.class")
@@ -412,15 +425,14 @@ final class InputMethodBindingController {
 
     @GuardedBy("ImfLock.class")
     private void removeCurrentToken() {
-        int curTokenDisplayId = mService.getCurTokenDisplayIdLocked();
-
         if (DEBUG) {
             Slog.v(TAG,
-                    "Removing window token: " + mCurToken + " for display: " + curTokenDisplayId);
+                    "Removing window token: " + mCurToken + " for display: " + mCurTokenDisplayId);
         }
         mWindowManagerInternal.removeWindowToken(mCurToken, true /* removeWindows */,
-                false /* animateExit */, curTokenDisplayId);
+                false /* animateExit */, mCurTokenDisplayId);
         mCurToken = null;
+        mCurTokenDisplayId = Display.INVALID_DISPLAY;
     }
 
     @GuardedBy("ImfLock.class")
@@ -445,7 +457,7 @@ final class InputMethodBindingController {
 
             final int displayIdToShowIme = mService.getDisplayIdToShowImeLocked();
             mCurToken = new Binder();
-            mService.setCurTokenDisplayIdLocked(displayIdToShowIme);
+            mCurTokenDisplayId = displayIdToShowIme;
             if (DEBUG) {
                 Slog.v(TAG, "Adding window token: " + mCurToken + " for display: "
                         + displayIdToShowIme);
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 46c5772453c7..61512e970dba 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -581,18 +581,10 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
      */
     @GuardedBy("ImfLock.class")
     int getCurTokenDisplayIdLocked() {
-        return mCurTokenDisplayId;
-    }
-
-    @GuardedBy("ImfLock.class")
-    void setCurTokenDisplayIdLocked(int curTokenDisplayId) {
-        mCurTokenDisplayId = curTokenDisplayId;
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getCurTokenDisplayId();
     }
 
-    @GuardedBy("ImfLock.class")
-    @MultiUserUnawareField
-    private int mCurTokenDisplayId = INVALID_DISPLAY;
-
     /**
      * The display ID of the input method indicates the fallback display which returned by
      * {@link #computeImeDisplayIdForTarget}.
@@ -2472,7 +2464,6 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
         mImeWindowVis = 0;
         mBackDisposition = InputMethodService.BACK_DISPOSITION_DEFAULT;
         updateSystemUiLocked(mImeWindowVis, mBackDisposition);
-        mCurTokenDisplayId = INVALID_DISPLAY;
         mAutofillController.onResetSystemUi();
     }