diff --git a/res/color-night-v31/popup_color_first.xml b/res/color-night-v31/popup_shade_first.xml
similarity index 100%
rename from res/color-night-v31/popup_color_first.xml
rename to res/color-night-v31/popup_shade_first.xml
diff --git a/res/color-night-v31/popup_color_second.xml b/res/color-night-v31/popup_shade_second.xml
similarity index 100%
rename from res/color-night-v31/popup_color_second.xml
rename to res/color-night-v31/popup_shade_second.xml
diff --git a/res/color-night-v31/popup_color_third.xml b/res/color-night-v31/popup_shade_third.xml
similarity index 100%
rename from res/color-night-v31/popup_color_third.xml
rename to res/color-night-v31/popup_shade_third.xml
diff --git a/res/color-v31/popup_color_first.xml b/res/color-v31/popup_shade_first.xml
similarity index 100%
rename from res/color-v31/popup_color_first.xml
rename to res/color-v31/popup_shade_first.xml
diff --git a/res/color-v31/popup_color_second.xml b/res/color-v31/popup_shade_second.xml
similarity index 100%
rename from res/color-v31/popup_color_second.xml
rename to res/color-v31/popup_shade_second.xml
diff --git a/res/color-v31/popup_color_third.xml b/res/color-v31/popup_shade_third.xml
similarity index 100%
rename from res/color-v31/popup_color_third.xml
rename to res/color-v31/popup_shade_third.xml
diff --git a/res/color/popup_color_second.xml b/res/color/popup_color_second.xml
deleted file mode 100644
index d9999a2..0000000
--- a/res/color/popup_color_second.xml
+++ /dev/null
@@ -1,18 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!-- Copyright (C) 2021 The Android Open Source Project
-
-     Licensed under the Apache License, Version 2.0 (the "License");
-     you may not use this file except in compliance with the License.
-     You may obtain a copy of the License at
-
-          http://www.apache.org/licenses/LICENSE-2.0
-
-     Unless required by applicable law or agreed to in writing, software
-     distributed under the License is distributed on an "AS IS" BASIS,
-     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-     See the License for the specific language governing permissions and
-     limitations under the License.
--->
-<selector xmlns:android="http://schemas.android.com/apk/res/android" >
-    <item android:color="?attr/popupColorPrimary" />
-</selector>
diff --git a/res/color/popup_color_first.xml b/res/color/popup_shade_first.xml
similarity index 93%
copy from res/color/popup_color_first.xml
copy to res/color/popup_shade_first.xml
index d9999a2..151190b 100644
--- a/res/color/popup_color_first.xml
+++ b/res/color/popup_shade_first.xml
@@ -14,5 +14,5 @@
      limitations under the License.
 -->
 <selector xmlns:android="http://schemas.android.com/apk/res/android" >
-    <item android:color="?attr/popupColorPrimary" />
+    <item android:color="?attr/popupShadeFirst" />
 </selector>
diff --git a/res/color/popup_color_first.xml b/res/color/popup_shade_second.xml
similarity index 93%
rename from res/color/popup_color_first.xml
rename to res/color/popup_shade_second.xml
index d9999a2..8660850 100644
--- a/res/color/popup_color_first.xml
+++ b/res/color/popup_shade_second.xml
@@ -14,5 +14,5 @@
      limitations under the License.
 -->
 <selector xmlns:android="http://schemas.android.com/apk/res/android" >
-    <item android:color="?attr/popupColorPrimary" />
+    <item android:color="?attr/popupShadeSecond" />
 </selector>
diff --git a/res/color/popup_color_third.xml b/res/color/popup_shade_third.xml
similarity index 93%
rename from res/color/popup_color_third.xml
rename to res/color/popup_shade_third.xml
index d7e9e79..9544728 100644
--- a/res/color/popup_color_third.xml
+++ b/res/color/popup_shade_third.xml
@@ -14,5 +14,5 @@
      limitations under the License.
 -->
 <selector xmlns:android="http://schemas.android.com/apk/res/android" >
-    <item android:color="?attr/popupColorPrimary" />
+    <item android:color="?attr/popupShadeThird" />
 </selector>
\ No newline at end of file
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index ed1db30..00cf31c 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -25,6 +25,9 @@
     <attr name="popupColorPrimary" format="color" />
     <attr name="popupColorSecondary" format="color" />
     <attr name="popupColorTertiary" format="color" />
+    <attr name="popupShadeFirst" format="color" />
+    <attr name="popupShadeSecond" format="color" />
+    <attr name="popupShadeThird" format="color" />
     <attr name="isMainColorDark" format="boolean" />
     <attr name="isWorkspaceDarkText" format="boolean" />
     <attr name="workspaceTextColor" format="color" />
diff --git a/res/values/colors.xml b/res/values/colors.xml
index 76e821d..76fd1d7 100644
--- a/res/values/colors.xml
+++ b/res/values/colors.xml
@@ -53,6 +53,13 @@
     <color name="popup_color_secondary_dark">#202124</color>
     <color name="popup_color_tertiary_dark">#757575</color> <!-- Gray 600 -->
 
+    <color name="popup_shade_first_light">#F9F9F9</color>
+    <color name="popup_shade_second_light">#F1F1F1</color>
+    <color name="popup_shade_third_light">#E2E2E2</color>
+    <color name="popup_shade_first_dark">#303030</color>
+    <color name="popup_shade_second_dark">#262626</color>
+    <color name="popup_shade_third_dark">#1B1B1B</color>
+
     <color name="popup_notification_dot_light">#FFF</color>
     <color name="popup_notification_dot_dark">#757575</color>
 
@@ -62,8 +69,8 @@
     <color name="folder_hint_text_color_light">#FFF</color>
     <color name="folder_hint_text_color_dark">#FF000000</color>
 
-    <color name="folder_background_light">#FFFFFF</color>
-    <color name="folder_background_dark">#FF3C4043</color>
+    <color name="folder_background_light">#F9F9F9</color>
+    <color name="folder_background_dark">#464746</color>
 
     <color name="folder_dot_color">?attr/colorPrimary</color>
 
diff --git a/res/values/styles.xml b/res/values/styles.xml
index dc99893..416c711 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -37,6 +37,9 @@
         <item name="popupColorPrimary">@color/popup_color_primary_light</item>
         <item name="popupColorSecondary">@color/popup_color_secondary_light</item>
         <item name="popupColorTertiary">@color/popup_color_tertiary_light</item>
+        <item name="popupShadeFirst">@color/popup_shade_first_light</item>
+        <item name="popupShadeSecond">@color/popup_shade_second_light</item>
+        <item name="popupShadeThird">@color/popup_shade_third_light</item>
         <item name="popupNotificationDotColor">@color/popup_notification_dot_light</item>
         <item name="isMainColorDark">false</item>
         <item name="isWorkspaceDarkText">false</item>
@@ -100,6 +103,9 @@
         <item name="popupColorSecondary">@color/popup_color_secondary_dark</item>
         <item name="popupColorTertiary">@color/popup_color_tertiary_dark</item>
         <item name="popupNotificationDotColor">@color/popup_notification_dot_dark</item>
+        <item name="popupShadeFirst">@color/popup_shade_first_dark</item>
+        <item name="popupShadeSecond">@color/popup_shade_second_dark</item>
+        <item name="popupShadeThird">@color/popup_shade_third_dark</item>
         <item name="widgetsTheme">@style/WidgetContainerTheme.Dark</item>
         <item name="folderDotColor">@color/folder_dot_color</item>
         <item name="folderFillColor">@color/folder_background_dark</item>
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
index c946c72..b9f183c 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -58,7 +58,7 @@
 
 @RunWith(RobolectricTestRunner.class)
 public final class WidgetsDiffReporterTest {
-    private static final String TEST_PACKAGE_PREFIX = "com.google.test";
+    private static final String TEST_PACKAGE_PREFIX = "com.android.test";
     private static final WidgetListBaseRowEntryComparator COMPARATOR =
             new WidgetListBaseRowEntryComparator();
 
@@ -241,6 +241,30 @@
         assertThat(currentList).containsExactlyElementsIn(newList);
     }
 
+    @Test
+    public void headersContentsMix_contentMaxSpanSizeModified_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, E content].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mContentE));
+        // GIVEN the new list has max span size in "E content" modified.
+        List<WidgetsListBaseEntry> newList = List.of(
+                mHeaderA,
+                mHeaderB,
+                new WidgetsListContentEntry(
+                        mContentE.mPkgItem,
+                        mContentE.mTitleSectionName,
+                        mContentE.mWidgets,
+                        mContentE.getMaxSpanSizeInCells() + 1));
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notify "E content" has been changed.
+        verify(mAdapter).notifyItemChanged(/* position= */ 2);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
 
     private WidgetsListHeaderEntry createWidgetsHeaderEntry(String packageName, String appName,
             int numOfWidgets) {
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index 2d22c45..106cac0 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -49,8 +49,10 @@
 
 @RunWith(RobolectricTestRunner.class)
 public final class WidgetsListContentEntryTest {
-    private static final String PACKAGE_NAME = "com.google.test";
-    private final PackageItemInfo mPackageItemInfo = new PackageItemInfo(PACKAGE_NAME);
+    private static final String PACKAGE_NAME = "com.android.test";
+    private static final String PACKAGE_NAME_2 = "com.android.test2";
+    private final PackageItemInfo mPackageItemInfo1 = new PackageItemInfo(PACKAGE_NAME);
+    private final PackageItemInfo mPackageItemInfo2 = new PackageItemInfo(PACKAGE_NAME_2);
     private final ComponentName mWidget1 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget1");
     private final ComponentName mWidget2 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget2");
     private final ComponentName mWidget3 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget3");
@@ -91,7 +93,7 @@
         WidgetItem widgetItem3 = createWidgetItem(mWidget3, /* spanX= */ 2, /* spanY= */ 3);
 
         // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
-        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo1,
                 /* titleSectionName= */ "T",
                 List.of(widgetItem1, widgetItem2, widgetItem3));
 
@@ -100,7 +102,7 @@
                 .containsExactly(widgetItem3, widgetItem1, widgetItem2)
                 .inOrder();
         assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
-        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo1);
     }
 
     @Test
@@ -114,7 +116,7 @@
         WidgetItem widgetItem3 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 2);
 
         // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
-        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo1,
                 /* titleSectionName= */ "T",
                 List.of(widgetItem1, widgetItem2, widgetItem3));
 
@@ -124,7 +126,7 @@
                 .containsExactly(widgetItem2, widgetItem3, widgetItem1)
                 .inOrder();
         assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
-        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo1);
     }
 
     @Test
@@ -140,7 +142,7 @@
         WidgetItem widgetItem4 = createWidgetItem(mWidget3, /* spanX= */ 2, /* spanY= */ 2);
 
         // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
-        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo1,
                 /* titleSectionName= */ "T",
                 List.of(widgetItem1, widgetItem2, widgetItem3, widgetItem4));
 
@@ -151,9 +153,96 @@
                 .containsExactly(widgetItem4, widgetItem2, widgetItem1, widgetItem3)
                 .inOrder();
         assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
-        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo1);
     }
 
+    @Test
+    public void equals_entriesWithDifferentPackageItemInfo_returnFalse() {
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry1 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry2 = new WidgetsListContentEntry(
+                mPackageItemInfo2,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+
+        assertThat(widgetsListRowEntry1.equals(widgetsListRowEntry2)).isFalse();
+    }
+
+    @Test
+    public void equals_entriesWithDifferentTitleSectionName_returnFalse() {
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry1 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry2 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "S",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+
+        assertThat(widgetsListRowEntry1.equals(widgetsListRowEntry2)).isFalse();
+    }
+
+    @Test
+    public void equals_entriesWithDifferentWidgetsList_returnFalse() {
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetItem widgetItem2 = createWidgetItem(mWidget2, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry1 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry2 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem2),
+                /* maxSpanSizeInCells= */ 3);
+
+        assertThat(widgetsListRowEntry1.equals(widgetsListRowEntry2)).isFalse();
+    }
+
+    @Test
+    public void equals_entriesWithDifferentMaxSpanSize_returnFalse() {
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry1 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry2 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 2);
+
+        assertThat(widgetsListRowEntry1.equals(widgetsListRowEntry2)).isFalse();
+    }
+
+    @Test
+    public void equals_entriesWithSameContents_returnTrue() {
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry1 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+        WidgetsListContentEntry widgetsListRowEntry2 = new WidgetsListContentEntry(
+                mPackageItemInfo1,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1),
+                /* maxSpanSizeInCells= */ 3);
+
+        assertThat(widgetsListRowEntry1.equals(widgetsListRowEntry2)).isTrue();
+    }
+
+
     private WidgetItem createWidgetItem(ComponentName componentName, int spanX, int spanY) {
         String label = mWidgetsToLabels.get(componentName);
         ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index a89fb3b..b7fe348 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -183,12 +183,12 @@
 
         if (isAboveAnotherSurface) {
             mColors = new int[] {
-                    getColorStateList(context, R.color.popup_color_first).getDefaultColor()};
+                    getColorStateList(context, R.color.popup_shade_first).getDefaultColor()};
         } else {
             mColors = new int[] {
-                    getColorStateList(context, R.color.popup_color_first).getDefaultColor(),
-                    getColorStateList(context, R.color.popup_color_second).getDefaultColor(),
-                    getColorStateList(context, R.color.popup_color_third).getDefaultColor()};
+                    getColorStateList(context, R.color.popup_shade_first).getDefaultColor(),
+                    getColorStateList(context, R.color.popup_shade_second).getDefaultColor(),
+                    getColorStateList(context, R.color.popup_shade_third).getDefaultColor()};
         }
     }
 
diff --git a/src/com/android/launcher3/views/TopRoundedCornerView.java b/src/com/android/launcher3/views/TopRoundedCornerView.java
index 5519df1..92cce92 100644
--- a/src/com/android/launcher3/views/TopRoundedCornerView.java
+++ b/src/com/android/launcher3/views/TopRoundedCornerView.java
@@ -17,12 +17,10 @@
 
 import android.content.Context;
 import android.graphics.Canvas;
-import android.graphics.Paint;
 import android.graphics.Path;
 import android.graphics.RectF;
 import android.util.AttributeSet;
 
-import com.android.launcher3.R;
 import com.android.launcher3.util.Themes;
 
 /**
@@ -34,41 +32,23 @@
     private final Path mClipPath = new Path();
     private float[] mRadii;
 
-    private final Paint mNavBarScrimPaint;
-    private int mNavBarScrimHeight = 0;
-
     public TopRoundedCornerView(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
 
         float radius = Themes.getDialogCornerRadius(context);
         mRadii = new float[] {radius, radius, radius, radius, 0, 0, 0, 0};
-
-        mNavBarScrimPaint = new Paint();
-        mNavBarScrimPaint.setColor(Themes.getAttrColor(context, R.attr.allAppsNavBarScrimColor));
     }
 
     public TopRoundedCornerView(Context context, AttributeSet attrs) {
         this(context, attrs, 0);
     }
 
-    public void setNavBarScrimHeight(int height) {
-        if (mNavBarScrimHeight != height) {
-            mNavBarScrimHeight = height;
-            invalidate();
-        }
-    }
-
     @Override
     public void draw(Canvas canvas) {
         canvas.save();
         canvas.clipPath(mClipPath);
         super.draw(canvas);
         canvas.restore();
-
-        if (mNavBarScrimHeight > 0) {
-            canvas.drawRect(0, getHeight() - mNavBarScrimHeight, getWidth(), getHeight(),
-                    mNavBarScrimPaint);
-        }
     }
 
     @Override
diff --git a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
index 0328cf6..73b17f1 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
@@ -26,14 +26,39 @@
  */
 public final class WidgetsListContentEntry extends WidgetsListBaseEntry {
 
+    private final int mMaxSpanSizeInCells;
+
+    /**
+     * Constructor for {@link WidgetsListContentEntry}.
+     *
+     * @param pkgItem package info associated with the entry
+     * @param titleSectionName title section name associated with the entry.
+     * @param items list of widgets for the package.
+     */
     public WidgetsListContentEntry(PackageItemInfo pkgItem, String titleSectionName,
             List<WidgetItem> items) {
+        this(pkgItem, titleSectionName, items, /* maxSpanSizeInCells= */ 0);
+    }
+
+    /**
+     * Constructor for {@link WidgetsListContentEntry}.
+     *
+     * @param pkgItem package info associated with the entry
+     * @param titleSectionName title section name associated with the entry.
+     * @param items list of widgets for the package.
+     * @param maxSpanSizeInCells the max horizontal span in cells that is allowed for grouping more
+     *                           than one widgets in a table row.
+     */
+    public WidgetsListContentEntry(PackageItemInfo pkgItem, String titleSectionName,
+            List<WidgetItem> items, int maxSpanSizeInCells) {
         super(pkgItem, titleSectionName, items);
+        mMaxSpanSizeInCells = maxSpanSizeInCells;
     }
 
     @Override
     public String toString() {
-        return "Content:" + mPkgItem.packageName + ":" + mWidgets.size();
+        return "Content:" + mPkgItem.packageName + ":" + mWidgets.size() + " maxSpanSizeInCells: "
+                + mMaxSpanSizeInCells;
     }
 
     @Override
@@ -42,11 +67,36 @@
         return RANK_WIDGETS_LIST_CONTENT;
     }
 
+    /**
+     * Returns a copy of this {@link WidgetsListContentEntry} with updated
+     * {@param maxSpanSizeInCells}.
+     *
+     * @param maxSpanSizeInCells the maximum horizontal span in cells that is allowed for grouping
+     *                           more than one widgets in a table row.
+     */
+    public WidgetsListContentEntry withMaxSpanSize(int maxSpanSizeInCells) {
+        if (mMaxSpanSizeInCells == maxSpanSizeInCells) return this;
+        return new WidgetsListContentEntry(
+                mPkgItem,
+                mTitleSectionName,
+                mWidgets,
+                /* maxSpanSizeInCells= */ maxSpanSizeInCells);
+    }
+
+    /**
+     * Returns the max horizontal span size in cells that is allowed for grouping more than one
+     * widget in a table row.
+     */
+    public int getMaxSpanSizeInCells() {
+        return mMaxSpanSizeInCells;
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (!(obj instanceof WidgetsListContentEntry)) return false;
         WidgetsListContentEntry otherEntry = (WidgetsListContentEntry) obj;
         return mWidgets.equals(otherEntry.mWidgets) && mPkgItem.equals(otherEntry.mPkgItem)
-                && mTitleSectionName.equals(otherEntry.mTitleSectionName);
+                && mTitleSectionName.equals(otherEntry.mTitleSectionName)
+                && mMaxSpanSizeInCells == otherEntry.mMaxSpanSizeInCells;
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
index dfe447a..99374f5 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
@@ -115,7 +115,7 @@
                 // or did the widget size and desc, span, etc change?
                 if (!isSamePackageItemInfo(orgRowEntry.mPkgItem, newRowEntry.mPkgItem)
                         || hasHeaderUpdated(orgRowEntry, newRowEntry)
-                        || hasWidgetsListChanged(orgRowEntry, newRowEntry)) {
+                        || hasWidgetsListContentChanged(orgRowEntry, newRowEntry)) {
                     index = currentEntries.indexOf(orgRowEntry);
                     currentEntries.set(index, newRowEntry);
                     mListener.notifyItemChanged(index);
@@ -158,17 +158,15 @@
 
     /**
      * Returns {@code true} if both {@code curRow} & {@code newRow} are
-     * {@link WidgetsListContentEntry}s with a different list of widgets.
+     * {@link WidgetsListContentEntry}s with a different list or arrangement of widgets.
      */
-    private boolean hasWidgetsListChanged(WidgetsListBaseEntry curRow,
+    private boolean hasWidgetsListContentChanged(WidgetsListBaseEntry curRow,
             WidgetsListBaseEntry newRow) {
         if (!(curRow instanceof WidgetsListContentEntry)
                 || !(newRow instanceof WidgetsListContentEntry)) {
             return false;
         }
-        WidgetsListContentEntry orgRowEntry = (WidgetsListContentEntry) curRow;
-        WidgetsListContentEntry newRowEntry = (WidgetsListContentEntry) newRow;
-        return !orgRowEntry.mWidgets.equals(newRowEntry.mWidgets);
+        return !curRow.equals(newRow);
     }
 
     /**
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 1a58bb0..0106ef5 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -323,7 +323,6 @@
             clearNavBarColor();
         }
 
-        ((TopRoundedCornerView) mContent).setNavBarScrimHeight(mInsets.bottom);
         requestLayout();
     }
 
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index b668c90..08a2263 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -111,6 +111,7 @@
     @Nullable private PackageUserKey mPendingClickHeader;
     private final int mShortcutPreviewPadding;
     private final int mSpacingBetweenEntries;
+    private int mMaxSpanSize = 4;
 
     private final WidgetPreviewLoadedCallback mPreviewLoadedCallback =
             ignored -> updateVisibleEntries();
@@ -249,10 +250,14 @@
                 .filter(entry -> (mFilter == null || mFilter.test(entry))
                         && mHeaderAndSelectedContentFilter.test(entry))
                 .map(entry -> {
-                    // Adjust the original entries to expand headers for the selected content.
                     if (entry instanceof WidgetsListBaseEntry.Header<?>
                             && matchesKey(entry, mWidgetsContentVisiblePackageUserKey)) {
+                        // Adjust the original entries to expand headers for the selected content.
                         return ((WidgetsListBaseEntry.Header<?>) entry).withWidgetListShown();
+                    } else if (entry instanceof WidgetsListContentEntry) {
+                        // Adjust the original content entries to accommodate for the current
+                        // maxSpanSize.
+                        return ((WidgetsListContentEntry) entry).withMaxSpanSize(mMaxSpanSize);
                     }
                     return entry;
                 })
@@ -491,20 +496,12 @@
     }
 
     /**
-     * Sets the max horizontal spans that are allowed for grouping more than one widgets in a table
-     * row.
-     *
-     * <p>If there is only one widget in a row, that widget horizontal span is allowed to exceed
-     * {@code maxHorizontalSpans}.
-     * <p>Let's say the max horizontal spans is set to 5. Widgets can be grouped in the same row if
-     * their total horizontal spans added don't exceed 5.
-     * Example 1: Row 1: 2x2, 2x3, 1x1. Total horizontal spans is 5. This is okay.
-     * Example 2: Row 1: 2x2, 4x3, 1x1. the total horizontal spans is 7. This is wrong.
-     *            4x3 and 1x1 should be moved to a new row.
-     * Example 3: Row 1: 6x4. This is okay because this is the only item in the row.
+     * Sets the max horizontal span in cells that is allowed for grouping more than one widget in a
+     * table row.
      */
     public void setMaxHorizontalSpansPerRow(int maxHorizontalSpans) {
-        mWidgetsListTableViewHolderBinder.setMaxSpansPerRow(maxHorizontalSpans);
+        mMaxSpanSize = maxHorizontalSpans;
+        updateVisibleEntries();
     }
 
     /**
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
index 7b52663..57dec14 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -49,7 +49,6 @@
     private static final boolean DEBUG = false;
     private static final String TAG = "WidgetsListRowViewHolderBinder";
 
-    private int mMaxSpansPerRow = 4;
     private final LayoutInflater mLayoutInflater;
     private final OnClickListener mIconClickListener;
     private final OnLongClickListener mIconLongClickListener;
@@ -82,10 +81,6 @@
         mApplyBitmapDeferred = applyBitmapDeferred;
     }
 
-    public void setMaxSpansPerRow(int maxSpansPerRow) {
-        mMaxSpansPerRow = maxSpansPerRow;
-    }
-
     @Override
     public WidgetsRowViewHolder newViewHolder(ViewGroup parent) {
         if (DEBUG) {
@@ -113,7 +108,8 @@
                 position == mWidgetsListAdapter.getItemCount() - 1 ? LAST : MIDDLE);
 
         List<ArrayList<WidgetItem>> widgetItemsTable =
-                WidgetsTableUtils.groupWidgetItemsIntoTable(entry.mWidgets, mMaxSpansPerRow);
+                WidgetsTableUtils.groupWidgetItemsIntoTable(
+                        entry.mWidgets, entry.getMaxSpanSizeInCells());
         recycleTableBeforeBinding(table, widgetItemsTable);
         // Bind the widget items.
         for (int i = 0; i < widgetItemsTable.size(); i++) {
diff --git a/src/com/android/launcher3/widget/util/WidgetsTableUtils.java b/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
index 7294a3a..54aaf93 100644
--- a/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
+++ b/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
@@ -56,6 +56,13 @@
      * 3. The order shortcuts are grouped together in the same row until their total horizontal
      *    spans exceed the {@code maxSpansPerRow} - 1.
      * 4. If there is only one widget in a row, its width may exceed the {@code maxSpansPerRow}.
+     *
+     * <p>Let's say the {@code maxSpansPerRow} is set to 6. Widgets can be grouped in the same row
+     * if their total horizontal spans added don't exceed 5.
+     * Example 1: Row 1: 2x2, 2x3, 1x1. Total horizontal spans is 5. This is okay.
+     * Example 2: Row 1: 2x2, 4x3, 1x1. the total horizontal spans is 7. This is wrong. 4x3 and 1x1
+     * should be moved to a new row.
+     * Example 3: Row 1: 6x4. This is okay because this is the only item in the row.
      */
     public static List<ArrayList<WidgetItem>> groupWidgetItemsIntoTable(
             List<WidgetItem> widgetItems, final int maxSpansPerRow) {
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 2712bc0..bf4eba0 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -65,6 +65,7 @@
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.FailureWatcher;
 import com.android.launcher3.util.rule.LauncherActivityRule;
+import com.android.launcher3.util.rule.ScreenRecordRule;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.util.rule.TestStabilityRule;
 
@@ -204,6 +205,9 @@
     public ShellCommandRule mDisableHeadsUpNotification =
             ShellCommandRule.disableHeadsUpNotification();
 
+    @Rule
+    public ScreenRecordRule mScreenRecordRule = new ScreenRecordRule();
+
     protected void clearPackageData(String pkg) throws IOException, InterruptedException {
         final CountDownLatch count = new CountDownLatch(2);
         final BroadcastReceiver broadcastReceiver = new BroadcastReceiver() {
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 06bc26a..4dd44f4 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -36,6 +36,7 @@
 import com.android.launcher3.tapl.AppIconMenuItem;
 import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.tapl.Workspace;
+import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.views.OptionsPopupView;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 import com.android.launcher3.widget.picker.WidgetsRecyclerView;
@@ -92,6 +93,7 @@
     }
 
     @Test
+    @ScreenRecord //b/187080582
     public void testDevicePressMenu() throws Exception {
         mDevice.pressMenu();
         mDevice.waitForIdle();
diff --git a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
index 822fefc..745dc22 100644
--- a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
@@ -42,6 +42,7 @@
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.Wait.Condition;
+import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.ShellCommandRule;
 
 import org.junit.Before;
@@ -77,6 +78,7 @@
     public void testEmpty() throws Throwable { /* needed while the broken tests are being fixed */ }
 
     @Test
+    @ScreenRecord  //b/192010616
     public void testPinWidgetNoConfig() throws Throwable {
         runTest("pinWidgetNoConfig", true, (info, view) -> info instanceof LauncherAppWidgetInfo &&
                 ((LauncherAppWidgetInfo) info).appWidgetId == mAppWidgetId &&
@@ -85,6 +87,7 @@
     }
 
     @Test
+    @ScreenRecord  //b/192005114
     public void testPinWidgetNoConfig_customPreview() throws Throwable {
         // Command to set custom preview
         Intent command = RequestPinItemActivity.getCommandIntent(
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
index 4c47947..dc59bdd 100644
--- a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -2,6 +2,8 @@
 
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
+import android.os.FileUtils;
+import android.os.ParcelFileDescriptor.AutoCloseInputStream;
 import android.util.Log;
 
 import androidx.test.uiautomator.UiDevice;
@@ -12,9 +14,12 @@
 import org.junit.rules.TestWatcher;
 import org.junit.runner.Description;
 
-import java.io.ByteArrayOutputStream;
 import java.io.File;
+import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
 
 public class FailureWatcher extends TestWatcher {
     private static final String TAG = "FailureWatcher";
@@ -26,20 +31,6 @@
         mLauncher = launcher;
     }
 
-    private static void dumpViewHierarchy(UiDevice device) {
-        final ByteArrayOutputStream stream = new ByteArrayOutputStream();
-        try {
-            device.dumpWindowHierarchy(stream);
-            stream.flush();
-            stream.close();
-            for (String line : stream.toString().split("\\r?\\n")) {
-                Log.e(TAG, line.trim());
-            }
-        } catch (IOException e) {
-            Log.e(TAG, "error dumping XML to logcat", e);
-        }
-    }
-
     @Override
     protected void succeeded(Description description) {
         super.succeeded(description);
@@ -53,22 +44,41 @@
 
     public static void onError(UiDevice device, Description description, Throwable e) {
         if (device == null) return;
-        final String pathname = getInstrumentation().getTargetContext().
-                getFilesDir().getPath() + "/TestScreenshot-" + description.getMethodName()
-                + ".png";
-        Log.e(TAG, "Failed test " + description.getMethodName() +
-                ", screenshot will be saved to " + pathname +
-                ", track trace is below, UI object dump is further below:\n" +
-                Log.getStackTraceString(e));
-        dumpViewHierarchy(device);
+        final File parentFile = getInstrumentation().getTargetContext().getFilesDir();
+        final File sceenshot = new File(parentFile,
+                "TestScreenshot-" + description.getMethodName() + ".png");
+        final File hierarchy = new File(parentFile,
+                "Hierarchy-" + description.getMethodName() + ".zip");
 
-        try {
-            final String dumpsysResult = device.executeShellCommand(
-                    "dumpsys activity service TouchInteractionService");
-            Log.d(TAG, "TouchInteractionService: " + dumpsysResult);
-        } catch (IOException ex) {
+        // Dump window hierarchy
+        try (ZipOutputStream out = new ZipOutputStream(new FileOutputStream(hierarchy))) {
+            out.putNextEntry(new ZipEntry("bugreport.txt"));
+            dumpStringCommand("dumpsys window windows", out);
+            dumpStringCommand("dumpsys package", out);
+            dumpStringCommand("dumpsys activity service TouchInteractionService", out);
+            out.closeEntry();
+
+            out.putNextEntry(new ZipEntry("visible_windows.zip"));
+            dumpCommand("cmd window dump-visible-window-views", out);
+            out.closeEntry();
+        } catch (IOException ex) { }
+
+        Log.e(TAG, "Failed test " + description.getMethodName()
+                + ",\nscreenshot will be saved to " + sceenshot
+                + ",\nUI dump at: " + hierarchy
+                + " (use go/web-hv to open the dump file)", e);
+        device.takeScreenshot(sceenshot);
+    }
+
+    private static void dumpStringCommand(String cmd, OutputStream out) throws IOException {
+        out.write(("\n\n" + cmd + "\n").getBytes());
+        dumpCommand(cmd, out);
+    }
+
+    private static void dumpCommand(String cmd, OutputStream out) throws IOException {
+        try (AutoCloseInputStream in = new AutoCloseInputStream(getInstrumentation()
+                .getUiAutomation().executeShellCommand(cmd))) {
+            FileUtils.copy(in, out);
         }
-
-        device.takeScreenshot(new File(pathname));
     }
 }
diff --git a/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
new file mode 100644
index 0000000..00b1cdd
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util.rule;
+
+import static androidx.test.InstrumentationRegistry.getInstrumentation;
+
+import android.app.Instrumentation;
+import android.app.UiAutomation;
+import android.os.ParcelFileDescriptor;
+import android.util.Log;
+
+import androidx.test.uiautomator.UiDevice;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.io.File;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Rule which captures a screen record for a test.
+ * After adding this rule to the test class, apply the annotation @ScreenRecord to individual tests
+ */
+public class ScreenRecordRule implements TestRule {
+
+    private static final String TAG = "ScreenRecordRule";
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        if (description.getAnnotation(ScreenRecord.class) == null) {
+            return base;
+        }
+
+        return new Statement() {
+            @Override
+            public void evaluate() throws Throwable {
+                Instrumentation inst = getInstrumentation();
+                UiAutomation automation = inst.getUiAutomation();
+                UiDevice device = UiDevice.getInstance(inst);
+
+                File outputFile = new File(inst.getTargetContext().getFilesDir(),
+                        "screenrecord-" + description.getMethodName() + ".mp4");
+                device.executeShellCommand("killall screenrecord");
+                ParcelFileDescriptor output =
+                        automation.executeShellCommand("screenrecord " + outputFile);
+                String screenRecordPid = device.executeShellCommand("pidof screenrecord");
+                try {
+                    base.evaluate();
+                } finally {
+                    device.executeShellCommand("kill -INT " + screenRecordPid);
+                    Log.e(TAG, "Screenrecord captured at: " + outputFile);
+                    output.close();
+                }
+            }
+        };
+    }
+
+    /**
+     * Interface to indicate that the test should capture screenrecord
+     */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.METHOD)
+    public @interface ScreenRecord { }
+}