diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java
index b8e83635de6ea..d6e1a82a68ff6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java
@@ -25,7 +25,6 @@ import static android.provider.Settings.Global.DEVELOPMENT_FORCE_DESKTOP_MODE_ON
 import static android.util.RotationUtils.rotateBounds;
 import static android.util.RotationUtils.rotateInsets;
 import static android.view.Display.FLAG_SHOULD_SHOW_SYSTEM_DECORATIONS;
-import static android.view.InsetsState.ITYPE_EXTRA_NAVIGATION_BAR;
 import static android.view.Surface.ROTATION_0;
 import static android.view.Surface.ROTATION_270;
 import static android.view.Surface.ROTATION_90;
@@ -46,9 +45,9 @@ import android.view.Display;
 import android.view.DisplayCutout;
 import android.view.DisplayInfo;
 import android.view.Gravity;
-import android.view.InsetsSource;
 import android.view.InsetsState;
 import android.view.Surface;
+import android.view.WindowInsets;
 
 import androidx.annotation.VisibleForTesting;
 
@@ -372,23 +371,20 @@ public class DisplayLayout {
 
         // Only navigation bar
         if (hasNavigationBar) {
-            final InsetsSource extraNavBar = insetsState.peekSource(ITYPE_EXTRA_NAVIGATION_BAR);
-            final boolean hasExtraNav = extraNavBar != null && extraNavBar.isVisible();
+            final Insets insets = insetsState.calculateInsets(
+                    insetsState.getDisplayFrame(),
+                    WindowInsets.Type.navigationBars(),
+                    false /* ignoreVisibility */);
+            outInsets.set(insets.left, insets.top, insets.right, insets.bottom);
             int position = navigationBarPosition(res, displayWidth, displayHeight, displayRotation);
             int navBarSize =
                     getNavigationBarSize(res, position, displayWidth > displayHeight, uiMode);
             if (position == NAV_BAR_BOTTOM) {
-                outInsets.bottom = hasExtraNav
-                        ? Math.max(navBarSize, extraNavBar.getFrame().height())
-                        : navBarSize;
+                outInsets.bottom = Math.max(outInsets.bottom , navBarSize);
             } else if (position == NAV_BAR_RIGHT) {
-                outInsets.right = hasExtraNav
-                        ? Math.max(navBarSize, extraNavBar.getFrame().width())
-                        : navBarSize;
+                outInsets.right = Math.max(outInsets.right , navBarSize);
             } else if (position == NAV_BAR_LEFT) {
-                outInsets.left = hasExtraNav
-                        ? Math.max(navBarSize, extraNavBar.getFrame().width())
-                        : navBarSize;
+                outInsets.left = Math.max(outInsets.left , navBarSize);
             }
         }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
index 05fd889521f27..e8784d77be1bf 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
@@ -18,7 +18,6 @@ package com.android.wm.shell.compatui;
 
 import static android.app.TaskInfo.CAMERA_COMPAT_CONTROL_HIDDEN;
 import static android.app.TaskInfo.CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED;
-import static android.view.InsetsState.ITYPE_EXTRA_NAVIGATION_BAR;
 import static android.view.WindowInsets.Type.navigationBars;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
@@ -333,7 +332,8 @@ public class CompatUIControllerTest extends ShellTestCase {
         mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
                 /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener);
         InsetsState insetsState = new InsetsState();
-        InsetsSource insetsSource = new InsetsSource(ITYPE_EXTRA_NAVIGATION_BAR, navigationBars());
+        InsetsSource insetsSource = new InsetsSource(
+                InsetsSource.createId(null, 0, navigationBars()), navigationBars());
         insetsSource.setFrame(0, 0, 1000, 1000);
         insetsState.addSource(insetsSource);
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
index f3f615d97a568..4de5298855650 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
@@ -20,7 +20,6 @@ import static android.app.TaskInfo.CAMERA_COMPAT_CONTROL_DISMISSED;
 import static android.app.TaskInfo.CAMERA_COMPAT_CONTROL_HIDDEN;
 import static android.app.TaskInfo.CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED;
 import static android.app.TaskInfo.CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
-import static android.view.InsetsState.ITYPE_EXTRA_NAVIGATION_BAR;
 import static android.view.WindowInsets.Type.navigationBars;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
@@ -338,8 +337,10 @@ public class CompatUIWindowManagerTest extends ShellTestCase {
         // Update if the insets change on the existing display layout
         clearInvocations(mWindowManager);
         InsetsState insetsState = new InsetsState();
-        InsetsSource insetsSource = new InsetsSource(ITYPE_EXTRA_NAVIGATION_BAR, navigationBars());
-        insetsSource.setFrame(0, 0, 1000, 1000);
+        insetsState.setDisplayFrame(new Rect(0, 0, 1000, 2000));
+        InsetsSource insetsSource = new InsetsSource(
+                InsetsSource.createId(null, 0, navigationBars()), navigationBars());
+        insetsSource.setFrame(0, 1800, 1000, 2000);
         insetsState.addSource(insetsSource);
         displayLayout.setInsets(mContext.getResources(), insetsState);
         mWindowManager.updateDisplayLayout(displayLayout);