Snap for 11571437 from b5c4820a5b to 24Q3-release

Change-Id: Ib82d9757f3ac51eaac8efa588b6109b949c8b44b

quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java

@@ -529,52 +529,26 @@ public class TaskbarActivityContext extends BaseTaskbarContext {
 
     /**
      * Creates {@link WindowManager.LayoutParams} for Taskbar, and also sets LP.paramsForRotation
-     * for taskbar showing as navigation bar
+     * for taskbar
      */
     private WindowManager.LayoutParams createAllWindowParams() {
         final int windowType =
                 ENABLE_TASKBAR_NAVBAR_UNIFICATION ? TYPE_NAVIGATION_BAR : TYPE_NAVIGATION_BAR_PANEL;
         WindowManager.LayoutParams windowLayoutParams =
                 createDefaultWindowLayoutParams(windowType, TaskbarActivityContext.WINDOW_TITLE);
-        if (!isPhoneButtonNavMode()) {
-            return windowLayoutParams;
-        }
 
-        // Provide WM layout params for all rotations to cache, see NavigationBar#getBarLayoutParams
-        int width = WindowManager.LayoutParams.MATCH_PARENT;
-        int height = WindowManager.LayoutParams.MATCH_PARENT;
-        int gravity = Gravity.BOTTOM;
         windowLayoutParams.paramsForRotation = new WindowManager.LayoutParams[4];
         for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
             WindowManager.LayoutParams lp =
                     createDefaultWindowLayoutParams(windowType,
                             TaskbarActivityContext.WINDOW_TITLE);
-            switch (rot) {
-                case Surface.ROTATION_0, Surface.ROTATION_180 -> {
-                    // Defaults are fine
-                    width = WindowManager.LayoutParams.MATCH_PARENT;
-                    height = mLastRequestedNonFullscreenSize;
-                    gravity = Gravity.BOTTOM;
-                }
-                case Surface.ROTATION_90 -> {
-                    width = mLastRequestedNonFullscreenSize;
-                    height = WindowManager.LayoutParams.MATCH_PARENT;
-                    gravity = Gravity.END;
-                }
-                case Surface.ROTATION_270 -> {
-                    width = mLastRequestedNonFullscreenSize;
-                    height = WindowManager.LayoutParams.MATCH_PARENT;
-                    gravity = Gravity.START;
-                }
-
+            if (isPhoneButtonNavMode()) {
+                populatePhoneButtonNavModeWindowLayoutParams(rot, lp);
             }
-            lp.width = width;
-            lp.height = height;
-            lp.gravity = gravity;
             windowLayoutParams.paramsForRotation[rot] = lp;
         }
 
-        // Override current layout params
+        // Override with current layout params
         WindowManager.LayoutParams currentParams =
                 windowLayoutParams.paramsForRotation[getDisplay().getRotation()];
         windowLayoutParams.width = currentParams.width;
@@ -584,6 +558,32 @@ public class TaskbarActivityContext extends BaseTaskbarContext {
         return windowLayoutParams;
     }
 
+    /**
+     * Update {@link WindowManager.LayoutParams} with values specific to phone and 3 button
+     * navigation users
+     */
+    private void populatePhoneButtonNavModeWindowLayoutParams(int rot,
+            WindowManager.LayoutParams lp) {
+        lp.width = WindowManager.LayoutParams.MATCH_PARENT;
+        lp.height = WindowManager.LayoutParams.MATCH_PARENT;
+        lp.gravity = Gravity.BOTTOM;
+
+        // Override with per-rotation specific values
+        switch (rot) {
+            case Surface.ROTATION_0, Surface.ROTATION_180 -> {
+                lp.height = mLastRequestedNonFullscreenSize;
+            }
+            case Surface.ROTATION_90 -> {
+                lp.width = mLastRequestedNonFullscreenSize;
+                lp.gravity = Gravity.END;
+            }
+            case Surface.ROTATION_270 -> {
+                lp.width = mLastRequestedNonFullscreenSize;
+                lp.gravity = Gravity.START;
+            }
+        }
+    }
+
     public void onConfigurationChanged(@Config int configChanges) {
         mControllers.onConfigurationChanged(configChanges);
         if (!mIsUserSetupComplete) {
@@ -944,8 +944,14 @@ public class TaskbarActivityContext extends BaseTaskbarContext {
         }
         if (landscapePhoneButtonNav) {
             mWindowLayoutParams.width = size;
+            for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
+                mWindowLayoutParams.paramsForRotation[rot].width = size;
+            }
         } else {
             mWindowLayoutParams.height = size;
+            for (int rot = Surface.ROTATION_0; rot <= Surface.ROTATION_270; rot++) {
+                mWindowLayoutParams.paramsForRotation[rot].height = size;
+            }
         }
         mControllers.taskbarInsetsController.onTaskbarOrBubblebarWindowHeightOrInsetsChanged();
         notifyUpdateLayoutParams();

quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt

@@ -118,11 +118,9 @@ class TaskbarInsetsController(val context: TaskbarActivityContext) : LoggableTas
                 getProvidedInsets(insetsRoundedCornerFlag)
             }
 
-        if (!context.isGestureNav) {
-            if (windowLayoutParams.paramsForRotation != null) {
-                for (layoutParams in windowLayoutParams.paramsForRotation) {
-                    layoutParams.providedInsets = getProvidedInsets(insetsRoundedCornerFlag)
-                }
+        if (windowLayoutParams.paramsForRotation != null) {
+            for (layoutParams in windowLayoutParams.paramsForRotation) {
+                layoutParams.providedInsets = getProvidedInsets(insetsRoundedCornerFlag)
             }
         }
 
@@ -156,19 +154,12 @@ class TaskbarInsetsController(val context: TaskbarActivityContext) : LoggableTas
             )
         }
 
-        val gravity = windowLayoutParams.gravity
-
         // Pre-calculate insets for different providers across different rotations for this gravity
         for (rotation in Surface.ROTATION_0..Surface.ROTATION_270) {
             // Add insets for navbar rotated params
-            if (windowLayoutParams.paramsForRotation != null) {
-                val layoutParams = windowLayoutParams.paramsForRotation[rotation]
-                for (provider in layoutParams.providedInsets) {
-                    setProviderInsets(provider, layoutParams.gravity, rotation)
-                }
-            }
-            for (provider in windowLayoutParams.providedInsets) {
-                setProviderInsets(provider, gravity, rotation)
+            val layoutParams = windowLayoutParams.paramsForRotation[rotation]
+            for (provider in layoutParams.providedInsets) {
+                setProviderInsets(provider, layoutParams.gravity, rotation)
             }
         }
         context.notifyUpdateLayoutParams()

quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java

@@ -900,12 +900,12 @@ public class TaskbarStashController implements TaskbarControllers.LoggableTaskba
         }
 
         // 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());

quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java

@@ -21,6 +21,7 @@ import android.app.PendingIntent;
 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;
@@ -157,6 +158,28 @@ public class ApiWrapper {
         }
     }
 
+    /**
+     * 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.
      */

quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.java

@@ -304,6 +304,12 @@ public class LandscapePagedViewHandler implements RecentsPagedOrientationHandler
         }
     }
 
+    @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,

quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java

@@ -209,6 +209,12 @@ public class PortraitPagedViewHandler extends DefaultPagedViewHandler implements
                 : thumbnailView.getMeasuredWidth()) - (2 * padding);
     }
 
+    @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,

quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.java

@@ -176,6 +176,9 @@ public interface RecentsPagedOrientationHandler extends PagedOrientationHandler
             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.

quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.java

@@ -113,6 +113,12 @@ public class SeascapePagedViewHandler extends LandscapePagedViewHandler {
         }
     }
 
+    @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) {

quickstep/src/com/android/quickstep/views/TaskMenuView.java

@@ -140,11 +140,9 @@ public class TaskMenuView extends AbstractFloatingView {
 
     @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 @@ public class TaskMenuView extends AbstractFloatingView {
      * 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) {

quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java

@@ -503,7 +503,6 @@ public class TaplTestsQuickstep extends AbstractQuickStepTest {
 
     @Test
     @PortraitLandscape
-    @ScreenRecord // b/326839375
     public void testOverviewDeadzones() throws Exception {
         startTestAppsWithCheck();
 

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"

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" />

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>

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>

src/com/android/launcher3/InvariantDeviceProfile.java

@@ -72,6 +72,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
 public class InvariantDeviceProfile {
@@ -577,6 +578,45 @@ public class InvariantDeviceProfile {
         return filteredProfiles;
     }
 
+    /**
+     * 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
      */

src/com/android/launcher3/Launcher.java

@@ -2365,7 +2365,8 @@ public class Launcher extends StatefulActivity<LauncherState>
      * 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.

src/com/android/launcher3/LauncherSettings.java

@@ -138,6 +138,11 @@ public class LauncherSettings {
          */
         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 @@ public class LauncherSettings {
                 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);
             }
         }

src/com/android/launcher3/allapps/PrivateProfileManager.java

@@ -23,15 +23,12 @@ import static com.android.launcher3.allapps.BaseAllAppsAdapter.VIEW_TYPE_PRIVATE
 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.BitmapInfo;
 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 @@ import java.util.function.Predicate;
  */
 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 @@ public class PrivateProfileManager extends UserProfileManager {
                 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 @@ public class PrivateProfileManager extends UserProfileManager {
     /** 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 @@ public class PrivateProfileManager extends UserProfileManager {
 
     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() {

src/com/android/launcher3/model/data/ItemInfo.java

@@ -94,6 +94,10 @@ public class ItemInfo {
      * {@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;
 

src/com/android/launcher3/model/data/ItemInfoWithIcon.java

@@ -120,11 +120,6 @@ public abstract class ItemInfoWithIcon extends ItemInfo {
      */
     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.

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;
+    }
+}

src/com/android/launcher3/provider/RestoreDbTask.java

@@ -50,8 +50,10 @@ import androidx.annotation.NonNull;
 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 @@ public class RestoreDbTask {
         // 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) {

src/com/android/launcher3/touch/ItemClickHandler.java

@@ -369,8 +369,8 @@ public class ItemClickHandler {
                 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());

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 @@ public class FloatingSurfaceView extends AbstractFloatingView implements
         if (mContract == null) {
             return;
         }
-        View icon = mLauncher.getFirstMatchForAppClose(-1,
+        View icon = mLauncher.getFirstMatchForAppClose(NO_MATCHING_ID,
                 mContract.componentName.getPackageName(), mContract.user,
                 false /* supportsAllAppsState */);
 

src/com/android/launcher3/widget/WidgetCell.java

@@ -57,6 +57,8 @@ import com.android.launcher3.model.data.ItemInfoWithIcon;
 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 @@ public class WidgetCell extends LinearLayout {
      * 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 @@ public class WidgetCell extends LinearLayout {
         mWidgetDims.setText(null);
         mWidgetDescription.setText(null);
         mWidgetDescription.setVisibility(GONE);
+        showDescription(true);
+        showDimensions(true);
 
         if (mActiveRequest != null) {
             mActiveRequest.cancel();
@@ -186,6 +190,7 @@ public class WidgetCell extends LinearLayout {
             mWidgetImageContainer.removeView(mAppWidgetHostViewPreview);
         }
         mAppWidgetHostViewPreview = null;
+        mPreviewContainerSize = new Size(0, 0);
         mAppWidgetHostViewScale = 1f;
         mPreviewContainerScale = 1f;
         mItem = null;
@@ -201,30 +206,21 @@ public class WidgetCell extends LinearLayout {
      * 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 @@ public class WidgetCell extends LinearLayout {
         }
     }
 
+    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,
@@ -383,6 +390,16 @@ public class WidgetCell extends LinearLayout {
         mWidgetDescription.setVisibility(show ? VISIBLE : GONE);
     }
 
+    /**
+     * 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 @@ public class WidgetCell extends LinearLayout {
     @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);
     }
 

src/com/android/launcher3/widget/WidgetImageView.java

@@ -82,15 +82,27 @@ public class WidgetImageView extends View {
     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;

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 @@ public class WidgetsBottomSheet extends BaseWidgetSheet {
                 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);

src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java

@@ -36,41 +36,49 @@ public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
             (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 @@ public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
 
     @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 @@ public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
         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 @@ public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
                 mPkgItem,
                 mTitleSectionName,
                 mWidgets,
+                mVisibleWidgetsCount,
                 mIsSearchEntry,
                 /* isWidgetListShown= */ true);
     }
@@ -122,7 +137,28 @@ public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
                 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 @@ public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
                 pkgItem,
                 titleSectionName,
                 items,
-                /* forSearch */ true,
+                /* isSearchEntry */ true,
                 /* isWidgetListShown= */ false);
     }
 }

src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java

@@ -93,18 +93,19 @@ public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots
      * @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 @@ public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots
      * @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 @@ public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots
         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 @@ public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots
 
         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 @@ public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots
             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 @@ public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots
      * 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 @@ public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots
         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. */

src/com/android/launcher3/widget/picker/WidgetsFullSheet.java

@@ -107,7 +107,8 @@ public class WidgetsFullSheet extends BaseWidgetSheet
             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 @@ public class WidgetsFullSheet extends BaseWidgetSheet
         }
 
         if (enableCategorizedWidgetSuggestions()) {
-            mHasRecommendedWidgets = mWidgetRecommendationsView.setRecommendations(
+            mRecommendedWidgetsCount = mWidgetRecommendationsView.setRecommendations(
                     mActivityContext.getPopupDataProvider().getCategorizedRecommendedWidgets(),
                     mDeviceProfile,
                     /* availableHeight= */ getMaxAvailableHeightForRecommendations(),
@@ -589,7 +590,7 @@ public class WidgetsFullSheet extends BaseWidgetSheet
                     /* cellPadding= */ mWidgetCellHorizontalPadding
             );
         } else {
-            mHasRecommendedWidgets = mWidgetRecommendationsView.setRecommendations(
+            mRecommendedWidgetsCount = mWidgetRecommendationsView.setRecommendations(
                     mActivityContext.getPopupDataProvider().getRecommendedWidgets(),
                     mDeviceProfile,
                     /* availableHeight= */ getMaxAvailableHeightForRecommendations(),
@@ -597,7 +598,8 @@ public class WidgetsFullSheet extends BaseWidgetSheet
                     /* cellPadding= */ mWidgetCellHorizontalPadding
             );
         }
-        mWidgetRecommendationsContainer.setVisibility(mHasRecommendedWidgets ? VISIBLE : GONE);
+        mWidgetRecommendationsContainer.setVisibility(
+                mRecommendedWidgetsCount > 0 ? VISIBLE : GONE);
     }
 
     @Px
@@ -790,7 +792,7 @@ public class WidgetsFullSheet extends BaseWidgetSheet
     }
 
     /** 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;
         }

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 @@ public final class WidgetsListTableViewHolderBinder
                 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 @@ public final class WidgetsListTableViewHolderBinder
                 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()) {

src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java

@@ -17,11 +17,13 @@ package com.android.launcher3.widget.picker;
 
 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.TableLayout;
 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 @@ public final class WidgetsRecommendationTableLayout extends TableLayout {
      * 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 @@ public final class WidgetsRecommendationTableLayout extends TableLayout {
         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);
+            }
+            if (totalHeight + rowHeight <= recommendationTableMaxHeight) {
+                totalHeight += rowHeight;
+                filteredRows.add(new ArrayList<>(widgetItems));
             }
-            totalHeight += rowHeight;
         }
 
-        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();
     }
 }

src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java

@@ -167,7 +167,7 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
     @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 class WidgetsTwoPaneSheet extends WidgetsFullSheet {
     public void onRecommendedWidgetsBound() {
         super.onRecommendedWidgetsBound();
 
-        if (mSuggestedWidgetsContainer == null && mHasRecommendedWidgets) {
+        if (mSuggestedWidgetsContainer == null && mRecommendedWidgetsCount > 0) {
             setupSuggestedWidgets(LayoutInflater.from(getContext()));
             mSuggestedWidgetsHeader.callOnClick();
         }
@@ -209,8 +209,9 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
         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 @@ public class WidgetsTwoPaneSheet extends WidgetsFullSheet {
     @Override
     @Px
     protected float getMaxTableHeight(@Px float noWidgetsViewHeight) {
-        return Float.MAX_VALUE;
+        return mContent.getMeasuredHeight() - measureHeightWithVerticalMargins(mHeaderTitle);
     }
 
     @Override

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
+        }
+    }
+}

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),
+    )

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 @@ public final class WidgetsTableUtils {
     /**
      * 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,13 +45,28 @@ public final class WidgetsTableUtils {
         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 @@ public final class WidgetsTableUtils {
             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;

src_ui_overrides/com/android/launcher3/uioverrides/ApiWrapper.java

@@ -107,6 +107,13 @@ public class ApiWrapper {
                         .authority(context.getPackageName()).build());
     }
 
+    /**
+     * 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.
      */

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()
+                }
+            }
+        }
+}

tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java

@@ -47,6 +47,7 @@ import androidx.test.runner.AndroidJUnit4;
 
 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 @@ public class PrivateProfileManagerTest {
     }
 
     @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 {

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."
+            }
+        }
+    }
+}

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))
+    }
+}

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)
+        }
+    }
+}

tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java

@@ -63,6 +63,7 @@ public final class WidgetsTableUtilsTest {
     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 @@ public final class WidgetsTableUtilsTest {
 
 
     @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 @@ public final class WidgetsTableUtilsTest {
                 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 @@ public final class WidgetsTableUtilsTest {
                 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 @@ public final class WidgetsTableUtilsTest {
     }
 
     @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 @@ public final class WidgetsTableUtilsTest {
                 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 @@ public final class WidgetsTableUtilsTest {
                 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 @@ public final class WidgetsTableUtilsTest {
                 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() {