diff --git a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
index 19a48db..5dc4613 100644
--- a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
+++ b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
@@ -15,11 +15,20 @@
  */
 package com.android.quickstep.util;
 
-import android.content.Context;
-import android.view.Display;
+import static android.view.Display.DEFAULT_DISPLAY;
 
+import android.content.Context;
+import android.util.ArrayMap;
+import android.view.Surface;
+import android.view.WindowManager;
+import android.view.WindowMetrics;
+
+import com.android.launcher3.util.WindowBounds;
+import com.android.launcher3.util.window.CachedDisplayInfo;
 import com.android.launcher3.util.window.WindowManagerProxy;
 
+import java.util.Set;
+
 /**
  * Extension of {@link WindowManagerProxy} with some assumption for the default system Launcher
  */
@@ -30,17 +39,23 @@
     }
 
     @Override
-    protected String getDisplayId(Display display) {
-        return display.getUniqueId();
+    public int getRotation(Context displayInfoContext) {
+        return displayInfoContext.getResources().getConfiguration().windowConfiguration
+                .getRotation();
     }
 
     @Override
-    public boolean isInternalDisplay(Display display) {
-        return display.getType() == Display.TYPE_INTERNAL;
-    }
-
-    @Override
-    public int getRotation(Context context) {
-        return context.getResources().getConfiguration().windowConfiguration.getRotation();
+    public ArrayMap<CachedDisplayInfo, WindowBounds[]> estimateInternalDisplayBounds(
+            Context displayInfoContext) {
+        ArrayMap<CachedDisplayInfo, WindowBounds[]> result = new ArrayMap<>();
+        WindowManager windowManager = displayInfoContext.getSystemService(WindowManager.class);
+        Set<WindowMetrics> possibleMaximumWindowMetrics =
+                windowManager.getPossibleMaximumWindowMetrics(DEFAULT_DISPLAY);
+        for (WindowMetrics windowMetrics : possibleMaximumWindowMetrics) {
+            CachedDisplayInfo info = getDisplayInfo(windowMetrics, Surface.ROTATION_0);
+            WindowBounds[] bounds = estimateWindowBounds(displayInfoContext, info);
+            result.put(info, bounds);
+        }
+        return result;
     }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
index 9e5d958..1c15e1e 100644
--- a/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
+++ b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
@@ -35,7 +35,6 @@
 import android.util.ArrayMap;
 import android.util.DisplayMetrics;
 import android.util.Size;
-import android.view.Display;
 import android.view.MotionEvent;
 import android.view.Surface;
 
@@ -290,15 +289,17 @@
     private DisplayController.Info createDisplayInfo(Size screenSize, int rotation) {
         Point displaySize = new Point(screenSize.getWidth(), screenSize.getHeight());
         RotationUtils.rotateSize(displaySize, rotation);
-        CachedDisplayInfo cdi = new CachedDisplayInfo(displaySize, rotation);
-        WindowBounds wm = new WindowBounds(
+        CachedDisplayInfo cachedDisplayInfo = new CachedDisplayInfo(displaySize, rotation);
+        WindowBounds windowBounds = new WindowBounds(
                 new Rect(0, 0, displaySize.x, displaySize.y),
                 new Rect());
         WindowManagerProxy wmProxy = mock(WindowManagerProxy.class);
-        doReturn(cdi).when(wmProxy).getDisplayInfo(any(), any());
-        doReturn(wm).when(wmProxy).getRealBounds(any(), any(), any());
+        doReturn(cachedDisplayInfo).when(wmProxy).getDisplayInfo(any());
+        doReturn(windowBounds).when(wmProxy).getRealBounds(any(), any());
+        ArrayMap<CachedDisplayInfo, WindowBounds[]> internalDisplayBounds = new ArrayMap<>();
+        doReturn(internalDisplayBounds).when(wmProxy).estimateInternalDisplayBounds(any());
         return new DisplayController.Info(
-                getApplicationContext(), mock(Display.class), wmProxy, new ArrayMap<>());
+                getApplicationContext(), wmProxy, new ArrayMap<>());
     }
 
     private float generateTouchRegionHeight(Size screenSize, int rotation) {
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
index 7d414f4..d43aafa 100644
--- a/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
+++ b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
@@ -23,8 +23,6 @@
 import android.graphics.Rect;
 import android.graphics.RectF;
 import android.util.ArrayMap;
-import android.util.Pair;
-import android.view.Display;
 import android.view.Surface;
 import android.view.SurfaceControl;
 
@@ -148,7 +146,7 @@
                 int rotation = mDisplaySize.x > mDisplaySize.y
                         ? Surface.ROTATION_90 : Surface.ROTATION_0;
                 CachedDisplayInfo cdi =
-                        new CachedDisplayInfo("test-display", mDisplaySize, rotation , new Rect());
+                        new CachedDisplayInfo(mDisplaySize, rotation, new Rect());
                 WindowBounds wm = new WindowBounds(
                         new Rect(0, 0, mDisplaySize.x, mDisplaySize.y),
                         mDisplayInsets);
@@ -164,15 +162,15 @@
                 }
 
                 WindowManagerProxy wmProxy = mock(WindowManagerProxy.class);
-                doReturn(cdi).when(wmProxy).getDisplayInfo(any(), any());
-                doReturn(wm).when(wmProxy).getRealBounds(any(), any(), any());
+                doReturn(cdi).when(wmProxy).getDisplayInfo(any());
+                doReturn(wm).when(wmProxy).getRealBounds(any(), any());
 
-                ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> perDisplayBoundsCache =
+                ArrayMap<CachedDisplayInfo, WindowBounds[]> perDisplayBoundsCache =
                         new ArrayMap<>();
-                perDisplayBoundsCache.put(cdi.id, Pair.create(cdi.normalize(), allBounds));
+                perDisplayBoundsCache.put(cdi.normalize(), allBounds);
 
                 DisplayController.Info mockInfo = new Info(
-                        helper.sandboxContext, mock(Display.class), wmProxy, perDisplayBoundsCache);
+                        helper.sandboxContext, wmProxy, perDisplayBoundsCache);
 
                 DisplayController controller =
                         DisplayController.INSTANCE.get(helper.sandboxContext);
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index db43b44..73ca4a3 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -234,7 +234,8 @@
                         /*allowDisabledGrid=*/false),
                 defaultDeviceType);
 
-        Info myInfo = new Info(context, display);
+        Context displayContext = context.createDisplayContext(display);
+        Info myInfo = new Info(displayContext);
         @DeviceType int deviceType = getDeviceType(myInfo);
         DisplayOption myDisplayOption = invDistWeightedInterpolate(
                 myInfo,
@@ -642,7 +643,7 @@
                             + "\nconfig: " + config
                             + "\ndisplayMetrics: " + res.getDisplayMetrics()
                             + "\nrotation: " + rotation
-                            + "\n" + stringWriter.toString(),
+                            + "\n" + stringWriter,
                     new Exception());
         }
         return getBestMatch(screenWidth, screenHeight, rotation);
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 9ed6bee..1a77674 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -41,7 +41,6 @@
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
-import android.util.Pair;
 import android.view.Display;
 
 import androidx.annotation.AnyThread;
@@ -117,8 +116,9 @@
                 getPackageFilter(TARGET_OVERLAY_PACKAGE, ACTION_OVERLAY_CHANGED));
 
         WindowManagerProxy wmProxy = WindowManagerProxy.INSTANCE.get(context);
-        mInfo = new Info(getDisplayInfoContext(display), display,
-                wmProxy, wmProxy.estimateInternalDisplayBounds(context));
+        Context displayInfoContext = getDisplayInfoContext(display);
+        mInfo = new Info(displayInfoContext, wmProxy,
+                wmProxy.estimateInternalDisplayBounds(displayInfoContext));
     }
 
     /**
@@ -216,18 +216,18 @@
         WindowManagerProxy wmProxy = WindowManagerProxy.INSTANCE.get(mContext);
         Info oldInfo = mInfo;
 
-        Context displayContext = getDisplayInfoContext(display);
-        Info newInfo = new Info(displayContext, display, wmProxy, oldInfo.mPerDisplayBounds);
+        Context displayInfoContext = getDisplayInfoContext(display);
+        Info newInfo = new Info(displayInfoContext, wmProxy, oldInfo.mPerDisplayBounds);
 
         if (newInfo.densityDpi != oldInfo.densityDpi || newInfo.fontScale != oldInfo.fontScale
                 || newInfo.navigationMode != oldInfo.navigationMode) {
             // Cache may not be valid anymore, recreate without cache
-            newInfo = new Info(displayContext, display, wmProxy,
-                    wmProxy.estimateInternalDisplayBounds(displayContext));
+            newInfo = new Info(displayInfoContext, wmProxy,
+                    wmProxy.estimateInternalDisplayBounds(displayInfoContext));
         }
 
         int change = 0;
-        if (!newInfo.displayId.equals(oldInfo.displayId)) {
+        if (!newInfo.normalizedDisplayInfo.equals(oldInfo.normalizedDisplayInfo)) {
             change |= CHANGE_ACTIVE_SCREEN;
         }
         if (newInfo.rotation != oldInfo.rotation) {
@@ -242,35 +242,16 @@
         if (!newInfo.supportedBounds.equals(oldInfo.supportedBounds)
                 || !newInfo.mPerDisplayBounds.equals(oldInfo.mPerDisplayBounds)) {
             change |= CHANGE_SUPPORTED_BOUNDS;
-
-            Point currentS = newInfo.currentSize;
-            Pair<CachedDisplayInfo, WindowBounds[]> cachedBounds =
-                    oldInfo.mPerDisplayBounds.get(newInfo.displayId);
-            Point expectedS = cachedBounds == null ? null : cachedBounds.first.size;
-            if (newInfo.supportedBounds.size() != oldInfo.supportedBounds.size()) {
-                Log.e("b/198965093",
-                        "Inconsistent number of displays"
-                                + "\ndisplay state: " + display.getState()
-                                + "\noldInfo.supportedBounds: " + oldInfo.supportedBounds
-                                + "\nnewInfo.supportedBounds: " + newInfo.supportedBounds);
-            }
-            if (expectedS != null
-                    && (Math.min(currentS.x, currentS.y) != Math.min(expectedS.x, expectedS.y)
-                    || Math.max(currentS.x, currentS.y) != Math.max(expectedS.x, expectedS.y))
-                    && display.getState() == Display.STATE_OFF) {
-                Log.e("b/198965093",
-                        "Display size changed while display is off, ignoring change");
-                return;
-            }
         }
         Log.d("b/198965093", "handleInfoChange"
-                + "\n\tchange: " + change
-                + "\n\tConfiguration diff: " + newInfo.mConfiguration.diff(oldInfo.mConfiguration));
+                + "\n\tchange: 0b" + Integer.toBinaryString(change)
+                + "\n\tConfiguration diff: 0x" + Integer.toHexString(
+                        newInfo.mConfiguration.diff(oldInfo.mConfiguration)));
 
         if (change != 0) {
             mInfo = newInfo;
             final int flags = change;
-            MAIN_EXECUTOR.execute(() -> notifyChange(displayContext, flags));
+            MAIN_EXECUTOR.execute(() -> notifyChange(displayInfoContext, flags));
         }
     }
 
@@ -288,8 +269,8 @@
     public static class Info {
 
         // Cached property
+        public final CachedDisplayInfo normalizedDisplayInfo;
         public final int rotation;
-        public final String displayId;
         public final Point currentSize;
         public final Rect cutout;
 
@@ -302,60 +283,71 @@
 
         public final Set<WindowBounds> supportedBounds = new ArraySet<>();
 
-        private final ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> mPerDisplayBounds =
+        private final ArrayMap<CachedDisplayInfo, WindowBounds[]> mPerDisplayBounds =
                 new ArrayMap<>();
 
         // TODO(b/198965093): Remove after investigation
         private Configuration mConfiguration;
 
-        public Info(Context context, Display display) {
+        public Info(Context displayInfoContext) {
             /* don't need system overrides for external displays */
-            this(context, display, new WindowManagerProxy(), new ArrayMap<>());
+            this(displayInfoContext, new WindowManagerProxy(), new ArrayMap<>());
         }
 
         // Used for testing
-        public Info(Context context, Display display,
+        public Info(Context displayInfoContext,
                 WindowManagerProxy wmProxy,
-                ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> perDisplayBoundsCache) {
-            CachedDisplayInfo displayInfo = wmProxy.getDisplayInfo(context, display);
+                ArrayMap<CachedDisplayInfo, WindowBounds[]> perDisplayBoundsCache) {
+            CachedDisplayInfo displayInfo = wmProxy.getDisplayInfo(displayInfoContext);
+            normalizedDisplayInfo = displayInfo.normalize();
             rotation = displayInfo.rotation;
             currentSize = displayInfo.size;
-            displayId = displayInfo.id;
             cutout = displayInfo.cutout;
 
-            Configuration config = context.getResources().getConfiguration();
+            Configuration config = displayInfoContext.getResources().getConfiguration();
             fontScale = config.fontScale;
             densityDpi = config.densityDpi;
             mScreenSizeDp = new PortraitSize(config.screenHeightDp, config.screenWidthDp);
-            navigationMode = parseNavigationMode(context);
+            navigationMode = parseNavigationMode(displayInfoContext);
 
             // TODO(b/198965093): Remove after investigation
             mConfiguration = config;
 
             mPerDisplayBounds.putAll(perDisplayBoundsCache);
-            Pair<CachedDisplayInfo, WindowBounds[]> cachedValue = mPerDisplayBounds.get(displayId);
+            WindowBounds[] cachedValue = mPerDisplayBounds.get(normalizedDisplayInfo);
 
-            WindowBounds realBounds = wmProxy.getRealBounds(context, display, displayInfo);
+            WindowBounds realBounds = wmProxy.getRealBounds(displayInfoContext, displayInfo);
             if (cachedValue == null) {
-                supportedBounds.add(realBounds);
-            } else {
+                // Unexpected normalizedDisplayInfo is found, recreate the cache
+                Log.e("b/198965093", "Unexpected normalizedDisplayInfo found, invalidating cache");
+                mPerDisplayBounds.clear();
+                mPerDisplayBounds.putAll(wmProxy.estimateInternalDisplayBounds(displayInfoContext));
+                cachedValue = mPerDisplayBounds.get(normalizedDisplayInfo);
+                if (cachedValue == null) {
+                    Log.e("b/198965093", "normalizedDisplayInfo not found in estimation: "
+                            + normalizedDisplayInfo);
+                    supportedBounds.add(realBounds);
+                }
+            }
+
+            if (cachedValue != null) {
                 // Verify that the real bounds are a match
-                WindowBounds expectedBounds = cachedValue.second[displayInfo.rotation];
+                WindowBounds expectedBounds = cachedValue[displayInfo.rotation];
                 if (!realBounds.equals(expectedBounds)) {
                     WindowBounds[] clone = new WindowBounds[4];
-                    System.arraycopy(cachedValue.second, 0, clone, 0, 4);
+                    System.arraycopy(cachedValue, 0, clone, 0, 4);
                     clone[displayInfo.rotation] = realBounds;
-                    cachedValue = Pair.create(displayInfo.normalize(), clone);
-                    mPerDisplayBounds.put(displayId, cachedValue);
+                    mPerDisplayBounds.put(normalizedDisplayInfo, clone);
                 }
             }
             mPerDisplayBounds.values().forEach(
-                    pair -> Collections.addAll(supportedBounds, pair.second));
+                    windowBounds -> Collections.addAll(supportedBounds, windowBounds));
             Log.e("b/198965093", "mConfiguration: " + mConfiguration);
             Log.d("b/198965093", "displayInfo: " + displayInfo);
             Log.d("b/198965093", "realBounds: " + realBounds);
-            mPerDisplayBounds.values().forEach(pair -> Log.d("b/198965093",
-                    "perDisplayBounds - " + pair.first + ": " + Arrays.deepToString(pair.second)));
+            Log.d("b/198965093", "normalizedDisplayInfo: " + normalizedDisplayInfo);
+            mPerDisplayBounds.forEach((key, value) -> Log.d("b/198965093",
+                    "perDisplayBounds - " + key + ": " + Arrays.deepToString(value)));
         }
 
         /**
@@ -383,14 +375,14 @@
     public void dump(PrintWriter pw) {
         Info info = mInfo;
         pw.println("DisplayController.Info:");
-        pw.println("  id=" + info.displayId);
+        pw.println("  normalizedDisplayInfo=" + info.normalizedDisplayInfo);
         pw.println("  rotation=" + info.rotation);
         pw.println("  fontScale=" + info.fontScale);
         pw.println("  densityDpi=" + info.densityDpi);
         pw.println("  navigationMode=" + info.navigationMode.name());
         pw.println("  currentSize=" + info.currentSize);
-        info.mPerDisplayBounds.values().forEach(pair -> pw.println(
-                "  perDisplayBounds - " + pair.first + ": " + Arrays.deepToString(pair.second)));
+        info.mPerDisplayBounds.forEach((key, value) -> pw.println(
+                "  perDisplayBounds - " + key + ": " + Arrays.deepToString(value)));
     }
 
     /**
diff --git a/src/com/android/launcher3/util/window/CachedDisplayInfo.java b/src/com/android/launcher3/util/window/CachedDisplayInfo.java
index 06b9829..23f37aa 100644
--- a/src/com/android/launcher3/util/window/CachedDisplayInfo.java
+++ b/src/com/android/launcher3/util/window/CachedDisplayInfo.java
@@ -30,7 +30,6 @@
  */
 public class CachedDisplayInfo {
 
-    public final String id;
     public final Point size;
     public final int rotation;
     public final Rect cutout;
@@ -40,11 +39,10 @@
     }
 
     public CachedDisplayInfo(Point size, int rotation) {
-        this("", size, rotation, new Rect());
+        this(size, rotation, new Rect());
     }
 
-    public CachedDisplayInfo(String id, Point size, int rotation, Rect cutout) {
-        this.id = id;
+    public CachedDisplayInfo(Point size, int rotation, Rect cutout) {
         this.size = size;
         this.rotation = rotation;
         this.cutout = cutout;
@@ -62,16 +60,15 @@
 
         Rect newCutout = new Rect(cutout);
         rotateRect(newCutout, deltaRotation(rotation, Surface.ROTATION_0));
-        return new CachedDisplayInfo(id, newSize, Surface.ROTATION_0, newCutout);
+        return new CachedDisplayInfo(newSize, Surface.ROTATION_0, newCutout);
     }
 
     @Override
     public String toString() {
         return "CachedDisplayInfo{"
-                + "id='" + id + '\''
-                + ", size=" + size
-                + ", rotation=" + rotation
+                + "size=" + size
                 + ", cutout=" + cutout
+                + ", rotation=" + rotation
                 + '}';
     }
 
@@ -80,13 +77,13 @@
         if (this == o) return true;
         if (!(o instanceof CachedDisplayInfo)) return false;
         CachedDisplayInfo that = (CachedDisplayInfo) o;
-        return rotation == that.rotation && Objects.equals(id, that.id)
-                && Objects.equals(size, that.size) && Objects.equals(cutout,
-                that.cutout);
+        return rotation == that.rotation
+                && Objects.equals(size, that.size)
+                && Objects.equals(cutout, that.cutout);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(id, size, rotation, cutout);
+        return Objects.hash(size, rotation, cutout);
     }
 }
diff --git a/src/com/android/launcher3/util/window/WindowManagerProxy.java b/src/com/android/launcher3/util/window/WindowManagerProxy.java
index 92f718e..d5a065a 100644
--- a/src/com/android/launcher3/util/window/WindowManagerProxy.java
+++ b/src/com/android/launcher3/util/window/WindowManagerProxy.java
@@ -16,7 +16,6 @@
 package com.android.launcher3.util.window;
 
 import static android.view.Display.DEFAULT_DISPLAY;
-import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
 
 import static com.android.launcher3.ResourceUtils.INVALID_RESOURCE_HANDLE;
 import static com.android.launcher3.ResourceUtils.NAVBAR_HEIGHT;
@@ -41,7 +40,6 @@
 import android.hardware.display.DisplayManager;
 import android.os.Build;
 import android.util.ArrayMap;
-import android.util.Pair;
 import android.view.Display;
 import android.view.DisplayCutout;
 import android.view.Surface;
@@ -88,20 +86,12 @@
      * Returns a map of normalized info of internal displays to estimated window bounds
      * for that display
      */
-    public ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> estimateInternalDisplayBounds(
-            Context context) {
-        Display[] displays = context.getSystemService(DisplayManager.class).getDisplays();
-        ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> result = new ArrayMap<>();
-        for (Display display : displays) {
-            if (isInternalDisplay(display)) {
-                Context displayContext = Utilities.ATLEAST_S
-                        ? context.createWindowContext(display, TYPE_APPLICATION, null)
-                        : context.createDisplayContext(display);
-                CachedDisplayInfo info = getDisplayInfo(displayContext, display).normalize();
-                WindowBounds[] bounds = estimateWindowBounds(context, info);
-                result.put(info.id, Pair.create(info, bounds));
-            }
-        }
+    public ArrayMap<CachedDisplayInfo, WindowBounds[]> estimateInternalDisplayBounds(
+            Context displayInfoContext) {
+        CachedDisplayInfo info = getDisplayInfo(displayInfoContext).normalize();
+        WindowBounds[] bounds = estimateWindowBounds(displayInfoContext, info);
+        ArrayMap<CachedDisplayInfo, WindowBounds[]> result = new ArrayMap<>();
+        result.put(info, bounds);
         return result;
     }
 
@@ -109,12 +99,11 @@
      * Returns the real bounds for the provided display after applying any insets normalization
      */
     @TargetApi(Build.VERSION_CODES.R)
-    public WindowBounds getRealBounds(Context windowContext,
-            Display display, CachedDisplayInfo info) {
+    public WindowBounds getRealBounds(Context displayInfoContext, CachedDisplayInfo info) {
         if (!Utilities.ATLEAST_R) {
             Point smallestSize = new Point();
             Point largestSize = new Point();
-            display.getCurrentSizeRange(smallestSize, largestSize);
+            getDisplay(displayInfoContext).getCurrentSizeRange(smallestSize, largestSize);
 
             if (info.size.y > info.size.x) {
                 // Portrait
@@ -122,17 +111,16 @@
                         info.rotation);
             } else {
                 // Landscape
-                new WindowBounds(info.size.x, info.size.y, largestSize.x, smallestSize.y,
+                return new WindowBounds(info.size.x, info.size.y, largestSize.x, smallestSize.y,
                         info.rotation);
             }
         }
 
-        WindowMetrics wm = windowContext.getSystemService(WindowManager.class)
+        WindowMetrics windowMetrics = displayInfoContext.getSystemService(WindowManager.class)
                 .getMaximumWindowMetrics();
-
         Rect insets = new Rect();
-        normalizeWindowInsets(windowContext, wm.getWindowInsets(), insets);
-        return new WindowBounds(wm.getBounds(), insets, info.rotation);
+        normalizeWindowInsets(displayInfoContext, windowMetrics.getWindowInsets(), insets);
+        return new WindowBounds(windowMetrics.getBounds(), insets, info.rotation);
     }
 
     /**
@@ -169,12 +157,9 @@
         insetsBuilder.setInsetsIgnoringVisibility(WindowInsets.Type.navigationBars(), newNavInsets);
 
         Insets statusBarInsets = oldInsets.getInsets(WindowInsets.Type.statusBars());
-
-
         int statusBarHeight = getDimenByName(systemRes,
                 (isPortrait) ? STATUS_BAR_HEIGHT_PORTRAIT : STATUS_BAR_HEIGHT_LANDSCAPE,
                 STATUS_BAR_HEIGHT);
-
         Insets newStatusBarInsets = Insets.of(
                 statusBarInsets.left,
                 Math.max(statusBarInsets.top, statusBarHeight),
@@ -202,21 +187,14 @@
     }
 
     /**
-     * Returns true if the display is an internal displays
-     */
-    protected boolean isInternalDisplay(Display display) {
-        return display.getDisplayId() == Display.DEFAULT_DISPLAY;
-    }
-
-    /**
      * Returns a list of possible WindowBounds for the display keyed on the 4 surface rotations
      */
-    public WindowBounds[] estimateWindowBounds(Context context, CachedDisplayInfo display) {
+    protected WindowBounds[] estimateWindowBounds(Context context, CachedDisplayInfo displayInfo) {
         int densityDpi = context.getResources().getConfiguration().densityDpi;
-        int rotation = display.rotation;
-        Rect safeCutout = display.cutout;
+        int rotation = displayInfo.rotation;
+        Rect safeCutout = displayInfo.cutout;
 
-        int minSize = Math.min(display.size.x, display.size.y);
+        int minSize = Math.min(displayInfo.size.x, displayInfo.size.y);
         int swDp = (int) dpiFromPx(minSize, densityDpi);
 
         Resources systemRes;
@@ -255,7 +233,7 @@
         Point tempSize = new Point();
         for (int i = 0; i < 4; i++) {
             int rotationChange = deltaRotation(rotation, i);
-            tempSize.set(display.size.x, display.size.y);
+            tempSize.set(displayInfo.size.x, displayInfo.size.y);
             rotateSize(tempSize, rotationChange);
             Rect bounds = new Rect(0, 0, tempSize.x, tempSize.y);
 
@@ -311,48 +289,58 @@
      * Returns a CachedDisplayInfo initialized for the current display
      */
     @TargetApi(Build.VERSION_CODES.S)
-    public CachedDisplayInfo getDisplayInfo(Context displayContext, Display display) {
-        int rotation = getRotation(displayContext);
-        Rect cutoutRect = new Rect();
-        Point size = new Point();
+    public CachedDisplayInfo getDisplayInfo(Context displayInfoContext) {
+        int rotation = getRotation(displayInfoContext);
         if (Utilities.ATLEAST_S) {
-            WindowMetrics wm = displayContext.getSystemService(WindowManager.class)
+            WindowMetrics windowMetrics = displayInfoContext.getSystemService(WindowManager.class)
                     .getMaximumWindowMetrics();
-            DisplayCutout cutout = wm.getWindowInsets().getDisplayCutout();
-            if (cutout != null) {
-                cutoutRect.set(cutout.getSafeInsetLeft(), cutout.getSafeInsetTop(),
-                        cutout.getSafeInsetRight(), cutout.getSafeInsetBottom());
-            }
-
-            size.set(wm.getBounds().right, wm.getBounds().bottom);
+            return getDisplayInfo(windowMetrics, rotation);
         } else {
+            Point size = new Point();
+            Display display = getDisplay(displayInfoContext);
             display.getRealSize(size);
+            Rect cutoutRect = new Rect();
+            return new CachedDisplayInfo(size, rotation, cutoutRect);
         }
-        return new CachedDisplayInfo(getDisplayId(display), size, rotation, cutoutRect);
     }
 
     /**
-     * Returns a unique ID representing the display
+     * Returns a CachedDisplayInfo initialized for the current display
      */
-    protected String getDisplayId(Display display) {
-        return Integer.toString(display.getDisplayId());
+    @TargetApi(Build.VERSION_CODES.S)
+    protected CachedDisplayInfo getDisplayInfo(WindowMetrics windowMetrics, int rotation) {
+        Point size = new Point(windowMetrics.getBounds().right, windowMetrics.getBounds().bottom);
+        Rect cutoutRect = new Rect();
+        DisplayCutout cutout = windowMetrics.getWindowInsets().getDisplayCutout();
+        if (cutout != null) {
+            cutoutRect.set(cutout.getSafeInsetLeft(), cutout.getSafeInsetTop(),
+                    cutout.getSafeInsetRight(), cutout.getSafeInsetBottom());
+        }
+        return new CachedDisplayInfo(size, rotation, cutoutRect);
     }
 
     /**
-     * Returns rotation of the display associated with the context.
+     * Returns rotation of the display associated with the context, or rotation of DEFAULT_DISPLAY
+     * if the context isn't associated with a display.
      */
-    public int getRotation(Context context) {
-        Display d = null;
+    public int getRotation(Context displayInfoContext) {
+        return getDisplay(displayInfoContext).getRotation();
+    }
+
+    /**
+     *
+     * Returns the display associated with the context, or DEFAULT_DISPLAY if the context isn't
+     * associated with a display.
+     */
+    protected Display getDisplay(Context displayInfoContext) {
         if (Utilities.ATLEAST_R) {
             try {
-                d = context.getDisplay();
+                return displayInfoContext.getDisplay();
             } catch (UnsupportedOperationException e) {
                 // Ignore
             }
         }
-        if (d == null) {
-            d = context.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY);
-        }
-        return d.getRotation();
+        return displayInfoContext.getSystemService(DisplayManager.class).getDisplay(
+                DEFAULT_DISPLAY);
     }
 }
diff --git a/tests/src/com/android/launcher3/deviceemulator/DisplayEmulator.java b/tests/src/com/android/launcher3/deviceemulator/DisplayEmulator.java
index 31468c5..e2ed65f 100644
--- a/tests/src/com/android/launcher3/deviceemulator/DisplayEmulator.java
+++ b/tests/src/com/android/launcher3/deviceemulator/DisplayEmulator.java
@@ -47,8 +47,7 @@
      * By changing the WindowManagerProxy we can override the window insets information
      **/
     private IWindowManager changeWindowManagerInstance(DeviceEmulationData deviceData) {
-        WindowManagerProxy.INSTANCE.initializeForTesting(
-                new TestWindowManagerProxy(mContext, deviceData));
+        WindowManagerProxy.INSTANCE.initializeForTesting(new TestWindowManagerProxy(deviceData));
         return WindowManagerGlobal.getWindowManagerService();
     }
 
@@ -57,8 +56,7 @@
         WindowManagerProxy original = WindowManagerProxy.INSTANCE.get(mContext);
         // Set up emulation
         final int userId = UserHandle.myUserId();
-        WindowManagerProxy.INSTANCE.initializeForTesting(
-                new TestWindowManagerProxy(mContext, device));
+        WindowManagerProxy.INSTANCE.initializeForTesting(new TestWindowManagerProxy(device));
         IWindowManager wm = changeWindowManagerInstance(device);
         // Change density twice to force display controller to reset its state
         wm.setForcedDisplayDensityForUser(Display.DEFAULT_DISPLAY, device.density / 2, userId);
diff --git a/tests/src/com/android/launcher3/deviceemulator/TestWindowManagerProxy.java b/tests/src/com/android/launcher3/deviceemulator/TestWindowManagerProxy.java
index cbea688..2d6bbcc 100644
--- a/tests/src/com/android/launcher3/deviceemulator/TestWindowManagerProxy.java
+++ b/tests/src/com/android/launcher3/deviceemulator/TestWindowManagerProxy.java
@@ -19,7 +19,6 @@
 import android.content.res.Resources;
 import android.graphics.Point;
 import android.graphics.Rect;
-import android.view.Display;
 import android.view.WindowInsets;
 
 import com.android.launcher3.deviceemulator.models.DeviceEmulationData;
@@ -32,17 +31,12 @@
 
     private final DeviceEmulationData mDevice;
 
-    public TestWindowManagerProxy(Context context, DeviceEmulationData device) {
+    public TestWindowManagerProxy(DeviceEmulationData device) {
         super(true);
         mDevice = device;
     }
 
     @Override
-    public boolean isInternalDisplay(Display display) {
-        return display.getDisplayId() == Display.DEFAULT_DISPLAY;
-    }
-
-    @Override
     protected int getDimenByName(Resources res, String resName) {
         Integer mock = mDevice.resourceOverrides.get(resName);
         return mock != null ? mock : super.getDimenByName(res, resName);
@@ -54,27 +48,25 @@
     }
 
     @Override
-    public CachedDisplayInfo getDisplayInfo(Context context, Display display) {
-        int rotation = display.getRotation();
+    public CachedDisplayInfo getDisplayInfo(Context displayInfoContext) {
+        int rotation = getRotation(displayInfoContext);
         Point size = new Point(mDevice.width, mDevice.height);
         RotationUtils.rotateSize(size, rotation);
         Rect cutout = new Rect(mDevice.cutout);
         RotationUtils.rotateRect(cutout, rotation);
-        return new CachedDisplayInfo(getDisplayId(display), size, rotation, cutout);
+        return new CachedDisplayInfo(size, rotation, cutout);
     }
 
     @Override
-    public WindowBounds getRealBounds(Context windowContext, Display display,
-            CachedDisplayInfo info) {
-        return estimateInternalDisplayBounds(windowContext)
-                .get(getDisplayId(display)).second[display.getRotation()];
+    public WindowBounds getRealBounds(Context displayInfoContext, CachedDisplayInfo info) {
+        return estimateInternalDisplayBounds(displayInfoContext).get(
+                getDisplayInfo(displayInfoContext))[getDisplay(displayInfoContext).getRotation()];
     }
 
     @Override
     public WindowInsets normalizeWindowInsets(Context context, WindowInsets oldInsets,
             Rect outInsets) {
-        outInsets.set(getRealBounds(context, context.getDisplay(),
-                getDisplayInfo(context, context.getDisplay())).insets);
+        outInsets.set(getRealBounds(context, getDisplayInfo(context)).insets);
         return oldInsets;
     }
 }