Merge "In the case the gird migration goes to a taller grid keep the grid configuration" into main
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 27a895c..cf9a68b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -944,10 +944,14 @@
         }
         if (landscapePhoneButtonNav) {
             mWindowLayoutParams.width = size;
-            mWindowLayoutParams.paramsForRotation[getDisplay().getRotation()].width = size;
+            for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
+                mWindowLayoutParams.paramsForRotation[rot].width = size;
+            }
         } else {
             mWindowLayoutParams.height = size;
-            mWindowLayoutParams.paramsForRotation[getDisplay().getRotation()].height = size;
+            for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
+                mWindowLayoutParams.paramsForRotation[rot].height = size;
+            }
         }
         mControllers.taskbarInsetsController.onTaskbarOrBubblebarWindowHeightOrInsetsChanged();
         notifyUpdateLayoutParams();
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 629c951..3d58464 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -900,12 +900,12 @@
         }
 
         // Only update the following flags when system gesture is not in progress.
+        updateStateForFlag(FLAG_STASHED_IN_TASKBAR_ALL_APPS, false);
         setStashedImeState();
     }
 
     private void setStashedImeState() {
         boolean shouldStashForIme = shouldStashForIme();
-        updateStateForFlag(FLAG_STASHED_IN_TASKBAR_ALL_APPS, false);
         if (hasAnyFlag(FLAG_STASHED_IN_APP_IME) != shouldStashForIme) {
             updateStateForFlag(FLAG_STASHED_IN_APP_IME, shouldStashForIme);
             applyState(TASKBAR_STASH_DURATION_FOR_IME, getTaskbarStashStartDelayForIme());
diff --git a/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java b/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
index dcc3b05..873dea8 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
@@ -21,6 +21,7 @@
 import android.app.Person;
 import android.content.Context;
 import android.content.Intent;
+import android.content.IntentSender;
 import android.content.pm.ActivityInfo;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
@@ -158,6 +159,28 @@
     }
 
     /**
+     * Returns an intent which can be used to open Private Space Settings.
+     */
+    public static Intent getPrivateSpaceSettingsIntent(Context context) {
+        if (android.os.Flags.allowPrivateProfile() && Flags.enablePrivateSpace()) {
+            LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
+            IntentSender intentSender = launcherApps.getPrivateSpaceSettingsIntent();
+            if (intentSender == null) {
+                return null;
+            }
+            StartActivityParams params = new StartActivityParams((PendingIntent) null, 0);
+            params.intentSender = intentSender;
+            ActivityOptions options = ActivityOptions.makeBasic()
+                    .setPendingIntentBackgroundActivityStartMode(ActivityOptions
+                            .MODE_BACKGROUND_ACTIVITY_START_ALLOWED);
+            params.options = options.toBundle();
+            params.requireActivityResult = false;
+            return ProxyActivityStarter.getLaunchIntent(context, params);
+        }
+        return null;
+    }
+
+    /**
      * Checks if an activity is flagged as non-resizeable.
      */
     public static boolean isNonResizeableActivity(LauncherActivityInfo lai) {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java b/quickstep/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java
index 72218bf..8c92c7d 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java
@@ -135,8 +135,7 @@
                 config.duration = Math.max(config.duration, scrollDuration);
 
                 // Sync scroll so that it ends before or at the same time as the taskbar animation.
-                if (DisplayController.isTransientTaskbar(mActivity)
-                        && mActivity.getDeviceProfile().isTaskbarPresent) {
+                if (mActivity.getDeviceProfile().isTaskbarPresent) {
                     config.duration = Math.min(config.duration, TASKBAR_TO_HOME_DURATION);
                 }
                 overview.snapToPage(DEFAULT_PAGE, Math.toIntExact(config.duration));
diff --git a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
index 36bdad4..62ce341 100644
--- a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
@@ -328,11 +328,15 @@
 
             // No "save app pair" menu item if:
             // - app pairs feature is not enabled
+            // - we are in 3p launcher
             // - the task in question is a single task
             // - at least one app in app pair is unpinnable
             // - the Overview Actions Button should be visible
-            if (!FeatureFlags.enableAppPairs() || !taskView.containsMultipleTasks()
-                    || hasUnpinnableApp || shouldShowActionsButtonInstead) {
+            if (!FeatureFlags.enableAppPairs()
+                    || !recentsView.supportsAppPairs()
+                    || !taskView.containsMultipleTasks()
+                    || hasUnpinnableApp
+                    || shouldShowActionsButtonInstead) {
                 return null;
             }
 
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 415f73f..7880124 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -40,7 +40,6 @@
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.MOTION_DOWN;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.MOTION_MOVE;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.MOTION_UP;
-import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.RECENTS_ANIMATION_START_PENDING;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SYSUI_PROXY;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_UNFOLD_ANIMATION_FORWARDER;
@@ -734,15 +733,17 @@
         // an ACTION_HOVER_ENTER will fire as well.
         boolean isHoverActionWithoutConsumer = enableCursorHoverStates()
                 && isHoverActionWithoutConsumer(event);
-        if (mTaskAnimationManager.isRecentsAnimationStartPending()
-                && (action == ACTION_DOWN || isHoverActionWithoutConsumer)) {
-            ActiveGestureLog.INSTANCE.addLog(
-                    new CompoundString("TIS.onInputEvent: ")
-                            .append("Cannot process input event: a recents animation has been ")
-                            .append("requested, but hasn't started."),
-                    RECENTS_ANIMATION_START_PENDING);
-            return;
-        }
+
+        // TODO(b/285636175): Uncomment this once WM can properly guarantee all animation callbacks
+//        if (mTaskAnimationManager.isRecentsAnimationStartPending()
+//                && (action == ACTION_DOWN || isHoverActionWithoutConsumer)) {
+//            ActiveGestureLog.INSTANCE.addLog(
+//                    new CompoundString("TIS.onInputEvent: ")
+//                            .append("Cannot process input event: a recents animation has been ")
+//                            .append("requested, but hasn't started."),
+//                    RECENTS_ANIMATION_START_PENDING);
+//            return;
+//        }
 
         SafeCloseable traceToken = TraceHelper.INSTANCE.allowIpcs("TIS.onInputEvent");
 
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackNavBarTouchController.java b/quickstep/src/com/android/quickstep/fallback/FallbackNavBarTouchController.java
index 8a87f63..69de3b0 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackNavBarTouchController.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackNavBarTouchController.java
@@ -45,8 +45,7 @@
             NavBarPosition navBarPosition = new NavBarPosition(sysUINavigationMode,
                     DisplayController.INSTANCE.get(mActivity).getInfo());
             mTriggerSwipeUpTracker = new TriggerSwipeUpTouchTracker(mActivity,
-                    true /* disableHorizontalSwipe */, navBarPosition,
-                    null /* onInterceptTouch */, this);
+                    true /* disableHorizontalSwipe */, navBarPosition, this);
         } else {
             mTriggerSwipeUpTracker = null;
         }
@@ -78,7 +77,4 @@
     public void onSwipeUp(boolean wasFling, PointF finalVelocity) {
         mActivity.<FallbackRecentsView>getOverviewPanel().startHome();
     }
-
-    @Override
-    public void onSwipeUpCancelled() {}
 }
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index 0ee50a4..32d8be9 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -302,4 +302,10 @@
     protected boolean canLaunchFullscreenTask() {
         return !mActivity.isInState(OVERVIEW_SPLIT_SELECT);
     }
+
+    /** Returns if app pairs are supported in this launcher. */
+    @Override
+    public boolean supportsAppPairs() {
+        return false;
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OverviewWithoutFocusInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OverviewWithoutFocusInputConsumer.java
index 41730bb..42e8694 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OverviewWithoutFocusInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OverviewWithoutFocusInputConsumer.java
@@ -51,7 +51,7 @@
         mGestureState = gestureState;
         mInputMonitor = inputMonitor;
         mTriggerSwipeUpTracker = new TriggerSwipeUpTouchTracker(context, disableHorizontalSwipe,
-                deviceState.getNavBarPosition(), this::onInterceptTouch, this);
+                deviceState.getNavBarPosition(), this);
     }
 
     @Override
@@ -69,7 +69,8 @@
         mTriggerSwipeUpTracker.onMotionEvent(ev);
     }
 
-    private void onInterceptTouch() {
+    @Override
+    public void onSwipeUpTouchIntercepted() {
         if (mInputMonitor != null) {
             TestLogging.recordEvent(TestProtocol.SEQUENCE_PILFER, "pilferPointers");
             mInputMonitor.pilferPointers();
@@ -93,7 +94,4 @@
                         .build())
                 .log(LAUNCHER_HOME_GESTURE);
     }
-
-    @Override
-    public void onSwipeUpCancelled() {}
 }
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/SysUiOverlayInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/SysUiOverlayInputConsumer.java
index 4806ac1..871d075 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/SysUiOverlayInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/SysUiOverlayInputConsumer.java
@@ -54,7 +54,7 @@
         mContext = context;
         mInputMonitor = inputMonitor;
         mTriggerSwipeUpTracker = new TriggerSwipeUpTouchTracker(context, true,
-                deviceState.getNavBarPosition(), this::onInterceptTouch, this);
+                deviceState.getNavBarPosition(), this);
     }
 
     @Override
@@ -72,7 +72,8 @@
         mTriggerSwipeUpTracker.onMotionEvent(ev);
     }
 
-    private void onInterceptTouch() {
+    @Override
+    public void onSwipeUpTouchIntercepted() {
         if (mInputMonitor != null) {
             TestLogging.recordEvent(TestProtocol.SEQUENCE_PILFER, "pilferPointers");
             mInputMonitor.pilferPointers();
@@ -88,9 +89,4 @@
             Log.e(TAG, "Exception calling closeSystemDialogs " + e.getMessage());
         }
     }
-
-    @Override
-    public void onSwipeUpCancelled() {
-
-    }
 }
diff --git a/quickstep/src/com/android/quickstep/interaction/NavBarGestureHandler.java b/quickstep/src/com/android/quickstep/interaction/NavBarGestureHandler.java
index c4a2216..c00f508 100644
--- a/quickstep/src/com/android/quickstep/interaction/NavBarGestureHandler.java
+++ b/quickstep/src/com/android/quickstep/interaction/NavBarGestureHandler.java
@@ -65,7 +65,7 @@
         mSwipeUpTouchTracker =
                 new TriggerSwipeUpTouchTracker(context, true /*disableHorizontalSwipe*/,
                         new NavBarPosition(NavigationMode.NO_BUTTON, displayInfo),
-                        null /*onInterceptTouch*/, this);
+                        this);
         mMotionPauseDetector = new MotionPauseDetector(context);
 
         final Resources resources = context.getResources();
diff --git a/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.java b/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.java
index 8648b56..f345aeb 100644
--- a/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.java
+++ b/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.java
@@ -305,6 +305,12 @@
     }
 
     @Override
+    public int getTaskMenuHeight(float taskInsetMargin, DeviceProfile deviceProfile,
+            float taskMenuX, float taskMenuY) {
+        return (int) (taskMenuX - taskInsetMargin);
+    }
+
+    @Override
     public void setTaskOptionsMenuLayoutOrientation(DeviceProfile deviceProfile,
             LinearLayout taskMenuLayout, int dividerSpacing,
             ShapeDrawable dividerDrawable) {
diff --git a/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java b/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java
index 60e6a25..5cd9776 100644
--- a/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java
+++ b/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java
@@ -210,6 +210,12 @@
     }
 
     @Override
+    public int getTaskMenuHeight(float taskInsetMargin, DeviceProfile deviceProfile,
+            float taskMenuX, float taskMenuY) {
+        return (int) (deviceProfile.availableHeightPx - taskInsetMargin - taskMenuY);
+    }
+
+    @Override
     public void setTaskOptionsMenuLayoutOrientation(DeviceProfile deviceProfile,
             LinearLayout taskMenuLayout, int dividerSpacing,
             ShapeDrawable dividerDrawable) {
diff --git a/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.java b/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.java
index 01c1225..4b65d53 100644
--- a/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.java
+++ b/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.java
@@ -176,6 +176,9 @@
             View taskMenuView, float taskInsetMargin, View taskViewIcon);
     int getTaskMenuWidth(View thumbnailView, DeviceProfile deviceProfile,
             @StagePosition int stagePosition);
+
+    int getTaskMenuHeight(float taskInsetMargin, DeviceProfile deviceProfile, float taskMenuX,
+            float taskMenuY);
     /**
      * Sets linear layout orientation for {@link com.android.launcher3.popup.SystemShortcut} items
      * inside task menu view.
diff --git a/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.java b/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.java
index a964639..89c678c 100644
--- a/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.java
+++ b/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.java
@@ -114,6 +114,12 @@
     }
 
     @Override
+    public int getTaskMenuHeight(float taskInsetMargin, DeviceProfile deviceProfile,
+            float taskMenuX, float taskMenuY) {
+        return (int) (deviceProfile.availableWidthPx - taskInsetMargin - taskMenuX);
+    }
+
+    @Override
     public void setSplitTaskSwipeRect(DeviceProfile dp, Rect outRect, SplitBounds splitInfo,
             int desiredStagePosition) {
         float topLeftTaskPercent = splitInfo.appsStackedVertically
diff --git a/quickstep/src/com/android/quickstep/util/TriggerSwipeUpTouchTracker.java b/quickstep/src/com/android/quickstep/util/TriggerSwipeUpTouchTracker.java
index 7bbde30..c63a58e 100644
--- a/quickstep/src/com/android/quickstep/util/TriggerSwipeUpTouchTracker.java
+++ b/quickstep/src/com/android/quickstep/util/TriggerSwipeUpTouchTracker.java
@@ -28,6 +28,8 @@
 import android.view.MotionEvent;
 import android.view.VelocityTracker;
 
+import androidx.annotation.NonNull;
+
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 
@@ -41,21 +43,20 @@
     private final float mMinFlingVelocity;
     private final boolean mDisableHorizontalSwipe;
     private final NavBarPosition mNavBarPosition;
-    private final Runnable mOnInterceptTouch;
+
+    @NonNull
     private final OnSwipeUpListener mOnSwipeUp;
 
     private boolean mInterceptedTouch;
     private VelocityTracker mVelocityTracker;
 
     public TriggerSwipeUpTouchTracker(Context context, boolean disableHorizontalSwipe,
-            NavBarPosition navBarPosition, Runnable onInterceptTouch,
-            OnSwipeUpListener onSwipeUp) {
+            NavBarPosition navBarPosition, @NonNull OnSwipeUpListener onSwipeUp) {
         mSquaredTouchSlop = Utilities.squaredTouchSlop(context);
         mMinFlingVelocity = context.getResources().getDimension(
                 R.dimen.quickstep_fling_threshold_speed);
         mNavBarPosition = navBarPosition;
         mDisableHorizontalSwipe = disableHorizontalSwipe;
-        mOnInterceptTouch = onInterceptTouch;
         mOnSwipeUp = onSwipeUp;
 
         init();
@@ -103,10 +104,7 @@
                         }
 
                         mInterceptedTouch = true;
-
-                        if (mOnInterceptTouch != null) {
-                            mOnInterceptTouch.run();
-                        }
+                        mOnSwipeUp.onSwipeUpTouchIntercepted();
                     }
                 }
                 break;
@@ -124,7 +122,8 @@
         }
     }
 
-    private void endTouchTracking() {
+    /** Finishes the tracking. All events after this call are ignored */
+    public void endTouchTracking() {
         if (mVelocityTracker != null) {
             mVelocityTracker.recycle();
             mVelocityTracker = null;
@@ -151,12 +150,10 @@
             isSwipeUp = squaredHypot(displacementX, displacementY) >= mSquaredTouchSlop;
         }
 
-        if (mOnSwipeUp != null) {
-            if (isSwipeUp) {
-                mOnSwipeUp.onSwipeUp(wasFling, new PointF(velocityX, velocityY));
-            } else {
-                mOnSwipeUp.onSwipeUpCancelled();
-            }
+        if (isSwipeUp) {
+            mOnSwipeUp.onSwipeUp(wasFling, new PointF(velocityX, velocityY));
+        } else {
+            mOnSwipeUp.onSwipeUpCancelled();
         }
     }
 
@@ -172,6 +169,9 @@
         void onSwipeUp(boolean wasFling, PointF finalVelocity);
 
         /** Called on touch up if a swipe up was not detected. */
-        void onSwipeUpCancelled();
+        default void onSwipeUpCancelled() { }
+
+        /** Called when the touch for swipe up is intercepted. */
+        default void onSwipeUpTouchIntercepted() { }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/OverviewActionsView.java b/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
index 7a1c49a..e0091a5 100644
--- a/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
+++ b/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
@@ -106,6 +106,7 @@
     public @interface AppPairButtonHiddenFlags { }
     public static final int FLAG_SINGLE_TASK_HIDE_APP_PAIR = 1 << 0;
     public static final int FLAG_SMALL_SCREEN_HIDE_APP_PAIR = 1 << 1;
+    public static final int FLAG_3P_LAUNCHER_HIDE_APP_PAIR = 1 << 2;
 
     private MultiValueAlpha mMultiValueAlpha;
 
@@ -255,6 +256,13 @@
     }
 
     /**
+     * Updates flags to hide and show actions buttons for 1p/3p launchers.
+     */
+    public void updateFor3pLauncher(boolean is3pLauncher) {
+        updateAppPairButtonHiddenFlags(FLAG_3P_LAUNCHER_HIDE_APP_PAIR, is3pLauncher);
+    }
+
+    /**
      * Updates the proper flags to indicate whether the "Screenshot" button should be hidden.
      *
      * @param flag   The flag to update.
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 2cbeb31..133749d 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -4022,6 +4022,8 @@
         mActionsView.updateForGroupedTask(isCurrentSplit);
         // Update flags to see if actions bar should show buttons for tablets or phones.
         mActionsView.updateForSmallScreen(!mActivity.getDeviceProfile().isTablet);
+        // Update flags for 1p/3p launchers
+        mActionsView.updateFor3pLauncher(!supportsAppPairs());
 
         if (isDesktopModeSupported()) {
             boolean isCurrentDesktop = getCurrentPageTaskView() instanceof DesktopTaskView;
@@ -4029,6 +4031,11 @@
         }
     }
 
+    /** Returns if app pairs are supported in this launcher. Overridden in subclasses. */
+    public boolean supportsAppPairs() {
+        return true;
+    }
+
     /**
      * Returns all the tasks in the top row, without the focused task
      */
@@ -5162,9 +5169,13 @@
     public float getMaxScaleForFullScreen() {
         if (enableGridOnlyOverview() && mActivity.getDeviceProfile().isTablet
                 && !mOverviewGridEnabled) {
+            if (mLastComputedCarouselTaskSize.isEmpty()) {
+                mSizeStrategy.calculateCarouselTaskSize(mActivity, mActivity.getDeviceProfile(),
+                        mLastComputedCarouselTaskSize, getPagedOrientationHandler());
+            }
             mTempRect.set(mLastComputedCarouselTaskSize);
         } else {
-            if (mLastComputedTaskSize.height() == 0 || mLastComputedTaskSize.width() == 0) {
+            if (mLastComputedTaskSize.isEmpty()) {
                 getTaskSize(mLastComputedTaskSize);
             }
             mTempRect.set(mLastComputedTaskSize);
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index 137455e..c9aad1a 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -140,11 +140,9 @@
 
     @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
-        if (!enableOverviewIconMenu()) {
-            int maxMenuHeight = calculateMaxHeight();
-            if (MeasureSpec.getSize(heightMeasureSpec) > maxMenuHeight) {
-                heightMeasureSpec = MeasureSpec.makeMeasureSpec(maxMenuHeight, MeasureSpec.AT_MOST);
-            }
+        int maxMenuHeight = calculateMaxHeight();
+        if (MeasureSpec.getSize(heightMeasureSpec) > maxMenuHeight) {
+            heightMeasureSpec = MeasureSpec.makeMeasureSpec(maxMenuHeight, MeasureSpec.AT_MOST);
         }
         super.onMeasure(widthMeasureSpec, heightMeasureSpec);
     }
@@ -416,10 +414,9 @@
      * with a margin on the top and bottom.
      */
     private int calculateMaxHeight() {
-        float taskBottom = mTaskView.getHeight() + mTaskView.getPersistentTranslationY();
         float taskInsetMargin = getResources().getDimension(R.dimen.task_card_margin);
-
-        return (int) (taskBottom - taskInsetMargin - getTranslationY());
+        return mTaskView.getPagedOrientationHandler().getTaskMenuHeight(taskInsetMargin,
+                mActivity.getDeviceProfile(), getTranslationX(), getTranslationY());
     }
 
     private void setOnClosingStartCallback(Runnable onClosingStartCallback) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 9b48082..14d7842 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -30,6 +30,8 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_ICON_TAP_OR_LONGPRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_LAUNCH_TAP;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_NOT_PINNABLE;
+import static com.android.launcher3.testing.shared.TestProtocol.SUCCESSFUL_GESTURE_MISMATCH_EVENTS;
+import static com.android.launcher3.testing.shared.TestProtocol.testLogD;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_BOTTOM_OR_RIGHT;
@@ -861,6 +863,7 @@
     @Nullable
     public RunnableList launchTaskAnimated() {
         if (mTask != null) {
+            testLogD(SUCCESSFUL_GESTURE_MISMATCH_EVENTS, "TaskView.launchTaskAnimated");
             TestLogging.recordEvent(
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
             ActivityOptionsWrapper opts =  mActivity.getActivityLaunchOptions(this, null);
@@ -909,6 +912,7 @@
      */
     public void launchTask(@NonNull Consumer<Boolean> callback, boolean isQuickswitch) {
         if (mTask != null) {
+            testLogD(SUCCESSFUL_GESTURE_MISMATCH_EVENTS, "TaskView.launchTaskAnimated");
             TestLogging.recordEvent(
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
 
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/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 37dde10..8702f70 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.model;
 
+import static android.content.pm.ApplicationInfo.CATEGORY_PRODUCTIVITY;
+import static android.content.pm.ApplicationInfo.FLAG_INSTALLED;
 import static android.os.Process.myUserHandle;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
@@ -37,6 +39,8 @@
 import android.appwidget.AppWidgetManager;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.LauncherApps;
 import android.os.UserHandle;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.text.TextUtils;
@@ -81,6 +85,8 @@
     private FakeBgDataModelCallback mCallback = new FakeBgDataModelCallback();
     private LauncherModelHelper mModelHelper;
     private UserHandle mUserHandle;
+    private LauncherApps mLauncherApps;
+
 
     @Before
     public void setup() throws Exception {
@@ -103,12 +109,18 @@
         allWidgets = Arrays.asList(mApp1Provider1, mApp1Provider2, mApp2Provider1,
                 mApp4Provider1, mApp4Provider2, mApp5Provider1);
 
+        mLauncherApps = mModelHelper.sandboxContext.spyService(LauncherApps.class);
         doAnswer(i -> {
             String pkg = i.getArgument(0);
-            return ApplicationInfoBuilder.newBuilder().setPackageName(pkg).setName(
-                    "App " + pkg).build();
-        }).when(mModelHelper.sandboxContext.getPackageManager())
-                .getApplicationInfo(anyString(), anyInt());
+            ApplicationInfo applicationInfo = ApplicationInfoBuilder.newBuilder()
+                    .setPackageName(pkg)
+                    .setName("App " + pkg)
+                    .build();
+            applicationInfo.category = CATEGORY_PRODUCTIVITY;
+            applicationInfo.flags = FLAG_INSTALLED;
+            return applicationInfo;
+        }).when(mLauncherApps).getApplicationInfo(anyString(), anyInt(), any());
+
         AppWidgetManager manager = mModelHelper.sandboxContext.spyService(AppWidgetManager.class);
         doReturn(allWidgets).when(manager).getInstalledProviders();
         doReturn(allWidgets).when(manager).getInstalledProvidersForProfile(eq(myUserHandle()));
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 213f58f..077ca60 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -76,17 +76,21 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.model.Statement;
 
+import java.io.IOException;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
 public class FallbackRecentsTest {
 
     private static final String FALLBACK_LAUNCHER_TITLE = "Test launcher";
+    private static final Pattern COMPONENT_INFO_REGEX = Pattern.compile("ComponentInfo\\{(.*)\\}");
 
     private final UiDevice mDevice;
     private final LauncherInstrumentation mLauncher;
@@ -253,7 +257,7 @@
     //@NavigationModeSwitch
     @Test
     @ScreenRecordRule.ScreenRecord // b/321775748
-    public void testOverview() {
+    public void testOverview() throws IOException {
         startAppFast(getAppPackageName());
         startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
         startTestActivity(2);
@@ -261,7 +265,10 @@
         Wait.atMost("Expected three apps in the task list",
                 () -> mLauncher.getRecentTasks().size() >= 3, DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
 
+        checkTestLauncher();
         BaseOverview overview = mLauncher.getLaunchedAppState().switchToOverview();
+        checkTestLauncher();
+
         executeOnRecents(recents -> {
             assertTrue("Don't have at least 3 tasks", getTaskCount(recents) >= 3);
         });
@@ -303,6 +310,17 @@
                 mOtherLauncherActivity.packageName).text(FALLBACK_LAUNCHER_TITLE)), WAIT_TIME_MS));
     }
 
+    private void checkTestLauncher() throws IOException {
+        final Matcher matcher = COMPONENT_INFO_REGEX.matcher(
+                mDevice.executeShellCommand("cmd shortcut get-default-launcher"));
+        assertTrue("Incorrect output from get-default-launcher", matcher.find());
+        assertEquals("Current Launcher activity is incorrect",
+                "com.google.android.apps.nexuslauncher.tests/com.android"
+                        + ".launcher3.testcomponent.TestLauncherActivity",
+                matcher.group(1)
+        );
+    }
+
     private int getCurrentOverviewPage(RecentsActivity recents) {
         return recents.<RecentsView>getOverviewPanel().getCurrentPage();
     }
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 45a9527..5bcf72a 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -503,7 +503,6 @@
 
     @Test
     @PortraitLandscape
-    @ScreenRecord // b/326839375
     public void testOverviewDeadzones() throws Exception {
         startTestAppsWithCheck();
 
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/res/values/strings.xml b/res/values/strings.xml
index aaef15b..7bf1c87 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -75,12 +75,14 @@
     <!-- Widget suggestions header title in the full widgets picker for large screen devices
     in landscape mode. [CHAR_LIMIT=50] -->
     <string name="suggested_widgets_header_title">Suggestions</string>
-    <string name="productivity_widget_recommendation_category_label">Your Daily Essentials</string>
-    <string name="news_widget_recommendation_category_label">News For You</string>
+    <string name="productivity_widget_recommendation_category_label">Essentials</string>
+    <string name="news_widget_recommendation_category_label">News &amp; magazines</string>
     <string name="social_and_entertainment_widget_recommendation_category_label">Your Chill Zone</string>
-    <string name="fitness_widget_recommendation_category_label">Reach Your Fitness Goals</string>
-    <string name="weather_widget_recommendation_category_label">Stay Ahead of the Weather</string>
-    <string name="others_widget_recommendation_category_label">You Might Also Like</string>
+    <string name="entertainment_widget_recommendation_category_label">Entertainment</string>
+    <string name="social_widget_recommendation_category_label">Social</string>
+    <string name="fitness_widget_recommendation_category_label">Health &amp; fitness</string>
+    <string name="weather_widget_recommendation_category_label">Weather</string>
+    <string name="others_widget_recommendation_category_label">Suggested for you</string>
     <!-- accessibilityPaneTitle for the right pane when showing suggested widgets. -->
     <string name="widget_picker_right_pane_accessibility_title"><xliff:g id="selected_header" example="Calendar">%1$s</xliff:g> widgets on right, search and options on left</string>
     <!-- Label for showing the number of widgets an app has in the full widgets picker.
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 4f071fe..a9cf2ff 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -916,7 +916,8 @@
     /** Applies the given progress level to the this icon's progress bar. */
     @Nullable
     public PreloadIconDrawable applyProgressLevel() {
-        if (!(getTag() instanceof ItemInfoWithIcon)) {
+        if (!(getTag() instanceof ItemInfoWithIcon)
+                || !((ItemInfoWithIcon) getTag()).isActiveArchive()) {
             return null;
         }
 
@@ -973,6 +974,7 @@
         return info.isDisabled() || info.isPendingDownload();
     }
 
+
     public void applyDotState(ItemInfo itemInfo, boolean animate) {
         if (mIcon instanceof FastBitmapDrawable) {
             boolean wasDotted = mDotInfo != null;
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 42d4d50..2e0f676 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -72,6 +72,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
 public class InvariantDeviceProfile {
@@ -578,6 +579,45 @@
     }
 
     /**
+     * Returns the GridOption associated to the given file name or null if the fileName is not
+     * supported.
+     * Ej, launcher.db -> "normal grid", launcher_4_by_4.db -> "practical grid"
+     */
+    public GridOption getGridOptionFromFileName(Context context, String fileName) {
+        return parseAllGridOptions(context).stream()
+                .filter(gridOption -> Objects.equals(gridOption.dbFile, fileName))
+                .findFirst()
+                .orElse(null);
+    }
+
+    /**
+     * Returns the name of the given size on the current device or empty string if the size is not
+     * supported. Ej. 4x4 -> normal, 5x4 -> practical, etc.
+     * (Note: the name of the grid can be different for the same grid size depending of
+     * the values of the InvariantDeviceProfile)
+     *
+     */
+    public String getGridNameFromSize(Context context, Point size) {
+        return parseAllGridOptions(context).stream()
+                .filter(gridOption -> gridOption.numColumns == size.x
+                        && gridOption.numRows == size.y)
+                .map(gridOption -> gridOption.name)
+                .findFirst()
+                .orElse("");
+    }
+
+    /**
+     * Returns the grid option for the given gridName on the current device (Note: the gridOption
+     * be different for the same gridName depending on the values of the InvariantDeviceProfile).
+     */
+    public GridOption getGridOptionFromName(Context context, String gridName) {
+        return parseAllGridOptions(context).stream()
+                .filter(gridOption -> Objects.equals(gridOption.name, gridName))
+                .findFirst()
+                .orElse(null);
+    }
+
+    /**
      * @return all the grid options that can be shown on the device
      */
     public List<GridOption> parseAllGridOptions(Context context) {
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/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 34ebaf2..84b8ba1 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -139,6 +139,11 @@
         public static final int ITEM_TYPE_SEARCH_ACTION = 9;
 
         /**
+         * Private space install app button.
+         */
+        public static final int ITEM_TYPE_PRIVATE_SPACE_INSTALL_APP_BUTTON = 11;
+
+        /**
          * The custom icon bitmap.
          * <P>Type: BLOB</P>
          */
@@ -206,6 +211,8 @@
                 case ITEM_TYPE_TASK: return "TASK";
                 case ITEM_TYPE_QSB: return "QSB";
                 case ITEM_TYPE_APP_PAIR: return "APP_PAIR";
+                case ITEM_TYPE_PRIVATE_SPACE_INSTALL_APP_BUTTON:
+                    return "PRIVATE_SPACE_INSTALL_APP_BUTTON";
                 default: return String.valueOf(type);
             }
         }
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 1ebd49e..e7bb1d0 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -23,15 +23,12 @@
 import static com.android.launcher3.allapps.SectionDecorationInfo.ROUND_NOTHING;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_PRIVATE_PROFILE_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_NOT_PINNABLE;
-import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_PRIVATE_SPACE_INSTALL_APP;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.SettingsCache.PRIVATE_SPACE_HIDE_WHEN_LOCKED_URI;
 
 import android.content.Context;
 import android.content.Intent;
-import android.content.pm.PackageManager;
-import android.content.pm.ResolveInfo;
 import android.os.UserHandle;
 import android.os.UserManager;
 
@@ -43,6 +40,7 @@
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.PrivateSpaceInstallAppButtonInfo;
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.uioverrides.ApiWrapper;
 import com.android.launcher3.util.Preconditions;
@@ -60,10 +58,6 @@
  */
 public class PrivateProfileManager extends UserProfileManager {
 
-    // TODO (b/324573634): Fix the intent string.
-    public static final Intent PRIVATE_SPACE_INTENT = new
-            Intent("com.android.settings.action.PRIVATE_SPACE_SETUP_FLOW");
-
     private final ActivityAllAppsContainerView<?> mAllApps;
     private final Predicate<UserHandle> mPrivateProfileMatcher;
     private Set<String> mPreInstalledSystemPackages = new HashSet<>();
@@ -105,13 +99,13 @@
                 context, com.android.launcher3.R.drawable.private_space_install_app_icon);
         BitmapInfo bitmapInfo = LauncherIcons.obtain(context).createIconBitmap(shortcut);
 
-        AppInfo itemInfo = new AppInfo();
+        PrivateSpaceInstallAppButtonInfo itemInfo = new PrivateSpaceInstallAppButtonInfo();
         itemInfo.title = context.getResources().getString(R.string.ps_add_button_label);
         itemInfo.intent = mAppInstallerIntent;
         itemInfo.bitmap = bitmapInfo;
         itemInfo.contentDescription = context.getResources().getString(
                 com.android.launcher3.R.string.ps_add_button_content_description);
-        itemInfo.runtimeStatusFlags |= FLAG_PRIVATE_SPACE_INSTALL_APP | FLAG_NOT_PINNABLE;
+        itemInfo.runtimeStatusFlags |= FLAG_NOT_PINNABLE;
 
         BaseAllAppsAdapter.AdapterItem item = new BaseAllAppsAdapter.AdapterItem(VIEW_TYPE_ICON);
         item.itemInfo = itemInfo;
@@ -162,7 +156,8 @@
     /** Opens the Private Space Settings Page. */
     public void openPrivateSpaceSettings() {
         if (mPrivateSpaceSettingsAvailable) {
-            mAllApps.getContext().startActivity(PRIVATE_SPACE_INTENT);
+            mAllApps.getContext()
+                    .startActivity(ApiWrapper.getPrivateSpaceSettingsIntent(mAllApps.getContext()));
         }
     }
 
@@ -194,9 +189,8 @@
 
     private void initializePrivateSpaceSettingsState() {
         Preconditions.assertNonUiThread();
-        ResolveInfo resolveInfo = mAllApps.getContext().getPackageManager()
-                .resolveActivity(PRIVATE_SPACE_INTENT, PackageManager.MATCH_SYSTEM_ONLY);
-        setPrivateSpaceSettingsAvailable(resolveInfo != null);
+        Intent psSettingsIntent = ApiWrapper.getPrivateSpaceSettingsIntent(mAllApps.getContext());
+        setPrivateSpaceSettingsAvailable(psSettingsIntent != null);
     }
 
     private void setPreInstalledSystemPackages() {
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index db27832..fe327d0 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -37,7 +37,6 @@
 import android.view.SurfaceControlViewHost;
 import android.view.SurfaceControlViewHost.SurfacePackage;
 import android.view.View;
-import android.view.WindowManager.LayoutParams;
 import android.view.animation.AccelerateDecelerateInterpolator;
 
 import androidx.annotation.NonNull;
@@ -215,7 +214,6 @@
             return new ContextThemeWrapper(context,
                     Themes.getActivityThemeRes(context));
         }
-        context = context.createWindowContext(LayoutParams.TYPE_APPLICATION_OVERLAY, null);
         LocalColorExtractor.newInstance(context)
                 .applyColorsOverride(context, mWallpaperColors);
         return new ContextThemeWrapper(context,
diff --git a/src/com/android/launcher3/model/WidgetItem.java b/src/com/android/launcher3/model/WidgetItem.java
index 1dd58c3..3f88717 100644
--- a/src/com/android/launcher3/model/WidgetItem.java
+++ b/src/com/android/launcher3/model/WidgetItem.java
@@ -123,6 +123,7 @@
         if (!Flags.enableGeneratedPreviews() || generatedPreviews == null) {
             return false;
         }
-        return generatedPreviews.contains(widgetCategory);
+        return generatedPreviews.contains(widgetCategory)
+                && generatedPreviews.get(widgetCategory) != null;
     }
 }
diff --git a/src/com/android/launcher3/model/data/ItemInfo.java b/src/com/android/launcher3/model/data/ItemInfo.java
index 55849c2..f7cff78 100644
--- a/src/com/android/launcher3/model/data/ItemInfo.java
+++ b/src/com/android/launcher3/model/data/ItemInfo.java
@@ -94,6 +94,10 @@
      * {@link Favorites#ITEM_TYPE_APP_PAIR},
      * {@link Favorites#ITEM_TYPE_APPWIDGET} or
      * {@link Favorites#ITEM_TYPE_CUSTOM_APPWIDGET}.
+     * {@link Favorites#ITEM_TYPE_TASK}.
+     * {@link Favorites#ITEM_TYPE_QSB}.
+     * {@link Favorites#ITEM_TYPE_SEARCH_ACTION}.
+     * {@link Favorites#ITEM_TYPE_PRIVATE_SPACE_INSTALL_APP_BUTTON}.
      */
     public int itemType;
 
diff --git a/src/com/android/launcher3/model/data/ItemInfoWithIcon.java b/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
index 352c363..9fbc6bf 100644
--- a/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
+++ b/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
@@ -121,11 +121,6 @@
     public static final int FLAG_ARCHIVED = 1 << 14;
 
     /**
-     * Flag indicating it's the Private Space Install App icon.
-     */
-    public static final int FLAG_PRIVATE_SPACE_INSTALL_APP = 1 << 15;
-
-    /**
      * Status associated with the system state of the underlying item. This is calculated every
      * time a new info is created and not persisted on the disk.
      */
@@ -160,10 +155,6 @@
      * and its install session is active
      */
     public boolean isPendingDownload() {
-        if (isArchived()) {
-            return this.getProgressLevel() == 0
-                    && (this.runtimeStatusFlags & FLAG_INSTALL_SESSION_ACTIVE) != 0;
-        }
         return getProgressLevel() == 0;
     }
 
@@ -177,6 +168,11 @@
         return (runtimeStatusFlags & FLAG_ARCHIVED) != 0;
     }
 
+    /** Returns true if the app is archived and has an active install session. */
+    public boolean isActiveArchive() {
+        return isArchived() && (runtimeStatusFlags & FLAG_INSTALL_SESSION_ACTIVE) != 0;
+    }
+
     /**
      * Indicates whether we're using a low res icon
      */
diff --git a/src/com/android/launcher3/model/data/PrivateSpaceInstallAppButtonInfo.java b/src/com/android/launcher3/model/data/PrivateSpaceInstallAppButtonInfo.java
new file mode 100644
index 0000000..1e7281d
--- /dev/null
+++ b/src/com/android/launcher3/model/data/PrivateSpaceInstallAppButtonInfo.java
@@ -0,0 +1,29 @@
+/*
+ * 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.model.data;
+
+import com.android.launcher3.LauncherSettings;
+
+/**
+ * Represents the Private Space Install App button in AllAppsView.
+ */
+public class PrivateSpaceInstallAppButtonInfo extends AppInfo {
+
+    public PrivateSpaceInstallAppButtonInfo() {
+        itemType = LauncherSettings.Favorites.ITEM_TYPE_PRIVATE_SPACE_INSTALL_APP_BUTTON;
+    }
+}
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 22bc13b..f2b7d18 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -50,8 +50,10 @@
 import androidx.annotation.VisibleForTesting;
 import androidx.annotation.WorkerThread;
 
+import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherFiles;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -121,7 +123,48 @@
         // executed again.
         LauncherPrefs.get(context).removeSync(RESTORE_DEVICE);
 
-        idp.reinitializeAfterRestore(context);
+        if (Flags.narrowGridRestore()) {
+            String oldPhoneFileName = idp.dbFile;
+            removeOldDBs(context, oldPhoneFileName);
+            trySettingPreviousGidAsCurrent(context, idp, oldPhoneFileName);
+        } else {
+            idp.reinitializeAfterRestore(context);
+        }
+    }
+
+    /**
+     * Try setting the gird used in the previous phone to the new one. If the current device doesn't
+     * support the previous grid option it will not be set.
+     */
+    private static void trySettingPreviousGidAsCurrent(Context context, InvariantDeviceProfile idp,
+            String oldPhoneDbFileName) {
+        InvariantDeviceProfile.GridOption gridOption = idp.getGridOptionFromFileName(context,
+                oldPhoneDbFileName);
+        if (gridOption != null) {
+            /*
+             * We do this because in some cases different devices have different names for grid
+             * options, in one device the grid option "normal" can be 4x4 while in other it
+             * could be "practical". Calling this changes the current device grid to the same
+             * we had in the other phone, in the case the current phone doesn't support the grid
+             * option we use the default and migrate the db to the default. Migration occurs on
+             * {@code GridSizeMigrationUtil#migrateGridIfNeeded}
+             */
+            idp.setCurrentGrid(context, gridOption.name);
+        }
+    }
+
+    /**
+     * Only keep the last database used on the previous device.
+     */
+    private static void removeOldDBs(Context context, String oldPhoneDbFileName) {
+        // At this point idp.dbFile contains the name of the dbFile from the previous phone
+        LauncherFiles.GRID_DB_FILES.stream()
+                .filter(dbName -> !dbName.equals(oldPhoneDbFileName))
+                .forEach(dbName -> {
+                    if (context.getDatabasePath(dbName).delete()) {
+                        FileLog.d(TAG, "Removed old grid db file: " + dbName);
+                    }
+                });
     }
 
     private static boolean performRestore(Context context, ModelDbController controller) {
diff --git a/src/com/android/launcher3/touch/ItemClickHandler.java b/src/com/android/launcher3/touch/ItemClickHandler.java
index 111931e..911568c 100644
--- a/src/com/android/launcher3/touch/ItemClickHandler.java
+++ b/src/com/android/launcher3/touch/ItemClickHandler.java
@@ -369,8 +369,8 @@
                 intent = ApiWrapper.getAppMarketActivityIntent(launcher,
                         itemInfoWithIcon.getTargetComponent().getPackageName(),
                         Process.myUserHandle());
-            } else if ((itemInfoWithIcon.runtimeStatusFlags
-                    & ItemInfoWithIcon.FLAG_PRIVATE_SPACE_INSTALL_APP) != 0) {
+            } else if (itemInfoWithIcon.itemType
+                    == LauncherSettings.Favorites.ITEM_TYPE_PRIVATE_SPACE_INSTALL_APP_BUTTON) {
                 intent = ApiWrapper.getAppMarketActivityIntent(launcher,
                         BuildConfig.APPLICATION_ID,
                         launcher.getAppsView().getPrivateProfileManager().getProfileUser());
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/WidgetRecommendationCategoryProvider.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
index 801b1f6..c3906a6 100644
--- a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
@@ -18,13 +18,13 @@
 
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
-import android.content.pm.PackageManager;
-import android.util.Log;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.R;
 import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.util.ResourceBasedOverride;
 
@@ -54,6 +54,7 @@
      * to display the recommendation grouped by categories.
      */
     @WorkerThread
+    @Nullable
     public WidgetRecommendationCategory getWidgetRecommendationCategory(Context context,
             WidgetItem item) {
         // This is a default implementation that uses application category to derive the category to
@@ -61,17 +62,16 @@
         // via the overridden WidgetRecommendationCategoryProvider resource.
 
         Preconditions.assertWorkerThread();
-        PackageManager pm = context.getPackageManager();
+        PackageManagerHelper pmHelper = new PackageManagerHelper(context);
         if (item.widgetInfo != null && item.widgetInfo.getComponent() != null) {
             String widgetComponentName = item.widgetInfo.getComponent().getClassName();
-            try {
-                int predictionCategory = pm.getApplicationInfo(
-                        item.widgetInfo.getComponent().getPackageName(), 0 /* flags */).category;
+            ApplicationInfo applicationInfo = pmHelper.getApplicationInfo(
+                    item.widgetInfo.getComponent().getPackageName(), item.widgetInfo.getUser(),
+                    0 /* flags */);
+            if (applicationInfo != null) {
+                int predictionCategory = applicationInfo.category;
                 return getCategoryFromApplicationCategory(context, predictionCategory,
                         widgetComponentName);
-            } catch (PackageManager.NameNotFoundException e) {
-                Log.e(TAG, "Failed to retrieve application category when determining the "
-                        + "widget category for " + widgetComponentName, e);
             }
         }
         return null;
@@ -112,17 +112,22 @@
                     R.string.news_widget_recommendation_category_label, /*order=*/1);
         }
 
-        if (applicationCategory == ApplicationInfo.CATEGORY_SOCIAL
-                || applicationCategory == ApplicationInfo.CATEGORY_AUDIO
+        if (applicationCategory == ApplicationInfo.CATEGORY_SOCIAL) {
+            return new WidgetRecommendationCategory(
+                    R.string.social_widget_recommendation_category_label,
+                    /*order=*/5);
+        }
+
+        if (applicationCategory == ApplicationInfo.CATEGORY_AUDIO
                 || applicationCategory == ApplicationInfo.CATEGORY_VIDEO
                 || applicationCategory == ApplicationInfo.CATEGORY_IMAGE) {
             return new WidgetRecommendationCategory(
-                    R.string.social_and_entertainment_widget_recommendation_category_label,
-                    /*order=*/4);
+                    R.string.entertainment_widget_recommendation_category_label,
+                    /*order=*/6);
         }
 
         return new WidgetRecommendationCategory(
-                R.string.others_widget_recommendation_category_label, /*order=*/5);
+                R.string.others_widget_recommendation_category_label, /*order=*/4);
     }
 
 }
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/src_ui_overrides/com/android/launcher3/uioverrides/ApiWrapper.java b/src_ui_overrides/com/android/launcher3/uioverrides/ApiWrapper.java
index efde7d8..90271c1 100644
--- a/src_ui_overrides/com/android/launcher3/uioverrides/ApiWrapper.java
+++ b/src_ui_overrides/com/android/launcher3/uioverrides/ApiWrapper.java
@@ -108,6 +108,13 @@
     }
 
     /**
+     * Returns an intent which can be used to open Private Space Settings.
+     */
+    public static Intent getPrivateSpaceSettingsIntent(Context context) {
+        return null;
+    }
+
+    /**
      * Checks if an activity is flagged as non-resizeable.
      */
     public static boolean isNonResizeableActivity(LauncherActivityInfo lai) {
diff --git a/tests/Android.bp b/tests/Android.bp
index e9111ea..24ae158 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -177,7 +177,7 @@
     name: "launcher-testing-shared",
     srcs: [
         "multivalentTests/shared/com/android/launcher3/testing/shared/**/*.java",
-        "multivalentTests/shared/com/android/launcher3/testing/shared/**/*.kt"
+        "multivalentTests/shared/com/android/launcher3/testing/shared/**/*.kt",
     ],
     resource_dirs: [],
     manifest: "multivalentTests/shared/AndroidManifest.xml",
@@ -225,8 +225,8 @@
     // multivalentTests directory is a shared folder for not only robolectric converted test
     // classes but also shared helper classes.
     srcs: [
-        "multivalentTests/src/com/android/launcher3/util/*.java",
-        "multivalentTests/src/com/android/launcher3/util/*.kt",
+        "multivalentTests/src/**/*.java",
+        "multivalentTests/src/**/*.kt",
 
         // Test util classes
         ":launcher-testing-helpers",
@@ -246,7 +246,8 @@
         "androidx.test.uiautomator_uiautomator",
         "androidx.core_core-animation-testing",
         "androidx.test.ext.junit",
-        "inline-mockito-robolectric-prebuilt",
+        "mockito-robolectric-prebuilt",
+        "mockito-kotlin2",
         "platform-parametric-runner-lib",
         "testables",
         "Launcher3TestResources",
diff --git a/tests/assets/databases/BackupAndRestore/launcher.db b/tests/assets/databases/BackupAndRestore/launcher.db
new file mode 100644
index 0000000..126d166
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher.db
Binary files differ
diff --git a/tests/assets/databases/BackupAndRestore/launcher_3_by_3.db b/tests/assets/databases/BackupAndRestore/launcher_3_by_3.db
new file mode 100644
index 0000000..6d8cd73
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher_3_by_3.db
Binary files differ
diff --git a/tests/assets/databases/BackupAndRestore/launcher_4_by_4.db b/tests/assets/databases/BackupAndRestore/launcher_4_by_4.db
new file mode 100644
index 0000000..00061dd
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher_4_by_4.db
Binary files differ
diff --git a/tests/assets/databases/BackupAndRestore/launcher_4_by_5.db b/tests/assets/databases/BackupAndRestore/launcher_4_by_5.db
new file mode 100644
index 0000000..e2e65aa
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher_4_by_5.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/multivalentTests/src/com/android/launcher3/util/rule/BackAndRestoreRule.kt b/tests/multivalentTests/src/com/android/launcher3/util/rule/BackAndRestoreRule.kt
new file mode 100644
index 0000000..da96939
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/rule/BackAndRestoreRule.kt
@@ -0,0 +1,119 @@
+/*
+ * 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.util.rule
+
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherPrefs
+import java.io.File
+import java.nio.file.Paths
+import kotlin.io.path.pathString
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * Removes all launcher's DBs from the device and copies the dbs in
+ * assets/databases/BackupAndRestore to the device. It also set's the needed LauncherPrefs variables
+ * needed to kickstart a backup and restore.
+ */
+class BackAndRestoreRule : TestRule {
+
+    private val phoneContext = getInstrumentation().targetContext
+
+    private fun dbBackUp() = File(phoneContext.dataDir.path, "/databasesBackUp")
+
+    private fun dbDirectory() = File(phoneContext.dataDir.path, "/databases")
+
+    private fun isWorkspaceDatabase(rawFileName: String): Boolean {
+        val fileName = Paths.get(rawFileName).fileName.pathString
+        return fileName.startsWith("launcher") && fileName.endsWith(".db")
+    }
+
+    fun getDatabaseFiles() = dbDirectory().listFiles().filter { isWorkspaceDatabase(it.name) }
+
+    /**
+     * Setting RESTORE_DEVICE would trigger a restore next time the Launcher starts, and we remove
+     * the widgets and apps ids to prevent issues when loading the database.
+     */
+    private fun setRestoreConstants() {
+        LauncherPrefs.get(phoneContext)
+            .put(LauncherPrefs.RESTORE_DEVICE.to(InvariantDeviceProfile.TYPE_MULTI_DISPLAY))
+        LauncherPrefs.get(phoneContext)
+            .remove(LauncherPrefs.OLD_APP_WIDGET_IDS, LauncherPrefs.APP_WIDGET_IDS)
+    }
+
+    private fun uploadDatabase(dbName: String) {
+        val file = File(File(getInstrumentation().targetContext.dataDir, "/databases"), dbName)
+        file.writeBytes(
+            getInstrumentation()
+                .context
+                .assets
+                .open("databases/BackupAndRestore/$dbName")
+                .readBytes()
+        )
+        file.setWritable(true, false)
+    }
+
+    private fun uploadDbs() {
+        uploadDatabase("launcher.db")
+        uploadDatabase("launcher_4_by_4.db")
+        uploadDatabase("launcher_4_by_5.db")
+        uploadDatabase("launcher_3_by_3.db")
+    }
+
+    private fun savePreviousState() {
+        dbBackUp().deleteRecursively()
+        if (!dbDirectory().renameTo(dbBackUp())) {
+            throw Exception("Unable to move databases to backup directory")
+        }
+        dbDirectory().mkdir()
+        if (!dbDirectory().exists()) {
+            throw Exception("Databases directory doesn't exists")
+        }
+    }
+
+    private fun restorePreviousState() {
+        dbDirectory().deleteRecursively()
+        if (!dbBackUp().renameTo(dbDirectory())) {
+            throw Exception("Unable to restore backup directory to databases directory")
+        }
+        dbBackUp().delete()
+    }
+
+    fun before() {
+        savePreviousState()
+        setRestoreConstants()
+        uploadDbs()
+    }
+
+    fun after() {
+        restorePreviousState()
+    }
+
+    override fun apply(base: Statement?, description: Description?): Statement =
+        object : Statement() {
+            override fun evaluate() {
+                before()
+                try {
+                    base?.evaluate()
+                } finally {
+                    after()
+                }
+            }
+        }
+}
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/PrivateProfileManagerTest.java b/tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java
index 0907f8f..eea4fe5 100644
--- a/tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java
+++ b/tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java
@@ -47,6 +47,7 @@
 
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.uioverrides.ApiWrapper;
 import com.android.launcher3.util.ActivityContextWrapper;
 import com.android.launcher3.util.UserIconInfo;
 import com.android.launcher3.util.rule.TestStabilityRule;
@@ -176,17 +177,15 @@
     }
 
     @Test
-    public void openPrivateSpaceSettings_triggersSecurityAndPrivacyIntent() {
-        Intent expectedIntent = PrivateProfileManager.PRIVATE_SPACE_INTENT;
+    public void openPrivateSpaceSettings_triggersCorrectIntent() {
+        Intent expectedIntent = ApiWrapper.getPrivateSpaceSettingsIntent(mContext);
         ArgumentCaptor<Intent> acIntent = ArgumentCaptor.forClass(Intent.class);
         mPrivateProfileManager.setPrivateSpaceSettingsAvailable(true);
 
         mPrivateProfileManager.openPrivateSpaceSettings();
 
         Mockito.verify(mContext).startActivity(acIntent.capture());
-        Intent actualIntent = acIntent.getValue();
-        assertEquals("Intent Action is different", expectedIntent.getAction(),
-                actualIntent.getAction());
+        assertEquals("Intent Action is different", expectedIntent, acIntent.getValue());
     }
 
     private static void awaitTasksCompleted() throws Exception {
diff --git a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
new file mode 100644
index 0000000..3e36bbb
--- /dev/null
+++ b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
@@ -0,0 +1,70 @@
+/*
+ * 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.backuprestore
+
+import android.platform.test.flag.junit.SetFlagsRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.MediumTest
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.Flags
+import com.android.launcher3.LauncherPrefs
+import com.android.launcher3.model.ModelDbController
+import com.android.launcher3.util.Executors.MODEL_EXECUTOR
+import com.android.launcher3.util.TestUtil
+import com.android.launcher3.util.rule.BackAndRestoreRule
+import com.android.launcher3.util.rule.setFlags
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Makes sure to test {@code RestoreDbTask#removeOldDBs}, we need to remove all the dbs that are not
+ * the last one used when we restore the device.
+ */
+@RunWith(AndroidJUnit4::class)
+@MediumTest
+class BackupAndRestoreDBSelectionTest {
+
+    @JvmField @Rule var backAndRestoreRule = BackAndRestoreRule()
+
+    @JvmField
+    @Rule
+    val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+
+    @Before
+    fun setUp() {
+        setFlagsRule.setFlags(true, Flags.FLAG_NARROW_GRID_RESTORE)
+    }
+
+    @Test
+    fun oldDatabasesNotPresentAfterRestore() {
+        val dbController = ModelDbController(getInstrumentation().targetContext)
+        dbController.tryMigrateDB(null)
+        TestUtil.runOnExecutorSync(MODEL_EXECUTOR) {
+            assert(backAndRestoreRule.getDatabaseFiles().size == 1) {
+                "There should only be one database after restoring, the last one used. Actual databases ${backAndRestoreRule.getDatabaseFiles()}"
+            }
+            assert(
+                !LauncherPrefs.get(getInstrumentation().targetContext)
+                    .has(LauncherPrefs.RESTORE_DEVICE)
+            ) {
+                "RESTORE_DEVICE shouldn't be present after a backup and restore."
+            }
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt b/tests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
index 11855e6..edd2652 100644
--- a/tests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
+++ b/tests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
@@ -111,6 +111,18 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(FLAG_ENABLE_GENERATED_PREVIEWS)
+    fun widgetItem_hasGeneratedPreview_nullPreview() {
+        appWidgetProviderInfo.generatedPreviewCategories =
+            WIDGET_CATEGORY_HOME_SCREEN or WIDGET_CATEGORY_KEYGUARD
+        createWidgetItem()
+        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_HOME_SCREEN)).isTrue()
+        // loadGeneratedPreview returns null for KEYGUARD, so this should still be false.
+        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_KEYGUARD)).isFalse()
+        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_SEARCHBOX)).isFalse()
+    }
+
+    @Test
     @RequiresFlagsDisabled(FLAG_ENABLE_GENERATED_PREVIEWS)
     fun widgetItem_hasGeneratedPreview_flagDisabled() {
         assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_HOME_SCREEN)).isFalse()
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/WidgetRecommendationCategoryProviderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
index c807771..60a4cd3 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
@@ -23,6 +23,7 @@
 import static android.content.pm.ApplicationInfo.CATEGORY_SOCIAL;
 import static android.content.pm.ApplicationInfo.CATEGORY_UNDEFINED;
 import static android.content.pm.ApplicationInfo.CATEGORY_VIDEO;
+import static android.content.pm.ApplicationInfo.FLAG_INSTALLED;
 
 import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
@@ -30,8 +31,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
@@ -41,7 +41,7 @@
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.pm.ApplicationInfo;
-import android.content.pm.PackageManager;
+import android.content.pm.LauncherApps;
 import android.os.Process;
 
 import androidx.test.core.content.pm.ApplicationInfoBuilder;
@@ -70,10 +70,25 @@
 public class WidgetRecommendationCategoryProviderTest {
     private static final String TEST_PACKAGE = "com.foo.test";
     private static final String TEST_APP_NAME = "foo";
-    public static final WidgetRecommendationCategory SOCIAL_AND_ENTERTAINMENT_CATEGORY =
+    private static final WidgetRecommendationCategory PRODUCTIVITY =
             new WidgetRecommendationCategory(
-                    R.string.social_and_entertainment_widget_recommendation_category_label,
-                    /*order=*/4);
+                    R.string.productivity_widget_recommendation_category_label,
+                    /*order=*/0);
+    private static final WidgetRecommendationCategory NEWS =
+            new WidgetRecommendationCategory(
+                    R.string.news_widget_recommendation_category_label, /*order=*/1);
+    private static final WidgetRecommendationCategory SUGGESTED_FOR_YOU =
+            new WidgetRecommendationCategory(
+                    R.string.others_widget_recommendation_category_label, /*order=*/4);
+    private static final WidgetRecommendationCategory SOCIAL =
+            new WidgetRecommendationCategory(
+                    R.string.social_widget_recommendation_category_label,
+                    /*order=*/5);
+    private static final WidgetRecommendationCategory ENTERTAINMENT =
+            new WidgetRecommendationCategory(
+                    R.string.entertainment_widget_recommendation_category_label,
+                    /*order=*/6);
+
     private final ApplicationInfo mTestAppInfo = ApplicationInfoBuilder.newBuilder().setPackageName(
             TEST_PACKAGE).setName(TEST_APP_NAME).build();
     private Context mContext;
@@ -82,7 +97,7 @@
 
     private WidgetItem mTestWidgetItem;
     @Mock
-    private PackageManager mPackageManager;
+    private LauncherApps mLauncherApps;
     private InvariantDeviceProfile mTestProfile;
 
     @Before
@@ -90,10 +105,12 @@
         MockitoAnnotations.initMocks(this);
         mContext = new ContextWrapper(getInstrumentation().getTargetContext()) {
             @Override
-            public PackageManager getPackageManager() {
-                return mPackageManager;
+            public Object getSystemService(String name) {
+                return LAUNCHER_APPS_SERVICE.equals(name) ? mLauncherApps : super.getSystemService(
+                        name);
             }
         };
+        mTestAppInfo.flags = FLAG_INSTALLED;
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
@@ -103,25 +120,22 @@
     @Test
     public void getWidgetRecommendationCategory_returnsMappedCategory() throws Exception {
         ImmutableMap<Integer, WidgetRecommendationCategory> testCategories = ImmutableMap.of(
-                CATEGORY_PRODUCTIVITY, new WidgetRecommendationCategory(
-                        R.string.productivity_widget_recommendation_category_label,
-                        /*order=*/
-                        0),
-                CATEGORY_NEWS, new WidgetRecommendationCategory(
-                        R.string.news_widget_recommendation_category_label, /*order=*/1),
-                CATEGORY_SOCIAL, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
-                CATEGORY_AUDIO, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
-                CATEGORY_IMAGE, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
-                CATEGORY_VIDEO, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
-                CATEGORY_UNDEFINED, new WidgetRecommendationCategory(
-                        R.string.others_widget_recommendation_category_label, /*order=*/5));
+                CATEGORY_PRODUCTIVITY, PRODUCTIVITY,
+                CATEGORY_NEWS, NEWS,
+                CATEGORY_SOCIAL, SOCIAL,
+                CATEGORY_AUDIO, ENTERTAINMENT,
+                CATEGORY_IMAGE, ENTERTAINMENT,
+                CATEGORY_VIDEO, ENTERTAINMENT,
+                CATEGORY_UNDEFINED, SUGGESTED_FOR_YOU);
 
         for (Map.Entry<Integer, WidgetRecommendationCategory> testCategory :
                 testCategories.entrySet()) {
 
             mTestAppInfo.category = testCategory.getKey();
-            when(mPackageManager.getApplicationInfo(anyString(), anyInt())).thenReturn(
-                    mTestAppInfo);
+            when(mLauncherApps.getApplicationInfo(/*packageName=*/ eq(TEST_PACKAGE),
+                    /*flags=*/ eq(0),
+                    /*user=*/ eq(Process.myUserHandle())))
+                    .thenReturn(mTestAppInfo);
 
             WidgetRecommendationCategory category = Executors.MODEL_EXECUTOR.submit(() ->
                     new WidgetRecommendationCategoryProvider().getWidgetRecommendationCategory(
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() {
diff --git a/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java b/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java
index b5414b7..3f96999 100644
--- a/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java
+++ b/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java
@@ -26,11 +26,11 @@
 import static com.android.launcher3.testing.shared.TestProtocol.REQUEST_TASKBAR_FROM_NAV_THRESHOLD;
 import static com.android.launcher3.testing.shared.TestProtocol.SUCCESSFUL_GESTURE_MISMATCH_EVENTS;
 import static com.android.launcher3.testing.shared.TestProtocol.TEST_INFO_RESPONSE_FIELD;
-import static com.android.launcher3.testing.shared.TestProtocol.testLogD;
 
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.SystemClock;
+import android.util.Log;
 import android.view.InputDevice;
 import android.view.MotionEvent;
 import android.view.ViewConfiguration;
@@ -141,7 +141,7 @@
 
             return new Taskbar(mLauncher);
         } finally {
-            testLogD(SUCCESSFUL_GESTURE_MISMATCH_EVENTS,
+            Log.d(SUCCESSFUL_GESTURE_MISMATCH_EVENTS,
                     "swipeUpToUnstashTaskbar: completed gesture");
             mLauncher.getTestInfo(REQUEST_DISABLE_BLOCK_TIMEOUT);
         }
diff --git a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
index afe5722..1cfbf09 100644
--- a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
+++ b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
@@ -19,8 +19,10 @@
 import static com.android.launcher3.tapl.OverviewTask.OverviewSplitTask.DEFAULT;
 import static com.android.launcher3.tapl.OverviewTask.OverviewSplitTask.SPLIT_BOTTOM_OR_RIGHT;
 import static com.android.launcher3.tapl.OverviewTask.OverviewSplitTask.SPLIT_TOP_OR_LEFT;
+import static com.android.launcher3.testing.shared.TestProtocol.SUCCESSFUL_GESTURE_MISMATCH_EVENTS;
 
 import android.graphics.Rect;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.test.uiautomator.By;
@@ -219,6 +221,7 @@
                     return new LaunchedAppState(mLauncher);
                 }
             } else {
+                Log.d(SUCCESSFUL_GESTURE_MISMATCH_EVENTS, "TaskView.launchTaskAnimated");
                 mLauncher.expectEvent(TestProtocol.SEQUENCE_MAIN, TASK_START_EVENT);
                 return new LaunchedAppState(mLauncher);
             }