diff --git a/api/test-current.txt b/api/test-current.txt
index d9b12cb52edb7..eb602692798de 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -3392,6 +3392,7 @@ package android.view.contentcapture {
 package android.view.inputmethod {
 
   public final class InputMethodManager {
+    method public int getDisplayId();
     method public boolean isInputMethodPickerShown();
   }
 
diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index d415387808dd4..83abf1a33b2d1 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -173,6 +173,6 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     }
 
     private InputMethodManager getImm() {
-        return mController.getViewRoot().mDisplayContext.getSystemService(InputMethodManager.class);
+        return mController.getViewRoot().mContext.getSystemService(InputMethodManager.class);
     }
 }
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index b813bc34c7633..4463e13ca5ee4 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -283,13 +283,7 @@ public final class ViewRootImpl implements ViewParent,
     @GuardedBy("mWindowCallbacks")
     final ArrayList<WindowCallbacks> mWindowCallbacks = new ArrayList<>();
     @UnsupportedAppUsage
-    final Context mContext;
-    /**
-     * TODO(b/116349163): Check if we can merge this into {@link #mContext}.
-     * @hide
-     */
-    @NonNull
-    public Context mDisplayContext;
+    public final Context mContext;
 
     @UnsupportedAppUsage
     final IWindowSession mWindowSession;
@@ -595,7 +589,6 @@ public final class ViewRootImpl implements ViewParent,
 
     public ViewRootImpl(Context context, Display display) {
         mContext = context;
-        mDisplayContext = context.createDisplayContext(display);
         mWindowSession = WindowManagerGlobal.getWindowSession();
         mDisplay = display;
         mBasePackageName = context.getBasePackageName();
@@ -1379,7 +1372,7 @@ public final class ViewRootImpl implements ViewParent,
         } else {
             mDisplay = preferredDisplay;
         }
-        mDisplayContext = mContext.createDisplayContext(mDisplay);
+        mContext.updateDisplay(mDisplay.getDisplayId());
     }
 
     void pokeDrawLockIfNeeded() {
@@ -2725,7 +2718,7 @@ public final class ViewRootImpl implements ViewParent,
                     .mayUseInputMethod(mWindowAttributes.flags);
             if (imTarget != mLastWasImTarget) {
                 mLastWasImTarget = imTarget;
-                InputMethodManager imm = mDisplayContext.getSystemService(InputMethodManager.class);
+                InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
                 if (imm != null && imTarget) {
                     imm.onPreWindowFocus(mView, hasWindowFocus);
                     imm.onPostWindowFocus(mView, mView.findFocus(),
@@ -2859,7 +2852,7 @@ public final class ViewRootImpl implements ViewParent,
             mLastWasImTarget = WindowManager.LayoutParams
                     .mayUseInputMethod(mWindowAttributes.flags);
 
-            InputMethodManager imm = mDisplayContext.getSystemService(InputMethodManager.class);
+            InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
             if (imm != null && mLastWasImTarget && !isInLocalFocusMode()) {
                 imm.onPreWindowFocus(mView, hasWindowFocus);
             }
@@ -4564,8 +4557,7 @@ public final class ViewRootImpl implements ViewParent,
                     enqueueInputEvent(event, null, 0, true);
                 } break;
                 case MSG_CHECK_FOCUS: {
-                    InputMethodManager imm =
-                            mDisplayContext.getSystemService(InputMethodManager.class);
+                    InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
                     if (imm != null) {
                         imm.checkFocus();
                     }
@@ -5110,7 +5102,7 @@ public final class ViewRootImpl implements ViewParent,
         @Override
         protected int onProcess(QueuedInputEvent q) {
             if (mLastWasImTarget && !isInLocalFocusMode()) {
-                InputMethodManager imm = mDisplayContext.getSystemService(InputMethodManager.class);
+                InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
                 if (imm != null) {
                     final InputEvent event = q.mEvent;
                     if (DEBUG_IMF) Log.v(mTag, "Sending input event to IME: " + event);
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index fd73856bf79b5..c7b1cdab4361a 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -510,7 +510,7 @@ public final class InputMethodManager {
             return null;
         }
         final InputMethodManager fallbackImm =
-                viewRootImpl.mDisplayContext.getSystemService(InputMethodManager.class);
+                viewRootImpl.mContext.getSystemService(InputMethodManager.class);
         if (fallbackImm == null) {
             Log.e(TAG, "b/117267690: Failed to get non-null fallback IMM. view=" + view);
             return null;
@@ -2887,6 +2887,16 @@ public final class InputMethodManager {
         }
     }
 
+    /**
+     * <p>This is used for CTS test only. Do not use this method outside of CTS package.<p/>
+     * @return the ID of this display which this {@link InputMethodManager} resides
+     * @hide
+     */
+    @TestApi
+    public int getDisplayId() {
+        return mDisplayId;
+    }
+
     void doDump(FileDescriptor fd, PrintWriter fout, String[] args) {
         final Printer p = new PrintWriterPrinter(fout);
         p.println("Input method client state for " + this + ":");