diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index 2ba1e016e03d2..1c811cb6e96cf 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -172,6 +172,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 7ad118e760d8d..fd3ecb48ee2ef 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(),
@@ -2855,7 +2848,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);
             }
@@ -4560,8 +4553,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();
                     }
@@ -5106,7 +5098,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 5e00425407ba8..2cfd6226351e2 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;