Merge "Suppress ViewInflationDuringSwipeUp test." into sc-dev
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index 1bf3627..33dc6cf 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -422,8 +422,8 @@
 
     @Override
     public Stream<SystemShortcut.Factory> getSupportedShortcuts() {
-        return Stream.concat(super.getSupportedShortcuts(),
-                Stream.of(WellbeingModel.SHORTCUT_FACTORY));
+        return Stream.concat(Stream.of(WellbeingModel.SHORTCUT_FACTORY),
+                super.getSupportedShortcuts());
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 65df237..31cf51c 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -34,6 +34,7 @@
 import static com.android.launcher3.anim.Interpolators.DEACCEL_1_7;
 import static com.android.launcher3.anim.Interpolators.EXAGGERATED_EASE;
 import static com.android.launcher3.anim.Interpolators.LINEAR;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_SCRIM_FOR_APP_LAUNCH;
 import static com.android.launcher3.config.FeatureFlags.KEYGUARD_ANIMATION;
 import static com.android.launcher3.config.FeatureFlags.SEPARATE_RECENTS_ACTIVITY;
 import static com.android.launcher3.dragndrop.DragLayer.ALPHA_INDEX_TRANSITIONS;
@@ -508,20 +509,23 @@
                 launcherAnimator.play(scaleAnim);
             });
 
-            int scrimColor = Themes.getAttrColor(mLauncher, R.attr.overviewScrimColor);
-            int scrimColorTrans = ColorUtils.setAlphaComponent(scrimColor, 0);
-            int[] colors = isAppOpening
-                    ? new int[] {scrimColorTrans, scrimColor}
-                    : new int[] {scrimColor, scrimColorTrans};
-            ScrimView scrimView = mLauncher.getScrimView();
-            if (scrimView.getBackground() instanceof ColorDrawable) {
-                scrimView.setBackgroundColor(colors[0]);
+            final boolean scrimEnabled = ENABLE_SCRIM_FOR_APP_LAUNCH.get();
+            if (scrimEnabled) {
+                int scrimColor = Themes.getAttrColor(mLauncher, R.attr.overviewScrimColor);
+                int scrimColorTrans = ColorUtils.setAlphaComponent(scrimColor, 0);
+                int[] colors = isAppOpening
+                        ? new int[]{scrimColorTrans, scrimColor}
+                        : new int[]{scrimColor, scrimColorTrans};
+                ScrimView scrimView = mLauncher.getScrimView();
+                if (scrimView.getBackground() instanceof ColorDrawable) {
+                    scrimView.setBackgroundColor(colors[0]);
 
-                ObjectAnimator scrim = ObjectAnimator.ofArgb(scrimView, VIEW_BACKGROUND_COLOR,
-                        colors);
-                scrim.setDuration(CONTENT_SCRIM_DURATION);
-                scrim.setInterpolator(DEACCEL_1_5);
-                launcherAnimator.play(scrim);
+                    ObjectAnimator scrim = ObjectAnimator.ofArgb(scrimView, VIEW_BACKGROUND_COLOR,
+                            colors);
+                    scrim.setDuration(CONTENT_SCRIM_DURATION);
+                    scrim.setInterpolator(DEACCEL_1_5);
+                    launcherAnimator.play(scrim);
+                }
             }
 
             // Pause page indicator animations as they lead to layer trashing.
@@ -532,7 +536,9 @@
                     SCALE_PROPERTY.set(view, 1f);
                     view.setLayerType(View.LAYER_TYPE_NONE, null);
                 });
-                scrimView.setBackgroundColor(Color.TRANSPARENT);
+                if (scrimEnabled) {
+                    mLauncher.getScrimView().setBackgroundColor(Color.TRANSPARENT);
+                }
                 mLauncher.getWorkspace().getPageIndicator().skipAnimationsToEnd();
             };
         }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index f0b02b3..ec9893c 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -192,7 +192,7 @@
     @Override
     public Stream<SystemShortcut.Factory> getSupportedShortcuts() {
         return Stream.concat(
-                super.getSupportedShortcuts(), Stream.of(mHotseatPredictionController));
+                Stream.of(mHotseatPredictionController), super.getSupportedShortcuts());
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 0ebaea2..4d47ef1 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -1381,11 +1381,17 @@
 
     /**
      * Cancels any running animation so that the active target can be overriden by a new swipe
-     * handle (in case of quick switch).
+     * handler (in case of quick switch).
      */
     private void cancelCurrentAnimation() {
         mCanceled = true;
         mCurrentShift.cancelAnimation();
+
+        // Cleanup when switching handlers
+        mInputConsumerProxy.unregisterCallback();
+        mActivityInitListener.unregister();
+        ActivityManagerWrapper.getInstance().unregisterTaskStackListener(mActivityRestartListener);
+        mTaskSnapshot = null;
     }
 
     private void invalidateHandler() {
diff --git a/res/drawable/ic_block_no_shadow.xml b/res/drawable/ic_block_no_shadow.xml
index be9aa07..6ac61f4 100644
--- a/res/drawable/ic_block_no_shadow.xml
+++ b/res/drawable/ic_block_no_shadow.xml
@@ -16,8 +16,8 @@
 <vector xmlns:android="http://schemas.android.com/apk/res/android"
     android:width="20dp"
     android:height="20dp"
-    android:viewportHeight="20.0"
-    android:viewportWidth="20.0"
+    android:viewportHeight="24.0"
+    android:viewportWidth="24.0"
     android:tint="?android:attr/textColorPrimary">
     <path
         android:fillColor="@android:color/white"
diff --git a/res/drawable/ic_drag_handle.xml b/res/drawable/ic_drag_handle.xml
index 0181ff1..9db75f4 100644
--- a/res/drawable/ic_drag_handle.xml
+++ b/res/drawable/ic_drag_handle.xml
@@ -19,7 +19,7 @@
         android:height="@dimen/deep_shortcut_drag_handle_size"
         android:viewportWidth="24.0"
         android:viewportHeight="24.0"
-        android:tint="?android:attr/textColorHint" >
+        android:tint="?android:attr/textColorPrimary" >
 
     <path
         android:pathData="M20,9H4v2h16V9z M4,15h16v-2H4V15z"
diff --git a/res/drawable/ic_uninstall_no_shadow.xml b/res/drawable/ic_uninstall_no_shadow.xml
index 14cecac..fbabdd2 100644
--- a/res/drawable/ic_uninstall_no_shadow.xml
+++ b/res/drawable/ic_uninstall_no_shadow.xml
@@ -16,8 +16,8 @@
 <vector xmlns:android="http://schemas.android.com/apk/res/android"
         android:width="20dp"
         android:height="20dp"
-        android:viewportWidth="20.0"
-        android:viewportHeight="20.0"
+        android:viewportWidth="24.0"
+        android:viewportHeight="24.0"
         android:tint="?android:attr/textColorPrimary" >
     <path
         android:fillColor="@android:color/white"
diff --git a/res/layout/system_shortcut.xml b/res/layout/system_shortcut.xml
index 6337fae..2cdf1f4 100644
--- a/res/layout/system_shortcut.xml
+++ b/res/layout/system_shortcut.xml
@@ -45,6 +45,6 @@
         android:layout_height="@dimen/system_shortcut_icon_size"
         android:layout_marginStart="@dimen/system_shortcut_margin_start"
         android:layout_gravity="start|center_vertical"
-        android:backgroundTint="?android:attr/textColorTertiary"/>
+        android:backgroundTint="?android:attr/textColorPrimary"/>
 
 </com.android.launcher3.shortcuts.DeepShortcutView>
diff --git a/res/layout/widget_shortcut_container.xml b/res/layout/widget_shortcut_container.xml
new file mode 100644
index 0000000..a4d8eb4
--- /dev/null
+++ b/res/layout/widget_shortcut_container.xml
@@ -0,0 +1,27 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2017 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.
+-->
+
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/system_shortcut_full"
+    android:layout_width="match_parent"
+    android:layout_height="@dimen/system_shortcut_header_height"
+    android:orientation="horizontal"
+    android:gravity="end|center_vertical"
+    android:elevation="@dimen/deep_shortcuts_elevation"
+    android:tag="@string/popup_container_iterate_children"
+    android:clipToPadding="true">
+</LinearLayout>
diff --git a/res/layout/widgets_table_container.xml b/res/layout/widgets_table_container.xml
index c6b70aa..4f70a05 100644
--- a/res/layout/widgets_table_container.xml
+++ b/res/layout/widgets_table_container.xml
@@ -13,7 +13,7 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<TableLayout
+<com.android.launcher3.widget.picker.WidgetsListTableView
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:id="@+id/widgets_table"
     android:layout_width="match_parent"
diff --git a/res/layout/work_apps_edu.xml b/res/layout/work_apps_edu.xml
index 2c58907..ffbf962 100644
--- a/res/layout/work_apps_edu.xml
+++ b/res/layout/work_apps_edu.xml
@@ -16,30 +16,38 @@
 <com.android.launcher3.allapps.WorkEduCard xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
-    android:padding="@dimen/work_edu_card_margin"
-    android:orientation="vertical"
-    android:background="@drawable/work_card"
     android:gravity="center">
 
-    <TextView
-        style="@style/PrimaryHeadline"
-        android:textColor="?android:attr/textColorPrimary"
-        android:id="@+id/work_apps_paused_title"
-        android:layout_width="wrap_content"
+    <LinearLayout
+        android:layout_width="match_parent"
         android:layout_height="wrap_content"
-        android:layout_marginTop="8dp"
-        android:layout_marginBottom="8dp"
-        android:text="@string/work_profile_edu_work_apps"
-        android:textAlignment="center"
-        android:textSize="20sp" />
+        android:orientation="vertical"
+        android:padding="@dimen/work_edu_card_margin"
+        android:background="@drawable/work_card"
+        android:layout_gravity="center_horizontal"
+        android:gravity="center"
+        android:id="@+id/wrapper">
 
-    <Button
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:id="@+id/action_btn"
-        android:textColor="?attr/workProfileOverlayTextColor"
-        android:text="@string/work_profile_edu_accept"
-        android:textAlignment="center"
-        android:background="@drawable/work_card_btn"
-        android:textSize="14sp" />
+        <TextView
+            style="@style/PrimaryHeadline"
+            android:textColor="?android:attr/textColorPrimary"
+            android:id="@+id/work_apps_paused_title"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_marginTop="8dp"
+            android:layout_marginBottom="8dp"
+            android:text="@string/work_profile_edu_work_apps"
+            android:textAlignment="center"
+            android:textSize="20sp" />
+
+        <Button
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:id="@+id/action_btn"
+            android:textColor="?attr/workProfileOverlayTextColor"
+            android:text="@string/work_profile_edu_accept"
+            android:textAlignment="center"
+            android:background="@drawable/work_card_btn"
+            android:textSize="14sp" />
+    </LinearLayout>
 </com.android.launcher3.allapps.WorkEduCard>
\ No newline at end of file
diff --git a/res/layout/work_apps_paused.xml b/res/layout/work_apps_paused.xml
index 7f1107f..14dcb8c 100644
--- a/res/layout/work_apps_paused.xml
+++ b/res/layout/work_apps_paused.xml
@@ -12,11 +12,10 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+<com.android.launcher3.allapps.WorkPausedCard xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="match_parent"
-    android:layout_height="match_parent"
-    android:background="?attr/allAppsScrimColor"
-    android:padding="48dp"
+    android:layout_height="wrap_content"
+    android:padding="@dimen/work_edu_card_margin"
     android:orientation="vertical"
     android:gravity="center">
 
@@ -39,5 +38,16 @@
         android:textColor="?attr/workProfileOverlayTextColor"
         android:text="@string/work_apps_paused_body"
         android:textAlignment="center"
+        android:layout_marginBottom="8dp"
         android:textSize="16sp" />
-</LinearLayout>
\ No newline at end of file
+
+    <Button
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:id="@+id/enable_work_apps"
+        android:textColor="?attr/workProfileOverlayTextColor"
+        android:text="@string/work_apps_enable_btn_text"
+        android:textAlignment="center"
+        android:background="@drawable/work_card_btn"
+        android:textSize="14sp" />
+</com.android.launcher3.allapps.WorkPausedCard>
\ No newline at end of file
diff --git a/res/values-v28/styles.xml b/res/values-v28/styles.xml
deleted file mode 100644
index 7df9ce5..0000000
--- a/res/values-v28/styles.xml
+++ /dev/null
@@ -1,23 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-/*
-* Copyright (C) 2019 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.
-*/
--->
-<resources>
-    <style name="TextHeadline" parent="@android:style/TextAppearance.DeviceDefault.DialogWindowTitle" >
-        <item name="android:textFontWeight">400</item>
-    </style>
-</resources>
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index e8d3212..6247dc2 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -48,7 +48,7 @@
     <dimen name="workspace_page_indicator_overlap_workspace">0dp</dimen>
 
 <!-- Drop target bar -->
-    <dimen name="dynamic_grid_drop_target_size">48dp</dimen>
+    <dimen name="dynamic_grid_drop_target_size">52dp</dimen>
     <dimen name="drop_target_vertical_gap">20dp</dimen>
 
 <!-- App Widget resize frame -->
@@ -94,11 +94,11 @@
     <dimen name="all_apps_background_canvas_height">475dp</dimen>
     <dimen name="all_apps_header_pill_height">48dp</dimen>
     <dimen name="all_apps_header_pill_corner_radius">18dp</dimen>
-    <dimen name="all_apps_header_pills_width">320dp</dimen>
+    <dimen name="all_apps_header_pills_width">348dp</dimen>
     <dimen name="all_apps_header_tab_height">48dp</dimen>
     <dimen name="all_apps_tabs_indicator_height">2dp</dimen>
     <dimen name="all_apps_header_top_padding">36dp</dimen>
-    <dimen name="all_apps_header_bottom_padding">16dp</dimen>
+    <dimen name="all_apps_header_bottom_padding">6dp</dimen>
     <dimen name="all_apps_work_profile_tab_footer_top_padding">16dp</dimen>
     <dimen name="all_apps_work_profile_tab_footer_bottom_padding">20dp</dimen>
     <dimen name="all_apps_tabs_vertical_padding">6dp</dimen>
@@ -140,7 +140,6 @@
 
     <dimen name="widget_list_top_bottom_corner_radius">28dp</dimen>
     <dimen name="widget_list_content_corner_radius">4dp</dimen>
-    <dimen name="widget_list_content_joined_corner_radius">0dp</dimen>
 
     <dimen name="widget_list_header_view_vertical_padding">20dp</dimen>
     <dimen name="widget_list_entry_bottom_margin">2dp</dimen>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 9823df3..60c0774 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -190,11 +190,11 @@
     <!-- Widgets: -->
     <skip />    
 
-    <!-- Text to show user in place of a gadget when we can't display it properly -->
-    <string name="gadget_error_text">Problem loading widget</string>
+    <!-- Error text that lets a user know that the widget can't load. -->
+    <string name="gadget_error_text">Can\'t load widget</string>
 
-    <!-- Text to show user in place of a gadget when it is not yet initialized. -->
-    <string name="gadget_setup_text">Setup</string>
+    <!-- Instructional text to encourage a user to finish setting up the widget. -->
+    <string name="gadget_setup_text">Tap to finish setup</string>
 
     <!-- Text to inform the user that they can't uninstall a system application -->
     <string name="uninstall_system_app_text">This is a system app and can\'t be uninstalled.</string>
diff --git a/res/values/styles.xml b/res/values/styles.xml
index 6f7d727..2119e58 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -214,7 +214,9 @@
         <item name="disabledIconAlpha">.54</item>
     </style>
 
-    <style name="BaseIconUnBounded" parent="@android:style/TextAppearance.DeviceDefault.DialogWindowTitle">
+    <style name="BaseIconRoot" parent="@android:style/TextAppearance.DeviceDefault.DialogWindowTitle"/>
+
+    <style name="BaseIconUnBounded" parent="BaseIconRoot">
         <item name="android:layout_width">match_parent</item>
         <item name="android:layout_height">match_parent</item>
         <item name="android:layout_gravity">center</item>
@@ -264,9 +266,8 @@
 
     <!-- Drop targets -->
     <style name="DropTargetButtonBase" parent="@android:style/TextAppearance.DeviceDefault">
-        <item name="android:drawablePadding">7.5dp</item>
-        <item name="android:paddingLeft">16dp</item>
-        <item name="android:paddingRight">16dp</item>
+        <item name="android:drawablePadding">8dp</item>
+        <item name="android:padding">16dp</item>
         <item name="android:textColor">?android:attr/textColorPrimary</item>
         <item name="android:textSize">@dimen/drop_target_text_size</item>
         <item name="android:singleLine">true</item>
@@ -276,7 +277,7 @@
 
     <style name="DropTargetButton" parent="DropTargetButtonBase" />
 
-    <style name="TextHeadline" parent="@android:style/TextAppearance.DeviceDefault" />
+    <style name="TextHeadline" parent="@android:style/TextAppearance.DeviceDefault.DialogWindowTitle" />
 
     <style name="PrimaryHeadline" parent="@android:style/TextAppearance.DeviceDefault.DialogWindowTitle"/>
 
diff --git a/robolectric_tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java b/robolectric_tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
index c7286e5..2d87957 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
@@ -36,6 +36,10 @@
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
 @RunWith(RobolectricTestRunner.class)
 public final class LauncherAppWidgetProviderInfoTest {
 
@@ -234,10 +238,11 @@
         Mockito.when(profile.getCellSize()).thenReturn(new Point(CELL_SIZE, CELL_SIZE));
 
         InvariantDeviceProfile idp = new InvariantDeviceProfile();
-        idp.supportedProfiles.add(profile);
+        List<DeviceProfile> supportedProfiles = new ArrayList<>(idp.supportedProfiles);
+        supportedProfiles.add(profile);
+        idp.supportedProfiles = Collections.unmodifiableList(supportedProfiles);
         idp.numColumns = NUM_OF_COLS;
         idp.numRows = NUM_OF_ROWS;
         return idp;
     }
-
 }
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index 84a03d5..4e2a508 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -107,7 +107,10 @@
                 /* iconClickListener= */ view -> {},
                 /* iconLongClickListener= */ view -> false);
         mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
-                LayoutInflater.from(mTestActivity), mOnHeaderClickListener, widgetsListAdapter);
+                LayoutInflater.from(mTestActivity),
+                mOnHeaderClickListener,
+                new WidgetsListDrawableFactory(mTestActivity),
+                widgetsListAdapter);
     }
 
     @After
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
index 075c58d..d6aea55 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
@@ -107,7 +107,10 @@
                 /* iconClickListener= */ view -> {},
                 /* iconLongClickListener= */ view -> false);
         mViewHolderBinder = new WidgetsListSearchHeaderViewHolderBinder(
-                LayoutInflater.from(mTestActivity), mOnHeaderClickListener, widgetsListAdapter);
+                LayoutInflater.from(mTestActivity),
+                mOnHeaderClickListener,
+                new WidgetsListDrawableFactory(mTestActivity),
+                widgetsListAdapter);
     }
 
     @After
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index 0c6e717..2f1326f 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -118,6 +118,7 @@
                 mOnIconClickListener,
                 mOnLongClickListener,
                 mWidgetPreviewLoader,
+                new WidgetsListDrawableFactory(mTestActivity),
                 widgetsListAdapter);
     }
 
diff --git a/src/com/android/launcher3/ButtonDropTarget.java b/src/com/android/launcher3/ButtonDropTarget.java
index 7db34a5..d1cb5b8 100644
--- a/src/com/android/launcher3/ButtonDropTarget.java
+++ b/src/com/android/launcher3/ButtonDropTarget.java
@@ -36,8 +36,6 @@
 import android.widget.PopupWindow;
 import android.widget.TextView;
 
-import androidx.appcompat.content.res.AppCompatResources;
-
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.dragndrop.DragController;
 import com.android.launcher3.dragndrop.DragLayer;
@@ -82,6 +80,8 @@
     private boolean mAccessibleDrag;
     /** An item must be dragged at least this many pixels before this drop target is enabled. */
     private final int mDragDistanceThreshold;
+    /** The size of the drawable shown in the drop target. */
+    private final int mDrawableSize;
 
     protected CharSequence mText;
     protected ColorStateList mOriginalTextColor;
@@ -103,6 +103,7 @@
 
         Resources resources = getResources();
         mDragDistanceThreshold = resources.getDimensionPixelSize(R.dimen.drag_distanceThreshold);
+        mDrawableSize = resources.getDimensionPixelSize(R.dimen.drop_target_text_size);
     }
 
     @Override
@@ -122,13 +123,9 @@
     protected void setDrawable(int resId) {
         // We do not set the drawable in the xml as that inflates two drawables corresponding to
         // drawableLeft and drawableStart.
-        if (mTextVisible) {
-            setCompoundDrawablesRelativeWithIntrinsicBounds(resId, 0, 0, 0);
-            mDrawable = getCompoundDrawablesRelative()[0];
-        } else {
-            setCompoundDrawablesRelativeWithIntrinsicBounds(0, resId, 0, 0);
-            mDrawable = getCompoundDrawablesRelative()[1];
-        }
+        mDrawable = getContext().getDrawable(resId).mutate();
+        mDrawable.setBounds(0, 0, mDrawableSize, mDrawableSize);
+        setCompoundDrawablesRelative(mDrawable, null, null, null);
     }
 
     public void setDropTargetBar(DropTargetBar dropTargetBar) {
@@ -329,11 +326,7 @@
         if (mTextVisible != isVisible || !TextUtils.equals(newText, getText())) {
             mTextVisible = isVisible;
             setText(newText);
-            if (mTextVisible) {
-                setCompoundDrawablesRelativeWithIntrinsicBounds(mDrawable, null, null, null);
-            } else {
-                setCompoundDrawablesRelativeWithIntrinsicBounds(null, mDrawable, null, null);
-            }
+            setCompoundDrawablesRelative(mDrawable, null, null, null);
         }
     }
 
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 318dde1..29a0223 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -140,7 +140,10 @@
     public int defaultLayoutId;
     int demoModeLayoutId;
 
-    public final List<DeviceProfile> supportedProfiles = new ArrayList<>();
+    /**
+     * An immutable list of supported profiles.
+     */
+    public List<DeviceProfile> supportedProfiles = Collections.EMPTY_LIST;
 
     @Nullable public DevicePaddings devicePaddings;
 
@@ -322,10 +325,10 @@
         // Supported overrides: numRows, numColumns, iconSize
         applyPartnerDeviceProfileOverrides(context, metrics);
 
-        supportedProfiles.clear();
+        final List<DeviceProfile> localSupportedProfiles = new ArrayList<>();
         defaultWallpaperSize = new Point(displayInfo.currentSize);
         for (WindowBounds bounds : displayInfo.supportedBounds) {
-            supportedProfiles.add(new DeviceProfile.Builder(context, this, displayInfo)
+            localSupportedProfiles.add(new DeviceProfile.Builder(context, this, displayInfo)
                     .setUseTwoPanels(isSplitDisplay)
                     .setWindowBounds(bounds).build());
 
@@ -343,6 +346,7 @@
             defaultWallpaperSize.x =
                     Math.max(defaultWallpaperSize.x, Math.round(parallaxFactor * displayWidth));
         }
+        supportedProfiles = Collections.unmodifiableList(localSupportedProfiles);
 
         ComponentName cn = new ComponentName(context.getPackageName(), getClass().getName());
         defaultWidgetPadding = AppWidgetHostView.getDefaultPaddingForWidget(context, cn, null);
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index a2d86bc..75f6278 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -122,7 +122,8 @@
 
     public static final boolean ATLEAST_R = Build.VERSION.SDK_INT >= Build.VERSION_CODES.R;
 
-    public static final boolean ATLEAST_S = BuildCompat.isAtLeastS();
+    public static final boolean ATLEAST_S = BuildCompat.isAtLeastS()
+            || Build.VERSION.SDK_INT >= Build.VERSION_CODES.S;
 
     /**
      * Set on a motion event dispatched from the nav bar. See {@link MotionEvent#setEdgeFlags(int)}.
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 2e1cc58..303bb01 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -2791,6 +2791,10 @@
                 removeWorkspaceItem(mDragInfo.cell);
             }
         } else if (mDragInfo != null) {
+            // When drag is cancelled, reattach content view back to its original parent.
+            if (mDragInfo.cell instanceof LauncherAppWidgetHostView) {
+                d.dragView.detachContentView(/* reattachToPreviousParent= */ true);
+            }
             final CellLayout cellLayout = mLauncher.getCellLayout(
                     mDragInfo.container, mDragInfo.screenId);
             if (cellLayout != null) {
diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index d749e02..39d1a00 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -21,9 +21,10 @@
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_CHANGE_PERMISSION;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_ENABLED;
 
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
 import android.content.Context;
-import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.graphics.Canvas;
 import android.graphics.Color;
@@ -82,10 +83,7 @@
         ScrimView.ScrimDrawingController {
 
     public static final float PULL_MULTIPLIER = .02f;
-    public static final float FLING_VELOCITY_MULTIPLIER = 2000f;
-
-    // Starts the springs after at least 25% of the animation has passed.
-    public static final float FLING_ANIMATION_THRESHOLD = 0.25f;
+    public static final float FLING_VELOCITY_MULTIPLIER = 1200f;
 
     private final Paint mHeaderPaint = new Paint(Paint.ANTI_ALIAS_FLAG);
 
@@ -494,15 +492,6 @@
         }
     }
 
-    @Override
-    protected void onConfigurationChanged(Configuration newConfig) {
-        super.onConfigurationChanged(newConfig);
-        View overlay = mAH[AdapterHolder.WORK].getOverlayView();
-        int v = newConfig.orientation == Configuration.ORIENTATION_LANDSCAPE ? GONE : VISIBLE;
-        overlay.findViewById(R.id.work_apps_paused_title).setVisibility(v);
-        overlay.findViewById(R.id.work_apps_paused_content).setVisibility(v);
-    }
-
     private void replaceRVContainer(boolean showTabs) {
         for (int i = 0; i < mAH.length; i++) {
             if (mAH[i].recyclerView != null) {
@@ -544,9 +533,6 @@
                     && mAllAppsStore.hasModelFlag(
                     FLAG_HAS_SHORTCUT_PERMISSION | FLAG_QUIET_MODE_CHANGE_PERMISSION));
         }
-        if (mSearchUiManager != null && mSearchUiManager.getEditText() != null) {
-            mSearchUiManager.getEditText().hideKeyboard();
-        }
     }
 
     // Used by tests only
@@ -658,20 +644,18 @@
     /**
      * Adds an update listener to {@param animator} that adds springs to the animation.
      */
-    public void addSpringFromFlingUpdateListener(ValueAnimator animator, float velocity) {
-        animator.addUpdateListener(new ValueAnimator.AnimatorUpdateListener() {
-            boolean shouldSpring = true;
-
+    public void addSpringFromFlingUpdateListener(ValueAnimator animator,
+            float velocity /* release velocity */,
+            float progress /* portion of the distance to travel*/) {
+        animator.addListener(new AnimatorListenerAdapter() {
             @Override
-            public void onAnimationUpdate(ValueAnimator valueAnimator) {
-                if (shouldSpring
-                        && valueAnimator.getAnimatedFraction() >= FLING_ANIMATION_THRESHOLD) {
-                    absorbSwipeUpVelocity(Math.max(100, Math.abs(
-                            Math.round(velocity * FLING_VELOCITY_MULTIPLIER))));
-                    // calculate the velocity of using the not user controlled interpolator
-                    // of when the container reach the end.
-                    shouldSpring = false;
-                }
+            public void onAnimationStart(Animator animator) {
+                float distance = (float) ((1 - progress) * getHeight()); // px
+                float settleVelocity = Math.min(0, distance
+                        / (AllAppsTransitionController.INTERP_COEFF * animator.getDuration())
+                        + velocity);
+                absorbSwipeUpVelocity(Math.max(1000, Math.abs(
+                        Math.round(settleVelocity * FLING_VELOCITY_MULTIPLIER))));
             }
         });
     }
@@ -704,7 +688,9 @@
         mHeaderPaint.setColor(mHeaderColor);
         mHeaderPaint.setAlpha((int) (getAlpha() * Color.alpha(mHeaderColor)));
         if (mHeaderPaint.getColor() != mScrimColor && mHeaderPaint.getColor() != 0) {
-            canvas.drawRect(0, 0, getWidth(), mSearchContainer.getTop() + getTranslationY(),
+            int bottom = mUsingTabs && mHeader.mHeaderCollapsed ? mHeader.getVisibleBottomBound()
+                    : mSearchContainer.getBottom();
+            canvas.drawRect(0, 0, getWidth(), bottom + getTranslationY(),
                     mHeaderPaint);
         }
     }
@@ -781,13 +767,6 @@
             mAH[AdapterHolder.MAIN].recyclerView.setVerticalFadingEdgeEnabled(!mUsingTabs
                     && verticalFadingEdge);
         }
-
-        private View getOverlayView() {
-            if (mOverlay == null) {
-                mOverlay = mLauncher.getLayoutInflater().inflate(R.layout.work_apps_paused, null);
-            }
-            return mOverlay;
-        }
     }
 
 
diff --git a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
index 5bbd02b..4ad694b 100644
--- a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
@@ -447,4 +447,14 @@
     public boolean hasOverlappingRendering() {
         return false;
     }
+
+    /**
+     * Returns distance between left and right app icons
+     */
+    public int getTabWidth() {
+        DeviceProfile grid = BaseDraggingActivity.fromContext(getContext()).getDeviceProfile();
+        int totalWidth = (grid.availableWidthPx - getPaddingLeft() - getPaddingRight());
+        int iconPadding = totalWidth / grid.numShownAllAppsColumns - grid.allAppsIconSizePx;
+        return totalWidth - iconPadding;
+    }
 }
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index 8ec8269..a0551f0 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -57,6 +57,8 @@
  */
 public class AllAppsTransitionController
         implements StateHandler<LauncherState>, OnDeviceProfileChangeListener {
+    // This constant should match the second derivative of the animator interpolator.
+    public static final float INTERP_COEFF = 1.7f;
     private static final float CONTENT_VISIBLE_MAX_THRESHOLD = 0.5f;
 
     public static final FloatProperty<AllAppsTransitionController> ALL_APPS_PROGRESS =
diff --git a/src/com/android/launcher3/allapps/FloatingHeaderView.java b/src/com/android/launcher3/allapps/FloatingHeaderView.java
index 450d2e2..f1381e9 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderView.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderView.java
@@ -72,8 +72,13 @@
                     }
 
                     int current = -mCurrentRV.getCurrentScrollY();
+                    boolean headerCollapsed = mHeaderCollapsed;
                     moved(current);
                     applyVerticalMove();
+                    if (headerCollapsed != mHeaderCollapsed) {
+                        AllAppsContainerView parent = (AllAppsContainerView) getParent();
+                        parent.invalidateHeader();
+                    }
                 }
             };
 
@@ -219,6 +224,8 @@
 
         mTabsHidden = tabsHidden;
         mTabLayout.setVisibility(tabsHidden ? View.GONE : View.VISIBLE);
+        mTabLayout.getLayoutParams().width =
+                mAH[AllAppsContainerView.AdapterHolder.MAIN].recyclerView.getTabWidth();
         mMainRV = setupRV(mMainRV, mAH[AllAppsContainerView.AdapterHolder.MAIN].recyclerView);
         mWorkRV = setupRV(mWorkRV, mAH[AllAppsContainerView.AdapterHolder.WORK].recyclerView);
         mParent = (ViewGroup) mMainRV.getParent();
@@ -429,6 +436,13 @@
         }
         return null;
     }
+
+    /**
+     * Returns visible height of FloatingHeaderView contents
+     */
+    public int getVisibleBottomBound() {
+        return getBottom() + mTranslationY;
+    }
 }
 
 
diff --git a/src/com/android/launcher3/allapps/WorkEduCard.java b/src/com/android/launcher3/allapps/WorkEduCard.java
index 29a42f8..9db7bf0 100644
--- a/src/com/android/launcher3/allapps/WorkEduCard.java
+++ b/src/com/android/launcher3/allapps/WorkEduCard.java
@@ -21,7 +21,7 @@
 import android.view.ViewGroup;
 import android.view.animation.Animation;
 import android.view.animation.AnimationUtils;
-import android.widget.LinearLayout;
+import android.widget.FrameLayout;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
@@ -29,7 +29,7 @@
 /**
  * Work profile toggle switch shown at the bottom of AllApps work tab
  */
-public class WorkEduCard extends LinearLayout implements View.OnClickListener,
+public class WorkEduCard extends FrameLayout implements View.OnClickListener,
         Animation.AnimationListener {
 
     private final Launcher mLauncher;
@@ -52,11 +52,24 @@
         mDismissAnim.setAnimationListener(this);
     }
 
+    @Override
+    protected void onAttachedToWindow() {
+        super.onAttachedToWindow();
+        mDismissAnim.reset();
+    }
+
+    @Override
+    protected void onDetachedFromWindow() {
+        super.onDetachedFromWindow();
+        mDismissAnim.cancel();
+    }
 
     @Override
     protected void onFinishInflate() {
         super.onFinishInflate();
         findViewById(R.id.action_btn).setOnClickListener(this);
+        MarginLayoutParams lp = ((MarginLayoutParams) findViewById(R.id.wrapper).getLayoutParams());
+        lp.width = mLauncher.getAppsView().getActiveRecyclerView().getTabWidth();
     }
 
     @Override
diff --git a/src/com/android/launcher3/allapps/WorkModeSwitch.java b/src/com/android/launcher3/allapps/WorkModeSwitch.java
index c22cd63..3680fb9 100644
--- a/src/com/android/launcher3/allapps/WorkModeSwitch.java
+++ b/src/com/android/launcher3/allapps/WorkModeSwitch.java
@@ -15,6 +15,7 @@
  */
 package com.android.launcher3.allapps;
 
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TURN_OFF_WORK_APPS_TAP;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
 import android.content.Context;
@@ -31,7 +32,7 @@
 import androidx.annotation.RequiresApi;
 
 import com.android.launcher3.Insettable;
-import com.android.launcher3.R;
+import com.android.launcher3.Launcher;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.pm.UserCache;
 
@@ -79,7 +80,9 @@
         clearAnimation();
         if (workTabVisible) {
             setEnabled(true);
-            setVisibility(VISIBLE);
+            if (mWorkEnabled) {
+                setVisibility(VISIBLE);
+            }
             setAlpha(0);
             animate().alpha(1).start();
         } else {
@@ -91,7 +94,9 @@
     public void onClick(View view) {
         if (Utilities.ATLEAST_P) {
             setEnabled(false);
-            UI_HELPER_EXECUTOR.post(() -> setToState(!mWorkEnabled));
+            Launcher.fromContext(getContext()).getStatsLogManager().logger().log(
+                    LAUNCHER_TURN_OFF_WORK_APPS_TAP);
+            UI_HELPER_EXECUTOR.post(() -> setWorkProfileEnabled(getContext(), false));
         }
     }
 
@@ -101,20 +106,18 @@
     public void updateCurrentState(boolean active) {
         mWorkEnabled = active;
         setEnabled(true);
-        setCompoundDrawablesRelativeWithIntrinsicBounds(
-                active ? R.drawable.ic_corp_off : R.drawable.ic_corp, 0, 0, 0);
-        setText(active ? R.string.work_apps_pause_btn_text : R.string.work_apps_enable_btn_text);
+        setVisibility(active ? VISIBLE : GONE);
     }
 
     @RequiresApi(Build.VERSION_CODES.P)
-    protected Boolean setToState(boolean toState) {
-        UserManager userManager = getContext().getSystemService(UserManager.class);
+    public static Boolean setWorkProfileEnabled(Context context, boolean enabled) {
+        UserManager userManager = context.getSystemService(UserManager.class);
         boolean showConfirm = false;
-        for (UserHandle userProfile : UserCache.INSTANCE.get(getContext()).getUserProfiles()) {
+        for (UserHandle userProfile : UserCache.INSTANCE.get(context).getUserProfiles()) {
             if (Process.myUserHandle().equals(userProfile)) {
                 continue;
             }
-            showConfirm |= !userManager.requestQuietModeEnabled(!toState, userProfile);
+            showConfirm |= !userManager.requestQuietModeEnabled(!enabled, userProfile);
         }
         return showConfirm;
     }
diff --git a/src/com/android/launcher3/allapps/WorkPausedCard.java b/src/com/android/launcher3/allapps/WorkPausedCard.java
new file mode 100644
index 0000000..7908b63
--- /dev/null
+++ b/src/com/android/launcher3/allapps/WorkPausedCard.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2020 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.allapps;
+
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TURN_ON_WORK_APPS_TAP;
+import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
+
+import android.content.Context;
+import android.content.res.Configuration;
+import android.util.AttributeSet;
+import android.view.View;
+import android.widget.Button;
+import android.widget.LinearLayout;
+
+import com.android.launcher3.Launcher;
+import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
+
+/**
+ * Work profile toggle switch shown at the bottom of AllApps work tab
+ */
+public class WorkPausedCard extends LinearLayout implements View.OnClickListener {
+
+    private final Launcher mLauncher;
+    private Button mBtn;
+
+    public WorkPausedCard(Context context) {
+        this(context, null, 0);
+    }
+
+    public WorkPausedCard(Context context, AttributeSet attrs) {
+        this(context, attrs, 0);
+    }
+
+    public WorkPausedCard(Context context, AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+        mLauncher = Launcher.getLauncher(getContext());
+    }
+
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        mBtn = findViewById(R.id.enable_work_apps);
+        mBtn.setOnClickListener(this);
+    }
+
+    @Override
+    public void onClick(View view) {
+        if (Utilities.ATLEAST_P) {
+            setEnabled(false);
+            mLauncher.getStatsLogManager().logger().log(LAUNCHER_TURN_ON_WORK_APPS_TAP);
+            UI_HELPER_EXECUTOR.post(() -> WorkModeSwitch.setWorkProfileEnabled(getContext(), true));
+        }
+    }
+
+    @Override
+    protected void onLayout(boolean changed, int l, int t, int r, int b) {
+        int orientation = getResources().getConfiguration().orientation;
+        getLayoutParams().height = orientation == Configuration.ORIENTATION_PORTRAIT
+                ? LayoutParams.MATCH_PARENT : LayoutParams.WRAP_CONTENT;
+        super.onLayout(changed, l, t, r, b);
+    }
+}
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 0e710b7..65b4c3b 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -220,6 +220,10 @@
             "ENABLE_TWO_PANEL_HOME", false,
             "Uses two panel on home screen. Only applicable on large screen devices.");
 
+    public static final BooleanFlag ENABLE_SCRIM_FOR_APP_LAUNCH = getDebugFlag(
+            "ENABLE_SCRIM_FOR_APP_LAUNCH", false,
+            "Enables scrim during app launch animation.");
+
     public static final BooleanFlag ENABLE_SPLIT_SELECT = getDebugFlag(
             "ENABLE_SPLIT_SELECT", false, "Uses new split screen selection overview UI");
 
diff --git a/src/com/android/launcher3/folder/ClippedFolderIconLayoutRule.java b/src/com/android/launcher3/folder/ClippedFolderIconLayoutRule.java
index 3457804..54967a99 100644
--- a/src/com/android/launcher3/folder/ClippedFolderIconLayoutRule.java
+++ b/src/com/android/launcher3/folder/ClippedFolderIconLayoutRule.java
@@ -34,7 +34,6 @@
         float totalScale = scaleForItem(curNumItems);
         float transX;
         float transY;
-        float overlayAlpha = 0;
 
         if (index == EXIT_INDEX) {
             // 0 1 * <-- Exit position (row 0, col 2)
@@ -55,10 +54,9 @@
         transY = mTmpPoint[1];
 
         if (params == null) {
-            params = new PreviewItemDrawingParams(transX, transY, totalScale, overlayAlpha);
+            params = new PreviewItemDrawingParams(transX, transY, totalScale);
         } else {
             params.update(transX, transY, totalScale);
-            params.overlayAlpha = overlayAlpha;
         }
         return params;
     }
diff --git a/src/com/android/launcher3/folder/FolderAnimationManager.java b/src/com/android/launcher3/folder/FolderAnimationManager.java
index 57b289e..4fe85be 100644
--- a/src/com/android/launcher3/folder/FolderAnimationManager.java
+++ b/src/com/android/launcher3/folder/FolderAnimationManager.java
@@ -79,7 +79,7 @@
     private final TimeInterpolator mLargeFolderPreviewItemOpenInterpolator;
     private final TimeInterpolator mLargeFolderPreviewItemCloseInterpolator;
 
-    private final PreviewItemDrawingParams mTmpParams = new PreviewItemDrawingParams(0, 0, 0, 0);
+    private final PreviewItemDrawingParams mTmpParams = new PreviewItemDrawingParams(0, 0, 0);
     private final FolderGridOrganizer mPreviewVerifier;
 
     private ObjectAnimator mBgColorAnimator;
diff --git a/src/com/android/launcher3/folder/FolderIcon.java b/src/com/android/launcher3/folder/FolderIcon.java
index ed2d0a9..5c11451 100644
--- a/src/com/android/launcher3/folder/FolderIcon.java
+++ b/src/com/android/launcher3/folder/FolderIcon.java
@@ -113,7 +113,7 @@
     FolderGridOrganizer mPreviewVerifier;
     ClippedFolderIconLayoutRule mPreviewLayoutRule;
     private PreviewItemManager mPreviewItemManager;
-    private PreviewItemDrawingParams mTmpParams = new PreviewItemDrawingParams(0, 0, 0, 0);
+    private PreviewItemDrawingParams mTmpParams = new PreviewItemDrawingParams(0, 0, 0);
     private List<WorkspaceItemInfo> mCurrentPreviewItems = new ArrayList<>();
 
     boolean mAnimating = false;
@@ -391,7 +391,7 @@
             to.offset(center[0] - animateView.getMeasuredWidth() / 2,
                     center[1] - animateView.getMeasuredHeight() / 2);
 
-            float finalAlpha = index < MAX_NUM_ITEMS_IN_PREVIEW ? 0.5f : 0f;
+            float finalAlpha = index < MAX_NUM_ITEMS_IN_PREVIEW ? 1f : 0f;
 
             float finalScale = scale * scaleRelativeToDragLayer;
 
@@ -402,15 +402,18 @@
                 finalScale *= containerScale;
             }
 
+            final int finalIndex = index;
             dragLayer.animateView(animateView, from, to, finalAlpha,
                     1, 1, finalScale, finalScale, DROP_IN_ANIMATION_DURATION,
                     Interpolators.DEACCEL_2, Interpolators.ACCEL_2,
-                    null, DragLayer.ANIMATION_END_DISAPPEAR, null);
+                    () -> {
+                        mPreviewItemManager.hidePreviewItem(finalIndex, false);
+                        mFolder.showItem(item);
+                    }, DragLayer.ANIMATION_END_DISAPPEAR, null);
 
             mFolder.hideItem(item);
 
             if (!itemAdded) mPreviewItemManager.hidePreviewItem(index, true);
-            final int finalIndex = index;
 
             FolderNameInfos nameInfos = new FolderNameInfos();
             if (FeatureFlags.FOLDER_NAME_SUGGEST.get()) {
@@ -430,8 +433,6 @@
     private void showFinalView(int finalIndex, final WorkspaceItemInfo item,
             FolderNameInfos nameInfos, InstanceId instanceId) {
         postDelayed(() -> {
-            mPreviewItemManager.hidePreviewItem(finalIndex, false);
-            mFolder.showItem(item);
             setLabelSuggestion(nameInfos, instanceId);
             invalidate();
         }, DROP_IN_ANIMATION_DURATION);
diff --git a/src/com/android/launcher3/folder/FolderPreviewItemAnim.java b/src/com/android/launcher3/folder/FolderPreviewItemAnim.java
index edfd2ba..e20bafb 100644
--- a/src/com/android/launcher3/folder/FolderPreviewItemAnim.java
+++ b/src/com/android/launcher3/folder/FolderPreviewItemAnim.java
@@ -45,7 +45,7 @@
             };
 
     private static final PreviewItemDrawingParams sTmpParams =
-            new PreviewItemDrawingParams(0, 0, 0, 0);
+            new PreviewItemDrawingParams(0, 0, 0);
     private static final float[] sTempParamsArray = new float[3];
 
     private final ObjectAnimator mAnimator;
diff --git a/src/com/android/launcher3/folder/PreviewItemDrawingParams.java b/src/com/android/launcher3/folder/PreviewItemDrawingParams.java
index 5746be8..58efdc1 100644
--- a/src/com/android/launcher3/folder/PreviewItemDrawingParams.java
+++ b/src/com/android/launcher3/folder/PreviewItemDrawingParams.java
@@ -27,17 +27,15 @@
     float transX;
     float transY;
     float scale;
-    float overlayAlpha;
     public FolderPreviewItemAnim anim;
     public boolean hidden;
     public Drawable drawable;
     public WorkspaceItemInfo item;
 
-    PreviewItemDrawingParams(float transX, float transY, float scale, float overlayAlpha) {
+    PreviewItemDrawingParams(float transX, float transY, float scale) {
         this.transX = transX;
         this.transY = transY;
         this.scale = scale;
-        this.overlayAlpha = overlayAlpha;
     }
 
     public void update(float transX, float transY, float scale) {
diff --git a/src/com/android/launcher3/folder/PreviewItemManager.java b/src/com/android/launcher3/folder/PreviewItemManager.java
index baafbc9..a6674fc 100644
--- a/src/com/android/launcher3/folder/PreviewItemManager.java
+++ b/src/com/android/launcher3/folder/PreviewItemManager.java
@@ -260,7 +260,7 @@
             params.remove(params.size() - 1);
         }
         while (items.size() > params.size()) {
-            params.add(new PreviewItemDrawingParams(0, 0, 0, 0));
+            params.add(new PreviewItemDrawingParams(0, 0, 0));
         }
 
         int numItemsInFirstPagePreview = page == 0 ? items.size() : MAX_NUM_ITEMS_IN_PREVIEW;
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index 345a2ac..ddff338 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -479,7 +479,13 @@
         LAUNCHER_THEMED_ICON_ENABLED(836),
 
         @UiEvent(doc = "User disabled themed icons option in wallpaper & style settings.")
-        LAUNCHER_THEMED_ICON_DISABLED(837)
+        LAUNCHER_THEMED_ICON_DISABLED(837),
+
+        @UiEvent(doc = "User tapped on 'Turn on work apps' button in all apps window.")
+        LAUNCHER_TURN_ON_WORK_APPS_TAP(838),
+
+        @UiEvent(doc = "User tapped on 'Turn off work apps' button in all apps window.")
+        LAUNCHER_TURN_OFF_WORK_APPS_TAP(839)
         ;
 
         // ADD MORE
diff --git a/src/com/android/launcher3/model/PackageUpdatedTask.java b/src/com/android/launcher3/model/PackageUpdatedTask.java
index 7bfa3ef..82b0f7c 100644
--- a/src/com/android/launcher3/model/PackageUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageUpdatedTask.java
@@ -29,6 +29,7 @@
 import android.os.UserManager;
 import android.util.Log;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -123,6 +124,14 @@
                         activitiesLists.put(
                                 packages[i], appsList.updatePackage(context, packages[i], mUser));
                         app.getWidgetCache().removePackage(packages[i], mUser);
+
+                        // The update may have changed which shortcuts/widgets are available.
+                        // Refresh the widgets for the package if we have an activity running.
+                        Launcher launcher = Launcher.ACTIVITY_TRACKER.getCreatedActivity();
+                        if (launcher != null) {
+                            launcher.refreshAndBindWidgetsForPackageUser(
+                                    new PackageUserKey(packages[i], mUser));
+                        }
                     }
                 }
                 // Since package was just updated, the target must be available now.
@@ -212,7 +221,8 @@
                         }
 
                         if (si.isPromise() && isNewApkAvailable) {
-                            boolean isTargetValid = true;
+                            boolean isTargetValid = !cn.getClassName().equals(
+                                    IconCache.EMPTY_CLASS_NAME);
                             if (si.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
                                 List<ShortcutInfo> shortcut =
                                         new ShortcutRequest(context, mUser)
@@ -225,7 +235,7 @@
                                     si.updateFromDeepShortcutInfo(shortcut.get(0), context);
                                     infoUpdated = true;
                                 }
-                            } else if (!cn.getClassName().equals(IconCache.EMPTY_CLASS_NAME)) {
+                            } else if (isTargetValid) {
                                 isTargetValid = context.getSystemService(LauncherApps.class)
                                         .isActivityEnabled(cn, mUser);
                             }
diff --git a/src/com/android/launcher3/popup/PopupContainerWithArrow.java b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
index 332390d..2ae12ac 100644
--- a/src/com/android/launcher3/popup/PopupContainerWithArrow.java
+++ b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
@@ -96,6 +96,8 @@
     private int mNumNotifications;
     private ViewGroup mNotificationContainer;
 
+    private ViewGroup mWidgetContainer;
+
     private ViewGroup mDeepShortcutContainer;
 
     private ViewGroup mSystemShortcutContainer;
@@ -234,7 +236,7 @@
 
     @Override
     protected List<View> getChildrenForColorExtraction() {
-        return Arrays.asList(mSystemShortcutContainer, mDeepShortcutContainer,
+        return Arrays.asList(mSystemShortcutContainer, mWidgetContainer, mDeepShortcutContainer,
                 mNotificationContainer);
     }
 
@@ -298,10 +300,24 @@
             updateHiddenShortcuts();
 
             if (!systemShortcuts.isEmpty()) {
-                mSystemShortcutContainer = inflateAndAdd(R.layout.system_shortcut_icons, this);
                 for (SystemShortcut shortcut : systemShortcuts) {
-                    initializeSystemShortcut(
-                            R.layout.system_shortcut_icon_only, mSystemShortcutContainer, shortcut);
+                    if (shortcut instanceof SystemShortcut.Widgets) {
+                        if (mWidgetContainer == null) {
+                            mWidgetContainer = inflateAndAdd(R.layout.widget_shortcut_container,
+                                    this);
+                        }
+                        initializeSystemShortcut(R.layout.system_shortcut, mWidgetContainer,
+                                shortcut);
+                    }
+                }
+                mSystemShortcutContainer = inflateAndAdd(R.layout.system_shortcut_icons, this);
+
+                for (SystemShortcut shortcut : systemShortcuts) {
+                    if (!(shortcut instanceof SystemShortcut.Widgets)) {
+                        initializeSystemShortcut(
+                                R.layout.system_shortcut_icon_only, mSystemShortcutContainer,
+                                shortcut);
+                    }
                 }
             }
         } else {
@@ -524,25 +540,34 @@
             mLauncher.getPopupDataProvider().setChangeListener(null);
         }
 
+        private View getWidgetsView(ViewGroup container) {
+            for (int i = container.getChildCount() - 1; i >= 0; --i) {
+                View systemShortcutView = container.getChildAt(i);
+                if (systemShortcutView.getTag() instanceof SystemShortcut.Widgets) {
+                    return systemShortcutView;
+                }
+            }
+            return null;
+        }
+
         @Override
         public void onWidgetsBound() {
             ItemInfo itemInfo = (ItemInfo) mOriginalIcon.getTag();
             SystemShortcut widgetInfo = SystemShortcut.WIDGETS.getShortcut(mLauncher, itemInfo);
-            View widgetsView = null;
-            int count = mSystemShortcutContainer.getChildCount();
-            for (int i = 0; i < count; i++) {
-                View systemShortcutView = mSystemShortcutContainer.getChildAt(i);
-                if (systemShortcutView.getTag() instanceof SystemShortcut.Widgets) {
-                    widgetsView = systemShortcutView;
-                    break;
-                }
+            View widgetsView = getWidgetsView(PopupContainerWithArrow.this);
+            if (widgetsView == null && mWidgetContainer != null) {
+                widgetsView = getWidgetsView(mWidgetContainer);
             }
 
             if (widgetInfo != null && widgetsView == null) {
                 // We didn't have any widgets cached but now there are some, so enable the shortcut.
                 if (mSystemShortcutContainer != PopupContainerWithArrow.this) {
-                    initializeSystemShortcut(R.layout.system_shortcut_icon_only,
-                            mSystemShortcutContainer, widgetInfo);
+                    if (mWidgetContainer == null) {
+                        mWidgetContainer = inflateAndAdd(R.layout.widget_shortcut_container,
+                                PopupContainerWithArrow.this);
+                    }
+                    initializeSystemShortcut(R.layout.system_shortcut, mWidgetContainer,
+                            widgetInfo);
                 } else {
                     // If using the expanded system shortcut (as opposed to just the icon), we need
                     // to reopen the container to ensure measurements etc. all work out. While this
@@ -554,8 +579,10 @@
                 }
             } else if (widgetInfo == null && widgetsView != null) {
                 // No widgets exist, but we previously added the shortcut so remove it.
-                if (mSystemShortcutContainer != PopupContainerWithArrow.this) {
-                    mSystemShortcutContainer.removeView(widgetsView);
+                if (mSystemShortcutContainer
+                        != PopupContainerWithArrow.this
+                        && mWidgetContainer != null) {
+                    mWidgetContainer.removeView(widgetsView);
                 } else {
                     close(false);
                     PopupContainerWithArrow.showForIcon(mOriginalIcon);
diff --git a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
index a086635..90f37f3 100644
--- a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
+++ b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
@@ -335,12 +335,10 @@
         mCurrentAnimation.dispatchOnStart();
         if (targetState == LauncherState.ALL_APPS && !UNSTABLE_SPRINGS.get()) {
             if (mAllAppsOvershootStarted) {
-
                 mLauncher.getAppsView().onRelease();
                 mAllAppsOvershootStarted = false;
-
             } else {
-                mLauncher.getAppsView().addSpringFromFlingUpdateListener(anim, velocity);
+                mLauncher.getAppsView().addSpringFromFlingUpdateListener(anim, velocity, progress);
             }
         }
         anim.start();
diff --git a/src/com/android/launcher3/util/UiThreadHelper.java b/src/com/android/launcher3/util/UiThreadHelper.java
index cb721e6..0f40179 100644
--- a/src/com/android/launcher3/util/UiThreadHelper.java
+++ b/src/com/android/launcher3/util/UiThreadHelper.java
@@ -21,6 +21,7 @@
 import android.annotation.SuppressLint;
 import android.app.Activity;
 import android.content.Context;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Message;
@@ -42,15 +43,30 @@
     private static final int MSG_HIDE_KEYBOARD = 1;
     private static final int MSG_SET_ORIENTATION = 2;
     private static final int MSG_RUN_COMMAND = 3;
+    private static final String STATS_LOGGER_KEY = "STATS_LOGGER_KEY";
 
     @SuppressLint("NewApi")
     public static void hideKeyboardAsync(ActivityContext activityContext, IBinder token) {
         View root = activityContext.getDragLayer();
 
-        Message.obtain(HANDLER.get(root.getContext()),
-                MSG_HIDE_KEYBOARD, token).sendToTarget();
-        Launcher.cast(activityContext).getStatsLogManager().logger().log(
-                LAUNCHER_ALLAPPS_KEYBOARD_CLOSED);
+        // Since the launcher context cannot be accessed directly from callback, adding secondary
+        // message to log keyboard close event asynchronously.
+        Bundle mHideKeyboardLoggerMsg = new Bundle();
+        mHideKeyboardLoggerMsg.putParcelable(
+                STATS_LOGGER_KEY,
+                Message.obtain(
+                        HANDLER.get(root.getContext()),
+                        () -> Launcher.cast(activityContext)
+                                .getStatsLogManager()
+                                .logger()
+                                .log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED)
+                )
+        );
+
+        Message mHideKeyboardMsg = Message.obtain(HANDLER.get(root.getContext()), MSG_HIDE_KEYBOARD,
+                token);
+        mHideKeyboardMsg.setData(mHideKeyboardLoggerMsg);
+        mHideKeyboardMsg.sendToTarget();
     }
 
     public static void setOrientationAsync(Activity activity, int orientation) {
@@ -81,7 +97,11 @@
         public boolean handleMessage(Message message) {
             switch (message.what) {
                 case MSG_HIDE_KEYBOARD:
-                    mIMM.hideSoftInputFromWindow((IBinder) message.obj, 0);
+                    if (mIMM.hideSoftInputFromWindow((IBinder) message.obj, 0)) {
+                        // log keyboard close event only when keyboard is actually closed
+                        ((Message) message.getData().getParcelable(STATS_LOGGER_KEY))
+                                .sendToTarget();
+                    }
                     return true;
                 case MSG_SET_ORIENTATION:
                     ((Activity) message.obj).setRequestedOrientation(message.arg1);
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index 3308eec..47a8914 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -97,9 +97,9 @@
 
     @Override
     public void updateAppWidget(RemoteViews remoteViews) {
-        super.updateAppWidget(remoteViews);
         WidgetManagerHelper widgetManagerHelper = new WidgetManagerHelper(getContext());
         if (widgetManagerHelper.isAppWidgetRestored(mInfo.appWidgetId)) {
+            super.updateAppWidget(remoteViews);
             reInflate();
         }
     }
@@ -149,7 +149,7 @@
             // The view displays three modes,
             //   1) App icon in the center
             //   2) Preload icon in the center
-            //   3) Setup icon in the center and app icon in the top right corner.
+            //   3) App icon in the center with a setup icon on the top left corner.
             if (mDisabledForSafeMode) {
                 FastBitmapDrawable disabledIcon = info.newIcon(getContext());
                 disabledIcon.setIsDisabled(true);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index e89aea7..6863c60 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -103,18 +103,25 @@
             OnClickListener iconClickListener, OnLongClickListener iconLongClickListener) {
         mLauncher = Launcher.getLauncher(context);
         mDiffReporter = new WidgetsDiffReporter(iconCache, this);
+        WidgetsListDrawableFactory listDrawableFactory = new WidgetsListDrawableFactory(context);
         mWidgetsListTableViewHolderBinder = new WidgetsListTableViewHolderBinder(context,
                 layoutInflater, iconClickListener, iconLongClickListener,
-                widgetPreviewLoader, /* listAdapter= */ this);
+                widgetPreviewLoader, listDrawableFactory, /* listAdapter= */ this);
         mViewHolderBinders.put(VIEW_TYPE_WIDGETS_LIST, mWidgetsListTableViewHolderBinder);
         mViewHolderBinders.put(
                 VIEW_TYPE_WIDGETS_HEADER,
                 new WidgetsListHeaderViewHolderBinder(
-                        layoutInflater, /* onHeaderClickListener= */this, /* listAdapter= */ this));
+                        layoutInflater,
+                        /* onHeaderClickListener= */ this,
+                        listDrawableFactory,
+                        /* listAdapter= */ this));
         mViewHolderBinders.put(
                 VIEW_TYPE_WIDGETS_SEARCH_HEADER,
                 new WidgetsListSearchHeaderViewHolderBinder(
-                        layoutInflater, /*onHeaderClickListener=*/ this, /* listAdapter= */ this));
+                        layoutInflater,
+                        /* onHeaderClickListener= */ this,
+                        listDrawableFactory,
+                        /* listAdapter= */ this));
     }
 
     @Override
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListDrawableFactory.java b/src/com/android/launcher3/widget/picker/WidgetsListDrawableFactory.java
new file mode 100644
index 0000000..c61e3a4
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListDrawableFactory.java
@@ -0,0 +1,113 @@
+/*
+ * 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.picker;
+
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.FIRST;
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.FIRST_EXPANDED;
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.LAST;
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.MIDDLE;
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.MIDDLE_EXPANDED;
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.SINGLE;
+
+import android.content.Context;
+import android.content.res.ColorStateList;
+import android.content.res.Resources;
+import android.graphics.drawable.Drawable;
+import android.graphics.drawable.GradientDrawable;
+import android.graphics.drawable.RippleDrawable;
+import android.graphics.drawable.StateListDrawable;
+
+import com.android.launcher3.R;
+import com.android.launcher3.util.Themes;
+
+/** Factory for creating drawables to use as background for list elements. */
+final class WidgetsListDrawableFactory {
+
+    private final float mTopBottomCornerRadius;
+    private final float mMiddleCornerRadius;
+    private final ColorStateList mSurfaceColor;
+    private final ColorStateList mRippleColor;
+
+    WidgetsListDrawableFactory(Context context) {
+        Resources res = context.getResources();
+        mTopBottomCornerRadius = res.getDimension(R.dimen.widget_list_top_bottom_corner_radius);
+        mMiddleCornerRadius = res.getDimension(R.dimen.widget_list_content_corner_radius);
+        mSurfaceColor = context.getColorStateList(R.color.surface);
+        mRippleColor = ColorStateList.valueOf(
+                Themes.getAttrColor(context, android.R.attr.colorControlHighlight));
+    }
+
+    /**
+     * Creates a drawable for widget header list items. This drawable supports all positions
+     * in {@link WidgetsListDrawableState}.
+     */
+    Drawable createHeaderBackgroundDrawable() {
+        StateListDrawable stateList = new StateListDrawable();
+        stateList.addState(
+                SINGLE.mStateSet,
+                createRoundedRectDrawable(mTopBottomCornerRadius, mTopBottomCornerRadius));
+        stateList.addState(
+                FIRST_EXPANDED.mStateSet,
+                createRoundedRectDrawable(mTopBottomCornerRadius, 0));
+        stateList.addState(
+                FIRST.mStateSet,
+                createRoundedRectDrawable(mTopBottomCornerRadius, mMiddleCornerRadius));
+        stateList.addState(
+                MIDDLE_EXPANDED.mStateSet,
+                createRoundedRectDrawable(mMiddleCornerRadius, 0));
+        stateList.addState(
+                MIDDLE.mStateSet,
+                createRoundedRectDrawable(mMiddleCornerRadius, mMiddleCornerRadius));
+        stateList.addState(
+                LAST.mStateSet,
+                createRoundedRectDrawable(mMiddleCornerRadius, mTopBottomCornerRadius));
+        return new RippleDrawable(mRippleColor, /* content= */ stateList, /* mask= */ stateList);
+    }
+
+    /**
+     * Creates a drawable for widget content list items. This state list supports the middle and
+     * last states.
+     */
+    Drawable createContentBackgroundDrawable() {
+        StateListDrawable stateList = new StateListDrawable();
+        stateList.addState(
+                MIDDLE.mStateSet,
+                createRoundedRectDrawable(0, mMiddleCornerRadius));
+        stateList.addState(
+                LAST.mStateSet,
+                createRoundedRectDrawable(0, mTopBottomCornerRadius));
+        return new RippleDrawable(mRippleColor, /* content= */ stateList, /* mask= */ stateList);
+    }
+
+    /** Creates a rounded-rect drawable with the specified radii. */
+    private Drawable createRoundedRectDrawable(float topRadius, float bottomRadius) {
+        GradientDrawable backgroundMask = new GradientDrawable();
+        backgroundMask.setColor(mSurfaceColor);
+        backgroundMask.setShape(GradientDrawable.RECTANGLE);
+        backgroundMask.setCornerRadii(
+                new float[]{
+                        topRadius,
+                        topRadius,
+                        topRadius,
+                        topRadius,
+                        bottomRadius,
+                        bottomRadius,
+                        bottomRadius,
+                        bottomRadius
+                });
+        return backgroundMask;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListDrawableState.java b/src/com/android/launcher3/widget/picker/WidgetsListDrawableState.java
new file mode 100644
index 0000000..94f292b
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListDrawableState.java
@@ -0,0 +1,44 @@
+/*
+ * 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.picker;
+
+/**
+ * Different possible list position states for an item in the widgets list to have. Note that only
+ * headers use the expanded state.
+ */
+enum WidgetsListDrawableState {
+    FIRST(new int[]{android.R.attr.state_first}),
+    FIRST_EXPANDED(new int[]{android.R.attr.state_first, android.R.attr.state_expanded}),
+    MIDDLE(new int[]{android.R.attr.state_middle}),
+    MIDDLE_EXPANDED(new int[]{android.R.attr.state_middle, android.R.attr.state_expanded}),
+    LAST(new int[]{android.R.attr.state_last}),
+    SINGLE(new int[]{android.R.attr.state_single});
+
+    final int[] mStateSet;
+
+    WidgetsListDrawableState(int[] stateSet) {
+        mStateSet = stateSet;
+    }
+
+    static WidgetsListDrawableState obtain(boolean isFirst, boolean isLast, boolean isExpanded) {
+        if (isFirst && isLast) return SINGLE;
+        if (isFirst && isExpanded) return FIRST_EXPANDED;
+        if (isFirst) return FIRST;
+        if (isLast) return LAST;
+        if (isExpanded) return MIDDLE_EXPANDED;
+        return MIDDLE;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListDrawables.java b/src/com/android/launcher3/widget/picker/WidgetsListDrawables.java
deleted file mode 100644
index b3bb544..0000000
--- a/src/com/android/launcher3/widget/picker/WidgetsListDrawables.java
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * 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.picker;
-
-import android.content.Context;
-import android.content.res.ColorStateList;
-import android.graphics.drawable.Drawable;
-import android.graphics.drawable.GradientDrawable;
-import android.graphics.drawable.RippleDrawable;
-
-import com.android.launcher3.R;
-import com.android.launcher3.util.Themes;
-
-/** Helper class for creating drawables to use as background for list elements. */
-final class WidgetsListDrawables {
-
-    private WidgetsListDrawables() {}
-
-    /** Creates a list background drawable with the specified radii. */
-    static Drawable createListBackgroundDrawable(
-            Context context,
-            float topRadius,
-            float bottomRadius) {
-        GradientDrawable backgroundMask = new GradientDrawable();
-        backgroundMask.setColor(context.getColorStateList(R.color.surface));
-        backgroundMask.setShape(GradientDrawable.RECTANGLE);
-
-        backgroundMask.setCornerRadii(
-                new float[]{
-                        topRadius,
-                        topRadius,
-                        topRadius,
-                        topRadius,
-                        bottomRadius,
-                        bottomRadius,
-                        bottomRadius,
-                        bottomRadius
-                });
-
-        return new RippleDrawable(
-                /* color= */ ColorStateList.valueOf(
-                        Themes.getAttrColor(context, android.R.attr.colorControlHighlight)),
-                /* content= */ backgroundMask,
-                /* mask= */ backgroundMask);
-    }
-
-}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
index fece359..cdab964 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -60,9 +60,6 @@
     @Nullable private Drawable mIconDrawable;
     private final int mIconSize;
     private final int mBottomMarginSize;
-    private final float mTopBottomCornerRadius;
-    private final float mMiddleCornerRadius;
-    private final float mJoinedCornerRadius;
 
     private ImageView mAppIcon;
     private TextView mTitle;
@@ -70,6 +67,7 @@
 
     private CheckBox mExpandToggle;
     private boolean mIsExpanded = false;
+    @Nullable private WidgetsListDrawableState mListDrawableState;
 
     public WidgetsListHeader(Context context) {
         this(context, /* attrs= */ null);
@@ -90,12 +88,6 @@
                 grid.iconSizePx);
         mBottomMarginSize =
                 getResources().getDimensionPixelSize(R.dimen.widget_list_entry_bottom_margin);
-        mTopBottomCornerRadius =
-                getResources().getDimension(R.dimen.widget_list_top_bottom_corner_radius);
-        mMiddleCornerRadius =
-                getResources().getDimension(R.dimen.widget_list_content_corner_radius);
-        mJoinedCornerRadius =
-                getResources().getDimension(R.dimen.widget_list_content_joined_corner_radius);
     }
 
     @Override
@@ -163,6 +155,14 @@
         }
     }
 
+    /** Sets the {@link WidgetsListDrawableState} and refreshes the background drawable. */
+    @UiThread
+    public void setListDrawableState(WidgetsListDrawableState state) {
+        if (state == mListDrawableState) return;
+        this.mListDrawableState = state;
+        refreshDrawableState();
+    }
+
     /** Apply app icon, labels and tag using a generic {@link WidgetsListHeaderEntry}. */
     @UiThread
     public void applyFromItemInfoWithIcon(WidgetsListHeaderEntry entry) {
@@ -263,20 +263,6 @@
         verifyHighRes();
     }
 
-    /** Updates the list to have a background drawable with the appropriate corner radii. */
-    @UiThread
-    public void updateListBackground(boolean isFirst, boolean isLast, boolean isExpanded) {
-        float topRadius = isFirst ? mTopBottomCornerRadius : mMiddleCornerRadius;
-        float bottomRadius = isLast
-                ? mTopBottomCornerRadius
-                : isExpanded
-                        ? mJoinedCornerRadius
-                        : mMiddleCornerRadius;
-        setBackground(
-                WidgetsListDrawables.createListBackgroundDrawable(
-                        getContext(), topRadius, bottomRadius));
-    }
-
     private void setTitles(WidgetsListSearchHeaderEntry entry) {
         mTitle.setText(entry.mPkgItem.title);
 
@@ -300,6 +286,17 @@
         }
     }
 
+    @Override
+    protected int[] onCreateDrawableState(int extraSpace) {
+        if (mListDrawableState == null) return super.onCreateDrawableState(extraSpace);
+        // Augment the state set from the super implementation with the custom states from
+        // mListDrawableState.
+        int[] drawableState =
+                super.onCreateDrawableState(extraSpace + mListDrawableState.mStateSet.length);
+        mergeDrawableStates(drawableState, mListDrawableState.mStateSet);
+        return drawableState;
+    }
+
     /** Verifies that the current icon is high-res otherwise posts a request to load the icon. */
     public void verifyHighRes() {
         if (mIconLoadRequest != null) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
index 22d6d22..2f8f1ba 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
@@ -30,13 +30,16 @@
         ViewHolderBinder<WidgetsListHeaderEntry, WidgetsListHeaderHolder> {
     private final LayoutInflater mLayoutInflater;
     private final OnHeaderClickListener mOnHeaderClickListener;
+    private final WidgetsListDrawableFactory mListDrawableFactory;
     private final WidgetsListAdapter mWidgetsListAdapter;
 
     public WidgetsListHeaderViewHolderBinder(LayoutInflater layoutInflater,
             OnHeaderClickListener onHeaderClickListener,
+            WidgetsListDrawableFactory listDrawableFactory,
             WidgetsListAdapter listAdapter) {
         mLayoutInflater = layoutInflater;
         mOnHeaderClickListener = onHeaderClickListener;
+        mListDrawableFactory = listDrawableFactory;
         mWidgetsListAdapter = listAdapter;
     }
 
@@ -44,7 +47,7 @@
     public WidgetsListHeaderHolder newViewHolder(ViewGroup parent) {
         WidgetsListHeader header = (WidgetsListHeader) mLayoutInflater.inflate(
                 R.layout.widgets_list_row_header, parent, false);
-
+        header.setBackground(mListDrawableFactory.createHeaderBackgroundDrawable());
         return new WidgetsListHeaderHolder(header);
     }
 
@@ -52,12 +55,13 @@
     public void bindViewHolder(WidgetsListHeaderHolder viewHolder, WidgetsListHeaderEntry data,
             int position) {
         WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
-        widgetsListHeader.updateListBackground(
-                /* isFirst= */ position == 0,
-                /* isLast= */ position == mWidgetsListAdapter.getItemCount() - 1,
-                /* isExpanded= */ data.isWidgetListShown());
         widgetsListHeader.applyFromItemInfoWithIcon(data);
         widgetsListHeader.setExpanded(data.isWidgetListShown());
+        widgetsListHeader.setListDrawableState(
+                WidgetsListDrawableState.obtain(
+                        /* isFirst= */ position == 0,
+                        /* isLast= */ position == mWidgetsListAdapter.getItemCount() - 1,
+                        /* isExpanded= */ data.isWidgetListShown()));
         widgetsListHeader.setOnExpandChangeListener(isExpanded ->
                 mOnHeaderClickListener.onHeaderClicked(
                         isExpanded,
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
index d5e03a4..31dd9ee 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
@@ -31,13 +31,16 @@
         ViewHolderBinder<WidgetsListSearchHeaderEntry, WidgetsListSearchHeaderHolder> {
     private final LayoutInflater mLayoutInflater;
     private final OnHeaderClickListener mOnHeaderClickListener;
+    private final WidgetsListDrawableFactory mListDrawableFactory;
     private final WidgetsListAdapter mWidgetsListAdapter;
 
     public WidgetsListSearchHeaderViewHolderBinder(LayoutInflater layoutInflater,
             OnHeaderClickListener onHeaderClickListener,
+            WidgetsListDrawableFactory listDrawableFactory,
             WidgetsListAdapter listAdapter) {
         mLayoutInflater = layoutInflater;
         mOnHeaderClickListener = onHeaderClickListener;
+        mListDrawableFactory = listDrawableFactory;
         mWidgetsListAdapter = listAdapter;
     }
 
@@ -45,7 +48,7 @@
     public WidgetsListSearchHeaderHolder newViewHolder(ViewGroup parent) {
         WidgetsListHeader header = (WidgetsListHeader) mLayoutInflater.inflate(
                 R.layout.widgets_list_row_header, parent, false);
-
+        header.setBackground(mListDrawableFactory.createHeaderBackgroundDrawable());
         return new WidgetsListSearchHeaderHolder(header);
     }
 
@@ -53,12 +56,13 @@
     public void bindViewHolder(WidgetsListSearchHeaderHolder viewHolder,
             WidgetsListSearchHeaderEntry data, int position) {
         WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
-        widgetsListHeader.updateListBackground(
-                /* isFirst= */ position == 0,
-                /* isLast= */ position == mWidgetsListAdapter.getItemCount() - 1,
-                /* isExpanded= */ data.isWidgetListShown());
         widgetsListHeader.applyFromItemInfoWithIcon(data);
         widgetsListHeader.setExpanded(data.isWidgetListShown());
+        widgetsListHeader.setListDrawableState(
+                WidgetsListDrawableState.obtain(
+                        /* isFirst= */ position == 0,
+                        /* isLast= */ position == mWidgetsListAdapter.getItemCount() - 1,
+                        /* isExpanded= */ data.isWidgetListShown()));
         widgetsListHeader.setOnExpandChangeListener(isExpanded ->
                 mOnHeaderClickListener.onHeaderClicked(isExpanded,
                         new PackageUserKey(data.mPkgItem.packageName, data.mPkgItem.user)));
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableView.java b/src/com/android/launcher3/widget/picker/WidgetsListTableView.java
new file mode 100644
index 0000000..d30e7b6
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableView.java
@@ -0,0 +1,59 @@
+/*
+ * 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.picker;
+
+import android.content.Context;
+import android.util.AttributeSet;
+import android.widget.TableLayout;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.UiThread;
+
+/**
+ * Extension of {@link TableLayout} to support the drawable states used by
+ * {@link WidgetsListDrawableState}.
+ */
+public class WidgetsListTableView extends TableLayout {
+
+    @Nullable private WidgetsListDrawableState mListDrawableState;
+
+    public WidgetsListTableView(Context context) {
+        super(context);
+    }
+
+    public WidgetsListTableView(Context context, AttributeSet attrs) {
+        super(context, attrs);
+    }
+
+    /** Sets the {@link WidgetsListDrawableState} and refreshes the background drawable. */
+    @UiThread
+    public void setListDrawableState(WidgetsListDrawableState state) {
+        if (state == mListDrawableState) return;
+        mListDrawableState = state;
+        refreshDrawableState();
+    }
+
+    @Override
+    protected int[] onCreateDrawableState(int extraSpace) {
+        if (mListDrawableState == null) return super.onCreateDrawableState(extraSpace);
+        // Augment the state set from the super implementation with the custom states from
+        // mListDrawableState.
+        int[] drawableState =
+                super.onCreateDrawableState(extraSpace + mListDrawableState.mStateSet.length);
+        mergeDrawableStates(drawableState, mListDrawableState.mStateSet);
+        return drawableState;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
index 8e310c5..7e8c55b 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -15,8 +15,10 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.LAST;
+import static com.android.launcher3.widget.picker.WidgetsListDrawableState.MIDDLE;
+
 import android.content.Context;
-import android.content.res.Resources;
 import android.util.Log;
 import android.view.Gravity;
 import android.view.LayoutInflater;
@@ -51,10 +53,8 @@
     private final OnClickListener mIconClickListener;
     private final OnLongClickListener mIconLongClickListener;
     private final WidgetPreviewLoader mWidgetPreviewLoader;
+    private final WidgetsListDrawableFactory mListDrawableFactory;
     private final WidgetsListAdapter mWidgetsListAdapter;
-    private final float mTopBottomCornerRadius;
-    private final float mMiddleCornerRadius;
-    private final float mJoinedCornerRadius;
     private boolean mApplyBitmapDeferred = false;
 
     public WidgetsListTableViewHolderBinder(
@@ -63,19 +63,14 @@
             OnClickListener iconClickListener,
             OnLongClickListener iconLongClickListener,
             WidgetPreviewLoader widgetPreviewLoader,
+            WidgetsListDrawableFactory listDrawableFactory,
             WidgetsListAdapter listAdapter) {
         mLayoutInflater = layoutInflater;
         mIconClickListener = iconClickListener;
         mIconLongClickListener = iconLongClickListener;
         mWidgetPreviewLoader = widgetPreviewLoader;
+        mListDrawableFactory = listDrawableFactory;
         mWidgetsListAdapter = listAdapter;
-        Resources resources = context.getResources();
-        mTopBottomCornerRadius =
-                resources.getDimension(R.dimen.widget_list_top_bottom_corner_radius);
-        mMiddleCornerRadius =
-                resources.getDimension(R.dimen.widget_list_content_corner_radius);
-        mJoinedCornerRadius =
-                resources.getDimension(R.dimen.widget_list_content_joined_corner_radius);
     }
 
     /**
@@ -97,28 +92,25 @@
             Log.v(TAG, "\nonCreateViewHolder");
         }
 
-        ViewGroup container = (ViewGroup) mLayoutInflater.inflate(
-                R.layout.widgets_table_container, parent, false);
-        return new WidgetsRowViewHolder(container);
+        WidgetsRowViewHolder viewHolder =
+                new WidgetsRowViewHolder(mLayoutInflater.inflate(
+                        R.layout.widgets_table_container, parent, false));
+        viewHolder.mTableContainer.setBackgroundDrawable(
+                mListDrawableFactory.createContentBackgroundDrawable());
+        return viewHolder;
     }
 
     @Override
     public void bindViewHolder(WidgetsRowViewHolder holder, WidgetsListContentEntry entry,
             int position) {
-        TableLayout table = holder.mTableContainer;
+        WidgetsListTableView table = holder.mTableContainer;
         if (DEBUG) {
             Log.d(TAG, String.format("onBindViewHolder [widget#=%d, table.getChildCount=%d]",
                     entry.mWidgets.size(), table.getChildCount()));
         }
 
-        // The content is always joined to an expanded header above.
-        float topRadius = mJoinedCornerRadius;
-        float bottomRadius = position == mWidgetsListAdapter.getItemCount() - 1
-                ? mTopBottomCornerRadius
-                : mMiddleCornerRadius;
-        table.setBackgroundDrawable(
-                WidgetsListDrawables.createListBackgroundDrawable(
-                        holder.itemView.getContext(), topRadius, bottomRadius));
+        table.setListDrawableState(
+                position == mWidgetsListAdapter.getItemCount() - 1 ? LAST : MIDDLE);
 
         List<ArrayList<WidgetItem>> widgetItemsTable =
                 WidgetsTableUtils.groupWidgetItemsIntoTable(entry.mWidgets, mMaxSpansPerRow);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
index aef1103..618e2cb 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
@@ -15,8 +15,7 @@
  */
 package com.android.launcher3.widget.picker;
 
-import android.view.ViewGroup;
-import android.widget.TableLayout;
+import android.view.View;
 
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
 
@@ -25,9 +24,9 @@
 /** A {@link ViewHolder} for showing widgets of an app in the full widget picker. */
 public final class WidgetsRowViewHolder extends ViewHolder {
 
-    public final TableLayout mTableContainer;
+    public final WidgetsListTableView mTableContainer;
 
-    public WidgetsRowViewHolder(ViewGroup v) {
+    public WidgetsRowViewHolder(View v) {
         super(v);
 
         mTableContainer = v.findViewById(R.id.widgets_table);
diff --git a/src/com/android/launcher3/widget/util/WidgetSizes.java b/src/com/android/launcher3/widget/util/WidgetSizes.java
index 5c8ea72..fcc0dde 100644
--- a/src/com/android/launcher3/widget/util/WidgetSizes.java
+++ b/src/com/android/launcher3/widget/util/WidgetSizes.java
@@ -99,11 +99,13 @@
      */
     public static Bundle getWidgetSizeOptions(Context context, ComponentName provider, int spanX,
             int spanY) {
+        boolean shouldInsetWidgets =
+                LauncherAppState.getIDP(context).getDeviceProfile(context).shouldInsetWidgets();
         ArrayList<SizeF> sizes = getWidgetSizes(context, spanX, spanY);
         Rect padding = getDefaultPaddingForWidget(context, provider, null);
         float density = context.getResources().getDisplayMetrics().density;
-        float xPaddingDips = (padding.left + padding.right) / density;
-        float yPaddingDips = (padding.top + padding.bottom) / density;
+        float xPaddingDips = shouldInsetWidgets ? (padding.left + padding.right) / density : 0;
+        float yPaddingDips = shouldInsetWidgets ? (padding.top + padding.bottom) / density : 0;
 
         ArrayList<SizeF> paddedSizes = sizes.stream()
                 .map(size -> new SizeF(
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index 51c5b12..a66b031 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -15,6 +15,7 @@
 import android.util.Log;
 
 import androidx.annotation.Nullable;
+import androidx.collection.ArrayMap;
 
 import com.android.launcher3.AppFilter;
 import com.android.launcher3.InvariantDeviceProfile;
@@ -37,6 +38,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -158,34 +160,29 @@
             Log.d(TAG, "addWidgetsAndShortcuts, widgetsShortcuts#=" + rawWidgetsShortcuts.size());
         }
 
-        // Temporary list for {@link PackageItemInfos} to avoid having to go through
+        // Temporary cache for {@link PackageItemInfos} to avoid having to go through
         // {@link mPackageItemInfos} to locate the key to be used for {@link #mWidgetsList}
-        HashMap<WidgetPackageOrCategoryKey, PackageItemInfo> tmpPackageItemInfos = new HashMap<>();
+        PackageItemInfoCache packageItemInfoCache = new PackageItemInfoCache();
 
-        // Clear the lists only if this is an update on all widgets and shortcuts. If packageUser
-        // isn't null, only updates the shortcuts and widgets for the app represented in
-        // packageUser.
         if (packageUser == null) {
+            // Clear the list if this is an update on all widgets and shortcuts.
             mWidgetsList.clear();
+        } else {
+            // Otherwise, only clear the widgets and shortcuts for the changed package.
+            mWidgetsList.remove(
+                    packageItemInfoCache.getOrCreate(new WidgetPackageOrCategoryKey(packageUser)));
         }
+
         // add and update.
         mWidgetsList.putAll(rawWidgetsShortcuts.stream()
                 .filter(new WidgetValidityCheck(app))
-                .collect(Collectors.groupingBy(item -> {
-                    WidgetPackageOrCategoryKey packageUserKey = getWidgetPackageOrCategoryKey(item);
-                    PackageItemInfo pInfo = tmpPackageItemInfos.get(packageUserKey);
-                    if (pInfo == null) {
-                        pInfo = new PackageItemInfo(item.componentName.getPackageName(),
-                                packageUserKey.mCategory);
-                        pInfo.user = item.user;
-                        tmpPackageItemInfos.put(packageUserKey,  pInfo);
-                    }
-                    return pInfo;
-                })));
+                .collect(Collectors.groupingBy(item ->
+                        packageItemInfoCache.getOrCreate(getWidgetPackageOrCategoryKey(item))
+                )));
 
         // Update each package entry
         IconCache iconCache = app.getIconCache();
-        for (PackageItemInfo p : tmpPackageItemInfos.values()) {
+        for (PackageItemInfo p : packageItemInfoCache.values()) {
             iconCache.getTitleAndIconForApp(p, true /* userLowResIcon */);
         }
     }
@@ -289,6 +286,10 @@
         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);
         }
@@ -310,4 +311,22 @@
             return mHashCode;
         }
     }
+
+    private static final class PackageItemInfoCache {
+        private final Map<WidgetPackageOrCategoryKey, PackageItemInfo> mMap = new ArrayMap<>();
+
+        PackageItemInfo getOrCreate(WidgetPackageOrCategoryKey key) {
+            PackageItemInfo pInfo = mMap.get(key);
+            if (pInfo == null) {
+                pInfo = new PackageItemInfo(key.mPackage, key.mCategory);
+                pInfo.user = key.mUser;
+                mMap.put(key,  pInfo);
+            }
+            return pInfo;
+        }
+
+        Collection<PackageItemInfo> values() {
+            return mMap.values();
+        }
+    }
 }
\ No newline at end of file
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 6f47df0..06bc26a 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -293,7 +293,7 @@
         try {
             final AppIconMenu menu = allApps.
                     getAppIcon(APP_NAME).
-                    openMenu();
+                    openDeepShortcutMenu();
 
             executeOnLauncher(
                     launcher -> assertTrue("Launcher internal state didn't switch to Showing Menu",
@@ -341,7 +341,7 @@
         try {
             final AppIconMenu menu = allApps
                     .getAppIcon(APP_NAME)
-                    .openMenu();
+                    .openDeepShortcutMenu();
             final AppIconMenuItem menuItem0 = menu.getMenuItem(0);
             final AppIconMenuItem menuItem2 = menu.getMenuItem(2);
 
diff --git a/tests/tapl/com/android/launcher3/tapl/AppIcon.java b/tests/tapl/com/android/launcher3/tapl/AppIcon.java
index 56723b1..57fd08a 100644
--- a/tests/tapl/com/android/launcher3/tapl/AppIcon.java
+++ b/tests/tapl/com/android/launcher3/tapl/AppIcon.java
@@ -51,6 +51,16 @@
         }
     }
 
+    /**
+     * Long-clicks the icon to open its menu, and looks at the deep shortcuts container only.
+     */
+    public AppIconMenu openDeepShortcutMenu() {
+        try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
+            return new AppIconMenu(mLauncher, mLauncher.clickAndGet(
+                    mObject, "deep_shortcuts_container", LONG_CLICK_EVENT));
+        }
+    }
+
     @Override
     protected void addExpectedEventsForLongClick() {
         mLauncher.expectEvent(TestProtocol.SEQUENCE_MAIN, LONG_CLICK_EVENT);