diff --git a/services/core/java/com/android/server/policy/WindowManagerPolicy.java b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
index 977b31e6d683c..1553966c9eca9 100644
--- a/services/core/java/com/android/server/policy/WindowManagerPolicy.java
+++ b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
@@ -496,6 +496,32 @@ public interface WindowManagerPolicy extends WindowManagerPolicyConstants {
      * @return int An arbitrary integer used to order windows, with lower numbers below higher ones.
      */
     default int getWindowLayerFromTypeLw(int type, boolean canAddInternalSystemWindow) {
+        return getWindowLayerFromTypeLw(type, canAddInternalSystemWindow,
+                false /* roundedCornerOverlay */);
+    }
+
+    /**
+     * Returns the layer assignment for the window type. Allows you to control how different
+     * kinds of windows are ordered on-screen.
+     *
+     * @param type The type of window being assigned.
+     * @param canAddInternalSystemWindow If the owner window associated with the type we are
+     *        evaluating can add internal system windows. I.e they have
+     *        {@link Manifest.permission#INTERNAL_SYSTEM_WINDOW}. If true, alert window
+     *        types {@link android.view.WindowManager.LayoutParams#isSystemAlertWindowType(int)}
+     *        can be assigned layers greater than the layer for
+     *        {@link android.view.WindowManager.LayoutParams#TYPE_APPLICATION_OVERLAY} Else, their
+     *        layers would be lesser.
+     * @param roundedCornerOverlay {#code true} to indicate that the owner window is rounded corner
+     *                             overlay.
+     * @return int An arbitrary integer used to order windows, with lower numbers below higher ones.
+     */
+    default int getWindowLayerFromTypeLw(int type, boolean canAddInternalSystemWindow,
+            boolean roundedCornerOverlay) {
+        // Always put the rounded corner layer to the top most.
+        if (roundedCornerOverlay && canAddInternalSystemWindow) {
+            return getMaxWindowLayer();
+        }
         if (type >= FIRST_APPLICATION_WINDOW && type <= LAST_APPLICATION_WINDOW) {
             return APPLICATION_LAYER;
         }
@@ -598,8 +624,17 @@ public interface WindowManagerPolicy extends WindowManagerPolicyConstants {
         }
     }
 
+    // TODO(b/155340867): consider to remove the logic after using pure Surface for rounded corner
+    //  overlay.
+    /**
+     * Returns the max window layer.
+     * <p>Note that the max window layer should be higher that the maximum value which reported
+     * by {@link #getWindowLayerFromTypeLw(int, boolean)} to contain rounded corner overlay.</p>
+     *
+     * @see WindowManager.LayoutParams#PRIVATE_FLAG_IS_ROUNDED_CORNERS_OVERLAY
+     */
     default int getMaxWindowLayer() {
-        return 35;
+        return 36;
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/DisplayAreaPolicyBuilder.java b/services/core/java/com/android/server/wm/DisplayAreaPolicyBuilder.java
index c8fadf62eb2f9..f505daa9b6283 100644
--- a/services/core/java/com/android/server/wm/DisplayAreaPolicyBuilder.java
+++ b/services/core/java/com/android/server/wm/DisplayAreaPolicyBuilder.java
@@ -365,7 +365,7 @@ class DisplayAreaPolicyBuilder {
          */
         private void build(@Nullable List<HierarchyBuilder> displayAreaGroupHierarchyBuilders) {
             final WindowManagerPolicy policy = mRoot.mWmService.mPolicy;
-            final int maxWindowLayerCount = policy.getMaxWindowLayer();
+            final int maxWindowLayerCount = policy.getMaxWindowLayer() + 1;
             final DisplayArea.Tokens[] displayAreaForLayer =
                     new DisplayArea.Tokens[maxWindowLayerCount];
             final Map<Feature, List<DisplayArea<WindowContainer>>> featureAreas =
@@ -560,6 +560,7 @@ class DisplayAreaPolicyBuilder {
             private final int mId;
             private final boolean[] mLayers;
             private NewDisplayAreaSupplier mNewDisplayAreaSupplier = DisplayArea::new;
+            private boolean mExcludeRoundedCorner = true;
 
             /**
              * Builds a new feature that applies to a set of window types as specified by the
@@ -578,7 +579,7 @@ class DisplayAreaPolicyBuilder {
                 mPolicy = policy;
                 mName = name;
                 mId = id;
-                mLayers = new boolean[mPolicy.getMaxWindowLayer()];
+                mLayers = new boolean[mPolicy.getMaxWindowLayer() + 1];
             }
 
             /**
@@ -633,7 +634,18 @@ class DisplayAreaPolicyBuilder {
                 return this;
             }
 
+            // TODO(b/155340867): consider to remove the logic after using pure Surface for rounded
+            // corner overlay.
+            Builder setExcludeRoundedCornerOverlay(boolean excludeRoundedCorner) {
+                mExcludeRoundedCorner = excludeRoundedCorner;
+                return this;
+            }
+
             Feature build() {
+                if (mExcludeRoundedCorner) {
+                    // Always put the rounded corner layer to the top most layer.
+                    mLayers[mPolicy.getMaxWindowLayer()] = false;
+                }
                 return new Feature(mName, mId, mLayers.clone(), mNewDisplayAreaSupplier);
             }
 
diff --git a/services/core/java/com/android/server/wm/RootDisplayArea.java b/services/core/java/com/android/server/wm/RootDisplayArea.java
index 393055eab111d..505af05ed642f 100644
--- a/services/core/java/com/android/server/wm/RootDisplayArea.java
+++ b/services/core/java/com/android/server/wm/RootDisplayArea.java
@@ -115,15 +115,11 @@ class RootDisplayArea extends DisplayArea<DisplayArea> {
     @Nullable
     DisplayArea.Tokens findAreaForToken(int windowType, boolean ownerCanManageAppTokens,
             boolean roundedCornerOverlay) {
-        // TODO(b/159767464): cover TYPE_INPUT_METHOD(_DIALOG) case here. mAreaForLayer doesn't
-        // contain IME container.
         int windowLayerFromType = mWmService.mPolicy.getWindowLayerFromTypeLw(windowType,
-                ownerCanManageAppTokens);
+                ownerCanManageAppTokens, roundedCornerOverlay);
         if (windowLayerFromType == APPLICATION_LAYER) {
             throw new IllegalArgumentException(
                     "There shouldn't be WindowToken on APPLICATION_LAYER");
-        } else if (roundedCornerOverlay) {
-            windowLayerFromType = mAreaForLayer.length - 1;
         }
         return mAreaForLayer[windowLayerFromType];
     }
diff --git a/services/core/java/com/android/server/wm/WindowToken.java b/services/core/java/com/android/server/wm/WindowToken.java
index 34875ed70889e..9e0fee3927c3c 100644
--- a/services/core/java/com/android/server/wm/WindowToken.java
+++ b/services/core/java/com/android/server/wm/WindowToken.java
@@ -827,15 +827,15 @@ class WindowToken extends WindowContainer<WindowState> {
      * system bars, or in other words extend outside of the "Decor Frame"
      */
     boolean canLayerAboveSystemBars() {
-        int layer = mWmService.mPolicy.getWindowLayerFromTypeLw(windowType,
-                mOwnerCanManageAppTokens);
+        int layer = getWindowLayerFromType();
         int navLayer = mWmService.mPolicy.getWindowLayerFromTypeLw(TYPE_NAVIGATION_BAR,
                 mOwnerCanManageAppTokens);
         return mOwnerCanManageAppTokens && (layer > navLayer);
     }
 
     int getWindowLayerFromType() {
-        return mWmService.mPolicy.getWindowLayerFromTypeLw(windowType, mOwnerCanManageAppTokens);
+        return mWmService.mPolicy.getWindowLayerFromTypeLw(windowType, mOwnerCanManageAppTokens,
+                mRoundedCornerOverlay);
     }
 
     int getOwnerUid() {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayAreaPolicyBuilderTest.java b/services/tests/wmtests/src/com/android/server/wm/DisplayAreaPolicyBuilderTest.java
index 8e78dce39ac7f..fb585ed124410 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayAreaPolicyBuilderTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayAreaPolicyBuilderTest.java
@@ -22,6 +22,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG;
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
+import static android.view.WindowManager.LayoutParams.TYPE_POINTER;
 import static android.view.WindowManager.LayoutParams.TYPE_PRESENTATION;
 import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
@@ -669,6 +670,15 @@ public class DisplayAreaPolicyBuilderTest {
         assertThat(token2.isDescendantOf(mGroupRoot2)).isTrue();
     }
 
+    @Test
+    public void testFeatureNotThrowArrayIndexOutOfBoundsException() {
+        final Feature feature1 = new Feature.Builder(mWms.mPolicy, "feature1",
+                FEATURE_VENDOR_FIRST + 5)
+                .all()
+                .except(TYPE_POINTER)
+                .build();
+    }
+
     private static Resources resourcesWithProvider(String provider) {
         Resources mock = mock(Resources.class);
         when(mock.getString(