diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
index f15d12b..8566e20 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
@@ -193,10 +193,14 @@
     }
 
     void closeQuickSwitchView() {
+        closeQuickSwitchView(true);
+    }
+
+    void closeQuickSwitchView(boolean animate) {
         if (mQuickSwitchViewController == null) {
             return;
         }
-        mQuickSwitchViewController.closeQuickSwitchView(true);
+        mQuickSwitchViewController.closeQuickSwitchView(animate);
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 5e325fb..0e1a6da 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -529,52 +529,26 @@
 
     /**
      * Creates {@link WindowManager.LayoutParams} for Taskbar, and also sets LP.paramsForRotation
-     * for taskbar showing as navigation bar
+     * for taskbar
      */
     private WindowManager.LayoutParams createAllWindowParams() {
         final int windowType =
                 ENABLE_TASKBAR_NAVBAR_UNIFICATION ? TYPE_NAVIGATION_BAR : TYPE_NAVIGATION_BAR_PANEL;
         WindowManager.LayoutParams windowLayoutParams =
                 createDefaultWindowLayoutParams(windowType, TaskbarActivityContext.WINDOW_TITLE);
-        if (!isPhoneButtonNavMode()) {
-            return windowLayoutParams;
-        }
 
-        // Provide WM layout params for all rotations to cache, see NavigationBar#getBarLayoutParams
-        int width = WindowManager.LayoutParams.MATCH_PARENT;
-        int height = WindowManager.LayoutParams.MATCH_PARENT;
-        int gravity = Gravity.BOTTOM;
         windowLayoutParams.paramsForRotation = new WindowManager.LayoutParams[4];
         for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
             WindowManager.LayoutParams lp =
                     createDefaultWindowLayoutParams(windowType,
                             TaskbarActivityContext.WINDOW_TITLE);
-            switch (rot) {
-                case Surface.ROTATION_0, Surface.ROTATION_180 -> {
-                    // Defaults are fine
-                    width = WindowManager.LayoutParams.MATCH_PARENT;
-                    height = mLastRequestedNonFullscreenSize;
-                    gravity = Gravity.BOTTOM;
-                }
-                case Surface.ROTATION_90 -> {
-                    width = mLastRequestedNonFullscreenSize;
-                    height = WindowManager.LayoutParams.MATCH_PARENT;
-                    gravity = Gravity.END;
-                }
-                case Surface.ROTATION_270 -> {
-                    width = mLastRequestedNonFullscreenSize;
-                    height = WindowManager.LayoutParams.MATCH_PARENT;
-                    gravity = Gravity.START;
-                }
-
+            if (isPhoneButtonNavMode()) {
+                populatePhoneButtonNavModeWindowLayoutParams(rot, lp);
             }
-            lp.width = width;
-            lp.height = height;
-            lp.gravity = gravity;
             windowLayoutParams.paramsForRotation[rot] = lp;
         }
 
-        // Override current layout params
+        // Override with current layout params
         WindowManager.LayoutParams currentParams =
                 windowLayoutParams.paramsForRotation[getDisplay().getRotation()];
         windowLayoutParams.width = currentParams.width;
@@ -584,6 +558,32 @@
         return windowLayoutParams;
     }
 
+    /**
+     * Update {@link WindowManager.LayoutParams} with values specific to phone and 3 button
+     * navigation users
+     */
+    private void populatePhoneButtonNavModeWindowLayoutParams(int rot,
+            WindowManager.LayoutParams lp) {
+        lp.width = WindowManager.LayoutParams.MATCH_PARENT;
+        lp.height = WindowManager.LayoutParams.MATCH_PARENT;
+        lp.gravity = Gravity.BOTTOM;
+
+        // Override with per-rotation specific values
+        switch (rot) {
+            case Surface.ROTATION_0, Surface.ROTATION_180 -> {
+                lp.height = mLastRequestedNonFullscreenSize;
+            }
+            case Surface.ROTATION_90 -> {
+                lp.width = mLastRequestedNonFullscreenSize;
+                lp.gravity = Gravity.END;
+            }
+            case Surface.ROTATION_270 -> {
+                lp.width = mLastRequestedNonFullscreenSize;
+                lp.gravity = Gravity.START;
+            }
+        }
+    }
+
     public void onConfigurationChanged(@Config int configChanges) {
         mControllers.onConfigurationChanged(configChanges);
         if (!mIsUserSetupComplete) {
@@ -944,8 +944,14 @@
         }
         if (landscapePhoneButtonNav) {
             mWindowLayoutParams.width = size;
+            for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
+                mWindowLayoutParams.paramsForRotation[rot].width = size;
+            }
         } else {
             mWindowLayoutParams.height = size;
+            for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
+                mWindowLayoutParams.paramsForRotation[rot].height = size;
+            }
         }
         mControllers.taskbarInsetsController.onTaskbarOrBubblebarWindowHeightOrInsetsChanged();
         notifyUpdateLayoutParams();
@@ -1551,4 +1557,9 @@
     public float getStashedTaskbarScale() {
         return mControllers.stashedHandleViewController.getStashedHandleHintScale().value;
     }
+
+    /** Closes the KeyboardQuickSwitchView without an animation if open. */
+    public void closeKeyboardQuickSwitchView() {
+        mControllers.keyboardQuickSwitchController.closeQuickSwitchView(false);
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
index aa457ca..567fad0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
@@ -118,11 +118,9 @@
                 getProvidedInsets(insetsRoundedCornerFlag)
             }
 
-        if (!context.isGestureNav) {
-            if (windowLayoutParams.paramsForRotation != null) {
-                for (layoutParams in windowLayoutParams.paramsForRotation) {
-                    layoutParams.providedInsets = getProvidedInsets(insetsRoundedCornerFlag)
-                }
+        if (windowLayoutParams.paramsForRotation != null) {
+            for (layoutParams in windowLayoutParams.paramsForRotation) {
+                layoutParams.providedInsets = getProvidedInsets(insetsRoundedCornerFlag)
             }
         }
 
@@ -156,19 +154,12 @@
             )
         }
 
-        val gravity = windowLayoutParams.gravity
-
         // Pre-calculate insets for different providers across different rotations for this gravity
         for (rotation in Surface.ROTATION_0..Surface.ROTATION_270) {
             // Add insets for navbar rotated params
-            if (windowLayoutParams.paramsForRotation != null) {
-                val layoutParams = windowLayoutParams.paramsForRotation[rotation]
-                for (provider in layoutParams.providedInsets) {
-                    setProviderInsets(provider, layoutParams.gravity, rotation)
-                }
-            }
-            for (provider in windowLayoutParams.providedInsets) {
-                setProviderInsets(provider, gravity, rotation)
+            val layoutParams = windowLayoutParams.paramsForRotation[rotation]
+            for (provider in layoutParams.providedInsets) {
+                setProviderInsets(provider, layoutParams.gravity, rotation)
             }
         }
         context.notifyUpdateLayoutParams()
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
index 65b5397..56c9a00 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
@@ -248,7 +248,15 @@
                 case TYPE_SHOW:
                     // already visible
                     return true;
+                case TYPE_KEYBOARD_INPUT: {
+                    if (visibleRecentsView.isHandlingTouch()) {
+                        return true;
+                    }
+                }
                 case TYPE_HIDE: {
+                    if (visibleRecentsView.isHandlingTouch()) {
+                        return true;
+                    }
                     mKeyboardTaskFocusIndex = INVALID_PAGE;
                     int currentPage = visibleRecentsView.getNextPage();
                     TaskView tv = (currentPage >= 0
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
index c56a621..d6a727f 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
@@ -33,6 +33,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DEVICE_DREAMING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DIALOG_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DISABLE_GESTURE_SPLIT_INVOCATION;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_HOME_DISABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_MAGNIFICATION_OVERLAP;
@@ -398,7 +399,8 @@
                 && (mSystemUiStateFlags & SYSUI_STATE_MAGNIFICATION_OVERLAP) == 0
                 && ((mSystemUiStateFlags & SYSUI_STATE_HOME_DISABLED) == 0
                         || (mSystemUiStateFlags & SYSUI_STATE_OVERVIEW_DISABLED) == 0)
-                && (mSystemUiStateFlags & SYSUI_STATE_DEVICE_DREAMING) == 0;
+                && (mSystemUiStateFlags & SYSUI_STATE_DEVICE_DREAMING) == 0
+                && (mSystemUiStateFlags & SYSUI_STATE_DISABLE_GESTURE_SPLIT_INVOCATION) == 0;
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 7880124..719c4f7 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -754,6 +754,10 @@
 
             boolean isOneHandedModeActive = mDeviceState.isOneHandedModeActive();
             boolean isInSwipeUpTouchRegion = mRotationTouchHelper.isInSwipeUpTouchRegion(event);
+            TaskbarActivityContext tac = mTaskbarManager.getCurrentActivityContext();
+            if (isInSwipeUpTouchRegion && tac != null) {
+                tac.closeKeyboardQuickSwitchView();
+            }
             if ((!isOneHandedModeActive && isInSwipeUpTouchRegion)
                     || isHoverActionWithoutConsumer) {
                 reasonString.append(!isOneHandedModeActive && isInSwipeUpTouchRegion
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 2575e61..34feee4 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -4231,30 +4231,31 @@
 
     @Override
     public boolean dispatchKeyEvent(KeyEvent event) {
-        if (event.getAction() == KeyEvent.ACTION_DOWN) {
-            switch (event.getKeyCode()) {
-                case KeyEvent.KEYCODE_TAB:
-                    return snapToPageRelative(event.isShiftPressed() ? -1 : 1, true /* cycle */,
-                            DIRECTION_TAB);
-                case KeyEvent.KEYCODE_DPAD_RIGHT:
-                    return snapToPageRelative(mIsRtl ? -1 : 1, true /* cycle */, DIRECTION_RIGHT);
-                case KeyEvent.KEYCODE_DPAD_LEFT:
-                    return snapToPageRelative(mIsRtl ? 1 : -1, true /* cycle */, DIRECTION_LEFT);
-                case KeyEvent.KEYCODE_DPAD_UP:
-                    return snapToPageRelative(1, false /* cycle */, DIRECTION_UP);
-                case KeyEvent.KEYCODE_DPAD_DOWN:
-                    return snapToPageRelative(1, false /* cycle */, DIRECTION_DOWN);
-                case KeyEvent.KEYCODE_DEL:
-                case KeyEvent.KEYCODE_FORWARD_DEL:
+        if (isHandlingTouch() || event.getAction() != KeyEvent.ACTION_DOWN) {
+            return super.dispatchKeyEvent(event);
+        }
+        switch (event.getKeyCode()) {
+            case KeyEvent.KEYCODE_TAB:
+                return snapToPageRelative(event.isShiftPressed() ? -1 : 1, true /* cycle */,
+                        DIRECTION_TAB);
+            case KeyEvent.KEYCODE_DPAD_RIGHT:
+                return snapToPageRelative(mIsRtl ? -1 : 1, true /* cycle */, DIRECTION_RIGHT);
+            case KeyEvent.KEYCODE_DPAD_LEFT:
+                return snapToPageRelative(mIsRtl ? 1 : -1, true /* cycle */, DIRECTION_LEFT);
+            case KeyEvent.KEYCODE_DPAD_UP:
+                return snapToPageRelative(1, false /* cycle */, DIRECTION_UP);
+            case KeyEvent.KEYCODE_DPAD_DOWN:
+                return snapToPageRelative(1, false /* cycle */, DIRECTION_DOWN);
+            case KeyEvent.KEYCODE_DEL:
+            case KeyEvent.KEYCODE_FORWARD_DEL:
+                dismissCurrentTask();
+                return true;
+            case KeyEvent.KEYCODE_NUMPAD_DOT:
+                if (event.isAltPressed()) {
+                    // Numpad DEL pressed while holding Alt.
                     dismissCurrentTask();
                     return true;
-                case KeyEvent.KEYCODE_NUMPAD_DOT:
-                    if (event.isAltPressed()) {
-                        // Numpad DEL pressed while holding Alt.
-                        dismissCurrentTask();
-                        return true;
-                    }
-            }
+                }
         }
         return super.dispatchKeyEvent(event);
     }
diff --git a/quickstep/tests/src/com/android/launcher3/model/AppEventProducerTest.java b/quickstep/tests/multivalentTests/src/com/android/launcher3/model/AppEventProducerTest.java
similarity index 100%
rename from quickstep/tests/src/com/android/launcher3/model/AppEventProducerTest.java
rename to quickstep/tests/multivalentTests/src/com/android/launcher3/model/AppEventProducerTest.java
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/RecentsHitboxExtenderTest.java b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/RecentsHitboxExtenderTest.java
similarity index 100%
rename from quickstep/tests/src/com/android/launcher3/taskbar/RecentsHitboxExtenderTest.java
rename to quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/RecentsHitboxExtenderTest.java
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
similarity index 100%
rename from quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
rename to quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationBarRotationContextTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/NavigationBarRotationContextTest.java
similarity index 100%
rename from quickstep/tests/src/com/android/quickstep/NavigationBarRotationContextTest.java
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/NavigationBarRotationContextTest.java
diff --git a/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
similarity index 100%
rename from quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaskGridNavHelperTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.java
similarity index 100%
rename from quickstep/tests/src/com/android/quickstep/util/TaskGridNavHelperTest.java
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.java
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaskKeyByLastActiveTimeCacheTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskKeyByLastActiveTimeCacheTest.java
similarity index 100%
rename from quickstep/tests/src/com/android/quickstep/util/TaskKeyByLastActiveTimeCacheTest.java
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskKeyByLastActiveTimeCacheTest.java
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
similarity index 100%
rename from quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
diff --git a/res/layout/widget_cell.xml b/res/layout/widget_cell.xml
index 55dd1de..4533873 100644
--- a/res/layout/widget_cell.xml
+++ b/res/layout/widget_cell.xml
@@ -17,7 +17,8 @@
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="0dp"
     android:layout_height="wrap_content"
-    android:paddingHorizontal="@dimen/widget_cell_horizontal_padding"
+    android:layout_marginStart="@dimen/widget_cell_horizontal_padding"
+    android:layout_marginEnd="@dimen/widget_cell_horizontal_padding"
     android:paddingVertical="@dimen/widget_cell_vertical_padding"
     android:layout_weight="1"
     android:orientation="vertical"
diff --git a/res/layout/widget_recommendations_table.xml b/res/layout/widget_recommendations_table.xml
index e3f0562..b53d2d5 100644
--- a/res/layout/widget_recommendations_table.xml
+++ b/res/layout/widget_recommendations_table.xml
@@ -17,5 +17,4 @@
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
-    android:paddingHorizontal="@dimen/widget_recommendations_table_horizontal_padding"
     android:paddingVertical="@dimen/widget_recommendations_table_vertical_padding" />
diff --git a/res/values-sw720dp/dimens.xml b/res/values-sw720dp/dimens.xml
index 3c79588..27aba6b 100644
--- a/res/values-sw720dp/dimens.xml
+++ b/res/values-sw720dp/dimens.xml
@@ -37,6 +37,7 @@
 
 <!-- Widget picker-->
     <dimen name="widget_list_horizontal_margin">30dp</dimen>
+    <dimen name="widget_cell_horizontal_padding">16dp</dimen>
 
     <!--  Folder spaces  -->
     <dimen name="folder_footer_horiz_padding">24dp</dimen>
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index a912e2d..97737fb 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -176,7 +176,7 @@
 
     <!-- Widget tray -->
     <dimen name="widget_cell_vertical_padding">8dp</dimen>
-    <dimen name="widget_cell_horizontal_padding">16dp</dimen>
+    <dimen name="widget_cell_horizontal_padding">8dp</dimen>
     <dimen name="widget_cell_font_size">14sp</dimen>
     <dimen name="widget_cell_app_icon_size">24dp</dimen>
     <dimen name="widget_cell_app_icon_padding">8dp</dimen>
@@ -187,7 +187,6 @@
     <dimen name="widget_picker_landscape_tablet_left_right_margin">117dp</dimen>
     <dimen name="widget_picker_two_panels_left_right_margin">0dp</dimen>
     <dimen name="widget_recommendations_table_vertical_padding">8dp</dimen>
-    <dimen name="widget_recommendations_table_horizontal_padding">16dp</dimen>
     <!-- Bottom margin for the search and recommended widgets container without work profile -->
     <dimen name="search_and_recommended_widgets_container_bottom_margin">16dp</dimen>
     <!-- Bottom margin for the search and recommended widgets container with work profile -->
@@ -198,7 +197,8 @@
 
     <dimen name="widget_list_header_view_vertical_padding">20dp</dimen>
     <dimen name="widget_list_entry_spacing">2dp</dimen>
-    <dimen name="widget_list_horizontal_margin">16dp</dimen>
+    <!-- Less margin on sides to let widgets table width be close to the workspace width. -->
+    <dimen name="widget_list_horizontal_margin">11dp</dimen>
     <!-- Margin applied to the recycler view with search bar & the list of widget apps below it. -->
     <dimen name="widget_list_left_pane_horizontal_margin">0dp</dimen>
     <dimen name="widget_list_horizontal_margin_two_pane">24dp</dimen>
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 8277c3e..72977ee 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2365,7 +2365,8 @@
      * Similar to {@link #getFirstMatch} but optimized to finding a suitable view for the app close
      * animation.
      *
-     * @param preferredItemId The id of the preferred item to match to if it exists.
+     * @param preferredItemId The id of the preferred item to match to if it exists,
+     *                        or ItemInfo#NO_MATCHING_ID if you want to not match by item id
      * @param packageName The package name of the app to match.
      * @param user The user of the app to match.
      * @param supportsAllAppsState If true and we are in All Apps state, looks for view in All Apps.
diff --git a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
index 01ea9fb..fbeab4e 100644
--- a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
@@ -189,6 +189,7 @@
     private float mBottomSheetAlpha = 1f;
     private boolean mForceBottomSheetVisible;
     private int mTabsProtectionAlpha;
+    private float mTotalHeaderProtectionHeight;
     @Nullable private AllAppsTransitionController mAllAppsTransitionController;
 
     private PrivateSpaceHeaderViewController mPrivateSpaceHeaderViewController;
@@ -1431,9 +1432,11 @@
                 mTmpPath.reset();
                 mTmpPath.addRoundRect(mTmpRectF, mBottomSheetCornerRadii, Direction.CW);
                 canvas.drawPath(mTmpPath, mHeaderPaint);
+                mTotalHeaderProtectionHeight = headerBottomWithScaleOnTablet;
             }
         } else {
             canvas.drawRect(0, 0, canvas.getWidth(), headerBottomWithScaleOnPhone, mHeaderPaint);
+            mTotalHeaderProtectionHeight = headerBottomWithScaleOnPhone;
         }
 
         // If tab exist (such as work profile), extend header with tab height
@@ -1463,10 +1466,19 @@
                     right,
                     tabBottomWithScale,
                     mHeaderPaint);
+            mTotalHeaderProtectionHeight = tabBottomWithScale;
         }
     }
 
     /**
+     * The height of the header protection is dynamically calculated during the time of drawing the
+     * header.
+     */
+    float getHeaderProtectionHeight() {
+        return mTotalHeaderProtectionHeight;
+    }
+
+    /**
      * redraws header protection
      */
     public void invalidateHeader() {
diff --git a/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewController.java b/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewController.java
index 6067454..fdc035e 100644
--- a/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewController.java
+++ b/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewController.java
@@ -41,15 +41,18 @@
 import android.widget.RelativeLayout;
 import android.widget.TextView;
 
+import androidx.annotation.VisibleForTesting;
 import androidx.recyclerview.widget.LinearSmoothScroller;
 import androidx.recyclerview.widget.RecyclerView;
 
 import com.android.app.animation.Interpolators;
+import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Flags;
 import com.android.launcher3.R;
 import com.android.launcher3.allapps.UserProfileManager.UserProfileState;
 import com.android.launcher3.anim.AnimatedPropertySetter;
 import com.android.launcher3.anim.PropertySetter;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.RecyclerViewFastScroller;
 
 import java.util.List;
@@ -59,7 +62,6 @@
  * {@link UserProfileState}
  */
 public class PrivateSpaceHeaderViewController {
-    private static final int EXPAND_SCROLL_DURATION = 2000;
     private static final int EXPAND_COLLAPSE_DURATION = 800;
     private static final int SETTINGS_OPACITY_DURATION = 160;
     private final ActivityAllAppsContainerView mAllApps;
@@ -174,23 +176,24 @@
                 && mAllApps.getActiveRecyclerView() == mainAdapterHolder.mRecyclerView) {
             // Animate the text and settings icon.
             updatePrivateStateAnimator(true, header);
-            mAllApps.getActiveRecyclerView().scrollToBottomWithMotion(EXPAND_SCROLL_DURATION);
+            DeviceProfile deviceProfile =
+                    ActivityContext.lookupContext(mAllApps.getContext()).getDeviceProfile();
+            AllAppsRecyclerView allAppsRecyclerView = mAllApps.getActiveRecyclerView();
+            scrollForViewToBeVisibleInContainer(allAppsRecyclerView,
+                    allAppsRecyclerView.getApps().getAdapterItems(),
+                    header.getHeight(), deviceProfile.allAppsCellHeightPx);
         }
     }
 
     /** Finds the private space header to scroll to and set the private space icons to GONE. */
     private void collapse() {
         AllAppsRecyclerView allAppsRecyclerView = mAllApps.getActiveRecyclerView();
-        for (int i = allAppsRecyclerView.getChildCount() - 1; i > 0; i--) {
-            int adapterPosition = allAppsRecyclerView.getChildAdapterPosition(
-                    allAppsRecyclerView.getChildAt(i));
-            List<BaseAllAppsAdapter.AdapterItem> allAppsAdapters = allAppsRecyclerView.getApps()
-                    .getAdapterItems();
-            if (adapterPosition < 0 || adapterPosition >= allAppsAdapters.size()) {
-                continue;
-            }
+        List<BaseAllAppsAdapter.AdapterItem> appListAdapterItems =
+                allAppsRecyclerView.getApps().getAdapterItems();
+        for (int i = appListAdapterItems.size() - 1; i > 0; i--) {
+            BaseAllAppsAdapter.AdapterItem currentItem = appListAdapterItems.get(i);
             // Scroll to the private space header.
-            if (allAppsAdapters.get(adapterPosition).viewType == VIEW_TYPE_PRIVATE_SPACE_HEADER) {
+            if (currentItem.viewType == VIEW_TYPE_PRIVATE_SPACE_HEADER) {
                 // Note: SmoothScroller is meant to be used once.
                 RecyclerView.SmoothScroller smoothScroller =
                         new LinearSmoothScroller(mAllApps.getContext()) {
@@ -198,7 +201,7 @@
                                 return LinearSmoothScroller.SNAP_TO_END;
                             }
                         };
-                smoothScroller.setTargetPosition(adapterPosition);
+                smoothScroller.setTargetPosition(i);
                 RecyclerView.LayoutManager layoutManager = allAppsRecyclerView.getLayoutManager();
                 if (layoutManager != null) {
                     layoutManager.startSmoothScroll(smoothScroller);
@@ -206,12 +209,67 @@
                 break;
             }
             // Make the private space apps gone to "collapse".
-            if (allAppsAdapters.get(adapterPosition).decorationInfo != null) {
-                allAppsRecyclerView.getChildAt(i).setVisibility(GONE);
+            if (currentItem.decorationInfo != null) {
+                RecyclerView.ViewHolder viewHolder =
+                        allAppsRecyclerView.findViewHolderForAdapterPosition(i);
+                if (viewHolder != null) {
+                    viewHolder.itemView.setVisibility(GONE);
+                }
             }
         }
     }
 
+    /**
+     * Upon expanding, only scroll to the item position in the adapter that allows the header to be
+     * visible.
+     */
+    @VisibleForTesting
+    public int scrollForViewToBeVisibleInContainer(
+            AllAppsRecyclerView allAppsRecyclerView,
+            List<BaseAllAppsAdapter.AdapterItem> appListAdapterItems,
+            int psHeaderHeight,
+            int allAppsCellHeight) {
+        int rowToExpandToWithRespectToHeader = -1;
+        int itemToScrollTo = -1;
+        // Looks for the item in the app list to scroll to so that the header is visible.
+        for (int i = 0; i < appListAdapterItems.size(); i++) {
+            BaseAllAppsAdapter.AdapterItem currentItem = appListAdapterItems.get(i);
+            if (currentItem.viewType == VIEW_TYPE_PRIVATE_SPACE_HEADER) {
+                itemToScrollTo = i;
+                continue;
+            }
+            if (itemToScrollTo != -1) {
+                if (rowToExpandToWithRespectToHeader == -1) {
+                    rowToExpandToWithRespectToHeader = currentItem.rowIndex;
+                }
+                int rowToScrollTo =
+                        (int) Math.floor((double) (mAllApps.getHeight() - psHeaderHeight
+                                - mAllApps.getHeaderProtectionHeight()) / allAppsCellHeight);
+                int currentRowDistance = currentItem.rowIndex - rowToExpandToWithRespectToHeader;
+                // rowToScrollTo - 1 since the item to scroll to is 0 indexed.
+                if (currentRowDistance == rowToScrollTo - 1) {
+                    itemToScrollTo = i;
+                    break;
+                }
+            }
+        }
+        if (itemToScrollTo != -1) {
+            // Note: SmoothScroller is meant to be used once.
+            RecyclerView.SmoothScroller smoothScroller =
+                    new LinearSmoothScroller(mAllApps.getContext()) {
+                        @Override protected int getVerticalSnapPreference() {
+                            return LinearSmoothScroller.SNAP_TO_ANY;
+                        }
+                    };
+            smoothScroller.setTargetPosition(itemToScrollTo);
+            RecyclerView.LayoutManager layoutManager = allAppsRecyclerView.getLayoutManager();
+            if (layoutManager != null) {
+                layoutManager.startSmoothScroll(smoothScroller);
+            }
+        }
+        return itemToScrollTo;
+    }
+
     PrivateProfileManager getPrivateProfileManager() {
         return mPrivateProfileManager;
     }
diff --git a/src/com/android/launcher3/model/GridSizeMigrationUtil.java b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
index 30d2cfb..c40484d 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationUtil.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
@@ -40,6 +40,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
@@ -122,6 +123,16 @@
         if (!needsToMigrate(srcDeviceState, destDeviceState)) {
             return true;
         }
+
+        if (Flags.gridMigrationFix()
+                && srcDeviceState.getColumns().equals(destDeviceState.getColumns())
+                && srcDeviceState.getRows() < destDeviceState.getRows()) {
+            // Only use this strategy when comparing the previous grid to the new grid and the
+            // columns are the same and the destination has more rows
+            copyTable(source, TABLE_NAME, target.getWritableDatabase(), TABLE_NAME, context);
+            destDeviceState.writeToPrefs(context);
+            return true;
+        }
         copyTable(source, TABLE_NAME, target.getWritableDatabase(), TMP_TABLE, context);
 
         HashSet<String> validPackages = getValidPackages(context);
diff --git a/src/com/android/launcher3/views/FloatingSurfaceView.java b/src/com/android/launcher3/views/FloatingSurfaceView.java
index c60e1a4..cab7982 100644
--- a/src/com/android/launcher3/views/FloatingSurfaceView.java
+++ b/src/com/android/launcher3/views/FloatingSurfaceView.java
@@ -15,6 +15,7 @@
  */
 package com.android.launcher3.views;
 
+import static com.android.launcher3.model.data.ItemInfo.NO_MATCHING_ID;
 import static com.android.launcher3.views.FloatingIconView.getLocationBoundsForView;
 import static com.android.launcher3.views.IconLabelDotView.setIconAndDotVisible;
 
@@ -159,7 +160,7 @@
         if (mContract == null) {
             return;
         }
-        View icon = mLauncher.getFirstMatchForAppClose(-1,
+        View icon = mLauncher.getFirstMatchForAppClose(NO_MATCHING_ID,
                 mContract.componentName.getPackageName(), mContract.user,
                 false /* supportsAllAppsState */);
 
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index f2f83c8..aaefe60 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -57,6 +57,8 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.CancellableTask;
 import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.picker.util.WidgetPreviewContainerSize;
+import com.android.launcher3.widget.util.WidgetSizes;
 
 import java.util.function.Consumer;
 
@@ -80,7 +82,7 @@
      * The requested scale of the preview container. It can be lower than this as well.
      */
     private float mPreviewContainerScale = 1f;
-
+    private Size mPreviewContainerSize = new Size(0, 0);
     private FrameLayout mWidgetImageContainer;
     private WidgetImageView mWidgetImage;
     private ImageView mWidgetBadge;
@@ -176,6 +178,8 @@
         mWidgetDims.setText(null);
         mWidgetDescription.setText(null);
         mWidgetDescription.setVisibility(GONE);
+        showDescription(true);
+        showDimensions(true);
 
         if (mActiveRequest != null) {
             mActiveRequest.cancel();
@@ -186,6 +190,7 @@
             mWidgetImageContainer.removeView(mAppWidgetHostViewPreview);
         }
         mAppWidgetHostViewPreview = null;
+        mPreviewContainerSize = new Size(0, 0);
         mAppWidgetHostViewScale = 1f;
         mPreviewContainerScale = 1f;
         mItem = null;
@@ -201,30 +206,21 @@
      * Applies the item to this view
      */
     public void applyFromCellItem(WidgetItem item) {
-        applyFromCellItem(item, 1f);
-    }
-
-    /**
-     * Applies the item to this view
-     */
-    public void applyFromCellItem(WidgetItem item, float previewScale) {
-        applyFromCellItem(item, previewScale, this::applyPreview, null);
+        applyFromCellItem(item, this::applyPreview, /*cachedPreview=*/null);
     }
 
     /**
      * Applies the item to this view
      * @param item item to apply
-     * @param previewScale factor to scale the preview
      * @param callback callback when preview is loaded in case the preview is being loaded or cached
      * @param cachedPreview previously cached preview bitmap is present
      */
-    public void applyFromCellItem(WidgetItem item, float previewScale,
-            @NonNull Consumer<Bitmap> callback, @Nullable Bitmap cachedPreview) {
-        mPreviewContainerScale = previewScale;
-
+    public void applyFromCellItem(WidgetItem item, @NonNull Consumer<Bitmap> callback,
+            @Nullable Bitmap cachedPreview) {
         Context context = getContext();
         mItem = item;
         mWidgetSize = getWidgetItemSizePx(getContext(), mActivity.getDeviceProfile(), mItem);
+        initPreviewContainerSizeAndScale();
 
         mWidgetName.setText(mItem.label);
         mWidgetName.setContentDescription(
@@ -278,6 +274,17 @@
         }
     }
 
+    private void initPreviewContainerSizeAndScale() {
+        WidgetPreviewContainerSize previewSize = WidgetPreviewContainerSize.Companion.forItem(mItem,
+                mActivity.getDeviceProfile());
+        mPreviewContainerSize = WidgetSizes.getWidgetSizePx(mActivity.getDeviceProfile(),
+                previewSize.spanX, previewSize.spanY);
+
+        float scaleX = (float) mPreviewContainerSize.getWidth() / mWidgetSize.getWidth();
+        float scaleY = (float) mPreviewContainerSize.getHeight() / mWidgetSize.getHeight();
+        mPreviewContainerScale = Math.min(scaleX, scaleY);
+    }
+
     private void setAppWidgetHostViewPreview(
             NavigableAppWidgetHostView appWidgetHostViewPreview,
             LauncherAppWidgetProviderInfo providerInfo,
@@ -384,6 +391,16 @@
     }
 
     /**
+     * Shows or hides the dimensions displayed below each widget.
+     *
+     * @param show a flag that shows the dimensions of the widget if {@code true}, hides it if
+     *             {@code false}.
+     */
+    public void showDimensions(boolean show) {
+        mWidgetDims.setVisibility(show ? VISIBLE : GONE);
+    }
+
+    /**
      * Set whether the app icon, for the app that provides the widget, should be shown next to the
      * title text of the widget.
      *
@@ -448,17 +465,22 @@
     @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
         ViewGroup.LayoutParams containerLp = mWidgetImageContainer.getLayoutParams();
-
-        mAppWidgetHostViewScale = mPreviewContainerScale;
         int maxWidth = MeasureSpec.getSize(widthMeasureSpec);
-        containerLp.width = Math.round(mWidgetSize.getWidth() * mAppWidgetHostViewScale);
+
+        // mPreviewContainerScale ensures the needed scaling with respect to original widget size.
+        mAppWidgetHostViewScale = mPreviewContainerScale;
+        containerLp.width = mPreviewContainerSize.getWidth();
+        containerLp.height = mPreviewContainerSize.getHeight();
+
+        // If we don't have enough available width, scale the preview container to fit.
         if (containerLp.width > maxWidth) {
             containerLp.width = maxWidth;
-            mAppWidgetHostViewScale = (float) containerLp.width / mWidgetSize.getWidth();
+            mAppWidgetHostViewScale = (float) containerLp.width / mPreviewContainerSize.getWidth();
+            containerLp.height = Math.round(
+                    mPreviewContainerSize.getHeight() * mAppWidgetHostViewScale);
         }
-        containerLp.height = Math.round(mWidgetSize.getHeight() * mAppWidgetHostViewScale);
-        // No need to call mWidgetImageContainer.setLayoutParams as we are in measure pass
 
+        // No need to call mWidgetImageContainer.setLayoutParams as we are in measure pass
         super.onMeasure(widthMeasureSpec, heightMeasureSpec);
     }
 
diff --git a/src/com/android/launcher3/widget/WidgetImageView.java b/src/com/android/launcher3/widget/WidgetImageView.java
index 11f4485..f0a23be 100644
--- a/src/com/android/launcher3/widget/WidgetImageView.java
+++ b/src/com/android/launcher3/widget/WidgetImageView.java
@@ -82,15 +82,27 @@
     private void updateDstRectF() {
         float myWidth = getWidth();
         float myHeight = getHeight();
-        float bitmapWidth = mDrawable.getIntrinsicWidth();
+        final float bitmapWidth = mDrawable.getIntrinsicWidth();
+        final float bitmapHeight = mDrawable.getIntrinsicHeight();
+        final float bitmapAspectRatio = bitmapWidth / bitmapHeight;
+        final float containerAspectRatio = myWidth / myHeight;
 
-        final float scale = bitmapWidth > myWidth ? myWidth / bitmapWidth : 1;
-        float scaledWidth = bitmapWidth * scale;
-        float scaledHeight = mDrawable.getIntrinsicHeight() * scale;
+        // Scale by width if image has larger aspect ratio than the container else by height; and
+        // avoid cropping the previews
+        final float scale = bitmapAspectRatio > containerAspectRatio ? myWidth / bitmapWidth
+                : myHeight / bitmapHeight;
 
-        mDstRectF.left = (myWidth - scaledWidth) / 2;
-        mDstRectF.right = (myWidth + scaledWidth) / 2;
+        final float scaledWidth = bitmapWidth * scale;
+        final float scaledHeight = bitmapHeight * scale;
 
+        // Avoid cropping by checking bounds after scaling.
+        if (scaledWidth > myWidth) {
+            mDstRectF.left = 0;
+            mDstRectF.right = scaledWidth;
+        } else {
+            mDstRectF.left = (myWidth - scaledWidth) / 2;
+            mDstRectF.right = (myWidth + scaledWidth) / 2;
+        }
         if (scaledHeight > myHeight) {
             mDstRectF.top = 0;
             mDstRectF.bottom = scaledHeight;
diff --git a/src/com/android/launcher3/widget/WidgetsBottomSheet.java b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
index ceb0072..ab1ad70 100644
--- a/src/com/android/launcher3/widget/WidgetsBottomSheet.java
+++ b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
@@ -16,7 +16,6 @@
 
 package com.android.launcher3.widget;
 
-import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_BOTTOM_WIDGETS_TRAY;
 
 import android.content.Context;
@@ -188,13 +187,9 @@
                 mWidgetCellHorizontalPadding)
                 .forEach(row -> {
                     TableRow tableRow = new TableRow(getContext());
-                    if (enableCategorizedWidgetSuggestions()) {
-                        // Vertically center align items, so that even if they don't fill bounds,
-                        // they can look organized when placed together in a row.
-                        tableRow.setGravity(Gravity.CENTER_VERTICAL);
-                    } else {
-                        tableRow.setGravity(Gravity.TOP);
-                    }
+                    // Vertically center align items, so that even if they don't fill bounds,
+                    // they can look organized when placed together in a row.
+                    tableRow.setGravity(Gravity.CENTER_VERTICAL);
                     row.forEach(widgetItem -> {
                         WidgetCell widget = addItemCell(tableRow);
                         widget.applyFromCellItem(widgetItem);
diff --git a/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
index 68f18ae..0d775c3 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
@@ -36,41 +36,49 @@
             (context, entry) -> entry.mWidgets.stream()
                     .map(item -> item.label).sorted().collect(Collectors.joining(", "));
 
-    private static final BiFunction<Context, WidgetsListHeaderEntry, String> SUBTITLE_DEFAULT =
-            (context, entry) -> {
-                List<WidgetItem> items = entry.mWidgets;
-                int wc = (int) items.stream().filter(item -> item.widgetInfo != null).count();
-                int sc = Math.max(0, items.size() - wc);
+    @Nullable
+    private static String buildWidgetsCountString(Context context, int wc, int sc) {
+        Resources resources = context.getResources();
+        if (wc == 0 && sc == 0) {
+            return null;
+        }
 
-                Resources resources = context.getResources();
-                if (wc == 0 && sc == 0) {
-                    return null;
-                }
-
-                String subtitle;
-                if (wc > 0 && sc > 0) {
-                    String widgetsCount = PluralMessageFormat.getIcuPluralString(context,
-                            R.string.widgets_count, wc);
-                    String shortcutsCount = PluralMessageFormat.getIcuPluralString(context,
-                            R.string.shortcuts_count, sc);
-                    subtitle = resources.getString(R.string.widgets_and_shortcuts_count,
-                            widgetsCount, shortcutsCount);
-                } else if (wc > 0) {
-                    subtitle = PluralMessageFormat.getIcuPluralString(context,
-                            R.string.widgets_count, wc);
-                } else {
-                    subtitle = PluralMessageFormat.getIcuPluralString(context,
-                            R.string.shortcuts_count, sc);
-                }
-                return subtitle;
-            };
+        String subtitle;
+        if (wc > 0 && sc > 0) {
+            String widgetsCount = PluralMessageFormat.getIcuPluralString(context,
+                    R.string.widgets_count, wc);
+            String shortcutsCount = PluralMessageFormat.getIcuPluralString(context,
+                    R.string.shortcuts_count, sc);
+            subtitle = resources.getString(R.string.widgets_and_shortcuts_count,
+                    widgetsCount, shortcutsCount);
+        } else if (wc > 0) {
+            subtitle = PluralMessageFormat.getIcuPluralString(context,
+                    R.string.widgets_count, wc);
+        } else {
+            subtitle = PluralMessageFormat.getIcuPluralString(context,
+                    R.string.shortcuts_count, sc);
+        }
+        return subtitle;
+    }
 
     private final boolean mIsWidgetListShown;
+    /** Selected widgets displayed */
+    private final int mVisibleWidgetsCount;
     private final boolean mIsSearchEntry;
 
     private WidgetsListHeaderEntry(PackageItemInfo pkgItem, String titleSectionName,
+            List<WidgetItem> items, int visibleWidgetsCount,
+            boolean isSearchEntry, boolean isWidgetListShown) {
+        super(pkgItem, titleSectionName, items);
+        mVisibleWidgetsCount = visibleWidgetsCount;
+        mIsSearchEntry = isSearchEntry;
+        mIsWidgetListShown = isWidgetListShown;
+    }
+
+    private WidgetsListHeaderEntry(PackageItemInfo pkgItem, String titleSectionName,
             List<WidgetItem> items, boolean isSearchEntry, boolean isWidgetListShown) {
         super(pkgItem, titleSectionName, items);
+        mVisibleWidgetsCount = (int) items.stream().filter(w -> w.widgetInfo != null).count();
         mIsSearchEntry = isSearchEntry;
         mIsWidgetListShown = isWidgetListShown;
     }
@@ -91,8 +99,13 @@
 
     @Nullable
     public String getSubtitle(Context context) {
-        return mIsSearchEntry
-                ? SUBTITLE_SEARCH.apply(context, this) : SUBTITLE_DEFAULT.apply(context, this);
+        if (mIsSearchEntry) {
+            return SUBTITLE_SEARCH.apply(context, this);
+        } else {
+            int shortcutsCount = Math.max(0,
+                    (int) mWidgets.stream().filter(WidgetItem::isShortcut).count());
+            return buildWidgetsCountString(context, mVisibleWidgetsCount, shortcutsCount);
+        }
     }
 
     @Override
@@ -102,6 +115,7 @@
         return mWidgets.equals(otherEntry.mWidgets) && mPkgItem.equals(otherEntry.mPkgItem)
                 && mTitleSectionName.equals(otherEntry.mTitleSectionName)
                 && mIsWidgetListShown == otherEntry.mIsWidgetListShown
+                && mVisibleWidgetsCount == otherEntry.mVisibleWidgetsCount
                 && mIsSearchEntry == otherEntry.mIsSearchEntry;
     }
 
@@ -112,6 +126,7 @@
                 mPkgItem,
                 mTitleSectionName,
                 mWidgets,
+                mVisibleWidgetsCount,
                 mIsSearchEntry,
                 /* isWidgetListShown= */ true);
     }
@@ -122,7 +137,28 @@
                 pkgItem,
                 titleSectionName,
                 items,
-                /* forSearch */ false,
+                /* isSearchEntry= */ false,
+                /* isWidgetListShown= */ false);
+    }
+
+    /**
+     * Creates a widget list holder for an header ("app" / "suggestions") which has widgets or/and
+     * shortcuts.
+     *
+     * @param pkgItem             package item info for the header section
+     * @param titleSectionName    title string for the header
+     * @param items               all items for the given header
+     * @param visibleWidgetsCount widgets count when only selected widgets are shown due to
+     *                            limited space.
+     */
+    public static WidgetsListHeaderEntry create(PackageItemInfo pkgItem, String titleSectionName,
+            List<WidgetItem> items, int visibleWidgetsCount) {
+        return new WidgetsListHeaderEntry(
+                pkgItem,
+                titleSectionName,
+                items,
+                visibleWidgetsCount,
+                /* isSearchEntry= */ false,
                 /* isWidgetListShown= */ false);
     }
 
@@ -132,7 +168,7 @@
                 pkgItem,
                 titleSectionName,
                 items,
-                /* forSearch */ true,
+                /* isSearchEntry */ true,
                 /* isWidgetListShown= */ false);
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java
index 426a3ae..811759d 100644
--- a/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java
@@ -93,18 +93,19 @@
      * @param availableWidth     width in px that the recommendations should display in
      * @param cellPadding        padding in px that should be applied to each widget in the
      *                           recommendations
-     * @return {@code false} if no recommendations could fit in the available space.
+     * @return number of recommendations that could fit in the available space.
      */
-    public boolean setRecommendations(
+    public int setRecommendations(
             List<WidgetItem> recommendedWidgets, DeviceProfile deviceProfile,
             final @Px float availableHeight, final @Px int availableWidth,
             final @Px int cellPadding) {
         this.mAvailableHeight = availableHeight;
-        removeAllViews();
+        clear();
 
-        maybeDisplayInTable(recommendedWidgets, deviceProfile, availableWidth, cellPadding);
+        int displayedWidgets = maybeDisplayInTable(recommendedWidgets, deviceProfile,
+                availableWidth, cellPadding);
         updateTitleAndIndicator();
-        return getChildCount() > 0;
+        return displayedWidgets;
     }
 
     /**
@@ -118,9 +119,9 @@
      * @param availableWidth  width in px that the recommendations should display in
      * @param cellPadding     padding in px that should be applied to each widget in the
      *                        recommendations
-     * @return {@code false} if no recommendations could fit in the available space.
+     * @return number of recommendations that could fit in the available space.
      */
-    public boolean setRecommendations(
+    public int setRecommendations(
             Map<WidgetRecommendationCategory, List<WidgetItem>> recommendations,
             DeviceProfile deviceProfile,
             final @Px float availableHeight, final @Px int availableWidth,
@@ -128,19 +129,23 @@
         this.mAvailableHeight = availableHeight;
         Context context = getContext();
         mPageIndicator.setPauseScroll(true, deviceProfile.isTwoPanels);
-        removeAllViews();
+        clear();
 
         int displayedCategories = 0;
+        int totalDisplayedWidgets = 0;
 
         // Render top MAX_CATEGORIES in separate tables. Each table becomes a page.
         for (Map.Entry<WidgetRecommendationCategory, List<WidgetItem>> entry :
                 new TreeMap<>(recommendations).entrySet()) {
             // If none of the recommendations for the category could fit in the mAvailableHeight, we
             // don't want to add that category; and we look for the next one.
-            if (maybeDisplayInTable(entry.getValue(), deviceProfile, availableWidth, cellPadding)) {
+            int displayedCount = maybeDisplayInTable(entry.getValue(), deviceProfile,
+                    availableWidth, cellPadding);
+            if (displayedCount > 0) {
                 mCategoryTitles.add(
                         context.getResources().getString(entry.getKey().categoryTitleRes));
                 displayedCategories++;
+                totalDisplayedWidgets += displayedCount;
             }
 
             if (displayedCategories == MAX_CATEGORIES) {
@@ -150,7 +155,12 @@
 
         updateTitleAndIndicator();
         mPageIndicator.setPauseScroll(false, deviceProfile.isTwoPanels);
-        return getChildCount() > 0;
+        return totalDisplayedWidgets;
+    }
+
+    private void clear() {
+        mCategoryTitles.clear();
+        removeAllViews();
     }
 
     /** Displays the page title and paging indicator if there are multiple pages. */
@@ -199,21 +209,8 @@
             for (int i = 0; i < getChildCount(); i++) {
                 View child = getChildAt(i);
                 measureChild(child, widthMeasureSpec, heightMeasureSpec);
-                if (mAvailableHeight == Float.MAX_VALUE) {
-                    // When we are not limited by height, use currentPage's height. This is the case
-                    // when the paged layout is placed in a scrollable container. We cannot use
-                    // height
-                    // of tallest child in such case, as it will display a scrollbar even for
-                    // smaller
-                    // pages that don't have more content.
-                    if (i == mCurrentPage) {
-                        int parentHeight = MeasureSpec.getSize(heightMeasureSpec);
-                        desiredHeight = Math.max(parentHeight, child.getMeasuredHeight());
-                    }
-                } else {
-                    // Use height of tallest child when we are limited to a certain height.
-                    desiredHeight = Math.max(desiredHeight, child.getMeasuredHeight());
-                }
+                // Use height of tallest child as we have limited height.
+                desiredHeight = Math.max(desiredHeight, child.getMeasuredHeight());
             }
 
             int finalHeight = resolveSizeAndState(desiredHeight, heightMeasureSpec, 0);
@@ -228,12 +225,14 @@
      * fits.
      * <p>Returns false if none of the recommendations could fit.</p>
      */
-    private boolean maybeDisplayInTable(List<WidgetItem> recommendedWidgets,
+    private int maybeDisplayInTable(List<WidgetItem> recommendedWidgets,
             DeviceProfile deviceProfile,
             final @Px int availableWidth, final @Px int cellPadding) {
         Context context = getContext();
         LayoutInflater inflater = LayoutInflater.from(context);
 
+        // Since we are limited by space, we don't sort recommendations - to show most relevant
+        // (if possible).
         List<ArrayList<WidgetItem>> rows = groupWidgetItemsUsingRowPxWithoutReordering(
                 recommendedWidgets,
                 context,
@@ -249,13 +248,13 @@
         recommendationsTable.setWidgetCellOnClickListener(mWidgetCellOnClickListener);
         recommendationsTable.setWidgetCellLongClickListener(mWidgetCellOnLongClickListener);
 
-        boolean displayedAtLeastOne = recommendationsTable.setRecommendedWidgets(rows,
+        int displayedCount = recommendationsTable.setRecommendedWidgets(rows,
                 deviceProfile, mAvailableHeight);
-        if (displayedAtLeastOne) {
+        if (displayedCount > 0) {
             addView(recommendationsTable);
         }
 
-        return displayedAtLeastOne;
+        return displayedCount;
     }
 
     /** Returns location of a widget cell for displaying the "touch and hold" education tip. */
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index c0f1070..848f6fa 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -107,7 +107,8 @@
             entry -> mCurrentUser.equals(entry.mPkgItem.user);
     private final Predicate<WidgetsListBaseEntry> mWorkWidgetsFilter;
     protected final boolean mHasWorkProfile;
-    protected boolean mHasRecommendedWidgets;
+    // Number of recommendations displayed
+    protected int mRecommendedWidgetsCount;
     protected final SparseArray<AdapterHolder> mAdapters = new SparseArray();
     @Nullable private ArrowTipView mLatestEducationalTip;
     private final OnLayoutChangeListener mLayoutChangeListenerToShowTips =
@@ -581,7 +582,7 @@
         }
 
         if (enableCategorizedWidgetSuggestions()) {
-            mHasRecommendedWidgets = mWidgetRecommendationsView.setRecommendations(
+            mRecommendedWidgetsCount = mWidgetRecommendationsView.setRecommendations(
                     mActivityContext.getPopupDataProvider().getCategorizedRecommendedWidgets(),
                     mDeviceProfile,
                     /* availableHeight= */ getMaxAvailableHeightForRecommendations(),
@@ -589,7 +590,7 @@
                     /* cellPadding= */ mWidgetCellHorizontalPadding
             );
         } else {
-            mHasRecommendedWidgets = mWidgetRecommendationsView.setRecommendations(
+            mRecommendedWidgetsCount = mWidgetRecommendationsView.setRecommendations(
                     mActivityContext.getPopupDataProvider().getRecommendedWidgets(),
                     mDeviceProfile,
                     /* availableHeight= */ getMaxAvailableHeightForRecommendations(),
@@ -597,7 +598,8 @@
                     /* cellPadding= */ mWidgetCellHorizontalPadding
             );
         }
-        mWidgetRecommendationsContainer.setVisibility(mHasRecommendedWidgets ? VISIBLE : GONE);
+        mWidgetRecommendationsContainer.setVisibility(
+                mRecommendedWidgetsCount > 0 ? VISIBLE : GONE);
     }
 
     @Px
@@ -790,7 +792,7 @@
     }
 
     /** private the height, in pixel, + the vertical margins of a given view. */
-    private static int measureHeightWithVerticalMargins(View view) {
+    protected static int measureHeightWithVerticalMargins(View view) {
         if (view.getVisibility() != VISIBLE) {
             return 0;
         }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
index ef3ccf0..36f8bf9 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -15,8 +15,6 @@
  */
 package com.android.launcher3.widget.picker;
 
-import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
-
 import android.content.Context;
 import android.graphics.Bitmap;
 import android.util.Log;
@@ -121,7 +119,7 @@
                 widget.setVisibility(View.VISIBLE);
 
                 // When preview loads, notify adapter to rebind the item and possibly animate
-                widget.applyFromCellItem(widgetItem, 1f,
+                widget.applyFromCellItem(widgetItem,
                         bitmap -> holder.onPreviewLoaded(Pair.create(widgetItem, bitmap)),
                         holder.previewCache.get(widgetItem));
                 widget.requestLayout();
@@ -150,13 +148,9 @@
                 tableRow = (TableRow) table.getChildAt(i);
             } else {
                 tableRow = new TableRow(table.getContext());
-                if (enableCategorizedWidgetSuggestions()) {
-                    // Vertically center align items, so that even if they don't fill bounds, they
-                    // can look organized when placed together in a row.
-                    tableRow.setGravity(Gravity.CENTER_VERTICAL);
-                } else {
-                    tableRow.setGravity(Gravity.TOP);
-                }
+                // Vertically center align items, so that even if they don't fill bounds, they
+                // can look organized when placed together in a row.
+                tableRow.setGravity(Gravity.CENTER_VERTICAL);
                 table.addView(tableRow);
             }
             if (tableRow.getChildCount() > widgetItems.size()) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
index 12564f4..76b8401 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
@@ -17,11 +17,13 @@
 
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.widget.util.WidgetSizes.getWidgetSizePx;
+import static com.android.launcher3.widget.util.WidgetsTableUtils.WIDGETS_TABLE_ROW_SIZE_COMPARATOR;
+
+import static java.lang.Math.max;
 
 import android.content.Context;
 import android.util.AttributeSet;
-import android.util.Log;
-import android.util.Size;
 import android.view.Gravity;
 import android.view.LayoutInflater;
 import android.view.View;
@@ -30,26 +32,23 @@
 import android.widget.TableRow;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.Px;
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.widget.WidgetCell;
-import com.android.launcher3.widget.util.WidgetSizes;
+import com.android.launcher3.widget.picker.util.WidgetPreviewContainerSize;
 
 import java.util.ArrayList;
 import java.util.List;
 
 /** A {@link TableLayout} for showing recommended widgets. */
 public final class WidgetsRecommendationTableLayout extends TableLayout {
-    private static final String TAG = "WidgetsRecommendationTableLayout";
-    private static final float DOWN_SCALE_RATIO = 0.9f;
-    private static final float MAX_DOWN_SCALE_RATIO = 0.5f;
     private final float mWidgetsRecommendationTableVerticalPadding;
     private final float mWidgetCellVerticalPadding;
     private final float mWidgetCellTextViewsHeight;
 
-    private float mRecommendationTableMaxHeight = Float.MAX_VALUE;
     @Nullable private OnLongClickListener mWidgetCellOnLongClickListener;
     @Nullable private OnClickListener mWidgetCellOnClickListener;
 
@@ -82,47 +81,40 @@
      * desired {@code recommendationTableMaxHeight}.
      *
      * <p>If the content can't fit {@code recommendationTableMaxHeight}, this view will remove a
-     * last row from the {@code recommendedWidgets} until it fits or only one row left. If the only
-     * row still doesn't fit, we scale down the preview image.
+     * last row from the {@code recommendedWidgets} until it fits or only one row left.
      *
      * <p>Returns {@code false} if none of the widgets could fit</p>
      */
-    public boolean setRecommendedWidgets(List<ArrayList<WidgetItem>> recommendedWidgets,
-            DeviceProfile deviceProfile,
-            float recommendationTableMaxHeight) {
-        mRecommendationTableMaxHeight = recommendationTableMaxHeight;
-        RecommendationTableData data = fitRecommendedWidgetsToTableSpace(/* previewScale= */ 1f,
-                deviceProfile,
-                recommendedWidgets);
-        bindData(data);
-        return !data.mRecommendationTable.isEmpty();
+    public int setRecommendedWidgets(List<ArrayList<WidgetItem>> recommendedWidgets,
+            DeviceProfile deviceProfile, float recommendationTableMaxHeight) {
+        List<ArrayList<WidgetItem>> rows = selectRowsThatFitInAvailableHeight(recommendedWidgets,
+                recommendationTableMaxHeight, deviceProfile);
+        bindData(rows);
+        return rows.stream().mapToInt(ArrayList::size).sum();
     }
 
-    private void bindData(RecommendationTableData data) {
-        if (data.mRecommendationTable.isEmpty()) {
+    private void bindData(List<ArrayList<WidgetItem>> recommendationTable) {
+        if (recommendationTable.isEmpty()) {
             setVisibility(GONE);
             return;
         }
 
         removeAllViews();
 
-        for (int i = 0; i < data.mRecommendationTable.size(); i++) {
-            List<WidgetItem> widgetItems = data.mRecommendationTable.get(i);
+        for (int i = 0; i < recommendationTable.size(); i++) {
+            List<WidgetItem> widgetItems = recommendationTable.get(i);
             TableRow tableRow = new TableRow(getContext());
-            if (enableCategorizedWidgetSuggestions()) {
-                // Vertically center align items, so that even if they don't fill bounds, they can
-                // look organized when placed together in a row.
-                tableRow.setGravity(Gravity.CENTER_VERTICAL);
-            } else {
-                tableRow.setGravity(Gravity.TOP);
-            }
+            // Vertically center align items, so that even if they don't fill bounds, they can
+            // look organized when placed together in a row.
+            tableRow.setGravity(Gravity.CENTER_VERTICAL);
             for (WidgetItem widgetItem : widgetItems) {
                 WidgetCell widgetCell = addItemCell(tableRow);
-                widgetCell.applyFromCellItem(widgetItem, data.mPreviewScale);
+                widgetCell.applyFromCellItem(widgetItem);
                 widgetCell.showAppIconInWidgetTitle(true);
                 widgetCell.showBadge();
                 if (enableCategorizedWidgetSuggestions()) {
                     widgetCell.showDescription(false);
+                    widgetCell.showDimensions(false);
                 }
             }
             addView(tableRow);
@@ -144,58 +136,32 @@
         return widget;
     }
 
-    private RecommendationTableData fitRecommendedWidgetsToTableSpace(
-            float previewScale,
-            DeviceProfile deviceProfile,
-            List<ArrayList<WidgetItem>> recommendedWidgetsInTable) {
-        if (previewScale < MAX_DOWN_SCALE_RATIO) {
-            Log.w(TAG, "Hide recommended widgets. Can't down scale previews to " + previewScale);
-            return new RecommendationTableData(List.of(), previewScale);
-        }
+    private List<ArrayList<WidgetItem>> selectRowsThatFitInAvailableHeight(
+            List<ArrayList<WidgetItem>> recommendedWidgets, @Px float recommendationTableMaxHeight,
+            DeviceProfile deviceProfile) {
+        List<ArrayList<WidgetItem>> filteredRows = new ArrayList<>();
         // A naive estimation of the widgets recommendation table height without inflation.
         float totalHeight = mWidgetsRecommendationTableVerticalPadding;
-        for (int i = 0; i < recommendedWidgetsInTable.size(); i++) {
-            List<WidgetItem> widgetItems = recommendedWidgetsInTable.get(i);
+
+        for (int i = 0; i < recommendedWidgets.size(); i++) {
+            List<WidgetItem> widgetItems = recommendedWidgets.get(i);
             float rowHeight = 0;
             for (int j = 0; j < widgetItems.size(); j++) {
                 WidgetItem widgetItem = widgetItems.get(j);
-                Size widgetSize = WidgetSizes.getWidgetItemSizePx(getContext(), deviceProfile,
-                        widgetItem);
-                float previewHeight = widgetSize.getHeight() * previewScale;
-                rowHeight = Math.max(rowHeight,
-                        previewHeight + mWidgetCellTextViewsHeight + mWidgetCellVerticalPadding);
+                WidgetPreviewContainerSize previewContainerSize =
+                        WidgetPreviewContainerSize.Companion.forItem(widgetItem, deviceProfile);
+                float widgetItemHeight = getWidgetSizePx(deviceProfile, previewContainerSize.spanX,
+                        previewContainerSize.spanY).getHeight();
+                rowHeight = max(rowHeight,
+                        widgetItemHeight + mWidgetCellTextViewsHeight + mWidgetCellVerticalPadding);
             }
-            totalHeight += rowHeight;
+            if (totalHeight + rowHeight <= recommendationTableMaxHeight) {
+                totalHeight += rowHeight;
+                filteredRows.add(new ArrayList<>(widgetItems));
+            }
         }
 
-        if (totalHeight < mRecommendationTableMaxHeight) {
-            return new RecommendationTableData(recommendedWidgetsInTable, previewScale);
-        }
-
-        if (recommendedWidgetsInTable.size() > 1) {
-            // We don't want to scale down widgets preview unless we really need to. Reduce the
-            // num of row by 1 to see if it fits.
-            return fitRecommendedWidgetsToTableSpace(
-                    previewScale,
-                    deviceProfile,
-                    recommendedWidgetsInTable.subList(/* fromIndex= */0,
-                            /* toIndex= */recommendedWidgetsInTable.size() - 1));
-        }
-
-        float nextPreviewScale = previewScale * DOWN_SCALE_RATIO;
-        return fitRecommendedWidgetsToTableSpace(nextPreviewScale, deviceProfile,
-                recommendedWidgetsInTable);
-    }
-
-    /** Data class for the widgets recommendation table and widgets preview scaling. */
-    private class RecommendationTableData {
-        private final List<ArrayList<WidgetItem>> mRecommendationTable;
-        private final float mPreviewScale;
-
-        RecommendationTableData(List<ArrayList<WidgetItem>> recommendationTable,
-                float previewScale) {
-            mRecommendationTable = recommendationTable;
-            mPreviewScale = previewScale;
-        }
+        // Perform re-ordering once we have filtered out recommendations that fit.
+        return filteredRows.stream().sorted(WIDGETS_TABLE_ROW_SIZE_COMPARATOR).toList();
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index 165b2fe..c3bb993 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -167,7 +167,7 @@
     @Override
     public void onWidgetsBound() {
         super.onWidgetsBound();
-        if (!mHasRecommendedWidgets && mSelectedHeader == null) {
+        if (mRecommendedWidgetsCount == 0 && mSelectedHeader == null) {
             mAdapters.get(mActivePage).mWidgetsListAdapter.selectFirstHeaderEntry();
             mAdapters.get(mActivePage).mWidgetsRecyclerView.scrollToTop();
         }
@@ -177,7 +177,7 @@
     public void onRecommendedWidgetsBound() {
         super.onRecommendedWidgetsBound();
 
-        if (mSuggestedWidgetsContainer == null && mHasRecommendedWidgets) {
+        if (mSuggestedWidgetsContainer == null && mRecommendedWidgetsCount > 0) {
             setupSuggestedWidgets(LayoutInflater.from(getContext()));
             mSuggestedWidgetsHeader.callOnClick();
         }
@@ -209,8 +209,9 @@
         packageItemInfo.title = suggestionsHeaderTitle;
         WidgetsListHeaderEntry widgetsListHeaderEntry = WidgetsListHeaderEntry.create(
                         packageItemInfo,
-                        suggestionsHeaderTitle,
-                        mActivityContext.getPopupDataProvider().getRecommendedWidgets())
+                        /*titleSectionName=*/ suggestionsHeaderTitle,
+                        /*items=*/ mActivityContext.getPopupDataProvider().getRecommendedWidgets(),
+                        /*visibleWidgetsCount=*/ mRecommendedWidgetsCount)
                 .withWidgetListShown();
 
         mSuggestedWidgetsHeader.applyFromItemInfoWithIcon(widgetsListHeaderEntry);
@@ -233,7 +234,7 @@
     @Override
     @Px
     protected float getMaxTableHeight(@Px float noWidgetsViewHeight) {
-        return Float.MAX_VALUE;
+        return mContent.getMeasuredHeight() - measureHeightWithVerticalMargins(mHeaderTitle);
     }
 
     @Override
diff --git a/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSize.kt b/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSize.kt
new file mode 100644
index 0000000..a0414ba
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSize.kt
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2024 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.widget.picker.util
+
+import com.android.launcher3.DeviceProfile
+import com.android.launcher3.model.WidgetItem
+import kotlin.math.abs
+
+/** Size of a preview container in terms of the grid spans. */
+data class WidgetPreviewContainerSize(@JvmField val spanX: Int, @JvmField val spanY: Int) {
+    companion object {
+        /**
+         * Returns the size of the preview container in which the given widget's preview should be
+         * displayed (by scaling it if necessary).
+         */
+        fun forItem(item: WidgetItem, dp: DeviceProfile): WidgetPreviewContainerSize {
+            val sizes =
+                if (dp.isTablet && !dp.isTwoPanels) {
+                    TABLET_WIDGET_PREVIEW_SIZES
+                } else {
+                    HANDHELD_WIDGET_PREVIEW_SIZES
+                }
+
+            for ((index, containerSize) in sizes.withIndex()) {
+                if (containerSize.spanX == item.spanX && containerSize.spanY == item.spanY) {
+                    return containerSize // Exact match!
+                }
+                if (containerSize.spanX <= item.spanX && containerSize.spanY <= item.spanY) {
+                    return findClosestFittingContainer(
+                        containerSizes = sizes.toList(),
+                        startIndex = index,
+                        item = item
+                    )
+                }
+            }
+            // Use largest container if no match found
+            return sizes.elementAt(0)
+        }
+
+        private fun findClosestFittingContainer(
+            containerSizes: List<WidgetPreviewContainerSize>,
+            startIndex: Int,
+            item: WidgetItem
+        ): WidgetPreviewContainerSize {
+            // Checks if it's a smaller container, but close enough to keep the down-scale minimal.
+            fun hasAcceptableSize(currentIndex: Int): Boolean {
+                val container = containerSizes[currentIndex]
+                val isSmallerThanItem =
+                    container.spanX <= item.spanX && container.spanY <= item.spanY
+                val isCloseToItemSize =
+                    (item.spanY - container.spanY <= 1) && (item.spanX - container.spanX <= 1)
+
+                return isSmallerThanItem && isCloseToItemSize
+            }
+
+            var currentIndex = startIndex
+            var match = containerSizes[currentIndex]
+            val itemCellSizeRatio = item.spanX.toFloat() / item.spanY
+            var lastCellSizeRatioDiff = Float.MAX_VALUE
+
+            // Look for a smaller container (up to an acceptable extent) with closest cell size
+            // ratio.
+            while (currentIndex <= containerSizes.lastIndex && hasAcceptableSize(currentIndex)) {
+                val current = containerSizes[currentIndex]
+                val currentCellSizeRatio = current.spanX.toFloat() / current.spanY
+                val currentCellSizeRatioDiff = abs(itemCellSizeRatio - currentCellSizeRatio)
+
+                if (currentCellSizeRatioDiff < lastCellSizeRatioDiff) {
+                    lastCellSizeRatioDiff = currentCellSizeRatioDiff
+                    match = containerSizes[currentIndex]
+                }
+                currentIndex++
+            }
+            return match
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizes.kt b/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizes.kt
new file mode 100644
index 0000000..a016676
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizes.kt
@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2024 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.widget.picker.util
+
+/**
+ * An ordered list of recommended sizes for the preview containers in handheld devices.
+ *
+ * Size of the preview container in which a widget's preview can be displayed.
+ */
+val HANDHELD_WIDGET_PREVIEW_SIZES: List<WidgetPreviewContainerSize> =
+    listOf(
+        WidgetPreviewContainerSize(spanX = 4, spanY = 3),
+        WidgetPreviewContainerSize(spanX = 4, spanY = 2),
+        WidgetPreviewContainerSize(spanX = 2, spanY = 3),
+        WidgetPreviewContainerSize(spanX = 2, spanY = 2),
+        WidgetPreviewContainerSize(spanX = 4, spanY = 1),
+        WidgetPreviewContainerSize(spanX = 2, spanY = 1),
+        WidgetPreviewContainerSize(spanX = 1, spanY = 1),
+    )
+
+/**
+ * An ordered list of recommended sizes for the preview containers in tablet devices (with larger
+ * grids).
+ *
+ * Size of the preview container in which a widget's preview can be displayed (by scaling the
+ * preview if necessary).
+ */
+val TABLET_WIDGET_PREVIEW_SIZES: List<WidgetPreviewContainerSize> =
+    listOf(
+        WidgetPreviewContainerSize(spanX = 3, spanY = 4),
+        WidgetPreviewContainerSize(spanX = 3, spanY = 3),
+        WidgetPreviewContainerSize(spanX = 3, spanY = 2),
+        WidgetPreviewContainerSize(spanX = 2, spanY = 3),
+        WidgetPreviewContainerSize(spanX = 2, spanY = 2),
+        WidgetPreviewContainerSize(spanX = 3, spanY = 1),
+        WidgetPreviewContainerSize(spanX = 2, spanY = 1),
+        WidgetPreviewContainerSize(spanX = 1, spanY = 1),
+    )
diff --git a/src/com/android/launcher3/widget/util/WidgetsTableUtils.java b/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
index 74d3062..5e0e203 100644
--- a/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
+++ b/src/com/android/launcher3/widget/util/WidgetsTableUtils.java
@@ -16,11 +16,13 @@
 package com.android.launcher3.widget.util;
 
 import android.content.Context;
+import android.util.Size;
 
 import androidx.annotation.Px;
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.widget.picker.util.WidgetPreviewContainerSize;
 
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -33,8 +35,8 @@
     /**
      * Groups widgets in the following order:
      * 1. Widgets always go before shortcuts.
-     * 2. Widgets with smaller horizontal spans will be shown first.
-     * 3. If widgets have the same horizontal spans, then widgets with a smaller vertical spans will
+     * 2. Widgets with smaller vertical spans will be shown first.
+     * 3. If widgets have the same vertical spans, then widgets with a smaller horizontal spans will
      *    go first.
      * 4. If both widgets have the same horizontal and vertical spans, they will use the same order
      *    from the given {@code widgetItems}.
@@ -43,14 +45,29 @@
         if (item.widgetInfo != null && otherItem.widgetInfo == null) return -1;
 
         if (item.widgetInfo == null && otherItem.widgetInfo != null) return 1;
-        if (item.spanX == otherItem.spanX) {
-            if (item.spanY == otherItem.spanY) return 0;
-            return item.spanY > otherItem.spanY ? 1 : -1;
+        if (item.spanY == otherItem.spanY) {
+            if (item.spanX == otherItem.spanX) return 0;
+            return item.spanX > otherItem.spanX ? 1 : -1;
         }
-        return item.spanX > otherItem.spanX ? 1 : -1;
+        return item.spanY > otherItem.spanY ? 1 : -1;
     };
 
     /**
+     * Comparator that enables displaying rows in increasing order of their size (totalW * H);
+     * except for shortcuts which always show at the bottom.
+     */
+    public static final Comparator<ArrayList<WidgetItem>> WIDGETS_TABLE_ROW_SIZE_COMPARATOR =
+            Comparator.comparingInt(row -> {
+                if (row.stream().anyMatch(WidgetItem::isShortcut)) {
+                    return Integer.MAX_VALUE;
+                } else {
+                    int rowWidth = row.stream().mapToInt(w -> w.spanX).sum();
+                    int rowHeight = row.get(0).spanY;
+                    return (rowWidth * rowHeight);
+                }
+            });
+
+    /**
      * Groups {@code widgetItems} items into a 2D array which matches their appearance in a UI
      * table. This takes liberty to rearrange widgets to make the table visually appealing.
      */
@@ -59,72 +76,70 @@
             final @Px int rowPx, final @Px int cellPadding) {
         List<WidgetItem> sortedWidgetItems = widgetItems.stream().sorted(WIDGET_SHORTCUT_COMPARATOR)
                 .collect(Collectors.toList());
-        return groupWidgetItemsUsingRowPxWithoutReordering(sortedWidgetItems, context, dp, rowPx,
+        List<ArrayList<WidgetItem>> rows = groupWidgetItemsUsingRowPxWithoutReordering(
+                sortedWidgetItems, context, dp, rowPx,
                 cellPadding);
+        return rows.stream().sorted(WIDGETS_TABLE_ROW_SIZE_COMPARATOR).toList();
     }
 
     /**
      * Groups {@code widgetItems} into a 2D array which matches their appearance in a UI table while
      * maintaining their order. This function is a variant of
-     * {@code groupWidgetItemsIntoTableWithoutReordering} in that this uses widget pixels for
-     * calculation.
+     * {@code groupWidgetItemsIntoTableWithoutReordering} in that this uses widget container's
+     * pixels for calculation.
      *
      * <p>Grouping:
      * 1. Widgets and shortcuts never group together in the same row.
-     * 2. The ordered widgets are grouped together in the same row until their individual occupying
-     *    pixels exceed the total allowed pixels for the cell.
+     * 2. Widgets are grouped together only if they have same preview container size.
+     * 3. Widgets are grouped together in the same row until the total of individual container sizes
+     *    exceed the total allowed pixels for the row.
      * 3. The ordered shortcuts are grouped together in the same row until their individual
      *    occupying pixels exceed the total allowed pixels for the cell.
      * 4. If there is only one widget in a row, its width may exceed the {@code rowPx}.
      *
-     * <p>Let's say the {@code rowPx} is set to 600 and we have 5 widgets. Widgets can be grouped
-     * in the same row if each of their individual occupying pixels does not exceed
-     * {@code rowPx} / 5 - 2 * {@code cellPadding}.
-     * Example 1: Row 1: 200x200, 200x300, 100x100. Average horizontal pixels is 200 and no widgets
-     * exceed that width. This is okay.
-     * Example 2: Row 1: 200x200, 400x300, 100x100. Average horizontal pixels is 200 and one widget
-     * exceed that width. This is not allowed.
-     * Example 3: Row 1: 700x400. This is okay because this is the only item in the row.
+     * <p>See WidgetTableUtilsTest
      */
     public static List<ArrayList<WidgetItem>> groupWidgetItemsUsingRowPxWithoutReordering(
             List<WidgetItem> widgetItems, Context context, final DeviceProfile dp,
             final @Px int rowPx, final @Px int cellPadding) {
-
         List<ArrayList<WidgetItem>> widgetItemsTable = new ArrayList<>();
         ArrayList<WidgetItem> widgetItemsAtRow = null;
+        // A row displays only items of same container size.
+        WidgetPreviewContainerSize containerSizeForRow = null;
+        @Px int currentRowWidth = 0;
+
         for (WidgetItem widgetItem : widgetItems) {
             if (widgetItemsAtRow == null) {
                 widgetItemsAtRow = new ArrayList<>();
                 widgetItemsTable.add(widgetItemsAtRow);
             }
             int numOfWidgetItems = widgetItemsAtRow.size();
-            @Px int individualSpan = (rowPx / (numOfWidgetItems + 1)) - (2 * cellPadding);
+
+            WidgetPreviewContainerSize containerSize =
+                    WidgetPreviewContainerSize.Companion.forItem(widgetItem, dp);
+            Size containerSizePx = WidgetSizes.getWidgetSizePx(dp, containerSize.spanX,
+                    containerSize.spanY);
+            @Px int containerWidth = containerSizePx.getWidth() + (2 * cellPadding);
+
             if (numOfWidgetItems == 0) {
                 widgetItemsAtRow.add(widgetItem);
-            } else if (
-                    // Since the size of the widget cell is determined by dividing the maximum span
-                    // pixels evenly, making sure that each widget would have enough span pixels to
-                    // show their contents.
-                    widgetItem.hasSameType(widgetItemsAtRow.get(numOfWidgetItems - 1))
-                    && widgetItemsAtRow.stream().allMatch(
-                            item -> WidgetSizes.getWidgetItemSizePx(context, dp, item)
-                                    .getWidth() <= individualSpan)
-                    && WidgetSizes.getWidgetItemSizePx(context, dp, widgetItem)
-                            .getWidth() <= individualSpan) {
+                containerSizeForRow = containerSize;
+                currentRowWidth = containerWidth;
+            } else if ((currentRowWidth + containerWidth) <= rowPx
+                    && widgetItem.hasSameType(widgetItemsAtRow.get(numOfWidgetItems - 1))
+                    && containerSize.equals(containerSizeForRow)) {
                 // Group items in the same row if
                 // 1. they are with the same type, i.e. a row can only have widgets or shortcuts but
                 //    never a mix of both.
-                // 2. Each widget will have horizontal cell span pixels that is at least as large as
-                //    it is required to fit in the horizontal content, unless the widget horizontal
-                //    span pixels is larger than the maximum allowed.
-                //    If an item has horizontal span pixels larger than the maximum allowed pixels
-                //    per row, we just place it in its own row regardless of the horizontal span
-                //    limit.
+                // 2. Each widget in the given row has same preview container size.
                 widgetItemsAtRow.add(widgetItem);
+                currentRowWidth += containerWidth;
             } else {
                 widgetItemsAtRow = new ArrayList<>();
                 widgetItemsTable.add(widgetItemsAtRow);
                 widgetItemsAtRow.add(widgetItem);
+                containerSizeForRow = containerSize;
+                currentRowWidth = containerWidth;
             }
         }
         return widgetItemsTable;
diff --git a/tests/assets/databases/GridMigrationTest/flagged_result5x5to5x8.db b/tests/assets/databases/GridMigrationTest/flagged_result5x5to5x8.db
new file mode 100644
index 0000000..8bea3ce
--- /dev/null
+++ b/tests/assets/databases/GridMigrationTest/flagged_result5x5to5x8.db
Binary files differ
diff --git a/tests/src/com/android/launcher3/celllayout/CellPosMapperTest.java b/tests/multivalentTests/src/com/android/launcher3/celllayout/CellPosMapperTest.java
similarity index 100%
rename from tests/src/com/android/launcher3/celllayout/CellPosMapperTest.java
rename to tests/multivalentTests/src/com/android/launcher3/celllayout/CellPosMapperTest.java
diff --git a/tests/src/com/android/launcher3/logging/FileLogTest.java b/tests/multivalentTests/src/com/android/launcher3/logging/FileLogTest.java
similarity index 100%
rename from tests/src/com/android/launcher3/logging/FileLogTest.java
rename to tests/multivalentTests/src/com/android/launcher3/logging/FileLogTest.java
diff --git a/tests/src/com/android/launcher3/model/data/ItemInfoWithIconTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/data/ItemInfoWithIconTest.kt
similarity index 100%
rename from tests/src/com/android/launcher3/model/data/ItemInfoWithIconTest.kt
rename to tests/multivalentTests/src/com/android/launcher3/model/data/ItemInfoWithIconTest.kt
diff --git a/tests/src/com/android/launcher3/popup/PopupPopulatorTest.java b/tests/multivalentTests/src/com/android/launcher3/popup/PopupPopulatorTest.java
similarity index 100%
rename from tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
rename to tests/multivalentTests/src/com/android/launcher3/popup/PopupPopulatorTest.java
diff --git a/tests/src/com/android/launcher3/util/window/WindowManagerProxyTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/window/WindowManagerProxyTest.kt
similarity index 100%
rename from tests/src/com/android/launcher3/util/window/WindowManagerProxyTest.kt
rename to tests/multivalentTests/src/com/android/launcher3/util/window/WindowManagerProxyTest.kt
diff --git a/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewControllerTest.java b/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewControllerTest.java
index 490cb47..043461d 100644
--- a/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewControllerTest.java
+++ b/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewControllerTest.java
@@ -18,6 +18,7 @@
 
 import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 
+import static com.android.launcher3.allapps.BaseAllAppsAdapter.VIEW_TYPE_PRIVATE_SPACE_HEADER;
 import static com.android.launcher3.allapps.UserProfileManager.STATE_DISABLED;
 import static com.android.launcher3.allapps.UserProfileManager.STATE_ENABLED;
 import static com.android.launcher3.allapps.UserProfileManager.STATE_TRANSITION;
@@ -25,13 +26,19 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.AdditionalAnswers.answer;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
 
+import android.content.ComponentName;
 import android.content.Context;
+import android.content.Intent;
 import android.graphics.Bitmap;
 import android.graphics.Canvas;
 import android.graphics.drawable.BitmapDrawable;
 import android.graphics.drawable.Drawable;
+import android.os.Process;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.widget.ImageButton;
@@ -44,6 +51,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.R;
+import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
 
 import org.junit.Before;
@@ -52,32 +60,53 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.ArrayList;
+import java.util.List;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class PrivateSpaceHeaderViewControllerTest {
 
+    private static final UserHandle MAIN_HANDLE = Process.myUserHandle();
+    private static final UserHandle PRIVATE_HANDLE = new UserHandle(11);
     private static final int CONTAINER_HEADER_ELEMENT_COUNT = 1;
     private static final int LOCK_UNLOCK_BUTTON_COUNT = 1;
     private static final int PS_SETTINGS_BUTTON_COUNT_VISIBLE = 1;
     private static final int PS_SETTINGS_BUTTON_COUNT_INVISIBLE = 0;
     private static final int PS_TRANSITION_IMAGE_COUNT = 1;
+    private static final int NUM_APP_COLS = 4;
+    private static final int NUM_PRIVATE_SPACE_APPS = 50;
+    private static final int ALL_APPS_HEIGHT = 10;
+    private static final int ALL_APPS_CELL_HEIGHT = 1;
+    private static final int PS_HEADER_HEIGHT = 1;
+    private static final int BIGGER_PS_HEADER_HEIGHT = 2;
+    private static final int SCROLL_NO_WHERE = -1;
+    private static final float HEADER_PROTECTION_HEIGHT = 1F;
 
     private Context mContext;
     private PrivateSpaceHeaderViewController mPsHeaderViewController;
     private RelativeLayout mPsHeaderLayout;
+    private AlphabeticalAppsList<?> mAlphabeticalAppsList;
     @Mock
     private PrivateProfileManager mPrivateProfileManager;
     @Mock
     private ActivityAllAppsContainerView mAllApps;
+    @Mock
+    private AllAppsStore<?> mAllAppsStore;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mContext = new ActivityContextWrapper(getApplicationContext());
+        when(mPrivateProfileManager.getItemInfoMatcher()).thenReturn(info ->
+                info != null && info.user.equals(PRIVATE_HANDLE));
         mPsHeaderViewController = new PrivateSpaceHeaderViewController(mAllApps,
                 mPrivateProfileManager);
         mPsHeaderLayout = (RelativeLayout) LayoutInflater.from(mContext).inflate(
                 R.layout.private_space_header, null);
+        mAlphabeticalAppsList = new AlphabeticalAppsList<>(mContext, mAllAppsStore,
+                null, mPrivateProfileManager);
+        mAlphabeticalAppsList.setNumAppsPerRowAllApps(NUM_APP_COLS);
     }
 
     @Test
@@ -223,6 +252,88 @@
         assertEquals(PS_TRANSITION_IMAGE_COUNT, totalLockUnlockButtonView);
     }
 
+    @Test
+    public void scrollForViewToBeVisibleInContainer_withHeader() {
+        when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
+        when(mPrivateProfileManager.addPrivateSpaceHeader(any()))
+                .thenAnswer(answer(this::addPrivateSpaceHeader));
+        when(mPrivateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
+        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps())
+                .thenReturn(iteminfo -> iteminfo.componentName == null
+                        || !iteminfo.componentName.getPackageName()
+                        .equals("com.android.launcher3.tests.camera"));
+        when(mAllApps.getContext()).thenReturn(mContext);
+        mAlphabeticalAppsList.updateItemFilter(info -> info != null
+                && info.user.equals(MAIN_HANDLE));
+        when(mAllApps.getHeight()).thenReturn(ALL_APPS_HEIGHT);
+        when(mAllApps.getHeaderProtectionHeight()).thenReturn(HEADER_PROTECTION_HEIGHT);
+        int rows = (int) (ALL_APPS_HEIGHT - PS_HEADER_HEIGHT - HEADER_PROTECTION_HEIGHT);
+        int position = rows * NUM_APP_COLS - (NUM_APP_COLS-1) + 1;
+
+        // The number of adapterItems should be the private space apps + one main app + header.
+        assertEquals(NUM_PRIVATE_SPACE_APPS + 1 + 1,
+                mAlphabeticalAppsList.getAdapterItems().size());
+        assertEquals(position,
+                mPsHeaderViewController.scrollForViewToBeVisibleInContainer(
+                        new AllAppsRecyclerView(mContext),
+                        mAlphabeticalAppsList.getAdapterItems(),
+                        PS_HEADER_HEIGHT,
+                        ALL_APPS_CELL_HEIGHT));
+    }
+
+    @Test
+    public void scrollForViewToBeVisibleInContainer_withHeaderAndLessAppRowSpace() {
+        when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
+        when(mPrivateProfileManager.addPrivateSpaceHeader(any()))
+                .thenAnswer(answer(this::addPrivateSpaceHeader));
+        when(mPrivateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
+        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps())
+                .thenReturn(iteminfo -> iteminfo.componentName == null
+                        || !iteminfo.componentName.getPackageName()
+                        .equals("com.android.launcher3.tests.camera"));
+        when(mAllApps.getContext()).thenReturn(mContext);
+        mAlphabeticalAppsList.updateItemFilter(info -> info != null
+                && info.user.equals(MAIN_HANDLE));
+        when(mAllApps.getHeight()).thenReturn(ALL_APPS_HEIGHT);
+        when(mAllApps.getHeaderProtectionHeight()).thenReturn(HEADER_PROTECTION_HEIGHT);
+        int rows = (int) (ALL_APPS_HEIGHT - BIGGER_PS_HEADER_HEIGHT - HEADER_PROTECTION_HEIGHT);
+        int position = rows * NUM_APP_COLS - (NUM_APP_COLS-1) + 1;
+
+        // The number of adapterItems should be the private space apps + one main app + header.
+        assertEquals(NUM_PRIVATE_SPACE_APPS + 1 + 1,
+                mAlphabeticalAppsList.getAdapterItems().size());
+        assertEquals(position,
+                mPsHeaderViewController.scrollForViewToBeVisibleInContainer(
+                        new AllAppsRecyclerView(mContext),
+                        mAlphabeticalAppsList.getAdapterItems(),
+                        BIGGER_PS_HEADER_HEIGHT,
+                        ALL_APPS_CELL_HEIGHT));
+    }
+
+    @Test
+    public void scrollForViewToBeVisibleInContainer_withNoHeader() {
+        when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
+        when(mPrivateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
+        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps())
+                .thenReturn(iteminfo -> iteminfo.componentName == null
+                        || !iteminfo.componentName.getPackageName()
+                        .equals("com.android.launcher3.tests.camera"));
+        when(mAllApps.getContext()).thenReturn(mContext);
+        mAlphabeticalAppsList.updateItemFilter(info -> info != null
+                && info.user.equals(MAIN_HANDLE));
+        when(mAllApps.getHeight()).thenReturn(ALL_APPS_HEIGHT);
+        when(mAllApps.getHeaderProtectionHeight()).thenReturn(HEADER_PROTECTION_HEIGHT);
+
+        // The number of adapterItems should be the private space apps + one main app.
+        assertEquals(NUM_PRIVATE_SPACE_APPS + 1,
+                mAlphabeticalAppsList.getAdapterItems().size());
+        assertEquals(SCROLL_NO_WHERE, mPsHeaderViewController.scrollForViewToBeVisibleInContainer(
+                new AllAppsRecyclerView(mContext),
+                mAlphabeticalAppsList.getAdapterItems(),
+                BIGGER_PS_HEADER_HEIGHT,
+                ALL_APPS_CELL_HEIGHT));
+    }
+
     private Bitmap getBitmap(Drawable drawable) {
         Bitmap result;
         if (drawable instanceof BitmapDrawable) {
@@ -249,4 +360,28 @@
     private static void awaitTasksCompleted() throws Exception {
         UI_HELPER_EXECUTOR.submit(() -> null).get();
     }
+
+    private int addPrivateSpaceHeader(List<BaseAllAppsAdapter.AdapterItem> adapterItemList) {
+        BaseAllAppsAdapter.AdapterItem privateSpaceHeader =
+                new BaseAllAppsAdapter.AdapterItem(VIEW_TYPE_PRIVATE_SPACE_HEADER);
+        adapterItemList.add(privateSpaceHeader);
+        return adapterItemList.size();
+    }
+
+    private AppInfo[] createAppInfoList() {
+        List<AppInfo> appInfos = new ArrayList<>();
+        ComponentName gmailComponentName = new ComponentName(mContext,
+                "com.android.launcher3.tests.Activity" + "Gmail");
+        AppInfo gmailAppInfo = new
+                AppInfo(gmailComponentName, "Gmail", MAIN_HANDLE, new Intent());
+        appInfos.add(gmailAppInfo);
+        ComponentName privateCameraComponentName = new ComponentName(
+                "com.android.launcher3.tests.camera", "CameraActivity");
+        for (int i = 0; i < NUM_PRIVATE_SPACE_APPS; i++) {
+            AppInfo privateCameraAppInfo = new AppInfo(privateCameraComponentName,
+                    "Private Camera " + i, PRIVATE_HANDLE, new Intent());
+            appInfos.add(privateCameraAppInfo);
+        }
+        return appInfos.toArray(AppInfo[]::new);
+    }
 }
diff --git a/tests/src/com/android/launcher3/model/GridMigrationTest.kt b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
index eb8604e..da4a208 100644
--- a/tests/src/com/android/launcher3/model/GridMigrationTest.kt
+++ b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
@@ -16,14 +16,18 @@
 
 package com.android.launcher3.model
 
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.Flags
 import com.android.launcher3.InvariantDeviceProfile.TYPE_PHONE
 import com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME
 import com.android.launcher3.celllayout.board.CellLayoutBoard
 import com.android.launcher3.pm.UserCache
 import com.android.launcher3.util.rule.TestToPhoneFileCopier
+import com.android.launcher3.util.rule.setFlags
+import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -65,11 +69,24 @@
 class GridMigrationTest {
     private val DB_FILE = "test_launcher.db"
 
+    @JvmField
+    @Rule
+    val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+
     // Copying the src db for all tests.
     @JvmField
     @Rule
     val fileCopier =
-        TestToPhoneFileCopier("databases/GridMigrationTest/$DB_FILE", "databases/$DB_FILE", true)
+        TestToPhoneFileCopier(
+            src = "databases/GridMigrationTest/$DB_FILE",
+            dest = "databases/$DB_FILE",
+            removeOnFinish = true
+        )
+
+    @Before
+    fun setup() {
+        setFlagsRule.setFlags(false, Flags.FLAG_GRID_MIGRATION_FIX)
+    }
 
     private fun migrate(src: GridMigrationData, dst: GridMigrationData) {
         GridSizeMigrationUtil.migrateGridIfNeeded(
@@ -86,8 +103,10 @@
      * same space in the db.
      */
     private fun validateDb(data: GridMigrationData) {
-        val cellLayoutBoard = CellLayoutBoard(data.gridState.columns, data.gridState.rows)
+        // The array size is just a big enough number to fit all the number of workspaces
+        val boards = Array(100) { CellLayoutBoard(data.gridState.columns, data.gridState.rows) }
         data.readEntries().forEach {
+            val cellLayoutBoard = boards[it.screenId]
             assert(cellLayoutBoard.isEmpty(it.cellX, it.cellY, it.spanX, it.spanY)) {
                 "Db has overlapping items"
             }
@@ -96,13 +115,13 @@
     }
 
     private fun compare(dst: GridMigrationData, target: GridMigrationData) {
-        val sortX = { it: GridSizeMigrationUtil.DbEntry -> it.cellX }
-        val sortY = { it: GridSizeMigrationUtil.DbEntry -> it.cellX }
+        val sort = compareBy<GridSizeMigrationUtil.DbEntry>({ it.cellX }, { it.cellY })
         val mapF = { it: GridSizeMigrationUtil.DbEntry ->
             EntryData(it.cellX, it.cellY, it.spanX, it.spanY, it.rank)
         }
-        val entriesDst = dst.readEntries().sortedBy(sortX).sortedBy(sortY).map(mapF)
-        val entriesTarget = target.readEntries().sortedBy(sortX).sortedBy(sortY).map(mapF)
+        val entriesDst = dst.readEntries().sortedWith(sort).map(mapF)
+        val entriesTarget = target.readEntries().sortedWith(sort).map(mapF)
+
         assert(entriesDst == entriesTarget) {
             "The elements on the dst database is not the same as in the target"
         }
@@ -128,9 +147,9 @@
     @Rule
     val result5x5to3x3 =
         TestToPhoneFileCopier(
-            "databases/GridMigrationTest/result5x5to3x3.db",
-            "databases/result5x5to3x3.db",
-            true
+            src = "databases/GridMigrationTest/result5x5to3x3.db",
+            dest = "databases/result5x5to3x3.db",
+            removeOnFinish = true
         )
 
     @Test
@@ -151,9 +170,9 @@
     @Rule
     val result5x5to4x7 =
         TestToPhoneFileCopier(
-            "databases/GridMigrationTest/result5x5to4x7.db",
-            "databases/result5x5to4x7.db",
-            true
+            src = "databases/GridMigrationTest/result5x5to4x7.db",
+            dest = "databases/result5x5to4x7.db",
+            removeOnFinish = true
         )
 
     @Test
@@ -174,9 +193,9 @@
     @Rule
     val result5x5to5x8 =
         TestToPhoneFileCopier(
-            "databases/GridMigrationTest/result5x5to5x8.db",
-            "databases/result5x5to5x8.db",
-            true
+            src = "databases/GridMigrationTest/result5x5to5x8.db",
+            dest = "databases/result5x5to5x8.db",
+            removeOnFinish = true
         )
 
     @Test
@@ -192,4 +211,32 @@
             target =
                 GridMigrationData("result5x5to5x8.db", DeviceGridState(5, 8, 5, TYPE_PHONE, ""))
         )
+
+    @JvmField
+    @Rule
+    val flaggedResult5x5to5x8 =
+        TestToPhoneFileCopier(
+            src = "databases/GridMigrationTest/flagged_result5x5to5x8.db",
+            dest = "databases/flagged_result5x5to5x8.db",
+            removeOnFinish = true
+        )
+
+    @Test
+    fun `flagged 5x5 to 5x8`() {
+        setFlagsRule.setFlags(true, Flags.FLAG_GRID_MIGRATION_FIX)
+        runTest(
+            src = GridMigrationData(DB_FILE, DeviceGridState(5, 5, 5, TYPE_PHONE, DB_FILE)),
+            dst =
+                GridMigrationData(
+                    null, // in memory db, to download a new db change null for the filename of the
+                    // db name to store it. Do not use existing names.
+                    DeviceGridState(5, 8, 5, TYPE_PHONE, "")
+                ),
+            target =
+                GridMigrationData(
+                    "flagged_result5x5to5x8.db",
+                    DeviceGridState(5, 8, 5, TYPE_PHONE, "")
+                )
+        )
+    }
 }
diff --git a/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt b/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt
index 72c4f16..d3516d1 100644
--- a/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt
+++ b/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt
@@ -49,7 +49,11 @@
         object : Statement() {
             override fun evaluate() {
                 before()
-                base.evaluate()
+                try {
+                    base.evaluate()
+                } finally {
+                    after()
+                }
             }
         }
 }
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetImageViewTest.kt b/tests/src/com/android/launcher3/widget/picker/WidgetImageViewTest.kt
new file mode 100644
index 0000000..6e751e0
--- /dev/null
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetImageViewTest.kt
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2024 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.widget.picker
+
+import android.content.Context
+import android.graphics.Rect
+import android.graphics.drawable.Drawable
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.MediumTest
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.widget.WidgetImageView
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.spy
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.whenever
+
+@MediumTest
+@RunWith(AndroidJUnit4::class)
+class WidgetImageViewTest {
+    private lateinit var context: Context
+    private lateinit var widgetImageView: WidgetImageView
+
+    @Mock private lateinit var testDrawable: Drawable
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        context = ActivityContextWrapper(ApplicationProvider.getApplicationContext())
+        widgetImageView = spy(WidgetImageView(context))
+    }
+
+    @Test
+    fun getBitmapBounds_aspectRatioLargerThanView_scaledByWidth() {
+        // view - 100 x 100
+        whenever(widgetImageView.width).thenReturn(100)
+        whenever(widgetImageView.height).thenReturn(100)
+        // bitmap - 200 x 100
+        whenever(testDrawable.intrinsicWidth).thenReturn(200)
+        whenever(testDrawable.intrinsicHeight).thenReturn(100)
+
+        widgetImageView.drawable = testDrawable
+        val bitmapBounds = widgetImageView.bitmapBounds
+
+        // new scaled width of bitmap is = 100, and height is scaled to 1/2 = 50
+        assertThat(bitmapBounds).isEqualTo(Rect(0, 25, 100, 75))
+    }
+
+    @Test
+    fun getBitmapBounds_aspectRatioSmallerThanView_scaledByHeight() {
+        // view - 100 x 100
+        whenever(widgetImageView.width).thenReturn(100)
+        whenever(widgetImageView.height).thenReturn(100)
+        // bitmap - 100 x 200
+        whenever(testDrawable.intrinsicWidth).thenReturn(100)
+        whenever(testDrawable.intrinsicHeight).thenReturn(200)
+        widgetImageView.drawable = testDrawable
+
+        val bitmapBounds = widgetImageView.bitmapBounds
+
+        // new scaled height of bitmap is = 100, and width is scaled to 1/2 = 50
+        assertThat(bitmapBounds).isEqualTo(Rect(25, 0, 75, 100))
+    }
+
+    @Test
+    fun getBitmapBounds_noScale_returnsOriginalDrawableBounds() {
+        // view - 200 x 100
+        whenever(widgetImageView.width).thenReturn(200)
+        whenever(widgetImageView.height).thenReturn(100)
+        // bitmap - 200 x 100
+        whenever(testDrawable.intrinsicWidth).thenReturn(200)
+        whenever(testDrawable.intrinsicHeight).thenReturn(100)
+
+        widgetImageView.drawable = testDrawable
+        val bitmapBounds = widgetImageView.bitmapBounds
+
+        // no scaling
+        assertThat(bitmapBounds).isEqualTo(Rect(0, 0, 200, 100))
+    }
+}
diff --git a/tests/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizesTest.kt b/tests/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizesTest.kt
new file mode 100644
index 0000000..040fbf5
--- /dev/null
+++ b/tests/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizesTest.kt
@@ -0,0 +1,154 @@
+/*
+ * Copyright (C) 2024 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.widget.picker.util
+
+import android.content.ComponentName
+import android.content.Context
+import android.graphics.Point
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.DeviceProfile
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.WidgetItem
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
+import com.google.common.truth.Truth.assertWithMessage
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class WidgetPreviewContainerSizesTest {
+    private lateinit var context: Context
+    private lateinit var deviceProfile: DeviceProfile
+    private lateinit var testInvariantProfile: InvariantDeviceProfile
+
+    @Mock private lateinit var iconCache: IconCache
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        context = ActivityContextWrapper(ApplicationProvider.getApplicationContext())
+        testInvariantProfile = LauncherAppState.getIDP(context)
+        deviceProfile = testInvariantProfile.getDeviceProfile(context).copy(context)
+    }
+
+    @Test
+    fun widgetPreviewContainerSize_forItem_returnsCorrectContainerSize() {
+        val testSizes = getTestSizes(deviceProfile)
+        val expectedPreviewContainers = testSizes.values.toList()
+
+        for ((index, widgetSize) in testSizes.keys.withIndex()) {
+            val widgetItem = createWidgetItem(widgetSize, context, testInvariantProfile, iconCache)
+
+            assertWithMessage("size for $widgetSize should be: ${expectedPreviewContainers[index]}")
+                .that(WidgetPreviewContainerSize.forItem(widgetItem, deviceProfile))
+                .isEqualTo(expectedPreviewContainers[index])
+        }
+    }
+
+    companion object {
+        private const val TEST_PACKAGE = "com.google.test"
+
+        private val HANDHELD_TEST_SIZES: Map<Point, WidgetPreviewContainerSize> =
+            mapOf(
+                // 1x1
+                Point(1, 1) to WidgetPreviewContainerSize(1, 1),
+                // 2x1
+                Point(2, 1) to WidgetPreviewContainerSize(2, 1),
+                Point(3, 1) to WidgetPreviewContainerSize(2, 1),
+                // 4x1
+                Point(4, 1) to WidgetPreviewContainerSize(4, 1),
+                // 2x2
+                Point(2, 2) to WidgetPreviewContainerSize(2, 2),
+                Point(3, 3) to WidgetPreviewContainerSize(2, 2),
+                Point(3, 2) to WidgetPreviewContainerSize(2, 2),
+                // 2x3
+                Point(2, 3) to WidgetPreviewContainerSize(2, 3),
+                Point(3, 4) to WidgetPreviewContainerSize(2, 3),
+                Point(3, 5) to WidgetPreviewContainerSize(2, 3),
+                // 4x2
+                Point(4, 2) to WidgetPreviewContainerSize(4, 2),
+                // 4x3
+                Point(4, 3) to WidgetPreviewContainerSize(4, 3),
+                Point(4, 4) to WidgetPreviewContainerSize(4, 3),
+            )
+
+        private val TABLET_TEST_SIZES: Map<Point, WidgetPreviewContainerSize> =
+            mapOf(
+                // 1x1
+                Point(1, 1) to WidgetPreviewContainerSize(1, 1),
+                // 2x1
+                Point(2, 1) to WidgetPreviewContainerSize(2, 1),
+                // 3x1
+                Point(3, 1) to WidgetPreviewContainerSize(3, 1),
+                Point(4, 1) to WidgetPreviewContainerSize(3, 1),
+                // 2x2
+                Point(2, 2) to WidgetPreviewContainerSize(2, 2),
+                // 2x3
+                Point(2, 3) to WidgetPreviewContainerSize(2, 3),
+                // 3x2
+                Point(3, 2) to WidgetPreviewContainerSize(3, 2),
+                Point(4, 2) to WidgetPreviewContainerSize(3, 2),
+                Point(5, 2) to WidgetPreviewContainerSize(3, 2),
+                // 3x3
+                Point(3, 3) to WidgetPreviewContainerSize(3, 3),
+                Point(4, 4) to WidgetPreviewContainerSize(3, 3),
+                // 3x4
+                Point(5, 4) to WidgetPreviewContainerSize(3, 4),
+                Point(3, 4) to WidgetPreviewContainerSize(3, 4),
+                Point(5, 5) to WidgetPreviewContainerSize(3, 4),
+                Point(6, 4) to WidgetPreviewContainerSize(3, 4),
+                Point(6, 5) to WidgetPreviewContainerSize(3, 4),
+            )
+
+        private fun getTestSizes(dp: DeviceProfile) =
+            if (dp.isTablet && !dp.isTwoPanels) {
+                TABLET_TEST_SIZES
+            } else {
+                HANDHELD_TEST_SIZES
+            }
+
+        private fun createWidgetItem(
+            widgetSize: Point,
+            context: Context,
+            invariantDeviceProfile: InvariantDeviceProfile,
+            iconCache: IconCache
+        ): WidgetItem {
+            val providerInfo =
+                createAppWidgetProviderInfo(
+                    ComponentName.createRelative(
+                        TEST_PACKAGE,
+                        /*cls=*/ ".WidgetProvider_" + widgetSize.x + "x" + widgetSize.y
+                    )
+                )
+            val widgetInfo =
+                LauncherAppWidgetProviderInfo.fromProviderInfo(context, providerInfo).apply {
+                    spanX = widgetSize.x
+                    spanY = widgetSize.y
+                }
+            return WidgetItem(widgetInfo, invariantDeviceProfile, iconCache, context)
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java b/tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
index 2c5a396..b2cb266 100644
--- a/tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
@@ -63,6 +63,7 @@
     private static final String TEST_PACKAGE = "com.google.test";
 
     private static final int SPACE_SIZE = 10;
+    // Note - actual widget size includes SPACE_SIZE (border) + cell padding.
     private static final int CELL_SIZE = 50;
     private static final int NUM_OF_COLS = 5;
     private static final int NUM_OF_ROWS = 5;
@@ -105,7 +106,7 @@
 
 
     @Test
-    public void groupWidgetItemsIntoTableWithReordering_widgetsOnly_maxSpanPxPerRow220_cellPadding0_shouldGroupWidgetsInTable() {
+    public void groupWithReordering_widgetsOnly_maxSpanPxPerRow220_cellPadding0() {
         List<WidgetItem> widgetItems = List.of(mWidget4x4, mWidget2x3, mWidget1x1, mWidget2x4,
                 mWidget2x2);
 
@@ -113,17 +114,20 @@
                 WidgetsTableUtils.groupWidgetItemsUsingRowPxWithReordering(widgetItems, mContext,
                         mTestDeviceProfile, 220, 0);
 
-        // Row 0: 1x1(50px), 2x2(110px)
-        // Row 1: 2x3(110px), 2x4(110px)
-        // Row 2: 4x4(230px)
-        assertThat(widgetItemInTable).hasSize(3);
-        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1, mWidget2x2);
-        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x3, mWidget2x4);
-        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget4x4);
+        // With reordering, rows displayed in order of increasing size.
+        // Row 0: 1x1(50px)
+        // Row 1: 2x2(in a 2x2 container - 110px)
+        // Row 2: 2x3(in a 2x3 container - 110px), 2x4(in a 2x3 container - 110px)
+        // Row 3: 4x4(in a 3x3 container in tablet - 170px; 4x3 on phone - 230px)
+        assertThat(widgetItemInTable).hasSize(4);
+        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x2);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget2x3, mWidget2x4);
+        assertThat(widgetItemInTable.get(3)).containsExactly(mWidget4x4);
     }
 
     @Test
-    public void groupWidgetItemsIntoTableWithReordering_widgetsOnly_maxSpanPxPerRow220_cellPadding10_shouldGroupWidgetsInTable() {
+    public void groupWithReordering_widgetsOnly_maxSpanPxPerRow220_cellPadding10() {
         List<WidgetItem> widgetItems = List.of(mWidget4x4, mWidget2x3, mWidget1x1, mWidget2x4,
                 mWidget2x2);
 
@@ -131,9 +135,13 @@
                 WidgetsTableUtils.groupWidgetItemsUsingRowPxWithReordering(widgetItems, mContext,
                         mTestDeviceProfile, 220, 10);
 
-        // Row 0: 1x1(50px), 2x2(110px)
-        // Row 1: 2x3(110px), 2x4(110px)
-        // Row 2: 4x4(230px)
+        // With reordering, but space taken up by cell padding, so, no grouping (even if 2x2 and 2x3
+        // use same preview container).
+        // Row 0: 1x1(50px)
+        // Row 1: 2x2(in a 2x2 container: 130px)
+        // Row 2: 2x3(in a 2x3 container: 130px)
+        // Row 3: 2x4(in a 2x3 container: 130px)
+        // Row 4: 4x4(in a 3x3 container in tablet - 190px; 4x3 on phone - 250px)
         assertThat(widgetItemInTable).hasSize(5);
         assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1);
         assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x2);
@@ -143,7 +151,29 @@
     }
 
     @Test
-    public void groupWidgetItemsIntoTableWithReordering_widgetsOnly_maxSpanPxPerRow350_cellPadding0_shouldGroupWidgetsInTable() {
+    public void groupWithReordering_widgetsOnly_maxSpanPxPerRow260_cellPadding10() {
+        List<WidgetItem> widgetItems = List.of(mWidget4x4, mWidget2x3, mWidget1x1, mWidget2x4,
+                mWidget2x2);
+
+        List<ArrayList<WidgetItem>> widgetItemInTable =
+                WidgetsTableUtils.groupWidgetItemsUsingRowPxWithReordering(widgetItems, mContext,
+                        mTestDeviceProfile, 260, 10);
+
+        // With reordering, even with cellPadding, enough space to group 2x3 and 2x4 (which also use
+        // same container)
+        // Row 0: 1x1(50px)
+        // Row 1: 2x2(in a 2x2 container: 130px)
+        // Row 2: 2x3(in a 2x3 container: 130px), 2x4(in a 2x3 container: 130px)
+        // Row 3: 4x4(in a 3x3 container in tablet - 190px; 4x3 on phone - 250px)
+        assertThat(widgetItemInTable).hasSize(4);
+        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x2);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget2x3, mWidget2x4);
+        assertThat(widgetItemInTable.get(3)).containsExactly(mWidget4x4);
+    }
+
+    @Test
+    public void groupWithReordering_widgetsOnly_maxSpanPxPerRow350_cellPadding0() {
         List<WidgetItem> widgetItems = List.of(mWidget4x4, mWidget2x3, mWidget1x1, mWidget2x4,
                 mWidget2x2);
 
@@ -151,17 +181,20 @@
                 WidgetsTableUtils.groupWidgetItemsUsingRowPxWithReordering(widgetItems, mContext,
                         mTestDeviceProfile, 350, 0);
 
-        // Row 0: 1x1(50px), 2x2(110px), 2x3(110px)
-        // Row 1: 2x4(110px)
-        // Row 2: 4x4(230px)
-        assertThat(widgetItemInTable).hasSize(3);
-        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1, mWidget2x2, mWidget2x3);
-        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x4);
-        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget4x4);
+        // With reordering, rows displayed in order of increasing size.
+        // Row 0: 1x1(50px)
+        // Row 1: 2x2(in a 2x2 container: 110px)
+        // Row 2: 2x3(in a 2x3 container: 110px), 2x4(in a 2x3 container: 110px)
+        // Row 3: 4x4(in a 3x3 container in tablet - 170px; 4x3 on phone - 230px)
+        assertThat(widgetItemInTable).hasSize(4);
+        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x2);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget2x3, mWidget2x4);
+        assertThat(widgetItemInTable.get(3)).containsExactly(mWidget4x4);
     }
 
     @Test
-    public void groupWidgetItemsIntoTableWithReordering_mixItems_maxSpanPxPerRow350_cellPadding0_shouldGroupWidgetsInTable() {
+    public void groupWithReordering_mixItems_maxSpanPxPerRow350_cellPadding0() {
         List<WidgetItem> widgetItems = List.of(mWidget4x4, mShortcut3, mWidget2x3, mShortcut1,
                 mWidget1x1, mShortcut2, mWidget2x4, mWidget2x2);
 
@@ -169,19 +202,22 @@
                 WidgetsTableUtils.groupWidgetItemsUsingRowPxWithReordering(widgetItems, mContext,
                         mTestDeviceProfile, 350, 0);
 
-        // Row 0: 1x1(50px), 2x2(110px), 2x3(110px)
-        // Row 1: 2x4(110px),
-        // Row 2: 4x4(230px)
-        // Row 3: shortcut3(50px), shortcut1(50px), shortcut2(50px)
-        assertThat(widgetItemInTable).hasSize(4);
-        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1, mWidget2x2, mWidget2x3);
-        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x4);
-        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget4x4);
-        assertThat(widgetItemInTable.get(3)).containsExactly(mShortcut3, mShortcut2, mShortcut1);
+        // With reordering - rows displays in order of increasing size:
+        // Row 0: 1x1(50px)
+        // Row 1: 2x2(110px)
+        // Row 2: 2x3 (in a 2x3 container 110px), 2x4 (in a 2x3 container 110px)
+        // Row 3: 4x4 (in a 3x3 container in tablet - 170px; 4x3 on phone - 230px)
+        // Row 4: shortcut3, shortcut1, shortcut2 (shortcuts are always displayed at bottom)
+        assertThat(widgetItemInTable).hasSize(5);
+        assertThat(widgetItemInTable.get(0)).containsExactly(mWidget1x1);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x2);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget2x3, mWidget2x4);
+        assertThat(widgetItemInTable.get(3)).containsExactly(mWidget4x4);
+        assertThat(widgetItemInTable.get(4)).containsExactly(mShortcut3, mShortcut2, mShortcut1);
     }
 
     @Test
-    public void groupWidgetItemsIntoTableWithoutReordering_maxSpanPxPerRow220_cellPadding0_shouldMaintainTheOrder() {
+    public void groupWithoutReordering_maxSpanPxPerRow220_cellPadding0() {
         List<WidgetItem> widgetItems =
                 List.of(mWidget4x4, mWidget2x3, mWidget1x1, mWidget2x4, mWidget2x2);
 
@@ -189,13 +225,19 @@
                 WidgetsTableUtils.groupWidgetItemsUsingRowPxWithoutReordering(widgetItems, mContext,
                         mTestDeviceProfile, 220, 0);
 
-        // Row 0: 4x4(230px)
-        // Row 1: 2x3(110px), 1x1(50px)
-        // Row 2: 2x4(110px), 2x2(110px)
-        assertThat(widgetItemInTable).hasSize(3);
+        // Without reordering, widgets are grouped only if the next one fits and uses same preview
+        // container:
+        // Row 0: 4x4(in a 3x3 container in tablet - 170px; 4x3 on phone - 230px)
+        // Row 1: 2x3(in a 2x3 container - 110px)
+        // Row 2: 1x1(50px)
+        // Row 3: 2x4(in a 2x3 container - 110px)
+        // Row 4: 2x2(in a 2x2 container - 110px)
+        assertThat(widgetItemInTable).hasSize(5);
         assertThat(widgetItemInTable.get(0)).containsExactly(mWidget4x4);
-        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x3, mWidget1x1);
-        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget2x4, mWidget2x2);
+        assertThat(widgetItemInTable.get(1)).containsExactly(mWidget2x3);
+        assertThat(widgetItemInTable.get(2)).containsExactly(mWidget1x1);
+        assertThat(widgetItemInTable.get(3)).containsExactly(mWidget2x4);
+        assertThat(widgetItemInTable.get(4)).containsExactly(mWidget2x2);
     }
 
     private void initDP() {