diff --git a/go/src/com/android/launcher3/model/WidgetsModel.java b/go/src/com/android/launcher3/model/WidgetsModel.java
index f8448da..1aa5d03 100644
--- a/go/src/com/android/launcher3/model/WidgetsModel.java
+++ b/go/src/com/android/launcher3/model/WidgetsModel.java
@@ -79,12 +79,13 @@
     }
 
     public WidgetItem getWidgetProviderInfoByProviderName(
-            ComponentName providerName) {
+            ComponentName providerName, UserHandle user) {
         return null;
     }
 
     /** Returns {@link PackageItemInfo} of a pending widget. */
-    public static PackageItemInfo newPendingItemInfo(ComponentName provider) {
-        return new PackageItemInfo(provider.getPackageName());
+    public static PackageItemInfo newPendingItemInfo(
+            Context context, ComponentName provider, UserHandle userHandle) {
+        return new PackageItemInfo(provider.getPackageName(), userHandle);
     }
 }
\ No newline at end of file
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 7b4bf62..ea8282f 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -795,7 +795,8 @@
             return true;
         }
 
-        if (!getRecentsView().isClearAllHidden()) {
+        if (!mActivity.getDeviceProfile().overviewShowAsGrid
+                && !getRecentsView().isClearAllHidden()) {
             getRecentsView().snapToPage(getRecentsView().indexOfChild(this));
             return false;
         } else {
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 9c20e1f..ca3ce78 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -177,7 +177,7 @@
         <attr name="minWidthDps" format="float" />
         <attr name="minHeightDps" format="float" />
 
-        <!-- These min cell values are only used if GridDisplayOption#isScalable is true-->
+        <!-- These min cell values are only used if GridDisplayOption#isScalable is true -->
         <attr name="minCellHeightDps" format="float" />
         <attr name="minCellWidthDps" format="float" />
         <!-- twoPanelPortraitMinCellWidthDps defaults to minCellHeightDps, if not specified -->
@@ -223,6 +223,12 @@
         <!-- allAppsIconTextSize defaults to iconTextSize, if not specified -->
         <attr name="allAppsIconTextSize" format="float" />
 
+        <!-- Margin on left and right of the workspace when GridDisplayOption#isScalable is true -->
+        <attr name="horizontalMargin" format="float"/>
+        <!-- twoPanelLandscapeHorizontalMargin defaults to horizontalMargin if not specified -->
+        <attr name="twoPanelLandscapeHorizontalMargin" format="float"/>
+        <!-- twoPanelPortraitHorizontalMargin defaults to horizontalMargin if not specified -->
+        <attr name="twoPanelPortraitHorizontalMargin" format="float"/>
     </declare-styleable>
 
     <declare-styleable name="CellLayout">
@@ -260,4 +266,18 @@
     <declare-styleable name="WidgetsListRowHeader">
         <attr name="appIconSize" format="dimension" />
     </declare-styleable>
+
+    <declare-styleable name="WidgetSections">
+        <!-- Component name of an app widget provider. -->
+        <attr name="provider" format="string" />
+        <!-- If true, keep the app widget under its app listing in addition to the widget category
+             in the widget picker. Defaults to false if not specified. -->
+        <attr name="alsoKeepInApp" format="boolean" />
+        <!-- The category of an app widget provider. Defaults to -1 if not specified. -->
+        <attr name="category" format="integer" />
+        <!-- The title name of a widget category. -->
+        <attr name="sectionTitle" format="reference" />
+        <!-- The icon drawable of a widget category. -->
+        <attr name="sectionDrawable" format="reference" />
+    </declare-styleable>
 </resources>
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 1cf49fc..b1294b4 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -29,9 +29,6 @@
     <dimen name="dynamic_grid_cell_layout_padding">5.5dp</dimen>
     <dimen name="dynamic_grid_cell_padding_x">8dp</dimen>
 
-    <dimen name="two_panels_home_side_padding_landscape">36dp</dimen>
-    <dimen name="two_panels_home_side_padding_portrait">9dp</dimen>
-
     <!-- Hotseat -->
     <dimen name="dynamic_grid_hotseat_top_padding">8dp</dimen>
     <dimen name="dynamic_grid_hotseat_bottom_padding">2dp</dimen>
@@ -41,7 +38,6 @@
     <dimen name="dynamic_grid_hotseat_side_padding">0dp</dimen>
 
     <!-- Scalable Grid -->
-    <dimen name="scalable_grid_left_right_margin">22dp</dimen>
     <dimen name="scalable_grid_qsb_bottom_margin">42dp</dimen>
 
     <!-- Workspace page indicator -->
diff --git a/res/xml/widget_sections.xml b/res/xml/widget_sections.xml
new file mode 100644
index 0000000..d755de6
--- /dev/null
+++ b/res/xml/widget_sections.xml
@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2021 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<widget-sections xmlns:launcher="http://schemas.android.com/apk/res-auto">
+    <section
+        launcher:category="0"
+        launcher:sectionDrawable="@drawable/ic_conversations_widget_category"
+        launcher:sectionTitle="@string/widget_category_conversations">
+        <widget launcher:provider="com.android.systemui/.people.widget.PeopleSpaceWidgetProvider" />
+    </section>
+</widget-sections>
\ No newline at end of file
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index 192fe30..07221fc 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -97,8 +97,8 @@
     private static final int PORTRAIT_TABLET_LEFT_RIGHT_PADDING_MULTIPLIER = 4;
 
     // Workspace
-    public final int desiredWorkspaceLeftRightOriginalPx;
-    public int desiredWorkspaceLeftRightMarginPx;
+    public final int desiredWorkspaceHorizontalMarginOriginalPx;
+    public int desiredWorkspaceHorizontalMarginPx;
     public final int cellLayoutBorderSpacingOriginalPx;
     public int cellLayoutBorderSpacingPx;
     public final int cellLayoutPaddingLeftRightPx;
@@ -270,11 +270,8 @@
 
         edgeMarginPx = res.getDimensionPixelSize(R.dimen.dynamic_grid_edge_margin);
 
-        desiredWorkspaceLeftRightMarginPx = isVerticalBarLayout() ? 0 : isScalableGrid
-                ? res.getDimensionPixelSize(R.dimen.scalable_grid_left_right_margin)
-                : res.getDimensionPixelSize(R.dimen.dynamic_grid_left_right_margin);
-        desiredWorkspaceLeftRightOriginalPx = desiredWorkspaceLeftRightMarginPx;
-
+        desiredWorkspaceHorizontalMarginPx = getHorizontalMarginPx(inv, res);
+        desiredWorkspaceHorizontalMarginOriginalPx = desiredWorkspaceHorizontalMarginPx;
 
         allAppsOpenVerticalTranslate = res.getDimensionPixelSize(
                 R.dimen.all_apps_open_vertical_translate);
@@ -446,6 +443,29 @@
                 new DotRenderer(allAppsIconSizePx, dotPath, DEFAULT_DOT_SIZE);
     }
 
+    private int getHorizontalMarginPx(InvariantDeviceProfile idp, Resources res) {
+        if (isVerticalBarLayout()) {
+            return 0;
+        }
+
+        int horizontalMarginPx;
+
+        if (isScalableGrid) {
+            if (isTwoPanels) {
+                if (isLandscape) {
+                    horizontalMarginPx = pxFromDp(idp.twoPanelLandscapeHorizontalMargin, mMetrics);
+                } else {
+                    horizontalMarginPx = pxFromDp(idp.twoPanelPortraitHorizontalMargin, mMetrics);
+                }
+            } else {
+                horizontalMarginPx = pxFromDp(idp.horizontalMargin, mMetrics);
+            }
+        } else {
+            horizontalMarginPx = res.getDimensionPixelSize(R.dimen.dynamic_grid_left_right_margin);
+        }
+        return horizontalMarginPx;
+    }
+
     private void updateHotseatIconSize(int hotseatIconSizePx) {
         // Ensure there is enough space for folder icons, which have a slightly larger radius.
         hotseatCellHeightPx = (int) Math.ceil(hotseatIconSizePx * ICON_OVERLAP_FACTOR);
@@ -554,7 +574,7 @@
             allAppsLeftRightPadding = Math.max(1, (availableWidthPx - usedWidth) / 2);
         } else {
             allAppsLeftRightPadding =
-                    desiredWorkspaceLeftRightMarginPx + cellLayoutPaddingLeftRightPx;
+                    desiredWorkspaceHorizontalMarginPx + cellLayoutPaddingLeftRightPx;
         }
     }
 
@@ -581,7 +601,7 @@
             int numColumns = isTwoPanels ? inv.numColumns * 2 : inv.numColumns;
             float usedWidth = (cellWidthPx * numColumns)
                     + (cellLayoutBorderSpacingPx * (numColumns - 1))
-                    + (desiredWorkspaceLeftRightMarginPx * 2);
+                    + (desiredWorkspaceHorizontalMarginPx * 2);
             // We do not subtract padding here, as we also scale the workspace padding if needed.
             scaleX = availableWidthPx / usedWidth;
             shouldScale = true;
@@ -647,7 +667,8 @@
             int cellContentHeight = iconSizePx + iconDrawablePaddingPx
                     + Utilities.calculateTextHeight(iconTextSizePx);
             cellYPaddingPx = Math.max(0, cellHeightPx - cellContentHeight) / 2;
-            desiredWorkspaceLeftRightMarginPx = (int) (desiredWorkspaceLeftRightOriginalPx * scale);
+            desiredWorkspaceHorizontalMarginPx =
+                    (int) (desiredWorkspaceHorizontalMarginOriginalPx * scale);
         } else {
             cellWidthPx = iconSizePx + iconDrawablePaddingPx;
             cellHeightPx = (int) Math.ceil(iconSizePx * ICON_OVERLAP_FACTOR)
@@ -851,9 +872,9 @@
             int paddingBottom = hotseatTop + workspacePageIndicatorHeight
                     + workspaceBottomPadding - mWorkspacePageIndicatorOverlapWorkspace;
 
-            padding.set(desiredWorkspaceLeftRightMarginPx,
+            padding.set(desiredWorkspaceHorizontalMarginPx,
                     (isScalableGrid ? workspaceTopPadding : edgeMarginPx),
-                    desiredWorkspaceLeftRightMarginPx,
+                    desiredWorkspaceHorizontalMarginPx,
                     paddingBottom);
         }
     }
@@ -1081,8 +1102,8 @@
 
         writer.println(prefix + pxToDpStr("cellLayoutBorderSpacingPx",
                 cellLayoutBorderSpacingPx));
-        writer.println(prefix + pxToDpStr("desiredWorkspaceLeftRightMarginPx",
-                desiredWorkspaceLeftRightMarginPx));
+        writer.println(prefix + pxToDpStr("desiredWorkspaceHorizontalMarginPx",
+                desiredWorkspaceHorizontalMarginPx));
 
         writer.println(prefix + pxToDpStr("allAppsIconSizePx", allAppsIconSizePx));
         writer.println(prefix + pxToDpStr("allAppsIconTextSizePx", allAppsIconTextSizePx));
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 8631040..92cc9e8 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -118,6 +118,9 @@
     public float twoPanelLandscapeMinCellHeightDps;
     public float twoPanelLandscapeMinCellWidthDps;
     public float borderSpacing;
+    public float horizontalMargin;
+    public float twoPanelLandscapeHorizontalMargin;
+    public float twoPanelPortraitHorizontalMargin;
 
     private SparseArray<TypedValue> mExtraAttrs;
 
@@ -289,6 +292,10 @@
         borderSpacing = displayOption.borderSpacing;
         allAppsCellSpacing = displayOption.allAppsCellSpacing;
 
+        horizontalMargin = displayOption.horizontalMargin;
+        twoPanelLandscapeHorizontalMargin = displayOption.twoPanelLandscapeHorizontalMargin;
+        twoPanelPortraitHorizontalMargin = displayOption.twoPanelPortraitHorizontalMargin;
+
         numShownHotseatIcons = closestProfile.numHotseatIcons;
         numDatabaseHotseatIcons = isSplitDisplay
                 ? closestProfile.numDatabaseHotseatIcons : closestProfile.numHotseatIcons;
@@ -728,6 +735,10 @@
         private float allAppsCellSpacing;
         private float borderSpacing;
 
+        private float horizontalMargin;
+        private float twoPanelLandscapeHorizontalMargin;
+        private float twoPanelPortraitHorizontalMargin;
+
         private final float[] iconSizes = new float[COUNT_TOTAL];
         private final float[] textSizes = new float[COUNT_TOTAL];
 
@@ -790,6 +801,14 @@
                     a.getFloat(R.styleable.ProfileDisplayOption_twoPanelLandscapeIconTextSize,
                             textSizes[INDEX_LANDSCAPE]);
 
+            horizontalMargin = a.getFloat(R.styleable.ProfileDisplayOption_horizontalMargin, 0);
+            twoPanelLandscapeHorizontalMargin = a.getFloat(
+                    R.styleable.ProfileDisplayOption_twoPanelLandscapeHorizontalMargin,
+                    horizontalMargin);
+            twoPanelPortraitHorizontalMargin = a.getFloat(
+                    R.styleable.ProfileDisplayOption_twoPanelPortraitHorizontalMargin,
+                    horizontalMargin);
+
             a.recycle();
         }
 
@@ -820,6 +839,9 @@
             twoPanelLandscapeMinCellWidthDps *= w;
             borderSpacing *= w;
             allAppsCellSpacing *= w;
+            horizontalMargin *= w;
+            twoPanelLandscapeHorizontalMargin *= w;
+            twoPanelPortraitHorizontalMargin *= w;
             return this;
         }
 
@@ -836,6 +858,9 @@
             twoPanelLandscapeMinCellWidthDps += p.twoPanelLandscapeMinCellWidthDps;
             borderSpacing += p.borderSpacing;
             allAppsCellSpacing += p.allAppsCellSpacing;
+            horizontalMargin += p.horizontalMargin;
+            twoPanelLandscapeHorizontalMargin += p.twoPanelLandscapeHorizontalMargin;
+            twoPanelPortraitHorizontalMargin += p.twoPanelPortraitHorizontalMargin;
             return this;
         }
     }
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 9d27bbd..0acafc0 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -147,6 +147,11 @@
             false,
             "Enable loading workspace icons in bulk.");
 
+    public static final BooleanFlag ENABLE_BULK_ALL_APPS_ICON_LOADING = getDebugFlag(
+            "ENABLE_BULK_ALL_APPS_ICON_LOADING",
+            false,
+            "Enable loading all apps icons in bulk.");
+
     // Keep as DeviceFlag for remote disable in emergency.
     public static final BooleanFlag ENABLE_OVERVIEW_SELECTIONS = new DeviceFlag(
             "ENABLE_OVERVIEW_SELECTIONS", true, "Show Select Mode button in Overview Actions");
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index e09b38b..468cb55 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -352,7 +352,7 @@
 
     private void inflateAndAddWidgets(LauncherAppWidgetInfo info, WidgetsModel widgetsModel) {
         WidgetItem widgetItem = widgetsModel.getWidgetProviderInfoByProviderName(
-                info.providerName);
+                info.providerName, info.user);
         if (widgetItem == null) {
             return;
         }
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 60d6e83..936eeb9 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -18,6 +18,7 @@
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
 
 import static java.util.stream.Collectors.groupingBy;
 
@@ -46,7 +47,6 @@
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherFiles;
-import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.icons.ComponentWithLabel.ComponentCachingLogic;
@@ -63,6 +63,8 @@
 import com.android.launcher3.util.InstantAppResolver;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.widget.WidgetSections;
+import com.android.launcher3.widget.WidgetSections.WidgetSection;
 
 import java.util.Collections;
 import java.util.List;
@@ -275,7 +277,8 @@
             getTitleAndIcon(appInfo, false);
             return appInfo.bitmap;
         } else {
-            PackageItemInfo pkgInfo = new PackageItemInfo(shortcutInfo.getPackage());
+            PackageItemInfo pkgInfo = new PackageItemInfo(shortcutInfo.getPackage(),
+                    shortcutInfo.getUserHandle());
             getTitleAndIconForApp(pkgInfo, false);
             return pkgInfo.bitmap;
         }
@@ -409,8 +412,10 @@
         CacheEntry entry = getEntryForPackageLocked(
                 infoInOut.packageName, infoInOut.user, useLowResIcon);
         applyCacheEntry(entry, infoInOut);
-        if (infoInOut.category == PackageItemInfo.CONVERSATIONS) {
-            infoInOut.title = mContext.getString(R.string.widget_category_conversations);
+        if (infoInOut.widgetCategory != NO_CATEGORY) {
+            WidgetSection widgetSection = WidgetSections.getWidgetSections(mContext)
+                    .get(infoInOut.widgetCategory);
+            infoInOut.title = mContext.getString(widgetSection.mSectionTitle);
             infoInOut.contentDescription = mPackageManager.getUserBadgedLabel(
                     infoInOut.title, infoInOut.user);
         }
diff --git a/src/com/android/launcher3/model/AllAppsList.java b/src/com/android/launcher3/model/AllAppsList.java
index 92b5885..dbed9a9 100644
--- a/src/com/android/launcher3/model/AllAppsList.java
+++ b/src/com/android/launcher3/model/AllAppsList.java
@@ -130,30 +130,54 @@
      * If the app is already in the list, doesn't add it.
      */
     public void add(AppInfo info, LauncherActivityInfo activityInfo) {
+        add(info, activityInfo, true);
+    }
+
+    public void add(AppInfo info, LauncherActivityInfo activityInfo, boolean loadIcon) {
         if (!mAppFilter.shouldShowApp(info.componentName)) {
             return;
         }
         if (findAppInfo(info.componentName, info.user) != null) {
             return;
         }
-        mIconCache.getTitleAndIcon(info, activityInfo, false /* useLowResIcon */);
-        info.sectionName = mIndex.computeSectionName(info.title);
+        if (loadIcon) {
+            mIconCache.getTitleAndIcon(info, activityInfo, false /* useLowResIcon */);
+            info.sectionName = mIndex.computeSectionName(info.title);
+        }
 
         data.add(info);
         mDataChanged = true;
     }
 
-    public void addPromiseApp(Context context, PackageInstallInfo installInfo) {
-        // only if not yet installed
-        if (!new PackageManagerHelper(context)
-                .isAppInstalled(installInfo.packageName, installInfo.user)) {
-            AppInfo info = new AppInfo(installInfo);
-            mIconCache.getTitleAndIcon(info, info.usingLowResIcon());
-            info.sectionName = mIndex.computeSectionName(info.title);
+    @Nullable
+    public AppInfo addPromiseApp(Context context, PackageInstallInfo installInfo) {
+        return addPromiseApp(context, installInfo, true);
+    }
 
-            data.add(info);
-            mDataChanged = true;
+    @Nullable
+    public AppInfo addPromiseApp(
+            Context context, PackageInstallInfo installInfo, boolean loadIcon) {
+        // only if not yet installed
+        if (new PackageManagerHelper(context)
+                .isAppInstalled(installInfo.packageName, installInfo.user)) {
+            return null;
         }
+        AppInfo promiseAppInfo = new AppInfo(installInfo);
+
+        if (loadIcon) {
+            mIconCache.getTitleAndIcon(promiseAppInfo, promiseAppInfo.usingLowResIcon());
+            promiseAppInfo.sectionName = mIndex.computeSectionName(promiseAppInfo.title);
+        }
+
+        data.add(promiseAppInfo);
+        mDataChanged = true;
+
+        return promiseAppInfo;
+    }
+
+    public void updateSectionName(AppInfo appInfo) {
+        appInfo.sectionName = mIndex.computeSectionName(appInfo.title);
+
     }
 
     /** Updates the given PackageInstallInfo's associated AppInfo's installation info. */
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
index 992a3f2..0fc4c2d 100644
--- a/src/com/android/launcher3/model/DeviceGridState.java
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -24,7 +24,6 @@
 import android.content.Context;
 import android.content.SharedPreferences;
 import android.text.TextUtils;
-import android.util.Log;
 
 import androidx.annotation.IntDef;
 
@@ -135,17 +134,8 @@
     public boolean isCompatible(DeviceGridState other) {
         if (this == other) return true;
         if (other == null) return false;
-        boolean isCompatible = mNumHotseat == other.mNumHotseat
+        return mNumHotseat == other.mNumHotseat
                 && deviceTypeCompatible(mDeviceType, other.mDeviceType)
                 && Objects.equals(mGridSizeString, other.mGridSizeString);
-        // TODO(b/198965093): Temporary fix for multi-display devices, ignore hotseat size changes
-        //  and type compatibility.
-        if ((mDeviceType == TYPE_MULTI_DISPLAY || other.mDeviceType == TYPE_MULTI_DISPLAY)
-                && !isCompatible && Objects.equals(mGridSizeString, other.mGridSizeString)) {
-            Log.d("b/198965093", "Hotseat and deice type compatibility ignored: " + this
-                    + ", other: " + other);
-            isCompatible = true;
-        }
-        return isCompatible;
     }
 }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 1249606..8a494ba 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -804,8 +804,9 @@
                                 if (appWidgetInfo.restoreStatus !=
                                         LauncherAppWidgetInfo.RESTORE_COMPLETED) {
                                     appWidgetInfo.pendingItemInfo = WidgetsModel.newPendingItemInfo(
-                                            appWidgetInfo.providerName);
-                                    appWidgetInfo.pendingItemInfo.user = appWidgetInfo.user;
+                                            mApp.getContext(),
+                                            appWidgetInfo.providerName,
+                                            appWidgetInfo.user);
                                     mIconCache.getTitleAndIconForApp(
                                             appWidgetInfo.pendingItemInfo, false);
                                 }
@@ -936,6 +937,8 @@
         List<LauncherActivityInfo> allActivityList = new ArrayList<>();
         // Clear the list of apps
         mBgAllAppsList.clear();
+
+        List<IconRequestInfo<AppInfo>> iconRequestInfos = new ArrayList<>();
         for (UserHandle user : profiles) {
             // Query for the set of apps
             final List<LauncherActivityInfo> apps = mLauncherApps.getActivityList(null, user);
@@ -948,18 +951,43 @@
             // Create the ApplicationInfos
             for (int i = 0; i < apps.size(); i++) {
                 LauncherActivityInfo app = apps.get(i);
-                // This builds the icon bitmaps.
-                mBgAllAppsList.add(new AppInfo(app, user, quietMode), app);
+                AppInfo appInfo = new AppInfo(app, user, quietMode);
+
+                iconRequestInfos.add(new IconRequestInfo<>(
+                        appInfo, app, /* useLowResIcon= */ false));
+                mBgAllAppsList.add(
+                        appInfo, app, !FeatureFlags.ENABLE_BULK_ALL_APPS_ICON_LOADING.get());
             }
             allActivityList.addAll(apps);
         }
 
+
         if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
             // get all active sessions and add them to the all apps list
             for (PackageInstaller.SessionInfo info :
                     mSessionHelper.getAllVerifiedSessions()) {
-                mBgAllAppsList.addPromiseApp(mApp.getContext(),
-                        PackageInstallInfo.fromInstallingState(info));
+                AppInfo promiseAppInfo = mBgAllAppsList.addPromiseApp(
+                        mApp.getContext(),
+                        PackageInstallInfo.fromInstallingState(info),
+                        !FeatureFlags.ENABLE_BULK_ALL_APPS_ICON_LOADING.get());
+
+                if (promiseAppInfo != null) {
+                    iconRequestInfos.add(new IconRequestInfo<>(
+                            promiseAppInfo,
+                            /* launcherActivityInfo= */ null,
+                            promiseAppInfo.usingLowResIcon()));
+                }
+            }
+        }
+
+        if (FeatureFlags.ENABLE_BULK_ALL_APPS_ICON_LOADING.get()) {
+            Trace.beginSection("LoadAllAppsIconsInBulk");
+            try {
+                mIconCache.getTitlesAndIconsInBulk(iconRequestInfos);
+                iconRequestInfos.forEach(iconRequestInfo ->
+                        mBgAllAppsList.updateSectionName(iconRequestInfo.itemInfo));
+            } finally {
+                Trace.endSection();
             }
         }
 
diff --git a/src/com/android/launcher3/model/data/IconRequestInfo.java b/src/com/android/launcher3/model/data/IconRequestInfo.java
index 2f566f6..5dc6a3b 100644
--- a/src/com/android/launcher3/model/data/IconRequestInfo.java
+++ b/src/com/android/launcher3/model/data/IconRequestInfo.java
@@ -50,6 +50,19 @@
     public IconRequestInfo(
             @NonNull T itemInfo,
             @Nullable LauncherActivityInfo launcherActivityInfo,
+            boolean useLowResIcon) {
+        this(
+                itemInfo,
+                launcherActivityInfo,
+                /* packageName= */ null,
+                /* resourceName= */ null,
+                /* iconBlob= */ null,
+                useLowResIcon);
+    }
+
+    public IconRequestInfo(
+            @NonNull T itemInfo,
+            @Nullable LauncherActivityInfo launcherActivityInfo,
             @Nullable String packageName,
             @Nullable String resourceName,
             @Nullable byte[] iconBlob,
diff --git a/src/com/android/launcher3/model/data/PackageItemInfo.java b/src/com/android/launcher3/model/data/PackageItemInfo.java
index a81fe6a..0055763 100644
--- a/src/com/android/launcher3/model/data/PackageItemInfo.java
+++ b/src/com/android/launcher3/model/data/PackageItemInfo.java
@@ -16,47 +16,41 @@
 
 package com.android.launcher3.model.data;
 
-import androidx.annotation.IntDef;
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
+import android.os.UserHandle;
 
 import com.android.launcher3.LauncherSettings;
 
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
 import java.util.Objects;
 
 /**
  * Represents a {@link Package} in the widget tray section.
  */
 public class PackageItemInfo extends ItemInfoWithIcon {
-    @Retention(RetentionPolicy.SOURCE)
-    @IntDef({NO_CATEGORY, CONVERSATIONS})
-    public @interface Category{}
-    /** The package is not categorized in the widget tray. */
-    public static final int NO_CATEGORY = 0;
-    /** The package is categorized to conversations widget in the widget tray. */
-    public static final int CONVERSATIONS = 1;
-
     /**
      * Package name of the {@link PackageItemInfo}.
      */
     public final String packageName;
 
     /** Represents a widget category shown in the widget tray section. */
-    @Category public final int category;
+    public final int widgetCategory;
 
-    public PackageItemInfo(String packageName) {
-        this(packageName, NO_CATEGORY);
+    public PackageItemInfo(String packageName, UserHandle user) {
+        this(packageName, NO_CATEGORY, user);
     }
 
-    public PackageItemInfo(String packageName, @Category int category) {
+    public PackageItemInfo(String packageName, int widgetCategory, UserHandle user) {
         this.packageName = packageName;
-        this.category = category;
+        this.widgetCategory = widgetCategory;
+        this.user = user;
         this.itemType = LauncherSettings.Favorites.ITEM_TYPE_NON_ACTIONABLE;
     }
 
     public PackageItemInfo(PackageItemInfo copy) {
         this.packageName = copy.packageName;
-        this.category = copy.category;
+        this.widgetCategory = copy.widgetCategory;
+        this.user = copy.user;
         this.itemType = LauncherSettings.Favorites.ITEM_TYPE_NON_ACTIONABLE;
     }
 
@@ -75,11 +69,13 @@
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         PackageItemInfo that = (PackageItemInfo) o;
-        return Objects.equals(packageName, that.packageName);
+        return Objects.equals(packageName, that.packageName)
+                && Objects.equals(user, that.user)
+                && widgetCategory == that.widgetCategory;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(packageName, user);
+        return Objects.hash(packageName, user, widgetCategory);
     }
 }
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDragLayer.java b/src/com/android/launcher3/secondarydisplay/SecondaryDragLayer.java
index f78f6dd..1820933 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDragLayer.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDragLayer.java
@@ -112,7 +112,7 @@
         for (int i = 0; i < count; i++) {
             final View child = getChildAt(i);
             if (child == mAppsView) {
-                int padding = 2 * (grid.desiredWorkspaceLeftRightMarginPx
+                int padding = 2 * (grid.desiredWorkspaceHorizontalMarginPx
                         + grid.cellLayoutPaddingLeftRightPx);
 
                 int maxWidth = grid.allAppsCellWidthPx * grid.numShownAllAppsColumns + padding;
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index d05ba3d..931ecb8 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -257,6 +257,14 @@
         }
         if (!newInfo.supportedBounds.equals(oldInfo.supportedBounds)) {
             change |= CHANGE_SUPPORTED_BOUNDS;
+
+            PortraitSize realSize = new PortraitSize(newInfo.currentSize.x, newInfo.currentSize.y);
+            PortraitSize expectedSize = oldInfo.mInternalDisplays.get(
+                    ApiWrapper.getUniqueId(display));
+            if (!realSize.equals(expectedSize) && display.getState() == Display.STATE_OFF) {
+                Log.e("b/198965093", "Display size changed while display is off, ignoring change");
+                return;
+            }
         }
 
         if (change != 0) {
diff --git a/src/com/android/launcher3/util/PackageUserKey.java b/src/com/android/launcher3/util/PackageUserKey.java
index 3a3b5a2..92d9737 100644
--- a/src/com/android/launcher3/util/PackageUserKey.java
+++ b/src/com/android/launcher3/util/PackageUserKey.java
@@ -1,19 +1,24 @@
 package com.android.launcher3.util;
 
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
 import android.os.UserHandle;
 import android.service.notification.StatusBarNotification;
+import android.text.TextUtils;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.PackageItemInfo;
 
-import java.util.Arrays;
+import java.util.Objects;
 
-/** Creates a hash key based on package name and user. */
+/** Creates a hash key based on package name, widget category, and user. */
 public class PackageUserKey {
 
     public String mPackageName;
+    public int mWidgetCategory;
     public UserHandle mUser;
     private int mHashCode;
 
@@ -27,14 +32,31 @@
         return new PackageUserKey(notification.getPackageName(), notification.getUser());
     }
 
+    /** Creates a {@link PackageUserKey} from {@link PackageItemInfo}. */
+    public static PackageUserKey fromPackageItemInfo(PackageItemInfo info) {
+        if (TextUtils.isEmpty(info.packageName) && info.widgetCategory != NO_CATEGORY) {
+            return new PackageUserKey(info.widgetCategory, info.user);
+        }
+        return new PackageUserKey(info.packageName, info.user);
+    }
+
     public PackageUserKey(String packageName, UserHandle user) {
         update(packageName, user);
     }
 
+    public PackageUserKey(int widgetCategory, UserHandle user) {
+        update(/* packageName= */ "", widgetCategory, user);
+    }
+
     public void update(String packageName, UserHandle user) {
+        update(packageName, NO_CATEGORY, user);
+    }
+
+    private void update(String packageName, int widgetCategory, UserHandle user) {
         mPackageName = packageName;
+        mWidgetCategory = widgetCategory;
         mUser = user;
-        mHashCode = Arrays.hashCode(new Object[] {packageName, user});
+        mHashCode = Objects.hash(packageName, widgetCategory, user);
     }
 
     /**
@@ -59,12 +81,14 @@
     public boolean equals(Object obj) {
         if (!(obj instanceof PackageUserKey)) return false;
         PackageUserKey otherKey = (PackageUserKey) obj;
-        return mPackageName.equals(otherKey.mPackageName) && mUser.equals(otherKey.mUser);
+        return Objects.equals(mPackageName, otherKey.mPackageName)
+                && mWidgetCategory == otherKey.mWidgetCategory
+                && Objects.equals(mUser, otherKey.mUser);
     }
 
     @NonNull
     @Override
     public String toString() {
-        return mPackageName + "#" + mUser;
+        return mPackageName + "#" + mUser + ",category=" + mWidgetCategory;
     }
 }
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index b6bb6aa..553ba13 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -17,7 +17,7 @@
 package com.android.launcher3.widget;
 
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
-import static com.android.launcher3.model.data.PackageItemInfo.CONVERSATIONS;
+import static com.android.launcher3.widget.WidgetSections.getWidgetSections;
 
 import android.content.Context;
 import android.graphics.Canvas;
@@ -89,8 +89,8 @@
         setOnClickListener(ItemClickHandler.INSTANCE);
 
         if (info.pendingItemInfo == null) {
-            info.pendingItemInfo = new PackageItemInfo(info.providerName.getPackageName());
-            info.pendingItemInfo.user = info.user;
+            info.pendingItemInfo = new PackageItemInfo(info.providerName.getPackageName(),
+                    info.user);
             cache.updateIconInBackground(this, info.pendingItemInfo);
         } else {
             reapplyItemInfo(info.pendingItemInfo);
@@ -338,10 +338,11 @@
      */
     @Nullable
     private Drawable getWidgetCategoryIcon() {
-        switch (mInfo.pendingItemInfo.category) {
-            case CONVERSATIONS:
-                return getContext().getDrawable(R.drawable.ic_conversations_widget_category);
+        if (mInfo.pendingItemInfo.widgetCategory == WidgetSections.NO_CATEGORY) {
+            return null;
         }
-        return null;
+        Context context = getContext();
+        return context.getDrawable(getWidgetSections(context).get(
+                mInfo.pendingItemInfo.widgetCategory).mSectionDrawable);
     }
 }
diff --git a/src/com/android/launcher3/widget/WidgetSections.java b/src/com/android/launcher3/widget/WidgetSections.java
new file mode 100644
index 0000000..c45b095
--- /dev/null
+++ b/src/com/android/launcher3/widget/WidgetSections.java
@@ -0,0 +1,134 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.widget;
+
+import static android.content.res.Resources.ID_NULL;
+
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.res.TypedArray;
+import android.content.res.XmlResourceParser;
+import android.util.ArrayMap;
+import android.util.AttributeSet;
+import android.util.SparseArray;
+import android.util.Xml;
+
+import androidx.annotation.DrawableRes;
+import androidx.annotation.StringRes;
+
+import com.android.launcher3.R;
+import com.android.launcher3.util.IntSet;
+
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
+import java.util.Map;
+
+/** A helper class to parse widget sections (categories) resource overlay. */
+public final class WidgetSections {
+    /** The package is not categorized in the widget tray. */
+    public static final int NO_CATEGORY = -1;
+
+    private static final String TAG_SECTION_NAME = "section";
+    private static final String TAG_WIDGET_NAME = "widget";
+
+    private static SparseArray<WidgetSection> sWidgetSections;
+    private static Map<ComponentName, IntSet> sWidgetsToCategories;
+
+    /** Returns a list of widget sections that are shown in the widget picker. */
+    public static synchronized SparseArray<WidgetSection> getWidgetSections(Context context) {
+        if (sWidgetSections != null) {
+            return sWidgetSections;
+        }
+        parseWidgetSectionsXml(context);
+        return sWidgetSections;
+    }
+
+    /** Returns a map which maps app widget providers to app widget categories. */
+    public static synchronized Map<ComponentName, IntSet> getWidgetsToCategory(
+            Context context) {
+        if (sWidgetsToCategories != null) {
+            return sWidgetsToCategories;
+        }
+        parseWidgetSectionsXml(context);
+        return sWidgetsToCategories;
+    }
+
+    private static synchronized void parseWidgetSectionsXml(Context context) {
+        SparseArray<WidgetSection> widgetSections = new SparseArray();
+        Map<ComponentName, IntSet> widgetsToCategories = new ArrayMap<>();
+        try (XmlResourceParser parser = context.getResources().getXml(R.xml.widget_sections)) {
+            final int depth = parser.getDepth();
+            int type;
+            while (((type = parser.next()) != XmlPullParser.END_TAG
+                    || parser.getDepth() > depth) && type != XmlPullParser.END_DOCUMENT) {
+                if ((type == XmlPullParser.START_TAG)
+                        && TAG_SECTION_NAME.equals(parser.getName())) {
+                    AttributeSet sectionAttributes = Xml.asAttributeSet(parser);
+                    WidgetSection section = new WidgetSection(context, sectionAttributes);
+                    final int sectionDepth = parser.getDepth();
+                    while (((type = parser.next()) != XmlPullParser.END_TAG
+                                    || parser.getDepth() > sectionDepth)
+                            && type != XmlPullParser.END_DOCUMENT) {
+                        if ((type == XmlPullParser.START_TAG)
+                                && TAG_WIDGET_NAME.equals(parser.getName())) {
+                            TypedArray a = context.obtainStyledAttributes(
+                                    Xml.asAttributeSet(parser), R.styleable.WidgetSections);
+                            ComponentName provider = ComponentName.unflattenFromString(
+                                    a.getString(R.styleable.WidgetSections_provider));
+                            boolean alsoKeepInApp = a.getBoolean(
+                                    R.styleable.WidgetSections_alsoKeepInApp,
+                                    /* defValue= */ false);
+                            final IntSet categories;
+                            if (widgetsToCategories.containsKey(provider)) {
+                                categories = widgetsToCategories.get(provider);
+                            } else {
+                                categories = new IntSet();
+                                widgetsToCategories.put(provider, categories);
+                            }
+                            if (alsoKeepInApp) {
+                                categories.add(NO_CATEGORY);
+                            }
+                            categories.add(section.mCategory);
+                        }
+                    }
+                    widgetSections.put(section.mCategory, section);
+                }
+            }
+            sWidgetSections = widgetSections;
+            sWidgetsToCategories = widgetsToCategories;
+        } catch (IOException | XmlPullParserException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /** A data class which contains a widget section's information. */
+    public static final class WidgetSection {
+        public final int mCategory;
+        @StringRes
+        public final int mSectionTitle;
+        @DrawableRes
+        public final int mSectionDrawable;
+
+        public WidgetSection(Context context, AttributeSet attrs) {
+            TypedArray a = context.obtainStyledAttributes(attrs, R.styleable.WidgetSections);
+            mCategory = a.getInt(R.styleable.WidgetSections_category, NO_CATEGORY);
+            mSectionTitle = a.getResourceId(R.styleable.WidgetSections_sectionTitle, ID_NULL);
+            mSectionDrawable = a.getResourceId(R.styleable.WidgetSections_sectionDrawable, ID_NULL);
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java b/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java
index e62425f..7f24905 100644
--- a/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.widget.model;
 
+import android.os.Process;
+
 import com.android.launcher3.model.data.PackageItemInfo;
 
 import java.util.Collections;
@@ -26,7 +28,9 @@
 public class WidgetListSpaceEntry extends WidgetsListBaseEntry {
 
     public WidgetListSpaceEntry() {
-        super(new PackageItemInfo(""), "", Collections.EMPTY_LIST);
+        super(new PackageItemInfo(/* packageName= */ "", Process.myUserHandle()),
+                /* titleSectionName= */ "",
+                Collections.EMPTY_LIST);
         mPkgItem.title = "";
     }
 
diff --git a/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java b/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
index 7372751..35f11bd 100644
--- a/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
+++ b/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
@@ -24,5 +24,5 @@
     /**
      * Calls when a header is clicked to show / hide widgets for a package.
      */
-    void onHeaderClicked(boolean showWidgets, PackageUserKey packageUserKey);
+    void onHeaderClicked(boolean showWidgets, PackageUserKey key);
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index d52134c..0e5a7d7 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -98,7 +98,7 @@
     private Predicate<WidgetsListBaseEntry> mHeaderAndSelectedContentFilter = entry ->
             entry instanceof WidgetsListHeaderEntry
                     || entry instanceof WidgetsListSearchHeaderEntry
-                    || new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)
+                    || PackageUserKey.fromPackageItemInfo(entry.mPkgItem)
                             .equals(mWidgetsContentVisiblePackageUserKey);
     @Nullable private Predicate<WidgetsListBaseEntry> mFilter = null;
     @Nullable private RecyclerView mRecyclerView;
@@ -252,10 +252,11 @@
         return entry instanceof WidgetsListBaseEntry.Header && matchesKey(entry, key);
     }
 
-    private static boolean matchesKey(
-            @NonNull WidgetsListBaseEntry entry, @Nullable PackageUserKey key) {
+    private static boolean matchesKey(@NonNull WidgetsListBaseEntry entry,
+            @Nullable PackageUserKey key) {
         if (key == null) return false;
         return entry.mPkgItem.packageName.equals(key.mPackageName)
+                && entry.mPkgItem.widgetCategory == key.mWidgetCategory
                 && entry.mPkgItem.user.equals(key.mUser);
     }
 
@@ -434,11 +435,10 @@
                         .filter(entry -> entry instanceof WidgetsListHeaderEntry)
                         .map(entry -> entry.mPkgItem)
                         .collect(Collectors.toMap(
-                                entry -> new PackageUserKey(entry.packageName, entry.user),
+                                entry -> PackageUserKey.fromPackageItemInfo(entry),
                                 entry -> entry));
         for (WidgetsListBaseEntry visibleEntry: mVisibleEntries) {
-            PackageUserKey key = new PackageUserKey(visibleEntry.mPkgItem.packageName,
-                    visibleEntry.mPkgItem.user);
+            PackageUserKey key = PackageUserKey.fromPackageItemInfo(visibleEntry.mPkgItem);
             PackageItemInfo packageItemInfo = packagesInfo.get(key);
             if (packageItemInfo != null
                     && !visibleEntry.mPkgItem.title.equals(packageItemInfo.title)) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
index ebd2d10..932e06d 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
 import android.content.Context;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
@@ -41,6 +43,8 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.PluralMessageFormat;
 import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.WidgetSections;
+import com.android.launcher3.widget.WidgetSections.WidgetSection;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
@@ -173,12 +177,12 @@
 
     private void setIcon(PackageItemInfo info) {
         Drawable icon;
-        switch (info.category) {
-            case PackageItemInfo.CONVERSATIONS:
-                icon = getContext().getDrawable(R.drawable.ic_conversations_widget_category);
-                break;
-            default:
-                icon = info.newIcon(getContext());
+        if (info.widgetCategory == NO_CATEGORY) {
+            icon = info.newIcon(getContext());
+        } else {
+            WidgetSection widgetSection = WidgetSections.getWidgetSections(getContext())
+                    .get(info.widgetCategory);
+            icon = getContext().getDrawable(widgetSection.mSectionDrawable);
         }
         applyDrawables(icon);
         mIconDrawable = icon;
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
index fadb637..c6a7285 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
@@ -62,9 +62,7 @@
                         (position & POSITION_LAST) != 0,
                         /* isExpanded= */ data.isWidgetListShown()));
         widgetsListHeader.setOnExpandChangeListener(isExpanded ->
-                mOnHeaderClickListener.onHeaderClicked(
-                        isExpanded,
-                        new PackageUserKey(data.mPkgItem.packageName, data.mPkgItem.user)
-                ));
+                mOnHeaderClickListener.onHeaderClicked(isExpanded,
+                        PackageUserKey.fromPackageItemInfo(data.mPkgItem)));
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
index bff43c1..2b27fc2 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
@@ -64,6 +64,6 @@
                         /* isExpanded= */ data.isWidgetListShown()));
         widgetsListHeader.setOnExpandChangeListener(isExpanded ->
                 mOnHeaderClickListener.onHeaderClicked(isExpanded,
-                        new PackageUserKey(data.mPkgItem.packageName, data.mPkgItem.user)));
+                        PackageUserKey.fromPackageItemInfo(data.mPkgItem)));
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index 12e9e1e..702f343 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -4,7 +4,10 @@
 import static android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_HIDE_FROM_PICKER;
 
 import static com.android.launcher3.pm.ShortcutConfigActivityInfo.queryList;
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
 
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.mapping;
 import static java.util.stream.Collectors.toList;
 
 import android.appwidget.AppWidgetProviderInfo;
@@ -13,6 +16,7 @@
 import android.content.pm.PackageManager;
 import android.os.UserHandle;
 import android.util.Log;
+import android.util.Pair;
 
 import androidx.annotation.Nullable;
 import androidx.collection.ArrayMap;
@@ -27,10 +31,12 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
+import com.android.launcher3.widget.WidgetSections;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
@@ -40,12 +46,12 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
 import java.util.function.Predicate;
-import java.util.stream.Collectors;
 
 /**
  * Widgets data model that is used by the adapters of the widget views and controllers.
@@ -61,9 +67,6 @@
     private static final String TAG = "WidgetsModel";
     private static final boolean DEBUG = false;
 
-    private static final ComponentName CONVERSATION_WIDGET = ComponentName.createRelative(
-            "com.android.systemui", ".people.widget.PeopleSpaceWidgetProvider");
-
     /* Map of widgets and shortcuts that are tracked per package. */
     private final Map<PackageItemInfo, List<WidgetItem>> mWidgetsList = new HashMap<>();
 
@@ -168,16 +171,15 @@
             mWidgetsList.clear();
         } else {
             // Otherwise, only clear the widgets and shortcuts for the changed package.
-            mWidgetsList.remove(
-                    packageItemInfoCache.getOrCreate(new WidgetPackageOrCategoryKey(packageUser)));
+            mWidgetsList.remove(packageItemInfoCache.getOrCreate(packageUser));
         }
 
         // add and update.
         mWidgetsList.putAll(rawWidgetsShortcuts.stream()
                 .filter(new WidgetValidityCheck(app))
-                .collect(Collectors.groupingBy(item ->
-                        packageItemInfoCache.getOrCreate(getWidgetPackageOrCategoryKey(item))
-                )));
+                .flatMap(widgetItem -> getPackageUserKeys(app.getContext(), widgetItem).stream()
+                        .map(key -> new Pair<>(packageItemInfoCache.getOrCreate(key), widgetItem)))
+                .collect(groupingBy(pair -> pair.first, mapping(pair -> pair.second, toList()))));
 
         // Update each package entry
         IconCache iconCache = app.getIconCache();
@@ -209,9 +211,9 @@
     }
 
     public WidgetItem getWidgetProviderInfoByProviderName(
-            ComponentName providerName) {
+            ComponentName providerName, UserHandle user) {
         List<WidgetItem> widgetsList = mWidgetsList.get(
-                new PackageItemInfo(providerName.getPackageName()));
+                new PackageItemInfo(providerName.getPackageName(), user));
         if (widgetsList == null) {
             return null;
         }
@@ -225,18 +227,40 @@
     }
 
     /** Returns {@link PackageItemInfo} of a pending widget. */
-    public static PackageItemInfo newPendingItemInfo(ComponentName provider) {
-        if (CONVERSATION_WIDGET.equals(provider)) {
-            return new PackageItemInfo(provider.getPackageName(), PackageItemInfo.CONVERSATIONS);
+    public static PackageItemInfo newPendingItemInfo(Context context, ComponentName provider,
+            UserHandle user) {
+        Map<ComponentName, IntSet> widgetsToCategories =
+                WidgetSections.getWidgetsToCategory(context);
+        if (widgetsToCategories.containsKey(provider)) {
+            Iterator<Integer> categoriesIterator = widgetsToCategories.get(provider).iterator();
+            int firstCategory = NO_CATEGORY;
+            while (categoriesIterator.hasNext() && firstCategory == NO_CATEGORY) {
+                firstCategory = categoriesIterator.next();
+            }
+            return new PackageItemInfo(provider.getPackageName(), firstCategory, user);
         }
-        return new PackageItemInfo(provider.getPackageName());
+        return new PackageItemInfo(provider.getPackageName(), user);
     }
 
-    private WidgetPackageOrCategoryKey getWidgetPackageOrCategoryKey(WidgetItem item) {
-        if (CONVERSATION_WIDGET.equals(item.componentName)) {
-            return new WidgetPackageOrCategoryKey(PackageItemInfo.CONVERSATIONS, item.user);
+    private List<PackageUserKey> getPackageUserKeys(Context context, WidgetItem item) {
+        Map<ComponentName, IntSet> widgetsToCategories =
+                WidgetSections.getWidgetsToCategory(context);
+        IntSet categories = widgetsToCategories.get(item.componentName);
+        if (categories == null || categories.isEmpty()) {
+            return Arrays.asList(
+                    new PackageUserKey(item.componentName.getPackageName(), item.user));
         }
-        return new WidgetPackageOrCategoryKey(item.componentName.getPackageName(), item.user);
+        List<PackageUserKey> packageUserKeys = new ArrayList<>();
+        categories.forEach(category -> {
+            if (category == NO_CATEGORY) {
+                packageUserKeys.add(
+                        new PackageUserKey(item.componentName.getPackageName(),
+                                item.user));
+            } else {
+                packageUserKeys.add(new PackageUserKey(category, item.user));
+            }
+        });
+        return packageUserKeys;
     }
 
     private static class WidgetValidityCheck implements Predicate<WidgetItem> {
@@ -279,53 +303,13 @@
         }
     }
 
-    /** A hash key for grouping widgets by package name or category. */
-    private static class WidgetPackageOrCategoryKey {
-        /**
-         * The package name of the widget provider.
-         *
-         * <p>This shouldn't be empty if {@link #mCategory} has a value,
-         * {@link PackageItemInfo#NO_CATEGORY}.
-         */
-        public final String mPackage;
-        /** A widget category. */
-        @PackageItemInfo.Category public final int mCategory;
-        public final UserHandle mUser;
-        private final int mHashCode;
-
-        WidgetPackageOrCategoryKey(PackageUserKey key) {
-            this(key.mPackageName, key.mUser);
-        }
-
-        WidgetPackageOrCategoryKey(String packageName, UserHandle user) {
-            this(packageName,  PackageItemInfo.NO_CATEGORY, user);
-        }
-
-        WidgetPackageOrCategoryKey(@PackageItemInfo.Category int category, UserHandle user) {
-            this("", category, user);
-        }
-
-        private WidgetPackageOrCategoryKey(String packageName,
-                @PackageItemInfo.Category int category, UserHandle user) {
-            mPackage = packageName;
-            mCategory = category;
-            mUser = user;
-            mHashCode = Arrays.hashCode(new Object[]{mPackage, mCategory, mUser});
-        }
-
-        @Override
-        public int hashCode() {
-            return mHashCode;
-        }
-    }
-
     private static final class PackageItemInfoCache {
-        private final Map<WidgetPackageOrCategoryKey, PackageItemInfo> mMap = new ArrayMap<>();
+        private final Map<PackageUserKey, PackageItemInfo> mMap = new ArrayMap<>();
 
-        PackageItemInfo getOrCreate(WidgetPackageOrCategoryKey key) {
+        PackageItemInfo getOrCreate(PackageUserKey key) {
             PackageItemInfo pInfo = mMap.get(key);
             if (pInfo == null) {
-                pInfo = new PackageItemInfo(key.mPackage, key.mCategory);
+                pInfo = new PackageItemInfo(key.mPackageName, key.mWidgetCategory, key.mUser);
                 pInfo.user = key.mUser;
                 mMap.put(key,  pInfo);
             }
diff --git a/tests/src/com/android/launcher3/util/PackageUserKeyTest.java b/tests/src/com/android/launcher3/util/PackageUserKeyTest.java
new file mode 100644
index 0000000..99490eb
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/PackageUserKeyTest.java
@@ -0,0 +1,162 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.util;
+
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.os.UserHandle;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.model.data.PackageItemInfo;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public final class PackageUserKeyTest {
+    @Rule
+    public ExpectedException exception = ExpectedException.none();
+
+    private static final String TEST_PACKAGE = "com.android.test.package";
+    private static final int CONVERSATIONS = 0;
+    private static final int WEATHER = 1;
+
+    @Test
+    public void fromPackageItemInfo_shouldCreateExpectedObject() {
+        PackageUserKey packageUserKey = PackageUserKey.fromPackageItemInfo(
+                new PackageItemInfo(TEST_PACKAGE, UserHandle.CURRENT));
+
+        assertThat(packageUserKey.mPackageName).isEqualTo(TEST_PACKAGE);
+        assertThat(packageUserKey.mWidgetCategory).isEqualTo(NO_CATEGORY);
+        assertThat(packageUserKey.mUser).isEqualTo(UserHandle.CURRENT);
+    }
+
+    @Test
+    public void constructor_packageNameAndUserHandle_shouldCreateExpectedObject() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.mPackageName).isEqualTo(TEST_PACKAGE);
+        assertThat(packageUserKey.mWidgetCategory).isEqualTo(NO_CATEGORY);
+        assertThat(packageUserKey.mUser).isEqualTo(UserHandle.CURRENT);
+    }
+
+    @Test
+    public void constructor_widgetCategoryAndUserHandle_shouldCreateExpectedObject() {
+        PackageUserKey packageUserKey = new PackageUserKey(CONVERSATIONS, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.mPackageName).isEqualTo("");
+        assertThat(packageUserKey.mWidgetCategory).isEqualTo(CONVERSATIONS);
+        assertThat(packageUserKey.mUser).isEqualTo(UserHandle.CURRENT);
+    }
+
+    @Test
+    public void equals_sameObject_shouldReturnTrue() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = packageUserKey;
+
+        assertThat(packageUserKey).isEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void equals_differentObjectSameContent_shouldReturnTrue() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void equals_compareAgainstNull_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isNotEqualTo(null);
+    }
+
+    @Test
+    public void equals_differentPackage_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE + "1",
+                UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isNotEqualTo(otherPackageUserKey);
+    }
+
+
+    @Test
+    public void equals_differentCategory_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(WEATHER, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(CONVERSATIONS, UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isNotEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void equals_differentUser_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(1));
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(2));
+
+        assertThat(packageUserKey).isNotEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void hashCode_sameObject_shouldBeTheSame() {
+        PackageUserKey packageUserKey = new PackageUserKey(WEATHER, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = packageUserKey;
+
+        assertThat(packageUserKey.hashCode()).isEqualTo(otherPackageUserKey.hashCode());
+    }
+
+    @Test
+    public void hashCode_differentObjectSameContent_shouldBeTheSame() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.hashCode()).isEqualTo(otherPackageUserKey.hashCode());
+    }
+
+    @Test
+    public void hashCode_differentPackage_shouldBeDifferent() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE + "1",
+                UserHandle.CURRENT);
+
+        assertThat(packageUserKey.hashCode()).isNotEqualTo(otherPackageUserKey.hashCode());
+    }
+
+
+    @Test
+    public void hashCode_differentCategory_shouldBeDifferent() {
+        PackageUserKey packageUserKey = new PackageUserKey(WEATHER, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(CONVERSATIONS, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.hashCode()).isNotEqualTo(otherPackageUserKey.hashCode());
+    }
+
+    @Test
+    public void hashCode_differentUser_shouldBeDifferent() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(1));
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(2));
+
+        assertThat(packageUserKey.hashCode()).isNotEqualTo(otherPackageUserKey.hashCode());
+    }
+}
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
index 6232938..b480a4c 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -288,9 +288,8 @@
 
     private PackageItemInfo createPackageItemInfo(String packageName, String appName,
             UserHandle userHandle) {
-        PackageItemInfo pInfo = new PackageItemInfo(packageName);
+        PackageItemInfo pInfo = new PackageItemInfo(packageName, userHandle);
         pInfo.title = appName;
-        pInfo.user = userHandle;
         pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
         return pInfo;
     }
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
index 44d6964..4e0bdda 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
@@ -215,14 +215,23 @@
     @Test
     public void setWidgetsOnSearch_expandedApp_shouldResetExpandedApp() {
         // GIVEN a list of widgets entries:
-        // [com.google.test0, com.google.test0 content,
-        //  com.google.test1, com.google.test1 content,
-        //  com.google.test2, com.google.test2 content]
-        // The visible widgets entries: [com.google.test0, com.google.test1, com.google.test2].
-        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(2);
+        // [Empty item
+        //  com.google.test0,
+        //  com.google.test0 content,
+        //  com.google.test1,
+        //  com.google.test1 content,
+        //  com.google.test2,
+        //  com.google.test2 content]
+        // The visible widgets entries:
+        // [Empty item,
+        // com.google.test0,
+        // com.google.test1,
+        // com.google.test2].
+        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(3);
         mAdapter.setWidgetsOnSearch(allEntries);
         // GIVEN com.google.test.1 header is expanded. The visible entries list becomes:
-        // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
+        // [Empty item, com.google.test0, com.google.test1, com.google.test1 content,
+        // com.google.test2]
         mAdapter.onHeaderClicked(/* showWidgets= */ true,
                 new PackageUserKey(TEST_PACKAGE_PLACEHOLDER + 1, mUserHandle));
         Mockito.reset(mListener);
@@ -231,7 +240,7 @@
         mAdapter.setWidgetsOnSearch(allEntries);
 
         // THEN expanded app is reset and the visible entries list becomes:
-        // [com.google.test0, com.google.test1, com.google.test2]
+        // [Empty item, com.google.test0, com.google.test1, com.google.test2]
         verify(mListener).onItemRangeChanged(eq(2), eq(1), isNull());
         verify(mListener).onItemRangeRemoved(/* positionStart= */ 3, /* itemCount= */ 1);
     }
@@ -257,9 +266,8 @@
 
             List<WidgetItem> widgetItems = generateWidgetItems(packageName, /* numOfWidgets= */ 1);
 
-            PackageItemInfo pInfo = new PackageItemInfo(packageName);
+            PackageItemInfo pInfo = new PackageItemInfo(packageName, widgetItems.get(0).user);
             pInfo.title = pInfo.packageName;
-            pInfo.user = widgetItems.get(0).user;
             pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
             result.add(new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems));
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index 969c12a..211318c 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -30,6 +30,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.widget.FrameLayout;
 import android.widget.TextView;
@@ -124,12 +125,12 @@
         widgetsListHeader.callOnClick();
 
         verify(mOnHeaderClickListener).onHeaderClicked(eq(true),
-                eq(new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)));
+                eq(PackageUserKey.fromPackageItemInfo(entry.mPkgItem)));
     }
 
     private WidgetsListHeaderEntry generateSampleAppHeader(String appName, String packageName,
             int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        PackageItemInfo appInfo = new PackageItemInfo(packageName, UserHandle.CURRENT);
         appInfo.title = appName;
         appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
index 453f4fb..66c2f36 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
@@ -30,6 +30,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.widget.FrameLayout;
 import android.widget.TextView;
@@ -124,12 +125,12 @@
         widgetsListHeader.callOnClick();
 
         verify(mOnHeaderClickListener).onHeaderClicked(eq(true),
-                eq(new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)));
+                eq(PackageUserKey.fromPackageItemInfo(entry.mPkgItem)));
     }
 
     private WidgetsListSearchHeaderEntry generateSampleSearchHeader(String appName,
             String packageName, int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        PackageItemInfo appInfo = new PackageItemInfo(packageName, UserHandle.CURRENT);
         appInfo.title = appName;
         appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index 5816b77..7ec4d20 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -28,6 +28,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.View.OnClickListener;
@@ -125,7 +126,7 @@
 
     private WidgetsListContentEntry generateSampleAppWithWidgets(String appName, String packageName,
             int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        PackageItemInfo appInfo = new PackageItemInfo(packageName, UserHandle.CURRENT);
         appInfo.title = appName;
         appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
diff --git a/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index 4b61b2c..d8f1f14 100644
--- a/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -26,6 +26,7 @@
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
+import android.os.UserHandle;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -53,8 +54,10 @@
 public final class WidgetsListContentEntryTest {
     private static final String PACKAGE_NAME = "com.android.test";
     private static final String PACKAGE_NAME_2 = "com.android.test2";
-    private final PackageItemInfo mPackageItemInfo1 = new PackageItemInfo(PACKAGE_NAME);
-    private final PackageItemInfo mPackageItemInfo2 = new PackageItemInfo(PACKAGE_NAME_2);
+    private final PackageItemInfo mPackageItemInfo1 = new PackageItemInfo(PACKAGE_NAME,
+            UserHandle.CURRENT);
+    private final PackageItemInfo mPackageItemInfo2 = new PackageItemInfo(PACKAGE_NAME_2,
+            UserHandle.CURRENT);
     private final ComponentName mWidget1 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget1");
     private final ComponentName mWidget2 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget2");
     private final ComponentName mWidget3 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget3");
diff --git a/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java b/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
index c862d6b..d812ab0 100644
--- a/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
@@ -189,9 +189,8 @@
 
     private PackageItemInfo createPackageItemInfo(String packageName, String appName,
             UserHandle userHandle) {
-        PackageItemInfo pInfo = new PackageItemInfo(packageName);
+        PackageItemInfo pInfo = new PackageItemInfo(packageName, userHandle);
         pInfo.title = appName;
-        pInfo.user = userHandle;
         pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
         return pInfo;
     }