diff --git a/Android.mk b/Android.mk
index 19ad328..127df79 100644
--- a/Android.mk
+++ b/Android.mk
@@ -145,7 +145,9 @@
     $(call all-java-files-under, quickstep/src) \
     $(call all-java-files-under, src_shortcuts_overrides)
 
-LOCAL_RESOURCE_DIR := $(LOCAL_PATH)/quickstep/res
+LOCAL_RESOURCE_DIR := \
+    $(LOCAL_PATH)/quickstep/res \
+    $(LOCAL_PATH)/quickstep/overview_ui_overrides/res
 LOCAL_PROGUARD_ENABLED := disabled
 
 
@@ -174,7 +176,9 @@
 LOCAL_OVERRIDES_PACKAGES := Home Launcher2 Launcher3
 LOCAL_REQUIRED_MODULES := privapp_whitelist_com.android.launcher3
 
-LOCAL_RESOURCE_DIR := $(LOCAL_PATH)/quickstep/res
+LOCAL_RESOURCE_DIR := \
+    $(LOCAL_PATH)/quickstep/res \
+    $(LOCAL_PATH)/quickstep/overview_ui_overrides/res
 
 LOCAL_FULL_LIBS_MANIFEST_FILES := \
     $(LOCAL_PATH)/quickstep/AndroidManifest-launcher.xml \
@@ -213,7 +217,8 @@
 LOCAL_RESOURCE_DIR := \
     $(LOCAL_PATH)/quickstep/res \
     $(LOCAL_PATH)/go/res \
-    $(LOCAL_PATH)/go/quickstep/res
+    $(LOCAL_PATH)/go/quickstep/res \
+    $(LOCAL_PATH)/go/quickstep/overview_ui_overrides/res
 
 LOCAL_PROGUARD_FLAG_FILES := proguard.flags
 LOCAL_PROGUARD_ENABLED := full
diff --git a/go/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml b/go/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
new file mode 100644
index 0000000..b438da3
--- /dev/null
+++ b/go/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
@@ -0,0 +1,98 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2021 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<com.android.quickstep.views.GoOverviewActionsView
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="@dimen/overview_actions_height"
+    android:layout_gravity="center_horizontal|bottom">
+
+    <LinearLayout
+        android:id="@+id/action_buttons"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center"
+        android:orientation="horizontal">
+
+        <Space
+            android:layout_width="0dp"
+            android:layout_height="1dp"
+            android:layout_weight="1" />
+
+        <Button
+            android:id="@+id/action_listen"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_listen"
+            android:drawablePadding="1dp"
+            android:text="@string/action_listen"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Button
+            android:id="@+id/action_translate"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_translate"
+            android:drawablePadding="1dp"
+            android:text="@string/action_translate"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Button
+            android:id="@+id/action_search"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_search"
+            android:drawablePadding="1dp"
+            android:text="@string/action_search"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Button
+            android:id="@+id/action_screenshot"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableTop="@drawable/ic_screenshot"
+            android:drawablePadding="1dp"
+            android:text="@string/action_screenshot"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor" />
+
+        <Space
+            android:layout_width="0dp"
+            android:layout_height="1dp"
+            android:layout_weight="1" />
+
+        <Button
+            android:id="@+id/action_share"
+            style="@style/OverviewActionButton"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:drawableStart="@drawable/ic_share"
+            android:text="@string/action_share"
+            android:theme="@style/ThemeControlHighlightWorkspaceColor"
+            android:visibility="gone" />
+
+        <Space
+            android:id="@+id/oav_three_button_space"
+            android:layout_width="0dp"
+            android:layout_height="1dp"
+            android:layout_weight="1"
+            android:visibility="gone" />
+    </LinearLayout>
+
+</com.android.quickstep.views.GoOverviewActionsView>
\ No newline at end of file
diff --git a/go/quickstep/overview_ui_overrides/res/values/config.xml b/go/quickstep/overview_ui_overrides/res/values/config.xml
new file mode 100644
index 0000000..ec21a01
--- /dev/null
+++ b/go/quickstep/overview_ui_overrides/res/values/config.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+/*
+* Copyright (C) 2021 The Android Open Source Project
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*      http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+-->
+<resources>
+    <string name="task_overlay_factory_class" translatable="false">
+        com.android.quickstep.TaskOverlayFactoryGo</string>
+</resources>
\ No newline at end of file
diff --git a/go/quickstep/res/drawable/ic_listen.xml b/go/quickstep/res/drawable/ic_listen.xml
new file mode 100644
index 0000000..a8e6c93
--- /dev/null
+++ b/go/quickstep/res/drawable/ic_listen.xml
@@ -0,0 +1,32 @@
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="28dp"
+    android:height="28dp"
+    android:viewportWidth="28"
+    android:viewportHeight="28">
+  <path
+      android:pathData="M10.5,15.17c2.58,0 4.67,-2.09 4.67,-4.67s-2.09,-4.67 -4.67,-4.67c-2.58,0 -4.67,2.09 -4.67,4.67S7.92,15.17 10.5,15.17zM10.5,8.17c1.28,0 2.33,1.05 2.33,2.33s-1.05,2.33 -2.33,2.33c-1.28,0 -2.33,-1.05 -2.33,-2.33S9.22,8.17 10.5,8.17z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M10.5,17.5c-3.11,0 -9.33,1.56 -9.33,4.67v2.33h18.67v-2.33C19.83,19.06 13.62,17.5 10.5,17.5zM3.5,22.17c0.26,-0.84 3.86,-2.33 7,-2.33c3.15,0 6.77,1.5 7,2.33H3.5z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M25.67,10.5c0,0.36 -0.02,0.71 -0.05,1.05c-0.01,0.15 -0.03,0.29 -0.05,0.43c-0.02,0.18 -0.05,0.36 -0.08,0.54c-0.04,0.2 -0.07,0.39 -0.12,0.58c-0.01,0.06 -0.03,0.11 -0.04,0.17c-0.59,2.34 -1.81,4.01 -2.52,4.82c-0.09,0.1 -0.18,0.2 -0.28,0.3c-0.17,0.18 -0.27,0.27 -0.27,0.27l-1.65,-1.63c1.34,-1.33 2.27,-3.07 2.6,-5.01c0.01,-0.08 0.02,-0.16 0.04,-0.24c0.06,-0.42 0.1,-0.85 0.1,-1.29c0,-0.44 -0.04,-0.88 -0.1,-1.3c-0.01,-0.06 -0.02,-0.13 -0.03,-0.19c-0.32,-1.95 -1.25,-3.7 -2.6,-5.04l1.65,-1.63c0,0 0.11,0.1 0.27,0.27c0.09,0.1 0.19,0.2 0.28,0.3c0.71,0.82 1.93,2.48 2.52,4.82c0.01,0.06 0.03,0.11 0.04,0.17c0.04,0.19 0.08,0.38 0.12,0.58c0.03,0.18 0.06,0.36 0.08,0.54c0.02,0.14 0.04,0.28 0.05,0.43C25.65,9.79 25.67,10.14 25.67,10.5z"
+      android:fillColor="#EA4335"/>
+  <path
+      android:pathData="M20.61,8.4C20.85,9.06 21,9.76 21,10.5s-0.15,1.44 -0.39,2.1c-0.28,0.77 -0.71,1.46 -1.25,2.05l-1.66,-1.64c0.56,-0.63 0.91,-1.44 0.95,-2.34c0,-0.06 0.02,-0.11 0.02,-0.17s-0.01,-0.11 -0.02,-0.17c-0.04,-0.9 -0.39,-1.71 -0.95,-2.34l1.66,-1.64C19.9,6.94 20.32,7.63 20.61,8.4z"
+      android:fillColor="#FBBC04"/>
+</vector>
diff --git a/go/quickstep/res/drawable/ic_search.xml b/go/quickstep/res/drawable/ic_search.xml
new file mode 100644
index 0000000..4307330
--- /dev/null
+++ b/go/quickstep/res/drawable/ic_search.xml
@@ -0,0 +1,32 @@
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="28dp"
+    android:height="28dp"
+    android:viewportWidth="28"
+    android:viewportHeight="28">
+  <path
+      android:pathData="M24.5,22.75l-6.84,-6.84c1,-1.35 1.59,-3.02 1.59,-4.83h-2.33c0,3.22 -2.62,5.83 -5.83,5.83v2.33c1.81,0 3.47,-0.6 4.83,-1.59l6.84,6.84L24.5,22.75z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M11.08,2.92v2.33c3.22,0 5.83,2.62 5.83,5.83h2.33C19.25,6.57 15.59,2.92 11.08,2.92z"
+      android:fillColor="#34A853"/>
+  <path
+      android:pathData="M5.25,11.08H2.92c0,4.51 3.66,8.17 8.17,8.17v-2.33C7.87,16.92 5.25,14.3 5.25,11.08z"
+      android:fillColor="#EA4335"/>
+  <path
+      android:pathData="M2.92,11.08h2.33c0,-3.22 2.62,-5.83 5.83,-5.83V2.92C6.57,2.92 2.92,6.57 2.92,11.08z"
+      android:fillColor="#FBBC04"/>
+</vector>
diff --git a/go/quickstep/res/drawable/ic_translate.xml b/go/quickstep/res/drawable/ic_translate.xml
new file mode 100644
index 0000000..1247807
--- /dev/null
+++ b/go/quickstep/res/drawable/ic_translate.xml
@@ -0,0 +1,32 @@
+<!-- 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="28dp"
+    android:height="28dp"
+    android:viewportWidth="28"
+    android:viewportHeight="28">
+  <path
+      android:pathData="M12.28,15.19l-0.07,-0.05c-0.61,-0.49 -1.15,-1.05 -1.65,-1.63c-1.05,-1.22 -1.88,-2.63 -2.39,-4.17H5.83c0.54,2.17 1.58,4.16 3.01,5.85l-5.93,5.23l1.75,1.75l5.91,-5.26c0.05,0.04 0.1,0.09 0.15,0.13l3.42,2.79l1.02,-2.33L12.28,15.19z"
+      android:fillColor="#FBBC04"/>
+  <path
+      android:pathData="M21.58,11.67h-2.33l-5.25,14h2.33l1.31,-3.5h5.54l1.32,3.5h2.33L21.58,11.67zM18.53,19.83l1.89,-5.05l1.89,5.05H18.53z"
+      android:fillColor="#4285F4"/>
+  <path
+      android:pathData="M11.67,2.33l-2.34,0l0,2.34l-8.16,0l0,2.33l10.5,0l0,-2.33z"
+      android:fillColor="#EA4335"/>
+  <path
+      android:pathData="M11.67,4.67V7H14c-0.61,2.42 -1.79,4.65 -3.44,6.5c0.5,0.59 1.04,1.15 1.65,1.63l0.07,0.05c2.03,-2.32 3.44,-5.14 4.05,-8.19h3.5V4.67H11.67z"
+      android:fillColor="#34A853"/>
+</vector>
diff --git a/go/quickstep/res/values/config.xml b/go/quickstep/res/values/config.xml
index f376774..9dca137 100644
--- a/go/quickstep/res/values/config.xml
+++ b/go/quickstep/res/values/config.xml
@@ -14,5 +14,11 @@
      limitations under the License.
 -->
 <resources>
+    <!-- The component to receive app sharing Intents -->
     <string name="app_sharing_component" translatable="false"/>
+    <!-- The package to receive Listen, Translate, and Search Intents -->
+    <string name="niu_actions_package" translatable="false"/>
+
+    <!-- Feature Flags -->
+    <bool name="enable_niu_actions">false</bool>
 </resources>
\ No newline at end of file
diff --git a/go/quickstep/res/values/strings.xml b/go/quickstep/res/values/strings.xml
index fdd8397..71e2f3a 100644
--- a/go/quickstep/res/values/strings.xml
+++ b/go/quickstep/res/values/strings.xml
@@ -3,4 +3,12 @@
 <resources xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
     <!-- Label for app share drop target. [CHAR_LIMIT=20] -->
     <string name="app_share_drop_target_label">Share App</string>
+
+    <!-- ******* Overview ******* -->
+    <!-- Label for a button that lets the user listen to the content of the current app. [CHAR_LIMIT=40] -->
+    <string name="action_listen">Listen</string>
+    <!-- Label for a button that translates a screenshot of the current app. [CHAR_LIMIT=40] -->
+    <string name="action_translate">Translate</string>
+    <!-- Label for a button that triggers Search on a screenshot of the current app. [CHAR_LIMIT=40] -->
+    <string name="action_search">Search</string>
 </resources>
diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
new file mode 100644
index 0000000..b102a39
--- /dev/null
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -0,0 +1,147 @@
+/*
+ * 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.quickstep;
+
+import static com.android.quickstep.views.OverviewActionsView.DISABLED_NO_THUMBNAIL;
+import static com.android.quickstep.views.OverviewActionsView.DISABLED_ROTATED;
+
+import android.annotation.SuppressLint;
+import android.content.Context;
+import android.content.Intent;
+import android.graphics.Matrix;
+import android.os.SystemClock;
+import android.text.TextUtils;
+
+import com.android.launcher3.R;
+import com.android.quickstep.views.OverviewActionsView;
+import com.android.quickstep.views.TaskThumbnailView;
+import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.model.ThumbnailData;
+
+/**
+ * Go-specific extension of the factory class that adds an overlay to TaskView
+ */
+public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
+    public static final String ACTION_LISTEN = "com.android.quickstep.ACTION_LISTEN";
+    public static final String ACTION_TRANSLATE = "com.android.quickstep.ACTION_TRANSLATE";
+    public static final String ACTION_SEARCH = "com.android.quickstep.ACTION_SEARCH";
+    public static final String ELAPSED_NANOS = "niu_actions_elapsed_realtime_nanos";
+
+    // Empty constructor required for ResourceBasedOverride
+    public TaskOverlayFactoryGo(Context context) {}
+
+    /**
+     * Create a new overlay instance for the given View
+     */
+    public TaskOverlayGo createOverlay(TaskThumbnailView thumbnailView) {
+        return new TaskOverlayGo(thumbnailView);
+    }
+
+    /**
+     * Overlay on each task handling Overview Action Buttons.
+     * @param <T> The type of View in which the overlay will be placed
+     */
+    public static final class TaskOverlayGo<T extends OverviewActionsView> extends TaskOverlay {
+
+        private String mPackageName;
+
+        private TaskOverlayGo(TaskThumbnailView taskThumbnailView) {
+            super(taskThumbnailView);
+        }
+
+        /**
+         * Called when the current task is interactive for the user
+         */
+        @Override
+        public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
+                boolean rotated) {
+            getActionsView().updateDisabledFlags(DISABLED_NO_THUMBNAIL, thumbnail == null);
+            mPackageName =
+                    mApplicationContext.getResources().getString(R.string.niu_actions_package);
+
+            if (thumbnail == null || TextUtils.isEmpty(mPackageName)) {
+                return;
+            }
+
+            getActionsView().updateDisabledFlags(DISABLED_ROTATED, rotated);
+            boolean isAllowedByPolicy = thumbnail.isRealSnapshot;
+            getActionsView().setCallbacks(new OverlayUICallbacksGoImpl(isAllowedByPolicy, task));
+        }
+
+        private void sendNIUIntent(String actionType) {
+            Intent intent = createNIUIntent(actionType);
+            mImageApi.shareAsDataWithExplicitIntent(/* crop */ null, intent);
+        }
+
+        private Intent createNIUIntent(String actionType) {
+            return new Intent(actionType)
+                    .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK)
+                    .addFlags(Intent.FLAG_ACTIVITY_NO_ANIMATION)
+                    .setPackage(mPackageName)
+                    .putExtra(ELAPSED_NANOS, SystemClock.elapsedRealtimeNanos());
+        }
+
+        protected class OverlayUICallbacksGoImpl extends OverlayUICallbacksImpl
+                implements OverlayUICallbacksGo {
+            public OverlayUICallbacksGoImpl(boolean isAllowedByPolicy, Task task) {
+                super(isAllowedByPolicy, task);
+            }
+
+            @SuppressLint("NewApi")
+            public void onListen() {
+                if (mIsAllowedByPolicy) {
+                    sendNIUIntent(ACTION_LISTEN);
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+
+            @SuppressLint("NewApi")
+            public void onTranslate() {
+                if (mIsAllowedByPolicy) {
+                    sendNIUIntent(ACTION_TRANSLATE);
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+
+            @SuppressLint("NewApi")
+            public void onSearch() {
+                if (mIsAllowedByPolicy) {
+                    sendNIUIntent(ACTION_SEARCH);
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+        }
+    }
+
+    /**
+     * Callbacks the Ui can generate. This is the only way for a Ui to call methods on the
+     * controller.
+     */
+    public interface OverlayUICallbacksGo extends OverlayUICallbacks {
+        /** User has requested to listen to the current content read aloud */
+        void onListen();
+
+        /** User has requested a translation of the current content */
+        void onTranslate();
+
+        /** User has requested a visual search of the current content */
+        void onSearch();
+    }
+}
diff --git a/go/quickstep/src/com/android/quickstep/views/GoOverviewActionsView.java b/go/quickstep/src/com/android/quickstep/views/GoOverviewActionsView.java
new file mode 100644
index 0000000..9997d16
--- /dev/null
+++ b/go/quickstep/src/com/android/quickstep/views/GoOverviewActionsView.java
@@ -0,0 +1,75 @@
+/*
+ * 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.quickstep.views;
+
+import android.content.Context;
+import android.util.AttributeSet;
+import android.view.View;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.R;
+import com.android.quickstep.TaskOverlayFactoryGo.OverlayUICallbacksGo;
+
+/**
+ * View for showing Go-specific action buttons in Overview
+ */
+public final class GoOverviewActionsView extends OverviewActionsView<OverlayUICallbacksGo> {
+    public GoOverviewActionsView(Context context) {
+        this(context, null);
+    }
+
+    public GoOverviewActionsView(Context context, @Nullable AttributeSet attrs) {
+        this(context, attrs, 0);
+    }
+
+    public GoOverviewActionsView(Context context, @Nullable AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+
+        if (getResources().getBoolean(R.bool.enable_niu_actions)) {
+            findViewById(R.id.action_listen).setOnClickListener(this);
+            findViewById(R.id.action_translate).setOnClickListener(this);
+            findViewById(R.id.action_search).setOnClickListener(this);
+        } else {
+            findViewById(R.id.action_listen).setVisibility(View.GONE);
+            findViewById(R.id.action_translate).setVisibility(View.GONE);
+            findViewById(R.id.action_search).setVisibility(View.GONE);
+        }
+    }
+
+    @Override
+    public void onClick(View view) {
+        super.onClick(view);
+
+        if (mCallbacks == null) {
+            return;
+        }
+        int id = view.getId();
+        if (id == R.id.action_listen) {
+            mCallbacks.onListen();
+        } else if (id == R.id.action_translate) {
+            mCallbacks.onTranslate();
+        } else if (id == R.id.action_search) {
+            mCallbacks.onSearch();
+        }
+    }
+}
diff --git a/go/src/com/android/launcher3/model/WidgetsModel.java b/go/src/com/android/launcher3/model/WidgetsModel.java
index 89b3831..a202095 100644
--- a/go/src/com/android/launcher3/model/WidgetsModel.java
+++ b/go/src/com/android/launcher3/model/WidgetsModel.java
@@ -25,7 +25,7 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.icons.ComponentWithLabelAndIcon;
 import com.android.launcher3.util.PackageUserKey;
-import com.android.launcher3.widget.WidgetListRowEntry;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -43,17 +43,17 @@
     public static final boolean GO_DISABLE_WIDGETS = true;
     public static final boolean GO_DISABLE_NOTIFICATION_DOTS = true;
 
-    private static final ArrayList<WidgetListRowEntry> EMPTY_WIDGET_LIST = new ArrayList<>();
+    private static final ArrayList<WidgetsListBaseEntry> EMPTY_WIDGET_LIST = new ArrayList<>();
 
     /**
-     * Returns a list of {@link WidgetListRowEntry}. All {@link WidgetItem} in a single row
-     * are sorted (based on label and user), but the overall list of {@link WidgetListRowEntry}s
-     * is not sorted. This list is sorted at the UI when using
-     * {@link com.android.launcher3.widget.WidgetsDiffReporter}
+     * Returns a list of {@link WidgetsListBaseEntry}. All {@link WidgetItem} in a single row are
+     * sorted (based on label and user), but the overall list of {@link WidgetsListBaseEntry}s is
+     * not sorted. This list is sorted at the UI when using
+     * {@link com.android.launcher3.widget.picker.WidgetsDiffReporter}
      *
-     * @see com.android.launcher3.widget.WidgetsListAdapter#setWidgets(ArrayList)
+     * @see com.android.launcher3.widget.picker.WidgetsListAdapter#setWidgets(List)
      */
-    public synchronized ArrayList<WidgetListRowEntry> getWidgetsList(Context context) {
+    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsList(Context context) {
         return EMPTY_WIDGET_LIST;
     }
 
diff --git a/quickstep/AndroidManifest.xml b/quickstep/AndroidManifest.xml
index 7431551..97f4a21 100644
--- a/quickstep/AndroidManifest.xml
+++ b/quickstep/AndroidManifest.xml
@@ -39,6 +39,7 @@
     <uses-permission android:name="android.permission.SET_ORIENTATION"/>
     <uses-permission android:name="android.permission.READ_FRAME_BUFFER"/>
     <uses-permission android:name="android.permission.MANAGE_ACCESSIBILITY"/>
+    <uses-permission android:name="android.permission.MONITOR_INPUT"/>
 
     <uses-permission android:name="${packageName}.permission.HOTSEAT_EDU" />
     <uses-permission android:name="android.permission.SYSTEM_APPLICATION_OVERLAY" />
diff --git a/quickstep/res/layout/overview_actions_container.xml b/quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
similarity index 100%
rename from quickstep/res/layout/overview_actions_container.xml
rename to quickstep/overview_ui_overrides/res/layout/overview_actions_container.xml
diff --git a/quickstep/res/layout/search_result_thumbnail.xml b/quickstep/overview_ui_overrides/res/values/config.xml
similarity index 64%
rename from quickstep/res/layout/search_result_thumbnail.xml
rename to quickstep/overview_ui_overrides/res/values/config.xml
index 5062b76..0f09439 100644
--- a/quickstep/res/layout/search_result_thumbnail.xml
+++ b/quickstep/overview_ui_overrides/res/values/config.xml
@@ -1,11 +1,11 @@
 <?xml version="1.0" encoding="utf-8"?>
-<!-- Copyright (C) 2020 The Android Open Source Project
+<!-- 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
+        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,
@@ -13,7 +13,6 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<com.android.launcher3.search.SearchResultThumbnailView
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="125dp"
-    android:layout_height="125dp"/>
\ No newline at end of file
+<resources>
+    <string name="task_overlay_factory_class" translatable="false"/>
+</resources>
\ No newline at end of file
diff --git a/quickstep/res/layout/fallback_search_input.xml b/quickstep/res/layout/fallback_search_input.xml
deleted file mode 100644
index 0fb2924..0000000
--- a/quickstep/res/layout/fallback_search_input.xml
+++ /dev/null
@@ -1,35 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-     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.
--->
-<com.android.launcher3.search.FallbackSearchInputView
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="match_parent"
-    android:layout_gravity="center_vertical"
-    android:layout_marginLeft="48dp"
-    android:layout_marginRight="48dp"
-    android:background="@android:color/transparent"
-    android:focusableInTouchMode="true"
-    android:gravity="start|center_vertical"
-    android:inputType="textNoSuggestions"
-    android:imeOptions="actionSearch|flagNoExtractUi"
-    android:maxLines="1"
-    android:privateImeOptions="bc_search"
-    android:scrollHorizontally="true"
-    android:singleLine="true"
-    android:textColor="?android:attr/textColorSecondary"
-    android:textColorHint="?android:attr/textColorTertiary"
-    android:textSize="16sp" />
\ No newline at end of file
diff --git a/quickstep/res/layout/search_edu_view.xml b/quickstep/res/layout/search_edu_view.xml
deleted file mode 100644
index d89f5c7..0000000
--- a/quickstep/res/layout/search_edu_view.xml
+++ /dev/null
@@ -1,73 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!-- Copyright (C) 2008 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.
--->
-<com.android.launcher3.search.DeviceSearchEdu
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="match_parent"
-    android:layout_gravity="center_horizontal"
-    android:orientation="vertical">
-
-
-    <FrameLayout
-        android:layout_height="wrap_content"
-        android:id="@+id/search_box_wrapper"
-        android:layout_width="match_parent">
-
-        <include
-            layout="@layout/fallback_search_input"
-            android:id="@+id/mock_search_box" />
-    </FrameLayout>
-
-    <LinearLayout
-        android:layout_height="wrap_content"
-        android:id="@+id/edu_wrapper"
-        android:padding="24dp"
-        android:layout_marginTop="40dp"
-        android:orientation="vertical"
-        android:layout_width="match_parent">
-
-        <TextView
-            style="@style/TextHeadline"
-            android:layout_width="match_parent"
-            android:gravity="center"
-            android:textSize="24sp"
-            android:textColor="?android:attr/textColorPrimary"
-            android:layout_height="wrap_content"
-            android:text="@string/search_edu_primary" />
-
-        <TextView
-            style="@style/TextHeadline"
-            android:layout_width="match_parent"
-            android:gravity="center"
-            android:textSize="18sp"
-            android:layout_marginTop="30dp"
-            android:textColor="?android:attr/textColorPrimary"
-            android:layout_height="wrap_content"
-            android:text="@string/search_edu_secondary" />
-
-        <Button
-            android:id="@+id/dismiss_edu"
-            android:layout_width="wrap_content"
-            android:layout_marginTop="@dimen/dynamic_grid_edge_margin"
-            android:background="?android:attr/selectableItemBackground"
-            android:layout_height="wrap_content"
-            android:textColor="?android:attr/textColorPrimary"
-            android:gravity="center"
-            android:layout_gravity="center"
-            android:text="@string/search_edu_dismiss" />
-
-    </LinearLayout>
-
-</com.android.launcher3.search.DeviceSearchEdu>
\ No newline at end of file
diff --git a/quickstep/res/layout/search_result_icon.xml b/quickstep/res/layout/search_result_icon.xml
deleted file mode 100644
index e1b6dfd..0000000
--- a/quickstep/res/layout/search_result_icon.xml
+++ /dev/null
@@ -1,20 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!-- Copyright (C) 2008 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.
--->
-
-<com.android.launcher3.search.SearchResultIcon xmlns:launcher="http://schemas.android.com/apk/res-auto"
-    style="@style/BaseIcon.AllApps"
-    launcher:iconDisplay="all_apps"
-    launcher:centerVertically="true" />
-
diff --git a/quickstep/res/layout/search_result_icon_row.xml b/quickstep/res/layout/search_result_icon_row.xml
deleted file mode 100644
index 084920a..0000000
--- a/quickstep/res/layout/search_result_icon_row.xml
+++ /dev/null
@@ -1,90 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!-- 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.
--->
-<com.android.launcher3.search.SearchResultIconRow xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:launcher="http://schemas.android.com/apk/res-auto"
-    android:orientation="horizontal"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:padding="@dimen/dynamic_grid_edge_margin">
-
-    <com.android.launcher3.search.SearchResultIcon
-        android:id="@+id/icon"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        launcher:iconDisplay="hero_app" />
-
-    <LinearLayout
-        android:layout_width="0dp"
-        android:layout_weight="1"
-        android:layout_height="wrap_content"
-        android:padding="@dimen/dynamic_grid_edge_margin"
-        android:orientation="vertical"
-        android:layout_gravity="center_vertical">
-
-        <TextView
-            android:id="@id/title"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:gravity="start|center_vertical"
-            android:maxLines="1"
-            android:textAlignment="viewStart"
-            android:textColor="?android:attr/textColorPrimary"
-            android:textSize="@dimen/search_hero_title_size" />
-
-        <TextView
-            android:layout_width="wrap_content"
-            android:id="@+id/subtitle"
-            android:maxLines="1"
-            android:textColor="?android:attr/textColorTertiary"
-            android:textSize="@dimen/search_hero_subtitle_size"
-            android:layout_height="wrap_content" />
-    </LinearLayout>
-
-    <com.android.launcher3.search.SearchResultIcon
-        android:id="@+id/shortcut_0"
-        style="@style/BaseIcon"
-        android:layout_width="@dimen/deep_shortcut_icon_size"
-        android:layout_height="match_parent"
-        android:layout_marginStart="4dp"
-        android:gravity="center"
-        launcher:iconDisplay="shortcut_popup"
-        android:textSize="@dimen/search_hero_inline_button_size"
-        launcher:iconSizeOverride="@dimen/deep_shortcut_icon_size"
-        launcher:layoutHorizontal="false" />
-
-    <com.android.launcher3.search.SearchResultIcon
-        android:id="@+id/shortcut_1"
-        style="@style/BaseIcon"
-        android:layout_width="@dimen/deep_shortcut_icon_size"
-        android:layout_height="match_parent"
-        android:layout_marginStart="4dp"
-        android:gravity="center"
-        launcher:iconDisplay="shortcut_popup"
-        android:textSize="@dimen/search_hero_inline_button_size"
-        launcher:iconSizeOverride="@dimen/deep_shortcut_icon_size"
-        launcher:layoutHorizontal="false" />
-
-    <com.android.launcher3.search.SearchResultIcon
-        android:id="@+id/shortcut_2"
-        style="@style/BaseIcon"
-        android:layout_width="@dimen/deep_shortcut_icon_size"
-        android:layout_height="match_parent"
-        android:layout_marginStart="4dp"
-        android:gravity="center"
-        launcher:iconDisplay="shortcut_popup"
-        android:textSize="@dimen/search_hero_inline_button_size"
-        launcher:iconSizeOverride="@dimen/deep_shortcut_icon_size"
-        launcher:layoutHorizontal="false" />
-</com.android.launcher3.search.SearchResultIconRow>
\ No newline at end of file
diff --git a/quickstep/res/layout/search_result_slice.xml b/quickstep/res/layout/search_result_slice.xml
deleted file mode 100644
index 4f884ff..0000000
--- a/quickstep/res/layout/search_result_slice.xml
+++ /dev/null
@@ -1,40 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!-- 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.
--->
-<com.android.launcher3.search.SearchResultIconSlice xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:launcher="http://schemas.android.com/apk/res-auto"
-    android:padding="@dimen/dynamic_grid_edge_margin"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content">
-
-    <FrameLayout
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content">
-
-        <com.android.launcher3.search.SearchResultIcon
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:id="@+id/icon"
-            launcher:iconDisplay="hero_app" />
-    </FrameLayout>
-
-    <androidx.slice.widget.SliceView
-        android:id="@+id/slice"
-        android:layout_height="wrap_content"
-        android:layout_weight="1"
-        android:layout_marginStart="@dimen/dynamic_grid_cell_padding_x"
-        android:layout_width="0dp" />
-
-</com.android.launcher3.search.SearchResultIconSlice>
-
diff --git a/quickstep/res/layout/search_result_small_icon_row.xml b/quickstep/res/layout/search_result_small_icon_row.xml
deleted file mode 100644
index 41856fe..0000000
--- a/quickstep/res/layout/search_result_small_icon_row.xml
+++ /dev/null
@@ -1,73 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!-- Copyright (C) 2021 The Android Open Source Project
-
-     Licensed under the Apache License, Version 2.0 (the "License");
-     you may not use this file except in compliance with the License.
-     You may obtain a copy of the License at
-
-          http://www.apache.org/licenses/LICENSE-2.0
-
-     Unless required by applicable law or agreed to in writing, software
-     distributed under the License is distributed on an "AS IS" BASIS,
-     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-     See the License for the specific language governing permissions and
-     limitations under the License.
--->
-<com.android.launcher3.search.SearchResultSmallIconRow
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:launcher="http://schemas.android.com/apk/res-auto"
-    android:orientation="horizontal"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:padding="@dimen/dynamic_grid_edge_margin">
-
-    <com.android.launcher3.search.SearchResultIcon
-        android:id="@+id/icon"
-        style="@style/BaseIcon"
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:drawablePadding="@dimen/dynamic_grid_icon_drawable_padding"
-        android:drawableTint="?android:attr/textColorPrimary"
-        android:padding="@dimen/dynamic_grid_edge_margin"
-        launcher:iconDisplay="hero_app"
-        launcher:iconSizeOverride="48dp"
-        launcher:matchTextInsetWithQuery="true"
-        launcher:layoutHorizontal="true" />
-
-    <LinearLayout
-        android:layout_width="0dp"
-        android:layout_weight="1"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center_vertical"
-        android:orientation="horizontal"
-        android:padding="@dimen/dynamic_grid_edge_margin" >
-
-        <TextView
-            android:id="@id/title"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:gravity="start|center_vertical"
-            android:maxLines="1"
-            android:paddingEnd="4dp"
-            android:textAlignment="viewStart"
-            android:textColor="?android:attr/textColorPrimary"
-            android:textSize="@dimen/search_hero_title_size" />
-        <TextView
-            android:id="@+id/delimeter"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:gravity="start|center_vertical"
-            android:maxLines="1"
-            android:text="\u2022"
-            android:textColor="?android:attr/textColorPrimary"
-            android:textSize="@dimen/search_hero_subtitle_size" />
-        <TextView
-            android:id="@+id/subtitle"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:maxLines="1"
-            android:paddingStart="4dp"
-            android:textColor="?android:attr/textColorTertiary"
-            android:textSize="@dimen/search_hero_subtitle_size" />
-    </LinearLayout>
-</com.android.launcher3.search.SearchResultSmallIconRow>
\ No newline at end of file
diff --git a/quickstep/res/layout/search_result_widget_live.xml b/quickstep/res/layout/search_result_widget_live.xml
deleted file mode 100644
index f2ac6cd..0000000
--- a/quickstep/res/layout/search_result_widget_live.xml
+++ /dev/null
@@ -1,43 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<com.android.launcher3.search.SearchResultWidget
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:launcher="http://schemas.android.com/apk/res-auto"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:orientation="vertical"
-    android:paddingVertical="@dimen/widget_section_vertical_padding"
-    android:layout_marginBottom="@dimen/widget_section_vertical_padding"
-    android:gravity="center"
-    >
-
-    <LinearLayout
-        android:layout_width="match_parent"
-        android:orientation="horizontal"
-        android:layout_height="wrap_content">
-
-        <com.android.launcher3.BubbleTextView
-            android:id="@+id/widget_provider"
-            android:layout_width="wrap_content"
-            android:layout_height="@dimen/widget_section_height"
-            android:drawablePadding="@dimen/widget_section_horizontal_padding"
-            android:focusable="true"
-            android:gravity="start|center_vertical"
-            android:paddingHorizontal="@dimen/widget_section_horizontal_padding"
-            android:paddingVertical="@dimen/widget_section_horizontal_padding"
-            android:singleLine="true"
-            android:textColor="?android:attr/textColorPrimary"
-            android:textSize="16sp"
-            android:textAlignment="viewStart"
-            launcher:iconDisplay="widget_section"
-            launcher:layoutHorizontal="true"
-            launcher:iconSizeOverride="@dimen/widget_section_icon_size" />
-
-        <TextView
-            android:id="@+id/widget_label"
-            android:layout_width="wrap_content"
-            android:layout_height="@dimen/widget_section_height"
-            android:textSize="16sp" />
-
-    </LinearLayout>
-
-</com.android.launcher3.search.SearchResultWidget>
\ No newline at end of file
diff --git a/quickstep/res/layout/search_result_widget_preview.xml b/quickstep/res/layout/search_result_widget_preview.xml
deleted file mode 100644
index 7af24a1..0000000
--- a/quickstep/res/layout/search_result_widget_preview.xml
+++ /dev/null
@@ -1,8 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<com.android.launcher3.search.SearchResultWidgetPreview xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:padding="@dimen/dynamic_grid_cell_padding_x"
-    android:layout_height="wrap_content">
-    <include layout="@layout/widget_cell" android:id="@+id/widget_cell"/>
-<!--    <include layout="@layout/widget_cell_content" />-->
-</com.android.launcher3.search.SearchResultWidgetPreview>
\ No newline at end of file
diff --git a/quickstep/res/layout/search_section_title.xml b/quickstep/res/layout/search_section_title.xml
deleted file mode 100644
index 5842e57..0000000
--- a/quickstep/res/layout/search_section_title.xml
+++ /dev/null
@@ -1,25 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!-- 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.
--->
-<com.android.launcher3.search.SearchSectionHeaderView xmlns:android="http://schemas.android.com/apk/res/android"
-    android:id="@+id/section_title"
-    android:layout_width="wrap_content"
-    android:layout_height="wrap_content"
-    style="@style/TextHeadline"
-    android:paddingStart="4dp"
-    android:paddingBottom="2dp"
-    android:paddingTop="12dp"
-    android:textColor="?android:attr/textColorPrimary"
-    android:textSize="18sp" />
\ No newline at end of file
diff --git a/quickstep/res/layout/task.xml b/quickstep/res/layout/task.xml
index c90706b..0f9a6aa 100644
--- a/quickstep/res/layout/task.xml
+++ b/quickstep/res/layout/task.xml
@@ -19,7 +19,6 @@
     android:layout_height="match_parent"
     android:clipChildren="false"
     android:defaultFocusHighlightEnabled="false"
-    android:elevation="4dp"
     android:focusable="true">
 
     <com.android.quickstep.views.TaskThumbnailView
diff --git a/quickstep/res/layout/task_menu.xml b/quickstep/res/layout/task_menu.xml
index 098b34f..3916ff9 100644
--- a/quickstep/res/layout/task_menu.xml
+++ b/quickstep/res/layout/task_menu.xml
@@ -24,20 +24,12 @@
     android:orientation="vertical"
     android:visibility="invisible">
 
-    <com.android.quickstep.views.IconView
-      android:id="@+id/task_icon"
-      android:layout_width="@dimen/task_thumbnail_icon_size"
-      android:layout_height="@dimen/task_thumbnail_icon_size"
-      android:layout_gravity="top|center_horizontal"
-      android:layout_marginBottom="@dimen/deep_shortcut_drawable_padding"
-      android:focusable="false"
-      android:importantForAccessibility="no" />
-
     <TextView
         android:id="@+id/task_name"
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
         android:gravity="center_horizontal"
+        android:layout_marginTop="16dp"
         android:layout_marginBottom="16dp"
         android:textSize="12sp"/>
 
diff --git a/quickstep/res/values/config.xml b/quickstep/res/values/config.xml
index 9ec303a..be66104 100644
--- a/quickstep/res/values/config.xml
+++ b/quickstep/res/values/config.xml
@@ -14,8 +14,6 @@
      limitations under the License.
 -->
 <resources>
-    <string name="task_overlay_factory_class" translatable="false"/>
-
     <string name="overscroll_plugin_factory_class" translatable="false" />
 
     <!-- Activities which block home gesture -->
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index 0f40775..68c3851 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -19,6 +19,7 @@
     <dimen name="task_thumbnail_top_margin">24dp</dimen>
     <dimen name="task_thumbnail_half_top_margin">12dp</dimen>
     <dimen name="task_thumbnail_icon_size">48dp</dimen>
+    <dimen name="task_icon_top_margin">-16dp</dimen>
     <!-- For screens without rounded corners -->
     <dimen name="task_corner_radius_small">2dp</dimen>
 
@@ -28,7 +29,7 @@
     <dimen name="overview_actions_bottom_margin_three_button">8dp</dimen>
     <dimen name="overview_actions_horizontal_margin">16dp</dimen>
 
-    <dimen name="recents_page_spacing">10dp</dimen>
+    <dimen name="recents_page_spacing">16dp</dimen>
     <dimen name="recents_clear_all_deadzone_vertical_margin">70dp</dimen>
 
     <!-- The speed in dp/s at which the user needs to be scrolling in recents such that we start
@@ -57,7 +58,7 @@
     <dimen name="task_card_menu_option_vertical_padding">8dp</dimen>
     <dimen name="task_card_menu_shadow_height">3dp</dimen>
     <dimen name="task_card_menu_horizontal_padding">0dp</dimen>
-    <dimen name="portrait_task_card_horz_space_big_overview">96dp</dimen>
+    <dimen name="portrait_task_card_horz_space_big_overview">132dp</dimen>
     <dimen name="portrait_modal_task_card_horz_space">60dp</dimen>
     <dimen name="landscape_task_card_horz_space">200dp</dimen>
     <dimen name="multi_window_task_card_horz_space">100dp</dimen>
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java b/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
index 5cb55ec..eca27b5 100644
--- a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
+++ b/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
@@ -43,8 +43,18 @@
 
 @RunWith(RobolectricTestRunner.class)
 public class OrientationTouchTransformerTest {
-    private static final int SIZE_WIDTH = 1080;
-    private static final int SIZE_HEIGHT = 2280;
+    static class ScreenSize {
+        int mHeight;
+        int mWidth;
+
+        ScreenSize(int height, int width) {
+            mHeight = height;
+            mWidth = width;
+        }
+    }
+
+    private static final ScreenSize NORMAL_SCREEN_SIZE = new ScreenSize(2280, 1080);
+    private static final ScreenSize LARGE_SCREEN_SIZE = new ScreenSize(3280, 1080);
     private static final float DENSITY_DISPLAY_METRICS = 3.0f;
 
     private OrientationTouchTransformer mTouchTransformer;
@@ -63,14 +73,16 @@
         DisplayMetrics mockDisplayMetrics = new DisplayMetrics();
         mockDisplayMetrics.density = DENSITY_DISPLAY_METRICS;
         when(mResources.getDisplayMetrics()).thenReturn(mockDisplayMetrics);
-        mInfo = createDisplayInfo(Surface.ROTATION_0);
+        mInfo = createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_0);
         mTouchTransformer = new OrientationTouchTransformer(mResources, NO_BUTTON, () -> 0);
     }
 
     @Test
     public void disabledMultipleRegions_shouldOverrideFirstRegion() {
-        float portraitRegionY = generateTouchRegionHeight(Surface.ROTATION_0) + 1;
-        float landscapeRegionY = generateTouchRegionHeight(Surface.ROTATION_90) + 1;
+        float portraitRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_0) + 1;
+        float landscapeRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_90) + 1;
 
         mTouchTransformer.createOrAddTouchRegion(mInfo);
         tapAndAssertTrue(100, portraitRegionY,
@@ -83,7 +95,8 @@
                 event -> mTouchTransformer.touchInAssistantRegion(event));
 
         // Override region
-        mTouchTransformer.createOrAddTouchRegion(createDisplayInfo(Surface.ROTATION_90));
+        mTouchTransformer
+            .createOrAddTouchRegion(createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_90));
         tapAndAssertFalse(100, portraitRegionY,
                 event -> mTouchTransformer.touchInValidSwipeRegions(event.getX(), event.getY()));
         tapAndAssertTrue(100, landscapeRegionY,
@@ -107,10 +120,13 @@
 
     @Test
     public void enableMultipleRegions_shouldOverrideFirstRegion() {
-        float portraitRegionY = generateTouchRegionHeight(Surface.ROTATION_0) + 1;
-        float landscapeRegionY = generateTouchRegionHeight(Surface.ROTATION_90) + 1;
+        float portraitRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_0) + 1;
+        float landscapeRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_90) + 1;
 
-        mTouchTransformer.createOrAddTouchRegion(createDisplayInfo(Surface.ROTATION_90));
+        mTouchTransformer
+            .createOrAddTouchRegion(createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_90));
         tapAndAssertFalse(100, portraitRegionY,
                 event -> mTouchTransformer.touchInValidSwipeRegions(event.getX(), event.getY()));
         tapAndAssertTrue(100, landscapeRegionY,
@@ -136,11 +152,14 @@
 
     @Test
     public void enableMultipleRegions_assistantTriggersInMostRecent() {
-        float portraitRegionY = generateTouchRegionHeight(Surface.ROTATION_0) + 1;
-        float landscapeRegionY = generateTouchRegionHeight(Surface.ROTATION_90) + 1;
+        float portraitRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_0) + 1;
+        float landscapeRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_90) + 1;
 
         mTouchTransformer.enableMultipleRegions(true, mInfo);
-        mTouchTransformer.createOrAddTouchRegion(createDisplayInfo(Surface.ROTATION_90));
+        mTouchTransformer
+            .createOrAddTouchRegion(createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_90));
         mTouchTransformer.createOrAddTouchRegion(mInfo);
         tapAndAssertTrue(0, portraitRegionY,
                 event -> mTouchTransformer.touchInAssistantRegion(event));
@@ -150,12 +169,15 @@
 
     @Test
     public void enableMultipleRegions_assistantTriggersInCurrentOrientationAfterDisable() {
-        float portraitRegionY = generateTouchRegionHeight(Surface.ROTATION_0) + 1;
-        float landscapeRegionY = generateTouchRegionHeight(Surface.ROTATION_90) + 1;
+        float portraitRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_0) + 1;
+        float landscapeRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_90) + 1;
 
         mTouchTransformer.enableMultipleRegions(true, mInfo);
         mTouchTransformer.createOrAddTouchRegion(mInfo);
-        mTouchTransformer.createOrAddTouchRegion(createDisplayInfo(Surface.ROTATION_90));
+        mTouchTransformer
+            .createOrAddTouchRegion(createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_90));
         mTouchTransformer.enableMultipleRegions(false, mInfo);
         tapAndAssertTrue(0, portraitRegionY,
                 event -> mTouchTransformer.touchInAssistantRegion(event));
@@ -164,6 +186,26 @@
     }
 
     @Test
+    public void assistantTriggersInCurrentScreenAfterScreenSizeChange() {
+        float smallerScreenPortraitRegionY =
+                generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_0) + 1;
+        float largerScreenPortraitRegionY =
+                generateTouchRegionHeight(LARGE_SCREEN_SIZE, Surface.ROTATION_0) + 1;
+
+        mTouchTransformer.enableMultipleRegions(false,
+                createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_0));
+        tapAndAssertTrue(0, smallerScreenPortraitRegionY,
+                event -> mTouchTransformer.touchInAssistantRegion(event));
+
+        mTouchTransformer
+            .enableMultipleRegions(false, createDisplayInfo(LARGE_SCREEN_SIZE, Surface.ROTATION_0));
+        tapAndAssertTrue(0, largerScreenPortraitRegionY,
+                event -> mTouchTransformer.touchInAssistantRegion(event));
+        tapAndAssertFalse(0, smallerScreenPortraitRegionY,
+                event -> mTouchTransformer.touchInAssistantRegion(event));
+    }
+
+    @Test
     public void applyTransform_taskNotFrozen_notInRegion() {
         mTouchTransformer.createOrAddTouchRegion(mInfo);
         tapAndAssertFalse(100, 100,
@@ -182,7 +224,7 @@
     public void applyTransform_taskFrozen_noRotate_inRegion() {
         mTouchTransformer.createOrAddTouchRegion(mInfo);
         mTouchTransformer.enableMultipleRegions(true, mInfo);
-        float y = generateTouchRegionHeight(Surface.ROTATION_0) + 1;
+        float y = generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_0) + 1;
         tapAndAssertTrue(100, y,
                 event -> mTouchTransformer.touchInValidSwipeRegions(event.getX(), event.getY()));
     }
@@ -190,15 +232,16 @@
     @Test
     public void applyTransform_taskNotFrozen_noRotate_inDefaultRegion() {
         mTouchTransformer.createOrAddTouchRegion(mInfo);
-        float y = generateTouchRegionHeight(Surface.ROTATION_0) + 1;
+        float y = generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_0) + 1;
         tapAndAssertTrue(100, y,
                 event -> mTouchTransformer.touchInValidSwipeRegions(event.getX(), event.getY()));
     }
 
     @Test
     public void applyTransform_taskNotFrozen_90Rotate_inRegion() {
-        mTouchTransformer.createOrAddTouchRegion(createDisplayInfo(Surface.ROTATION_90));
-        float y = generateTouchRegionHeight(Surface.ROTATION_90) + 1;
+        mTouchTransformer
+            .createOrAddTouchRegion(createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_90));
+        float y = generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_90) + 1;
         tapAndAssertTrue(100, y,
                 event -> mTouchTransformer.touchInValidSwipeRegions(event.getX(), event.getY()));
     }
@@ -210,9 +253,10 @@
     public void applyTransform_taskNotFrozen_90Rotate_inTwoRegions() {
         mTouchTransformer.createOrAddTouchRegion(mInfo);
         mTouchTransformer.enableMultipleRegions(true, mInfo);
-        mTouchTransformer.createOrAddTouchRegion(createDisplayInfo(Surface.ROTATION_90));
+        mTouchTransformer
+            .createOrAddTouchRegion(createDisplayInfo(NORMAL_SCREEN_SIZE, Surface.ROTATION_90));
         // Landscape point
-        float y1 = generateTouchRegionHeight(Surface.ROTATION_90) + 1;
+        float y1 = generateTouchRegionHeight(NORMAL_SCREEN_SIZE, Surface.ROTATION_90) + 1;
         MotionEvent inRegion1_down = generateMotionEvent(MotionEvent.ACTION_DOWN, 10, y1);
         MotionEvent inRegion1_up = generateMotionEvent(MotionEvent.ACTION_UP, 10, y1);
         // Portrait point in landscape orientation axis
@@ -231,18 +275,18 @@
         assertTrue(mTouchTransformer.touchInValidSwipeRegions(inRegion2.getX(), inRegion2.getY()));
     }
 
-    private DisplayController.Info createDisplayInfo(int rotation) {
-        Point p = new Point(SIZE_WIDTH, SIZE_HEIGHT);
+    private DisplayController.Info createDisplayInfo(ScreenSize screenSize, int rotation) {
+        Point p = new Point(screenSize.mWidth, screenSize.mHeight);
         if (rotation == Surface.ROTATION_90 || rotation == Surface.ROTATION_270) {
-            p = new Point(SIZE_HEIGHT, SIZE_WIDTH);
+            p = new Point(screenSize.mHeight, screenSize.mWidth);
         }
         return new DisplayController.Info(0, rotation, 0, p, p, p, null);
     }
 
-    private float generateTouchRegionHeight(int rotation) {
-        float height = SIZE_HEIGHT;
+    private float generateTouchRegionHeight(ScreenSize screenSize, int rotation) {
+        float height = screenSize.mHeight;
         if (rotation == Surface.ROTATION_90 || rotation == Surface.ROTATION_270) {
-            height = SIZE_WIDTH;
+            height = screenSize.mWidth;
         }
         return height - ResourceUtils.DEFAULT_NAVBAR_VALUE * DENSITY_DISPLAY_METRICS;
     }
diff --git a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java
index 5e81fba..0156e8f 100644
--- a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java
+++ b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java
@@ -74,11 +74,13 @@
     private static final int FLAG_UPDATE_PAUSED = 1 << 0;
     private static final int FLAG_DRAG_IN_PROGRESS = 1 << 1;
     private static final int FLAG_FILL_IN_PROGRESS = 1 << 2;
+    private static final int FLAG_REMOVING_PREDICTED_ICON = 1 << 3;
 
     private int mHotSeatItemsCount;
 
     private QuickstepLauncher mLauncher;
     private final Hotseat mHotseat;
+    private final Runnable mUpdateFillIfNotLoading = this::updateFillIfNotLoading;
 
     private List<ItemInfo> mPredictedItems = Collections.emptyList();
 
@@ -132,7 +134,8 @@
     private void onHotseatHierarchyChanged() {
         if (mPauseFlags == 0 && !mLauncher.isWorkspaceLoading()) {
             // Post update after a single frame to avoid layout within layout
-            MAIN_EXECUTOR.getHandler().post(this::updateFillIfNotLoading);
+            MAIN_EXECUTOR.getHandler().removeCallbacks(mUpdateFillIfNotLoading);
+            MAIN_EXECUTOR.getHandler().post(mUpdateFillIfNotLoading);
         }
     }
 
@@ -374,7 +377,7 @@
                 continue;
             }
             if (dragObject.dragSource == this && icon.equals(dragObject.originalView)) {
-                mHotseat.removeView(icon);
+                removeIconWithoutNotify(icon);
                 continue;
             }
             int rank = ((WorkspaceItemInfo) icon.getTag()).rank;
@@ -386,7 +389,7 @@
                 @Override
                 public void onAnimationSuccess(Animator animator) {
                     if (icon.getParent() != null) {
-                        mHotseat.removeView(icon);
+                        removeIconWithoutNotify(icon);
                     }
                 }
             });
@@ -395,6 +398,17 @@
         mIconRemoveAnimators.start();
     }
 
+    /**
+     * Removes icon while suppressing any extra tasks performed on view-hierarchy changes.
+     * This avoids recursive/redundant updates as the control updates the UI anyway after
+     * it's animation.
+     */
+    private void removeIconWithoutNotify(PredictedAppIcon icon) {
+        mPauseFlags |= FLAG_REMOVING_PREDICTED_ICON;
+        mHotseat.removeView(icon);
+        mPauseFlags &= ~FLAG_REMOVING_PREDICTED_ICON;
+    }
+
     @Override
     public void onDragStart(DropTarget.DragObject dragObject, DragOptions options) {
         removePredictedApps(mOutlineDrawings, dragObject);
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index c3f5c00..225823e 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -84,7 +84,7 @@
     private final InvariantDeviceProfile mIDP;
     private final AppEventProducer mAppEventProducer;
 
-    private boolean mActive = false;
+    protected boolean mActive = false;
 
     public QuickstepModelDelegate(Context context) {
         mAppEventProducer = new AppEventProducer(context, this::onAppTargetEvent);
@@ -200,7 +200,6 @@
                         .setPredictedTargetCount(mIDP.numHotseatIcons)
                         .setExtras(convertDataModelToAppTargetBundle(context, mDataModel))
                         .build()));
-
     }
 
     private void registerPredictor(PredictorState state, AppPredictor predictor) {
@@ -236,14 +235,14 @@
     static class PredictorState {
 
         public final FixedContainerItems items;
-        public final PersistedItemArray storage;
+        public final PersistedItemArray<ItemInfo> storage;
         public AppPredictor predictor;
 
         private List<AppTarget> mLastTargets;
 
         PredictorState(int container, String storageName) {
             items = new FixedContainerItems(container);
-            storage = new PersistedItemArray(storageName);
+            storage = new PersistedItemArray<>(storageName);
             mLastTargets = Collections.emptyList();
         }
 
@@ -255,7 +254,7 @@
         }
 
         /**
-         * Sets the new targets and returns true if it was different than before.
+         * Sets the new targets and returns true if it was the same as before.
          */
         boolean setTargets(List<AppTarget> newTargets) {
             List<AppTarget> oldTargets = mLastTargets;
@@ -289,7 +288,7 @@
         return true;
     }
 
-    private static class WorkspaceItemFactory implements PersistedItemArray.ItemFactory {
+    private static class WorkspaceItemFactory implements PersistedItemArray.ItemFactory<ItemInfo> {
 
         private final LauncherAppState mAppState;
         private final UserManagerState mUMS;
diff --git a/quickstep/src/com/android/launcher3/search/DeviceSearchAdapterProvider.java b/quickstep/src/com/android/launcher3/search/DeviceSearchAdapterProvider.java
deleted file mode 100644
index d415c36..0000000
--- a/quickstep/src/com/android/launcher3/search/DeviceSearchAdapterProvider.java
+++ /dev/null
@@ -1,152 +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.search;
-
-import static com.android.launcher3.allapps.AllAppsGridAdapter.VIEW_TYPE_ALL_APPS_DIVIDER;
-import static com.android.launcher3.allapps.AllAppsGridAdapter.VIEW_TYPE_ICON;
-
-import android.app.search.SearchTarget;
-import android.util.Log;
-import android.util.SparseIntArray;
-import android.view.LayoutInflater;
-import android.view.View;
-import android.view.ViewGroup;
-
-import com.android.app.search.LayoutType;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.R;
-import com.android.launcher3.allapps.AllAppsContainerView;
-import com.android.launcher3.allapps.AllAppsGridAdapter;
-import com.android.launcher3.allapps.search.SearchAdapterProvider;
-import com.android.launcher3.config.FeatureFlags;
-
-/**
- * Provides views for on-device search results
- */
-public class DeviceSearchAdapterProvider extends SearchAdapterProvider {
-
-    public static final int VIEW_TYPE_SEARCH_CORPUS_TITLE = 1 << 5;
-    public static final int VIEW_TYPE_SEARCH_SLICE = 1 << 7;
-    public static final int VIEW_TYPE_SEARCH_ICON = (1 << 8) | VIEW_TYPE_ICON;
-    public static final int VIEW_TYPE_SEARCH_ICON_ROW = (1 << 9);
-    public static final int VIEW_TYPE_SEARCH_SMALL_ICON_ROW = (1 << 10);
-    public static final int VIEW_TYPE_SEARCH_THUMBNAIL = 1 << 12;
-    public static final int VIEW_TYPE_SEARCH_WIDGET_LIVE = 1 << 15;
-    public static final int VIEW_TYPE_SEARCH_WIDGET_PREVIEW = 1 << 16;
-
-    private static final String TAG = "SearchServiceAdapter";
-
-    private final AllAppsContainerView mAppsView;
-    private final SparseIntArray mViewTypeToLayoutMap = new SparseIntArray();
-
-    public DeviceSearchAdapterProvider(Launcher launcher, AllAppsContainerView appsView) {
-        super(launcher);
-        mAppsView = appsView;
-
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_CORPUS_TITLE, R.layout.search_section_title);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_ICON, R.layout.search_result_icon);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_ICON_ROW, R.layout.search_result_icon_row);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_SMALL_ICON_ROW, R.layout.search_result_small_icon_row);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_SLICE, R.layout.search_result_slice);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_THUMBNAIL, R.layout.search_result_thumbnail);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_WIDGET_LIVE, R.layout.search_result_widget_live);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_SEARCH_WIDGET_PREVIEW,
-                R.layout.search_result_widget_preview);
-        mViewTypeToLayoutMap.put(VIEW_TYPE_ALL_APPS_DIVIDER, R.layout.all_apps_divider);
-    }
-
-    @Override
-    public void onBindView(AllAppsGridAdapter.ViewHolder holder, int position) {
-        SearchAdapterItem item = (SearchAdapterItem) mAppsView.getApps().getAdapterItems().get(
-                position);
-        SearchTargetHandler
-                payloadResultView =
-                (SearchTargetHandler) holder.itemView;
-        payloadResultView.apply(item.getSearchTarget(), item.getInlineItems());
-    }
-
-    @Override
-    public boolean isSearchView(int viewType) {
-        return mViewTypeToLayoutMap.get(viewType, -1) != -1;
-    }
-
-    @Override
-    public AllAppsGridAdapter.ViewHolder onCreateViewHolder(LayoutInflater inflater,
-            ViewGroup parent, int viewType) {
-        return new AllAppsGridAdapter.ViewHolder(inflater.inflate(
-                mViewTypeToLayoutMap.get(viewType), parent, false));
-    }
-
-    @Override
-    public int getGridSpanSize(int viewType, int appsPerRow) {
-        if (viewType == VIEW_TYPE_SEARCH_THUMBNAIL
-                || viewType == VIEW_TYPE_SEARCH_WIDGET_PREVIEW) {
-            return appsPerRow;
-        }
-        return super.getGridSpanSize(viewType, appsPerRow);
-    }
-
-
-    @Override
-    public boolean onAdapterItemSelected(AllAppsGridAdapter.AdapterItem adapterItem, View view) {
-        if (view instanceof SearchTargetHandler) {
-            return ((SearchTargetHandler) view).quickSelect();
-        }
-        return false;
-    }
-
-    /**
-     * Determines what view type should be used to present search target.
-     * Returns -1 if viewType is not found or if required field is not present
-     * to render the viewType.
-     */
-    public int getViewTypeForSearchTarget(SearchTarget t) {
-        switch (t.getLayoutType()) {
-            case LayoutType.TEXT_HEADER:
-                return VIEW_TYPE_SEARCH_CORPUS_TITLE;
-            case LayoutType.ICON_SINGLE_VERTICAL_TEXT:
-                return VIEW_TYPE_SEARCH_ICON;
-            case LayoutType.ICON_SLICE:
-                if (t.getSliceUri() != null) {
-                    return VIEW_TYPE_SEARCH_SLICE;
-                }
-                Log.w(TAG, "LayoutType.ICON_SLICE target doesn't contain sliceUri.");
-                break;
-            case LayoutType.ICON_DOUBLE_HORIZONTAL_TEXT:
-            case LayoutType.ICON_SINGLE_HORIZONTAL_TEXT:
-            case LayoutType.ICON_DOUBLE_HORIZONTAL_TEXT_BUTTON:
-            case LayoutType.ICON_HORIZONTAL_TEXT:
-                return VIEW_TYPE_SEARCH_ICON_ROW;
-            case LayoutType.SMALL_ICON_HORIZONTAL_TEXT:
-                return VIEW_TYPE_SEARCH_SMALL_ICON_ROW;
-            case LayoutType.THUMBNAIL:
-                if (t.getSearchAction() != null) {
-                    return VIEW_TYPE_SEARCH_THUMBNAIL;
-                }
-                Log.w(TAG, "LayoutType.THUMBNAIL target doesn't contain searchAction.");
-                break;
-            case LayoutType.WIDGET_PREVIEW:
-                return VIEW_TYPE_SEARCH_WIDGET_PREVIEW;
-            case LayoutType.WIDGET_LIVE:
-                return VIEW_TYPE_SEARCH_WIDGET_LIVE;
-            case LayoutType.DIVIDER:
-                return VIEW_TYPE_ALL_APPS_DIVIDER;
-        }
-
-        return -1;
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/DeviceSearchEdu.java b/quickstep/src/com/android/launcher3/search/DeviceSearchEdu.java
deleted file mode 100644
index 016ec1b..0000000
--- a/quickstep/src/com/android/launcher3/search/DeviceSearchEdu.java
+++ /dev/null
@@ -1,220 +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.search;
-
-import static com.android.launcher3.util.OnboardingPrefs.SEARCH_EDU_SEEN;
-
-import android.animation.PropertyValuesHolder;
-import android.content.Context;
-import android.graphics.Rect;
-import android.text.Editable;
-import android.text.TextWatcher;
-import android.util.AttributeSet;
-import android.view.KeyEvent;
-import android.view.LayoutInflater;
-import android.view.View;
-import android.view.inputmethod.InputMethodManager;
-import android.widget.EditText;
-import android.widget.TextView;
-
-import androidx.core.graphics.ColorUtils;
-
-import com.android.launcher3.AbstractFloatingView;
-import com.android.launcher3.Insettable;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherState;
-import com.android.launcher3.R;
-import com.android.launcher3.anim.Interpolators;
-import com.android.launcher3.statemanager.StateManager;
-import com.android.launcher3.util.Themes;
-import com.android.launcher3.views.AbstractSlideInView;
-
-/**
- * Feature education for on-device Search. Shown the first time user opens AllApps Search
- */
-public class DeviceSearchEdu extends AbstractSlideInView implements
-        StateManager.StateListener<LauncherState>, TextWatcher, Insettable,
-        TextView.OnEditorActionListener {
-
-    private static final long ANIMATION_DURATION = 350;
-    private static final int ANIMATION_CONTENT_TRANSLATION = 200;
-
-    private EditText mEduInput;
-
-    private View mInputWrapper;
-    private EditText mSearchInput;
-
-    private boolean mSwitchFocusOnDismiss;
-
-
-    public DeviceSearchEdu(Context context) {
-        this(context, null, 0);
-    }
-
-    public DeviceSearchEdu(Context context, AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public DeviceSearchEdu(Context context, AttributeSet attrs, int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-    }
-
-
-    private void dismiss() {
-        handleClose(true);
-        mLauncher.getOnboardingPrefs().markChecked(SEARCH_EDU_SEEN);
-    }
-
-    @Override
-    protected void handleClose(boolean animate) {
-        handleClose(animate, ANIMATION_DURATION);
-        mLauncher.getAllAppsController().getInsetController().show();
-        mLauncher.getStateManager().removeStateListener(this);
-    }
-
-    @Override
-    protected boolean isOfType(int type) {
-        return false;
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        mSearchInput = mLauncher.getAppsView().getSearchUiManager().getEditText();
-        mInputWrapper = findViewById(R.id.search_box_wrapper);
-        mContent = findViewById(R.id.edu_wrapper);
-
-        mEduInput = findViewById(R.id.mock_search_box);
-        mEduInput.setHint(R.string.all_apps_on_device_search_bar_hint);
-        mEduInput.addTextChangedListener(this);
-        if (mSearchInput != null) {
-            mEduInput.getLayoutParams().height = mSearchInput.getHeight();
-            mEduInput.setOnEditorActionListener(this);
-        } else {
-            mEduInput.setVisibility(INVISIBLE);
-        }
-
-        findViewById(R.id.dismiss_edu).setOnClickListener((view) -> {
-            mSwitchFocusOnDismiss = true;
-            dismiss();
-        });
-    }
-
-    private void showInternal() {
-        mLauncher.getStateManager().addStateListener(this);
-        AbstractFloatingView.closeAllOpenViews(mLauncher);
-        attachToContainer();
-        if (mSearchInput != null) {
-            Rect r = mLauncher.getViewBounds(mSearchInput);
-            mEduInput.requestFocus();
-            InputMethodManager imm = mLauncher.getSystemService(InputMethodManager.class);
-            imm.showSoftInput(mEduInput, InputMethodManager.SHOW_IMPLICIT);
-            ((LayoutParams) mInputWrapper.getLayoutParams()).setMargins(0, r.top, 0, 0);
-        }
-        animateOpen();
-    }
-
-    @Override
-    protected int getScrimColor(Context context) {
-        return ColorUtils.setAlphaComponent(Themes.getAttrColor(context, R.attr.allAppsScrimColor),
-                230);
-    }
-
-    protected void setTranslationShift(float translationShift) {
-        mTranslationShift = translationShift;
-        mContent.setAlpha(getBoxedProgress(1 - mTranslationShift, .25f, 1));
-        mContent.setTranslationY(ANIMATION_CONTENT_TRANSLATION * translationShift);
-        if (mColorScrim != null) {
-            mColorScrim.setAlpha(getBoxedProgress(1 - mTranslationShift, 0, .75f));
-        }
-    }
-
-    /**
-     * Given input [0-1], returns progress within bounds [min,max] allowing for staged animations
-     */
-    private float getBoxedProgress(float input, float min, float max) {
-        if (input < min) return 0;
-        if (input > max) return 1;
-        return (input - min) / (max - min);
-    }
-
-    private void animateOpen() {
-        if (mIsOpen || mOpenCloseAnimator.isRunning()) {
-            return;
-        }
-        mIsOpen = true;
-        mOpenCloseAnimator.setValues(
-                PropertyValuesHolder.ofFloat(TRANSLATION_SHIFT, TRANSLATION_SHIFT_OPENED));
-        mOpenCloseAnimator.setInterpolator(Interpolators.FAST_OUT_SLOW_IN);
-        mOpenCloseAnimator.setDuration(ANIMATION_DURATION);
-        mOpenCloseAnimator.start();
-    }
-
-    /**
-     * Show On-device search education view.
-     */
-    public static void show(Launcher launcher) {
-        LayoutInflater layoutInflater = LayoutInflater.from(launcher);
-        ((DeviceSearchEdu) layoutInflater.inflate(
-                R.layout.search_edu_view, launcher.getDragLayer(),
-                false)).showInternal();
-    }
-
-    @Override
-    public void onStateTransitionStart(LauncherState toState) {
-        dismiss();
-    }
-
-    @Override
-    protected void onCloseComplete() {
-        super.onCloseComplete();
-        if (mSearchInput != null && mSwitchFocusOnDismiss) {
-            mSearchInput.requestFocus();
-            mSearchInput.setSelection(mSearchInput.getText().length());
-        }
-    }
-
-    @Override
-    public void afterTextChanged(Editable editable) {
-        //Does nothing
-    }
-
-    @Override
-    public void beforeTextChanged(CharSequence charSequence, int i, int i1, int i2) {
-        //Does nothing
-    }
-
-    @Override
-    public void onTextChanged(CharSequence charSequence, int i, int i1, int i2) {
-        if (mSearchInput != null) {
-            mSearchInput.setText(charSequence.toString());
-            mSwitchFocusOnDismiss = true;
-            dismiss();
-        }
-    }
-
-    @Override
-    public void setInsets(Rect insets) {
-
-    }
-
-    @Override
-    public boolean onEditorAction(TextView textView, int i, KeyEvent keyEvent) {
-        mSearchInput.onEditorAction(i);
-        dismiss();
-        return true;
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/FallbackSearchInputView.java b/quickstep/src/com/android/launcher3/search/FallbackSearchInputView.java
deleted file mode 100644
index 7b10622..0000000
--- a/quickstep/src/com/android/launcher3/search/FallbackSearchInputView.java
+++ /dev/null
@@ -1,140 +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.search;
-
-import static com.android.launcher3.LauncherState.ALL_APPS;
-
-import android.content.Context;
-import android.graphics.Rect;
-import android.util.AttributeSet;
-
-import com.android.launcher3.ExtendedEditText;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.allapps.AllAppsContainerView;
-import com.android.launcher3.allapps.AllAppsGridAdapter.AdapterItem;
-import com.android.launcher3.allapps.AllAppsStore;
-import com.android.launcher3.allapps.AlphabeticalAppsList;
-import com.android.launcher3.allapps.FloatingHeaderView;
-import com.android.launcher3.allapps.search.AllAppsSearchBarController;
-import com.android.launcher3.allapps.search.SearchAlgorithm;
-
-import java.util.ArrayList;
-
-/**
- * A search view shown in all apps for on device search
- */
-public class FallbackSearchInputView extends ExtendedEditText
-        implements AllAppsSearchBarController.Callbacks, AllAppsStore.OnUpdateListener {
-
-    private final AllAppsSearchBarController mSearchBarController;
-
-    private AlphabeticalAppsList mApps;
-    private Runnable mOnResultsChanged;
-    private AllAppsContainerView mAppsView;
-
-    public FallbackSearchInputView(Context context) {
-        this(context, null);
-    }
-
-    public FallbackSearchInputView(Context context, AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public FallbackSearchInputView(Context context, AttributeSet attrs, int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-        mSearchBarController = new AllAppsSearchBarController();
-    }
-
-    @Override
-    protected void onAttachedToWindow() {
-        super.onAttachedToWindow();
-        Launcher.getLauncher(getContext()).getAppsView().getAppsStore().addUpdateListener(this);
-    }
-
-    @Override
-    protected void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-        Launcher.getLauncher(getContext()).getAppsView().getAppsStore().removeUpdateListener(this);
-    }
-
-    /**
-     * Initializes SearchInput
-     */
-    public void initialize(AllAppsContainerView appsView, SearchAlgorithm algo, Runnable changed) {
-        mOnResultsChanged = changed;
-        mApps = appsView.getApps();
-        mAppsView = appsView;
-        mSearchBarController.initialize(algo, this, Launcher.getLauncher(getContext()), this);
-    }
-
-    @Override
-    public void onSearchResult(String query, ArrayList<AdapterItem> items) {
-        if (mApps != null && getParent() != null) {
-            mApps.setSearchResults(items);
-            notifyResultChanged();
-            collapseAppsViewHeader(true);
-            mAppsView.setLastSearchQuery(query);
-        }
-    }
-
-    @Override
-    public void onAppendSearchResult(String query, ArrayList<AdapterItem> items) {
-        if (mApps != null && getParent() != null) {
-            mApps.appendSearchResults(items);
-            notifyResultChanged();
-        }
-    }
-
-    @Override
-    public void clearSearchResult() {
-        if (getParent() != null && mApps != null) {
-            mApps.setSearchResults(null);
-            notifyResultChanged();
-            collapseAppsViewHeader(false);
-            mAppsView.onClearSearchResult();
-        }
-    }
-
-    @Override
-    public void onAppsUpdated() {
-        mSearchBarController.refreshSearchResult();
-    }
-
-    private void collapseAppsViewHeader(boolean collapse) {
-        FloatingHeaderView header = mAppsView.getFloatingHeaderView();
-        if (header != null) {
-            header.setCollapsed(collapse);
-        }
-    }
-
-    private void notifyResultChanged() {
-        if (mOnResultsChanged != null) {
-            mOnResultsChanged.run();
-        }
-        mAppsView.onSearchResultsChanged();
-    }
-
-    @Override
-    protected void onFocusChanged(boolean focused, int direction, Rect previouslyFocusedRect) {
-        // TODO: Consider animating the state transition here
-        if (focused) {
-            // Getting focus will open the keyboard. Go to the all-apps state, so that the input
-            // box is at the top and there is enough space below to show search results.
-            Launcher.getLauncher(getContext()).getStateManager().goToState(ALL_APPS, false);
-        }
-        super.onFocusChanged(focused, direction, previouslyFocusedRect);
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchAdapterItem.java b/quickstep/src/com/android/launcher3/search/SearchAdapterItem.java
deleted file mode 100644
index 8983c4f..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchAdapterItem.java
+++ /dev/null
@@ -1,66 +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.search;
-
-import static com.android.launcher3.search.DeviceSearchAdapterProvider.VIEW_TYPE_SEARCH_ICON;
-import static com.android.launcher3.search.DeviceSearchAdapterProvider.VIEW_TYPE_SEARCH_ICON_ROW;
-import static com.android.launcher3.search.DeviceSearchAdapterProvider.VIEW_TYPE_SEARCH_SMALL_ICON_ROW;
-import static com.android.launcher3.search.DeviceSearchAdapterProvider.VIEW_TYPE_SEARCH_SLICE;
-import static com.android.launcher3.search.DeviceSearchAdapterProvider.VIEW_TYPE_SEARCH_THUMBNAIL;
-import static com.android.launcher3.search.DeviceSearchAdapterProvider.VIEW_TYPE_SEARCH_WIDGET_LIVE;
-import static com.android.launcher3.search.DeviceSearchAdapterProvider.VIEW_TYPE_SEARCH_WIDGET_PREVIEW;
-
-import android.app.search.SearchTarget;
-
-import com.android.launcher3.allapps.AllAppsGridAdapter;
-
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * Extension of AdapterItem that contains an extra payload specific to item
- */
-public class SearchAdapterItem extends AllAppsGridAdapter.AdapterItem {
-    private SearchTarget mSearchTarget;
-    private List<SearchTarget> mInlineItems = new ArrayList<>();
-
-
-    private static final int AVAILABLE_FOR_ACCESSIBILITY = VIEW_TYPE_SEARCH_SLICE
-            | VIEW_TYPE_SEARCH_THUMBNAIL
-            | VIEW_TYPE_SEARCH_ICON_ROW
-            | VIEW_TYPE_SEARCH_ICON
-            | VIEW_TYPE_SEARCH_SMALL_ICON_ROW
-            | VIEW_TYPE_SEARCH_WIDGET_PREVIEW
-            | VIEW_TYPE_SEARCH_WIDGET_LIVE;
-
-    public SearchAdapterItem(SearchTarget searchTarget, int type) {
-        mSearchTarget = searchTarget;
-        viewType = type;
-    }
-
-    public SearchTarget getSearchTarget() {
-        return mSearchTarget;
-    }
-
-    public List<SearchTarget> getInlineItems() {
-        return mInlineItems;
-    }
-    @Override
-    protected boolean isCountedForAccessibility() {
-        return (AVAILABLE_FOR_ACCESSIBILITY & viewType) == viewType;
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultIcon.java b/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
deleted file mode 100644
index 4c44479..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
+++ /dev/null
@@ -1,325 +0,0 @@
-/*
- * 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.search;
-
-import static com.android.launcher3.LauncherSettings.Favorites.EXTENDED_CONTAINERS;
-import static com.android.launcher3.model.data.SearchActionItemInfo.FLAG_BADGE_WITH_PACKAGE;
-import static com.android.launcher3.model.data.SearchActionItemInfo.FLAG_PRIMARY_ICON_FROM_TITLE;
-import static com.android.launcher3.search.SearchTargetUtil.BUNDLE_EXTRA_PRIMARY_ICON_FROM_TITLE;
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-
-import android.app.search.SearchAction;
-import android.app.search.SearchTarget;
-import android.app.search.SearchTargetEvent;
-import android.content.ComponentName;
-import android.content.Context;
-import android.content.pm.ShortcutInfo;
-import android.graphics.Bitmap;
-import android.graphics.BitmapFactory;
-import android.graphics.Paint;
-import android.graphics.PorterDuff;
-import android.graphics.PorterDuffXfermode;
-import android.graphics.Rect;
-import android.graphics.RectF;
-import android.graphics.drawable.Icon;
-import android.os.Bundle;
-import android.os.UserHandle;
-import android.util.AttributeSet;
-import android.view.View;
-import android.view.ViewGroup;
-
-import com.android.app.search.ResultType;
-import com.android.launcher3.BubbleTextView;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.allapps.AllAppsStore;
-import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.BitmapRenderer;
-import com.android.launcher3.icons.LauncherIcons;
-import com.android.launcher3.logger.LauncherAtomExtensions.DeviceSearchResultContainer;
-import com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers;
-import com.android.launcher3.model.data.AppInfo;
-import com.android.launcher3.model.data.ItemInfoWithIcon;
-import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.model.data.SearchActionItemInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.touch.ItemClickHandler;
-import com.android.launcher3.touch.ItemLongClickListener;
-import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.util.Themes;
-
-import java.io.IOException;
-import java.net.URL;
-import java.net.URLConnection;
-import java.util.List;
-import java.util.function.Consumer;
-
-/**
- * A {@link BubbleTextView} representing a single cell result in AllApps
- */
-public class SearchResultIcon extends BubbleTextView implements
-        SearchTargetHandler, View.OnClickListener,
-        View.OnLongClickListener {
-
-    //Play store thumbnail process workaround
-    private final Rect mTempRect = new Rect();
-    private final Paint mIconPaint = new Paint();
-    private static final int BITMAP_CROP_MASK_COLOR = 0xff424242;
-
-    private final Launcher mLauncher;
-    private final SearchSessionTracker mSearchSessionTracker;
-    private String mTargetId;
-    private Consumer<ItemInfoWithIcon> mOnItemInfoChanged;
-
-
-    public SearchResultIcon(Context context) {
-        this(context, null, 0);
-    }
-
-    public SearchResultIcon(Context context, AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public SearchResultIcon(Context context, AttributeSet attrs, int defStyle) {
-        super(context, attrs, defStyle);
-        mLauncher = Launcher.getLauncher(getContext());
-        mSearchSessionTracker = SearchSessionTracker.getInstance(getContext());
-    }
-
-    private boolean mLongPressSupported;
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        setLongPressTimeoutFactor(1f);
-        setOnFocusChangeListener(mLauncher.getFocusHandler());
-        setOnClickListener(this);
-        setOnLongClickListener(this);
-        setLayoutParams(new ViewGroup.LayoutParams(ViewGroup.LayoutParams.MATCH_PARENT,
-                mLauncher.getDeviceProfile().allAppsCellHeightPx));
-    }
-
-    /**
-     * Applies {@link SearchTarget} to view. registers a consumer after a corresponding {@link
-     * ItemInfoWithIcon} is created
-     */
-    public void apply(SearchTarget searchTarget, List<SearchTarget> inlineItems,
-            Consumer<ItemInfoWithIcon> cb) {
-        mOnItemInfoChanged = cb;
-        apply(searchTarget, inlineItems);
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        mTargetId = parentTarget.getId();
-        if (parentTarget.getShortcutInfo() != null) {
-            prepareUsingShortcutInfo(parentTarget.getShortcutInfo());
-            mLongPressSupported = true;
-        } else if (parentTarget.getSearchAction() != null) {
-            prepareUsingSearchAction(parentTarget);
-            mLongPressSupported = false;
-        } else {
-            String className = parentTarget.getExtras().getString(SearchTargetUtil.EXTRA_CLASS);
-            prepareUsingApp(new ComponentName(parentTarget.getPackageName(), className),
-                    parentTarget.getUserHandle());
-            mLongPressSupported = true;
-        }
-    }
-
-    private void prepareUsingSearchAction(SearchTarget searchTarget) {
-        SearchAction searchAction = searchTarget.getSearchAction();
-        Bundle extras = searchAction.getExtras();
-
-        SearchActionItemInfo itemInfo = new SearchActionItemInfo(searchAction.getIcon(),
-                searchTarget.getPackageName(), searchTarget.getUserHandle(),
-                searchAction.getTitle()) {
-            @Override
-            protected ExtendedContainers getExtendedContainer() {
-                return ExtendedContainers.newBuilder()
-                        .setDeviceSearchResultContainer(buildDeviceSearchResultContainer()).build();
-            }
-        };
-        itemInfo.setIntent(searchAction.getIntent());
-        itemInfo.setPendingIntent(searchAction.getPendingIntent());
-
-        //TODO: remove this after flags are introduced in SearchAction. Settings results require
-        // startActivityForResult
-        boolean isSettingsResult = searchTarget.getResultType() == ResultType.SETTING;
-        if ((extras != null && extras.getBoolean(
-                SearchTargetUtil.BUNDLE_EXTRA_SHOULD_START_FOR_RESULT))
-                || isSettingsResult) {
-            itemInfo.setFlags(SearchActionItemInfo.FLAG_SHOULD_START_FOR_RESULT);
-        } else if (extras != null && extras.getBoolean(
-                SearchTargetUtil.BUNDLE_EXTRA_SHOULD_START)) {
-            itemInfo.setFlags(SearchActionItemInfo.FLAG_SHOULD_START);
-        }
-        if (extras != null && extras.getBoolean(
-                SearchTargetUtil.BUNDLE_EXTRA_BADGE_WITH_PACKAGE)) {
-            itemInfo.setFlags(FLAG_BADGE_WITH_PACKAGE);
-        }
-        if (extras != null && extras.getBoolean(BUNDLE_EXTRA_PRIMARY_ICON_FROM_TITLE)) {
-            itemInfo.setFlags(FLAG_PRIMARY_ICON_FROM_TITLE);
-        }
-
-        notifyItemInfoChanged(itemInfo);
-        LauncherAppState appState = LauncherAppState.getInstance(mLauncher);
-        MODEL_EXECUTOR.post(() -> {
-            try (LauncherIcons li = LauncherIcons.obtain(getContext())) {
-                Icon icon = searchTarget.getSearchAction().getIcon();
-                BitmapInfo pkgBitmap = getPackageBitmap(appState, searchTarget);
-                if (itemInfo.hasFlags(FLAG_PRIMARY_ICON_FROM_TITLE)) {
-                    // create a bitmap with first char if FLAG_PRIMARY_ICON_FROM_TITLE is set
-                    itemInfo.bitmap = li.createIconBitmap(String.valueOf(itemInfo.title.charAt(0)),
-                            pkgBitmap.color);
-                } else if (icon == null) {
-                    // Use default icon from package name
-                    itemInfo.bitmap = pkgBitmap;
-                } else {
-                    boolean isPlayResult = searchTarget.getResultType() == ResultType.PLAY;
-                    if (isPlayResult) {
-                        Bitmap b = getPlayResultBitmap(searchAction.getIcon());
-                        itemInfo.bitmap = b == null
-                                ? BitmapInfo.LOW_RES_INFO : BitmapInfo.fromBitmap(b);
-                    } else {
-                        itemInfo.bitmap = li.createBadgedIconBitmap(icon.loadDrawable(getContext()),
-                                itemInfo.user, false);
-                    }
-                }
-
-                // badge with package name
-                if (itemInfo.hasFlags(FLAG_BADGE_WITH_PACKAGE) && itemInfo.bitmap != pkgBitmap) {
-                    itemInfo.bitmap = li.badgeBitmap(itemInfo.bitmap.icon, pkgBitmap);
-                }
-            }
-            MAIN_EXECUTOR.post(() -> applyFromSearchActionItemInfo(itemInfo));
-        });
-    }
-
-    private static BitmapInfo getPackageBitmap(LauncherAppState appState, SearchTarget target) {
-        PackageItemInfo pkgInfo = new PackageItemInfo(target.getPackageName());
-        pkgInfo.user = target.getUserHandle();
-        appState.getIconCache().getTitleAndIconForApp(pkgInfo, false);
-        return pkgInfo.bitmap;
-    }
-
-    private Bitmap getPlayResultBitmap(Icon icon) {
-        try {
-            int iconSize = getIconSize();
-            URL url = new URL(icon.getUri().toString());
-            URLConnection con = url.openConnection();
-            con.addRequestProperty("Cache-Control", "max-age: 0");
-            con.setUseCaches(true);
-            Bitmap bitmap = BitmapFactory.decodeStream(con.getInputStream());
-            return getRoundedBitmap(Bitmap.createScaledBitmap(bitmap, iconSize, iconSize, false));
-        } catch (IOException e) {
-            e.printStackTrace();
-        }
-        return null;
-    }
-
-    private Bitmap getRoundedBitmap(Bitmap bitmap) {
-        final int iconSize = bitmap.getWidth();
-        final float radius = Themes.getDialogCornerRadius(getContext());
-
-        return BitmapRenderer.createHardwareBitmap(iconSize, iconSize, (canvas) -> {
-            mTempRect.set(0, 0, iconSize, iconSize);
-            final RectF rectF = new RectF(mTempRect);
-
-            mIconPaint.setAntiAlias(true);
-            mIconPaint.reset();
-            canvas.drawARGB(0, 0, 0, 0);
-            mIconPaint.setColor(BITMAP_CROP_MASK_COLOR);
-            canvas.drawRoundRect(rectF, radius, radius, mIconPaint);
-
-            mIconPaint.setXfermode(new PorterDuffXfermode(PorterDuff.Mode.SRC_IN));
-            canvas.drawBitmap(bitmap, mTempRect, mTempRect, mIconPaint);
-        });
-    }
-
-    private void prepareUsingApp(ComponentName componentName, UserHandle userHandle) {
-        AllAppsStore appsStore = mLauncher.getAppsView().getAppsStore();
-        AppInfo appInfo = new AppInfo(
-                appsStore.getApp(new ComponentKey(componentName, userHandle))) {
-            @Override
-            protected ExtendedContainers getExtendedContainer() {
-                return ExtendedContainers.newBuilder()
-                        .setDeviceSearchResultContainer(buildDeviceSearchResultContainer()).build();
-            }
-        };
-        appInfo.container = EXTENDED_CONTAINERS;
-        if (appInfo == null) {
-            setVisibility(GONE);
-            return;
-        }
-        applyFromApplicationInfo(appInfo);
-        notifyItemInfoChanged(appInfo);
-    }
-
-    private void prepareUsingShortcutInfo(ShortcutInfo shortcutInfo) {
-        WorkspaceItemInfo workspaceItemInfo = new WorkspaceItemInfo(shortcutInfo, getContext()) {
-            @Override
-            protected ExtendedContainers getExtendedContainer() {
-                return ExtendedContainers.newBuilder()
-                        .setDeviceSearchResultContainer(buildDeviceSearchResultContainer()).build();
-            }
-        };
-        workspaceItemInfo.container = EXTENDED_CONTAINERS;
-        notifyItemInfoChanged(workspaceItemInfo);
-        LauncherAppState launcherAppState = LauncherAppState.getInstance(getContext());
-        MODEL_EXECUTOR.execute(() -> {
-            launcherAppState.getIconCache().getShortcutIcon(workspaceItemInfo, shortcutInfo);
-            MAIN_EXECUTOR.post(() -> applyFromWorkspaceItem(workspaceItemInfo));
-        });
-    }
-
-    @Override
-    public boolean quickSelect() {
-        this.performClick();
-        notifyEvent(mLauncher, mTargetId, SearchTargetEvent.ACTION_LAUNCH_KEYBOARD_FOCUS);
-        return true;
-    }
-
-    @Override
-    public void onClick(View view) {
-        ItemClickHandler.INSTANCE.onClick(view);
-        notifyEvent(mLauncher, mTargetId, SearchTargetEvent.ACTION_LAUNCH_TOUCH);
-    }
-
-    @Override
-    public boolean onLongClick(View view) {
-        if (!mLongPressSupported) {
-            return false;
-        }
-        notifyEvent(mLauncher, mTargetId, SearchTargetEvent.ACTION_LONGPRESS);
-        return ItemLongClickListener.INSTANCE_ALL_APPS.onLongClick(this);
-    }
-
-
-    private void notifyItemInfoChanged(ItemInfoWithIcon itemInfoWithIcon) {
-        if (mOnItemInfoChanged != null) {
-            mOnItemInfoChanged.accept(itemInfoWithIcon);
-            mOnItemInfoChanged = null;
-        }
-    }
-
-    private DeviceSearchResultContainer buildDeviceSearchResultContainer() {
-        return mSearchSessionTracker.getQueryLength()
-                .map(queryLength -> DeviceSearchResultContainer.newBuilder()
-                        .setQueryLength(queryLength))
-                .orElse(DeviceSearchResultContainer.newBuilder()).build();
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultIconRow.java b/quickstep/src/com/android/launcher3/search/SearchResultIconRow.java
deleted file mode 100644
index 12a1a1c..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchResultIconRow.java
+++ /dev/null
@@ -1,158 +0,0 @@
-/*
- * 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.search;
-
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-
-import android.app.search.SearchTarget;
-import android.content.Context;
-import android.content.pm.ShortcutInfo;
-import android.text.TextUtils;
-import android.util.AttributeSet;
-import android.widget.LinearLayout;
-import android.widget.TextView;
-
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.R;
-import com.android.launcher3.model.data.ItemInfoWithIcon;
-import com.android.launcher3.model.data.PackageItemInfo;
-
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * A full width representation of {@link SearchResultIcon} with a secondary label and inline
- * SearchTargets
- */
-public class SearchResultIconRow extends LinearLayout implements SearchTargetHandler {
-
-    public static final int MAX_INLINE_ITEMS = 3;
-
-    protected final Launcher mLauncher;
-    protected final SearchResultIcon[] mInlineIcons = new SearchResultIcon[MAX_INLINE_ITEMS];
-    private SearchResultIcon mResultIcon;
-
-    private final LauncherAppState mLauncherAppState;
-    private TextView mTitleView;
-    private TextView mSubTitleView;
-
-    private PackageItemInfo mProviderInfo;
-
-    public SearchResultIconRow(Context context) {
-        this(context, null, 0);
-    }
-
-    public SearchResultIconRow(Context context,
-            @Nullable AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public SearchResultIconRow(Context context, @Nullable AttributeSet attrs, int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-        mLauncher = Launcher.getLauncher(getContext());
-        mLauncherAppState = LauncherAppState.getInstance(getContext());
-    }
-
-    protected int getIconSize() {
-        return mLauncher.getDeviceProfile().allAppsIconSizePx;
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        int iconSize = getIconSize();
-
-        mResultIcon = findViewById(R.id.icon);
-
-        mTitleView = findViewById(R.id.title);
-        mSubTitleView = findViewById(R.id.subtitle);
-        mSubTitleView.setVisibility(GONE);
-
-        mResultIcon.getLayoutParams().height = iconSize;
-        mResultIcon.getLayoutParams().width = iconSize;
-        mResultIcon.setTextVisibility(false);
-
-        mInlineIcons[0] = findViewById(R.id.shortcut_0);
-        mInlineIcons[1] = findViewById(R.id.shortcut_1);
-        mInlineIcons[2] = findViewById(R.id.shortcut_2);
-        for (SearchResultIcon inlineIcon : mInlineIcons) {
-            inlineIcon.getLayoutParams().width = getIconSize();
-        }
-        setOnClickListener(mResultIcon);
-        setOnLongClickListener(mResultIcon);
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        mResultIcon.apply(parentTarget, children, this::onItemInfoCreated);
-
-        showSubtitleIfNeeded(null);
-
-        if (parentTarget.getShortcutInfo() != null) {
-            updateWithShortcutInfo(parentTarget.getShortcutInfo());
-        } else if (parentTarget.getSearchAction() != null) {
-            showSubtitleIfNeeded(parentTarget.getSearchAction().getSubtitle());
-        }
-        showInlineItems(children);
-    }
-
-    @Override
-    public boolean quickSelect() {
-        this.performClick();
-        return true;
-    }
-
-    private void updateWithShortcutInfo(ShortcutInfo shortcutInfo) {
-        PackageItemInfo packageItemInfo = new PackageItemInfo(shortcutInfo.getPackage());
-        if (packageItemInfo.equals(mProviderInfo)) return;
-        MODEL_EXECUTOR.post(() -> {
-            mLauncherAppState.getIconCache().getTitleAndIconForApp(packageItemInfo, true);
-            MAIN_EXECUTOR.post(() -> {
-                showSubtitleIfNeeded(packageItemInfo.title);
-                mProviderInfo = packageItemInfo;
-            });
-        });
-    }
-
-    protected void showSubtitleIfNeeded(CharSequence subTitle) {
-        if (!TextUtils.isEmpty(subTitle)) {
-            mSubTitleView.setText(subTitle);
-            mSubTitleView.setVisibility(VISIBLE);
-        } else {
-            mSubTitleView.setVisibility(GONE);
-        }
-    }
-
-    protected void showInlineItems(List<SearchTarget> children) {
-        for (int i = 0; i < MAX_INLINE_ITEMS; i++) {
-            if (i < children.size()) {
-                mInlineIcons[i].apply(children.get(i), new ArrayList<>());
-                mInlineIcons[i].setVisibility(VISIBLE);
-            } else {
-                mInlineIcons[i].setVisibility(GONE);
-            }
-        }
-    }
-
-    protected void onItemInfoCreated(ItemInfoWithIcon info) {
-        setTag(info);
-        mTitleView.setText(info.title);
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultIconSlice.java b/quickstep/src/com/android/launcher3/search/SearchResultIconSlice.java
deleted file mode 100644
index 4bf3432..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchResultIconSlice.java
+++ /dev/null
@@ -1,122 +0,0 @@
-/*
- * 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.search;
-
-import android.app.search.SearchTarget;
-import android.app.search.SearchTargetEvent;
-import android.content.Context;
-import android.util.AttributeSet;
-import android.widget.LinearLayout;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-import androidx.slice.SliceItem;
-import androidx.slice.widget.EventInfo;
-import androidx.slice.widget.SliceView;
-
-import com.android.launcher3.Launcher;
-import com.android.launcher3.R;
-import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.util.SafeCloseable;
-
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * A slice view wrapper with settings app icon at start
- */
-public class SearchResultIconSlice extends LinearLayout implements SearchTargetHandler,
-        SliceView.OnSliceActionListener {
-
-    private final Launcher mLauncher;
-
-    private SliceView mSliceView;
-    private SearchResultIcon mIcon;
-    private SafeCloseable mSliceSession;
-    private String mTargetId;
-
-    public SearchResultIconSlice(Context context) {
-        this(context, null, 0);
-    }
-
-    public SearchResultIconSlice(Context context,
-            @Nullable AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public SearchResultIconSlice(Context context, @Nullable AttributeSet attrs,
-            int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-        mLauncher = Launcher.getLauncher(getContext());
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        mSliceView = findViewById(R.id.slice);
-        mIcon = findViewById(R.id.icon);
-        mIcon.setTextVisibility(false);
-        int iconSize = mLauncher.getDeviceProfile().iconSizePx;
-        mIcon.getLayoutParams().height = iconSize;
-        mIcon.getLayoutParams().width = iconSize;
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        mTargetId = parentTarget.getId();
-        reset();
-        updateIcon(parentTarget, children);
-        mSliceSession = mLauncher.getLiveSearchManager()
-                .addObserver(parentTarget.getSliceUri(), mSliceView);
-    }
-
-    private void updateIcon(SearchTarget parentTarget, List<SearchTarget> children) {
-        if (children.size() == 1) {
-            mIcon.apply(children.get(0), new ArrayList<>());
-        } else {
-            PackageItemInfo pkgItem = new PackageItemInfo(parentTarget.getPackageName());
-            pkgItem.user = parentTarget.getUserHandle();
-            if (!pkgItem.equals(mIcon.getTag())) {
-                // The icon will load and apply high res icon automatically
-                mIcon.applyFromItemInfoWithIcon(pkgItem);
-            }
-        }
-    }
-
-    @Override
-    protected void onAttachedToWindow() {
-        super.onAttachedToWindow();
-        mSliceView.setOnSliceActionListener(this);
-    }
-
-    @Override
-    protected void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-        reset();
-    }
-
-    private void reset() {
-        mSliceView.setOnSliceActionListener(null);
-        if (mSliceSession != null) {
-            mSliceSession.close();
-        }
-    }
-
-    @Override
-    public void onSliceAction(@NonNull EventInfo eventInfo, @NonNull SliceItem sliceItem) {
-        notifyEvent(mLauncher, mTargetId, SearchTargetEvent.ACTION_TAP);
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultSmallIconRow.java b/quickstep/src/com/android/launcher3/search/SearchResultSmallIconRow.java
deleted file mode 100644
index ca8aa81..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchResultSmallIconRow.java
+++ /dev/null
@@ -1,142 +0,0 @@
-/*
- * 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.search;
-
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-
-import android.app.search.SearchTarget;
-import android.content.Context;
-import android.content.pm.ShortcutInfo;
-import android.text.TextUtils;
-import android.util.AttributeSet;
-import android.widget.LinearLayout;
-import android.widget.TextView;
-
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.R;
-import com.android.launcher3.model.data.ItemInfoWithIcon;
-import com.android.launcher3.model.data.PackageItemInfo;
-
-import java.util.List;
-
-/**
- * A full width representation of {@link SearchResultIcon} with a secondary label and inline
- * SearchTargets
- */
-public class SearchResultSmallIconRow extends LinearLayout implements SearchTargetHandler {
-
-    protected final Launcher mLauncher;
-    private final LauncherAppState mLauncherAppState;
-    protected SearchResultIcon mResultIcon;
-
-    private TextView mTitleView;
-    private TextView mDelimeterView;
-    private TextView mSubTitleView;
-
-    private PackageItemInfo mProviderInfo;
-
-    public SearchResultSmallIconRow(Context context) {
-        this(context, null, 0);
-    }
-
-    public SearchResultSmallIconRow(Context context,
-            @Nullable AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public SearchResultSmallIconRow(Context context,
-            @Nullable AttributeSet attrs, int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-        mLauncher = Launcher.getLauncher(getContext());
-        mLauncherAppState = LauncherAppState.getInstance(getContext());
-    }
-
-    protected int getIconSize() {
-        return mLauncher.getDeviceProfile().allAppsIconSizePx;
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        int iconSize = getIconSize();
-
-        mResultIcon = findViewById(R.id.icon);
-
-        mTitleView = findViewById(R.id.title);
-        mDelimeterView = findViewById(R.id.delimeter);
-        mDelimeterView.setVisibility(GONE);
-        mSubTitleView = findViewById(R.id.subtitle);
-        mSubTitleView.setVisibility(GONE);
-
-        mResultIcon.getLayoutParams().height = iconSize;
-        mResultIcon.getLayoutParams().width = iconSize;
-        mResultIcon.setTextVisibility(false);
-
-        setOnClickListener(mResultIcon);
-        setOnLongClickListener(mResultIcon);
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        mResultIcon.apply(parentTarget, children, this::onItemInfoCreated);
-
-        showSubtitleIfNeeded(null);
-
-        if (parentTarget.getShortcutInfo() != null) {
-            updateWithShortcutInfo(parentTarget.getShortcutInfo());
-        } else if (parentTarget.getSearchAction() != null) {
-            showSubtitleIfNeeded(parentTarget.getSearchAction().getSubtitle());
-        }
-    }
-
-    @Override
-    public boolean quickSelect() {
-        this.performClick();
-        return true;
-    }
-
-    private void updateWithShortcutInfo(ShortcutInfo shortcutInfo) {
-        PackageItemInfo packageItemInfo = new PackageItemInfo(shortcutInfo.getPackage());
-        if (packageItemInfo.equals(mProviderInfo)) return;
-        MODEL_EXECUTOR.post(() -> {
-            mLauncherAppState.getIconCache().getTitleAndIconForApp(packageItemInfo, true);
-            MAIN_EXECUTOR.post(() -> {
-                showSubtitleIfNeeded(packageItemInfo.title);
-                mProviderInfo = packageItemInfo;
-            });
-        });
-    }
-
-    protected void showSubtitleIfNeeded(CharSequence subTitle) {
-        if (!TextUtils.isEmpty(subTitle)) {
-            mSubTitleView.setText(subTitle);
-            mSubTitleView.setVisibility(VISIBLE);
-            mDelimeterView.setVisibility(VISIBLE);
-
-        } else {
-            mSubTitleView.setVisibility(GONE);
-        }
-    }
-
-    protected void onItemInfoCreated(ItemInfoWithIcon info) {
-        setTag(info);
-        mTitleView.setText(info.title);
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultThumbnailView.java b/quickstep/src/com/android/launcher3/search/SearchResultThumbnailView.java
deleted file mode 100644
index 8803c98..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchResultThumbnailView.java
+++ /dev/null
@@ -1,102 +0,0 @@
-/*
- * 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.search;
-
-
-import android.app.search.SearchTarget;
-import android.app.search.SearchTargetEvent;
-import android.content.Context;
-import android.graphics.Bitmap;
-import android.graphics.drawable.BitmapDrawable;
-import android.util.AttributeSet;
-import android.view.View;
-
-import androidx.core.graphics.drawable.RoundedBitmapDrawable;
-import androidx.core.graphics.drawable.RoundedBitmapDrawableFactory;
-
-import com.android.launcher3.Launcher;
-import com.android.launcher3.model.data.SearchActionItemInfo;
-import com.android.launcher3.touch.ItemClickHandler;
-import com.android.launcher3.util.Themes;
-
-import java.util.List;
-
-/**
- * A view representing a high confidence app search result that includes shortcuts
- */
-public class SearchResultThumbnailView extends androidx.appcompat.widget.AppCompatImageView
-        implements SearchTargetHandler, View.OnClickListener {
-
-    private SearchTarget mSearchTarget;
-
-    public SearchResultThumbnailView(Context context) {
-        super(context);
-    }
-
-    public SearchResultThumbnailView(Context context, AttributeSet attrs) {
-        super(context, attrs);
-    }
-
-    public SearchResultThumbnailView(Context context, AttributeSet attrs, int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        setOnFocusChangeListener(Launcher.getLauncher(getContext()).getFocusHandler());
-        setOnClickListener(this);
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        mSearchTarget = parentTarget;
-        Bitmap bitmap;
-
-        SearchActionItemInfo itemInfo = new SearchActionItemInfo(
-                parentTarget.getSearchAction().getIcon(),
-                parentTarget.getPackageName(),
-                parentTarget.getUserHandle(),
-                parentTarget.getSearchAction().getTitle());
-        itemInfo.setIntent(parentTarget.getSearchAction().getIntent());
-        itemInfo.setPendingIntent(parentTarget.getSearchAction().getPendingIntent());
-
-        bitmap = ((BitmapDrawable) itemInfo.getIcon()
-                .loadDrawable(getContext())).getBitmap();
-        // crop
-        if (bitmap.getWidth() < bitmap.getHeight()) {
-            bitmap = Bitmap.createBitmap(bitmap, 0,
-                    bitmap.getHeight() / 2 - bitmap.getWidth() / 2,
-                    bitmap.getWidth(), bitmap.getWidth());
-        } else {
-            bitmap = Bitmap.createBitmap(bitmap, bitmap.getWidth() / 2 - bitmap.getHeight() / 2,
-                    0,
-                    bitmap.getHeight(), bitmap.getHeight());
-        }
-        setTag(itemInfo);
-
-        RoundedBitmapDrawable drawable = RoundedBitmapDrawableFactory.create(null, bitmap);
-        drawable.setCornerRadius(Themes.getDialogCornerRadius(getContext()));
-        setImageDrawable(drawable);
-    }
-
-    @Override
-    public void onClick(View view) {
-        ItemClickHandler.onClickSearchAction(Launcher.getLauncher(getContext()),
-                (SearchActionItemInfo) view.getTag());
-        notifyEvent(getContext(), mSearchTarget.getId(), SearchTargetEvent.ACTION_LAUNCH_TOUCH);
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultWidget.java b/quickstep/src/com/android/launcher3/search/SearchResultWidget.java
deleted file mode 100644
index e22f6ab..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchResultWidget.java
+++ /dev/null
@@ -1,228 +0,0 @@
-/*
- * 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.search;
-
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-
-import android.app.search.SearchTarget;
-import android.app.search.SearchTargetEvent;
-import android.appwidget.AppWidgetHostView;
-import android.appwidget.AppWidgetProviderInfo;
-import android.content.Context;
-import android.graphics.Rect;
-import android.util.AttributeSet;
-import android.view.GestureDetector;
-import android.view.MotionEvent;
-import android.view.View;
-import android.widget.LinearLayout;
-import android.widget.TextView;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.AppWidgetResizeFrame;
-import com.android.launcher3.BubbleTextView;
-import com.android.launcher3.CheckLongPressHelper;
-import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.R;
-import com.android.launcher3.Utilities;
-import com.android.launcher3.allapps.search.SearchWidgetInfoContainer;
-import com.android.launcher3.dragndrop.DraggableView;
-import com.android.launcher3.icons.cache.HandlerRunnable;
-import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.touch.ItemLongClickListener;
-import com.android.launcher3.widget.PendingAddWidgetInfo;
-
-import java.util.List;
-
-/**
- * displays live version of a widget upon receiving {@link AppWidgetProviderInfo} from Search
- * provider
- */
-public class SearchResultWidget extends LinearLayout implements SearchTargetHandler, DraggableView,
-        View.OnLongClickListener {
-
-
-    private final Rect mWidgetOffset = new Rect();
-
-    private final Launcher mLauncher;
-    private final CheckLongPressHelper mLongPressHelper;
-    private final GestureDetector mClickDetector;
-    private final AppWidgetHostView mHostView;
-    private final float mScaleToFit;
-
-    private SearchWidgetInfoContainer mInfoContainer;
-    private HandlerRunnable mLabelRequest;
-    private BubbleTextView mWidgetProvider;
-    private TextView mWidgetLabel;
-
-    public SearchResultWidget(@NonNull Context context) {
-        this(context, null, 0);
-    }
-
-    public SearchResultWidget(@NonNull Context context,
-            @Nullable AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public SearchResultWidget(@NonNull Context context, @Nullable AttributeSet attrs,
-            int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-        mLauncher = Launcher.getLauncher(context);
-        mHostView = new AppWidgetHostView(context);
-        DeviceProfile grid = mLauncher.getDeviceProfile();
-        mScaleToFit = Math.min(grid.appWidgetScale.x, grid.appWidgetScale.y);
-
-        // detect tap event on widget container for search target event reporting
-        mClickDetector = new GestureDetector(context,
-                new ClickListener(
-                        () -> reportEvent(SearchTargetEvent.ACTION_LAUNCH_TOUCH)));
-        mLongPressHelper = new CheckLongPressHelper(this);
-        mLongPressHelper.setLongPressTimeoutFactor(1);
-        setOnLongClickListener(this);
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        mWidgetProvider = findViewById(R.id.widget_provider);
-        mWidgetLabel = findViewById(R.id.widget_label);
-        addView(mHostView);
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        AppWidgetProviderInfo providerInfo = parentTarget.getAppWidgetProviderInfo();
-        removeListener();
-
-        showWidgetInfo(providerInfo);
-        mInfoContainer = mLauncher.getLiveSearchManager().getPlaceHolderWidget(providerInfo);
-        if (mInfoContainer == null) {
-            setVisibility(GONE);
-            return;
-        }
-        setVisibility(VISIBLE);
-        mInfoContainer.attachWidget(mHostView);
-        PendingAddWidgetInfo info = (PendingAddWidgetInfo) mHostView.getTag();
-        int[] size = mLauncher.getWorkspace().estimateItemSize(info);
-        mHostView.getLayoutParams().width = size[0];
-        mHostView.getLayoutParams().height = size[1];
-        AppWidgetResizeFrame.updateWidgetSizeRanges(mHostView, mLauncher, info.spanX,
-                info.spanY);
-        mHostView.requestLayout();
-        setTag(info);
-    }
-
-    private void showWidgetInfo(AppWidgetProviderInfo providerInfo) {
-        PackageItemInfo pinfo = new PackageItemInfo(providerInfo.provider.getPackageName());
-        pinfo.user = providerInfo.getProfile();
-        mWidgetProvider.applyFromItemInfoWithIcon(pinfo);
-
-        mLabelRequest = new HandlerRunnable<>(
-                MODEL_EXECUTOR.getHandler(),
-                () -> LauncherAppState.getInstance(mLauncher).getIconCache()
-                        .getTitleNoCache(LauncherAppWidgetProviderInfo
-                                .fromProviderInfo(mLauncher, providerInfo)),
-                MAIN_EXECUTOR,
-                mWidgetLabel::setText);
-        Utilities.postAsyncCallback(MODEL_EXECUTOR.getHandler(), mLabelRequest);
-    }
-
-    /**
-     * Stops hostView from getting updates on a widget provider
-     */
-    public void removeListener() {
-        if (mInfoContainer != null) {
-            mInfoContainer.detachWidget(mHostView);
-            mInfoContainer = null;
-        }
-        if (mLabelRequest != null) {
-            mLabelRequest.cancel();
-            mLabelRequest = null;
-        }
-    }
-
-    @Override
-    protected void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-        removeListener();
-    }
-
-    private void reportEvent(int eventType) {
-        SearchSessionTracker.INSTANCE.get(getContext()).notifyEvent(
-                new SearchTargetEvent.Builder("search target id", eventType).build());
-    }
-
-    @Override
-    public boolean onInterceptTouchEvent(MotionEvent ev) {
-        mLongPressHelper.onTouchEvent(ev);
-        mClickDetector.onTouchEvent(ev);
-        return mLongPressHelper.hasPerformedLongPress();
-    }
-
-    @Override
-    public boolean onTouchEvent(MotionEvent ev) {
-        mLongPressHelper.onTouchEvent(ev);
-        return true;
-    }
-
-    @Override
-    public void cancelLongPress() {
-        super.cancelLongPress();
-        mLongPressHelper.cancelLongPress();
-    }
-
-    @Override
-    public int getViewType() {
-        return DraggableView.DRAGGABLE_WIDGET;
-    }
-
-    @Override
-    public void getSourceVisualDragBounds(Rect bounds) {
-        mHostView.getHitRect(mWidgetOffset);
-        int width = (int) (mHostView.getMeasuredWidth() * mScaleToFit);
-        int height = (int) (mHostView.getMeasuredHeight() * mScaleToFit);
-        bounds.set(mWidgetOffset.left,
-                mWidgetOffset.top,
-                width + mWidgetOffset.left,
-                height + mWidgetOffset.top);
-    }
-
-    @Override
-    public boolean onLongClick(View view) {
-        ItemLongClickListener.INSTANCE_ALL_APPS.onLongClick(view);
-        return false;
-    }
-
-
-    static class ClickListener extends GestureDetector.SimpleOnGestureListener {
-        private final Runnable mCb;
-
-        ClickListener(Runnable cb) {
-            mCb = cb;
-        }
-
-        @Override
-        public boolean onSingleTapConfirmed(MotionEvent e) {
-            mCb.run();
-            return super.onSingleTapConfirmed(e);
-        }
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchResultWidgetPreview.java b/quickstep/src/com/android/launcher3/search/SearchResultWidgetPreview.java
deleted file mode 100644
index 5bf1908..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchResultWidgetPreview.java
+++ /dev/null
@@ -1,130 +0,0 @@
-/*
- * 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.search;
-
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-
-import android.app.search.SearchTarget;
-import android.app.search.SearchTargetEvent;
-import android.appwidget.AppWidgetProviderInfo;
-import android.content.Context;
-import android.graphics.Point;
-import android.util.AttributeSet;
-import android.view.View;
-import android.widget.LinearLayout;
-import android.widget.Toast;
-
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.R;
-import com.android.launcher3.dragndrop.DragOptions;
-import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.touch.ItemLongClickListener;
-import com.android.launcher3.widget.BaseWidgetSheet;
-import com.android.launcher3.widget.PendingItemDragHelper;
-import com.android.launcher3.widget.WidgetCell;
-import com.android.launcher3.widget.WidgetImageView;
-
-import java.util.List;
-
-/**
- * displays preview of a widget upon receiving {@link AppWidgetProviderInfo} from Search provider
- */
-public class SearchResultWidgetPreview extends LinearLayout implements SearchTargetHandler,
-        View.OnClickListener, View.OnLongClickListener {
-
-    private final Launcher mLauncher;
-    private final LauncherAppState mAppState;
-    private WidgetCell mWidgetCell;
-    private Toast mWidgetToast;
-
-    private String mTargetId;
-    public SearchResultWidgetPreview(Context context) {
-        this(context, null, 0);
-    }
-
-    public SearchResultWidgetPreview(Context context,
-            @Nullable AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public SearchResultWidgetPreview(Context context, @Nullable AttributeSet attrs,
-            int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-        mLauncher = Launcher.getLauncher(context);
-        mAppState = LauncherAppState.getInstance(context);
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        mWidgetCell = findViewById(R.id.widget_cell);
-        mWidgetCell.setOnLongClickListener(this);
-        mWidgetCell.setOnClickListener(this);
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        mTargetId = parentTarget.getId();
-        AppWidgetProviderInfo providerInfo = parentTarget.getAppWidgetProviderInfo();
-        LauncherAppWidgetProviderInfo pInfo = LauncherAppWidgetProviderInfo.fromProviderInfo(
-                getContext(), providerInfo);
-        MODEL_EXECUTOR.post(() -> {
-            WidgetItem widgetItem = new WidgetItem(pInfo, mLauncher.getDeviceProfile().inv,
-                    mAppState.getIconCache());
-            MAIN_EXECUTOR.post(() -> {
-                mWidgetCell.applyFromCellItem(widgetItem, mAppState.getWidgetCache());
-                mWidgetCell.ensurePreview();
-            });
-        });
-    }
-
-    @Override
-    public boolean onLongClick(View view) {
-        view.cancelLongPress();
-        if (!ItemLongClickListener.canStartDrag(mLauncher)) return false;
-        if (mWidgetCell.getTag() == null) return false;
-
-        WidgetImageView imageView = mWidgetCell.getWidgetView();
-        if (imageView.getBitmap() == null) {
-            return false;
-        }
-
-        int[] loc = new int[2];
-        mLauncher.getDragLayer().getLocationInDragLayer(imageView, loc);
-
-        new PendingItemDragHelper(mWidgetCell).startDrag(
-                imageView.getBitmapBounds(), imageView.getBitmap().getWidth(), imageView.getWidth(),
-                new Point(loc[0], loc[1]), mLauncher.getAppsView(), new DragOptions());
-        reportEvent(SearchTargetEvent.ACTION_LONGPRESS);
-        return true;
-    }
-
-    @Override
-    public void onClick(View view) {
-        mWidgetToast = BaseWidgetSheet.showWidgetToast(getContext(), mWidgetToast);
-        reportEvent(SearchTargetEvent.ACTION_LAUNCH_TOUCH);
-    }
-
-    private void reportEvent(int eventType) {
-        SearchSessionTracker.INSTANCE.get(getContext()).notifyEvent(
-                new SearchTargetEvent.Builder(mTargetId, eventType).build());
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchSectionHeaderView.java b/quickstep/src/com/android/launcher3/search/SearchSectionHeaderView.java
deleted file mode 100644
index 9276841..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchSectionHeaderView.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * 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.search;
-
-import android.app.search.SearchTarget;
-import android.content.Context;
-import android.util.AttributeSet;
-import android.widget.TextView;
-
-import androidx.annotation.Nullable;
-
-import java.util.List;
-
-/**
- * Header text view that shows a title for a given section in All apps search
- */
-public class SearchSectionHeaderView extends TextView implements SearchTargetHandler {
-
-    public SearchSectionHeaderView(Context context) {
-        super(context);
-    }
-
-    public SearchSectionHeaderView(Context context,
-            @Nullable AttributeSet attrs) {
-        super(context, attrs);
-    }
-
-    public SearchSectionHeaderView(Context context, @Nullable AttributeSet attrs, int styleAttr) {
-        super(context, attrs, styleAttr);
-    }
-
-    @Override
-    public void apply(SearchTarget parentTarget, List<SearchTarget> children) {
-        setText(parentTarget.getSearchAction().getTitle());
-        setVisibility(VISIBLE);
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchServicePipeline.java b/quickstep/src/com/android/launcher3/search/SearchServicePipeline.java
deleted file mode 100644
index fac6ba7..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchServicePipeline.java
+++ /dev/null
@@ -1,153 +0,0 @@
-/*
- * 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.search;
-
-import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
-
-import android.app.search.Query;
-import android.app.search.SearchContext;
-import android.app.search.SearchSession;
-import android.app.search.SearchTarget;
-import android.app.search.SearchUiManager;
-import android.content.Context;
-import android.os.CancellationSignal;
-import android.text.TextUtils;
-import android.util.Log;
-
-import androidx.annotation.WorkerThread;
-
-import com.android.app.search.ResultType;
-import com.android.launcher3.allapps.AllAppsGridAdapter;
-import com.android.launcher3.allapps.AllAppsSectionDecorator;
-import com.android.launcher3.allapps.search.SearchPipeline;
-import com.android.launcher3.allapps.search.SearchSectionInfo;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.function.Consumer;
-
-/**
- * Search pipeline utilizing {@link android.app.search.SearchUiManager}
- */
-public class SearchServicePipeline implements SearchPipeline {
-    private static final int SUPPORTED_RESULT_TYPES =
-            ResultType.APPLICATION | ResultType.SHORTCUT | ResultType.PLAY | ResultType.PEOPLE
-                    | ResultType.SETTING;
-
-    private static final boolean DEBUG = true;
-    private static final int REQUEST_TIMEOUT = 200;
-    private static final String TAG = "SearchServicePipeline";
-
-    private final Context mContext;
-    private final SearchSession mSession;
-    private final DeviceSearchAdapterProvider mAdapterProvider;
-
-    private boolean mCanceled = false;
-
-
-    public SearchServicePipeline(Context context, DeviceSearchAdapterProvider adapterProvider) {
-        mContext = context;
-        mAdapterProvider = adapterProvider;
-        SearchUiManager manager = context.getSystemService(SearchUiManager.class);
-        mSession = manager.createSearchSession(
-                new SearchContext(SUPPORTED_RESULT_TYPES, REQUEST_TIMEOUT, null));
-        SearchSessionTracker.getInstance(context).setSearchSession(mSession);
-    }
-
-    @WorkerThread
-    @Override
-    public void query(String input, Consumer<ArrayList<AllAppsGridAdapter.AdapterItem>> callback,
-            CancellationSignal cancellationSignal) {
-        mCanceled = false;
-        Query query = new Query(input, System.currentTimeMillis(), null);
-        mSession.query(query, UI_HELPER_EXECUTOR, targets -> {
-            if (!mCanceled) {
-                if (DEBUG) {
-                    printSearchTargets(input, targets);
-                }
-                SearchSessionTracker.getInstance(mContext).setQuery(query);
-                callback.accept(this.onResult(targets));
-            }
-            Log.w(TAG, "Ignoring results due to cancel signal");
-        });
-    }
-
-    /**
-     * Given A list of search Targets, pairs a group of search targets to a AdapterItem that can
-     * be inflated in AllAppsRecyclerView
-     */
-    private ArrayList<AllAppsGridAdapter.AdapterItem> onResult(List<SearchTarget> searchTargets) {
-        HashMap<String, SearchAdapterItem> adapterMap = new LinkedHashMap<>();
-        List<SearchTarget> unmappedChildren = new ArrayList<>();
-        SearchSectionInfo section = new SearchSectionInfo();
-        section.setDecorationHandler(
-                new AllAppsSectionDecorator.SectionDecorationHandler(mContext, true));
-        for (SearchTarget target : searchTargets) {
-            if (!TextUtils.isEmpty(target.getParentId())) {
-                if (!addChildToParent(target, adapterMap)) {
-                    unmappedChildren.add(target);
-                }
-                continue;
-            }
-            int viewType = mAdapterProvider.getViewTypeForSearchTarget(target);
-            if (viewType != -1) {
-                SearchAdapterItem adapterItem = new SearchAdapterItem(target, viewType);
-                adapterItem.searchSectionInfo = section;
-                adapterMap.put(target.getId(), adapterItem);
-            }
-        }
-        for (SearchTarget s : unmappedChildren) {
-            if (!addChildToParent(s, adapterMap)) {
-                Log.w(TAG,
-                        "Unable to pair child " + s.getId() + " to parent " + s.getParentId());
-            }
-        }
-        return new ArrayList<>(adapterMap.values());
-    }
-
-    private void printSearchTargets(String query, List<SearchTarget> results) {
-        Log.d(TAG, " query=" + query + " size=" + results.size());
-        for (SearchTarget s : results) {
-            Log.d(TAG, "layoutType=" + s.getLayoutType() + " resultType=" + s.getResultType());
-        }
-    }
-
-    /**
-     * Adds a child SearchTarget to a collection of searchTarget children with a shared parentId.
-     * Returns false if no parent searchTarget with id=$parentId does not exists.
-     */
-    private boolean addChildToParent(SearchTarget target, HashMap<String, SearchAdapterItem> map) {
-        if (!map.containsKey(target.getParentId())) return false;
-        map.get(target.getParentId()).getInlineItems().add(target);
-        return true;
-    }
-
-    /**
-     * Unregister callbacks and destroy search session
-     */
-    public void destroy() {
-        mSession.destroy();
-    }
-
-    /**
-     * Cancels current ongoing search request.
-     */
-    public void cancel() {
-        mCanceled = true;
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchSessionTracker.java b/quickstep/src/com/android/launcher3/search/SearchSessionTracker.java
deleted file mode 100644
index 97a72bb..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchSessionTracker.java
+++ /dev/null
@@ -1,81 +0,0 @@
-/*
- * 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.
- */
-package com.android.launcher3.search;
-
-import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
-
-import android.app.search.Query;
-import android.app.search.SearchSession;
-import android.app.search.SearchTargetEvent;
-import android.content.Context;
-import android.util.Log;
-
-import androidx.annotation.Nullable;
-import androidx.annotation.WorkerThread;
-
-import com.android.launcher3.util.MainThreadInitializedObject;
-
-import java.util.Optional;
-
-/**
- * A singleton class to track and report search events back to SearchSession
- */
-public class SearchSessionTracker {
-
-    private static final String TAG = "SearchSessionTracker";
-
-    @Nullable
-    private SearchSession mSession;
-    private Query mQuery;
-
-    public static final MainThreadInitializedObject<SearchSessionTracker> INSTANCE =
-            new MainThreadInitializedObject<>(SearchSessionTracker::new);
-
-    private SearchSessionTracker(Context context) {
-    }
-
-    /**
-     * Returns instance of SearchSessionTracker
-     */
-    public static SearchSessionTracker getInstance(Context context) {
-        return SearchSessionTracker.INSTANCE.get(context);
-    }
-
-    @WorkerThread
-    public void setSearchSession(SearchSession session) {
-        mSession = session;
-    }
-
-    @WorkerThread
-    public void setQuery(Query query) {
-        mQuery = query;
-    }
-
-    public Optional<Integer> getQueryLength() {
-        return Optional.ofNullable(mQuery).map(Query::getInput).map(String::length);
-    }
-
-    /**
-     * Send the user event handling back to the {@link SearchSession} object.
-     */
-    public void notifyEvent(SearchTargetEvent event) {
-        if (mSession == null) {
-            Log.d(TAG, "Dropping event " + event.getTargetId());
-        }
-        Log.d(TAG, "notifyEvent:" + mQuery.getInput() + ":" + event.getTargetId());
-        UI_HELPER_EXECUTOR.post(() -> mSession.notifyEvent(mQuery, event));
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchTargetHandler.java b/quickstep/src/com/android/launcher3/search/SearchTargetHandler.java
deleted file mode 100644
index acf6f8a..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchTargetHandler.java
+++ /dev/null
@@ -1,46 +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.search;
-
-import android.app.search.SearchTarget;
-import android.app.search.SearchTargetEvent;
-import android.content.Context;
-
-import java.util.List;
-
-/**
- * An interface for supporting dynamic search results
- */
-public interface SearchTargetHandler {
-
-    /**
-     * Update view using values from {@link SearchTarget}
-     */
-    default void apply(SearchTarget parentTarget, List<SearchTarget> children) { }
-
-    /**
-     * Handle IME quick select event. returns whether event was handled.
-     */
-    default boolean quickSelect() {
-        return false;
-    }
-
-    default void notifyEvent(Context context, String id, int eventType) {
-        SearchTargetEvent.Builder builder = new SearchTargetEvent.Builder(id, eventType);
-        SearchSessionTracker.getInstance(context).notifyEvent(builder.build());
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/search/SearchTargetUtil.java b/quickstep/src/com/android/launcher3/search/SearchTargetUtil.java
deleted file mode 100644
index ede3b9d..0000000
--- a/quickstep/src/com/android/launcher3/search/SearchTargetUtil.java
+++ /dev/null
@@ -1,198 +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.search;
-
-import static com.android.app.search.LayoutType.DIVIDER;
-import static com.android.app.search.LayoutType.ICON_HORIZONTAL_TEXT;
-import static com.android.app.search.LayoutType.SMALL_ICON_HORIZONTAL_TEXT;
-import static com.android.app.search.LayoutType.THUMBNAIL;
-import static com.android.app.search.ResultType.ACTION;
-import static com.android.app.search.ResultType.SCREENSHOT;
-import static com.android.app.search.ResultType.SUGGEST;
-
-import android.app.PendingIntent;
-import android.app.search.SearchAction;
-import android.app.search.SearchTarget;
-import android.content.Context;
-import android.content.Intent;
-import android.content.pm.ShortcutInfo;
-import android.graphics.Bitmap;
-import android.graphics.Canvas;
-import android.graphics.drawable.Icon;
-import android.net.Uri;
-import android.os.Bundle;
-import android.os.Process;
-import android.os.UserHandle;
-
-import com.android.app.search.ResultType;
-
-public class SearchTargetUtil {
-
-    public static final String BUNDLE_EXTRA_SHOULD_START = "should_start";
-    public static final String BUNDLE_EXTRA_SHOULD_START_FOR_RESULT = "should_start_for_result";
-    public static final String BUNDLE_EXTRA_BADGE_WITH_PACKAGE = "badge_with_package";
-    public static final String BUNDLE_EXTRA_PRIMARY_ICON_FROM_TITLE = "primary_icon_from_title";
-
-    public static final String EXTRA_CLASS = "class";
-
-    private static final String TITLE = " title: weather ";
-    private static final String SUBTITLE = " subtitle: 68 degrees ";
-    private static final String PACKAGE2 = "com.google.android.gm";
-    private static final UserHandle USERHANDLE = Process.myUserHandle();
-
-
-    /**
-     * Generate SearchTargetUtil for ICON_HORIZONTAL_TEXT layout type.
-     *
-     * targets.add(SearchTargetUtil.generateIconDoubleHorizontalText_SearchAction(
-     * mContext, "red", Color.RED));
-     * targets.add(SearchTargetUtil.generateIconDoubleHorizontalText_SearchAction(
-     * mContext, "yellow", Color.YELLOW));
-     */
-    public static SearchTarget generateIcoHorizontalText_usingSearchAction(
-            Context context, String id, int color) {
-        SearchTarget.Builder builder =
-                new SearchTarget.Builder(ACTION, ICON_HORIZONTAL_TEXT, id)
-                        .setPackageName(PACKAGE2) /* required */
-                        .setUserHandle(USERHANDLE); /* required */
-
-        Intent intent = new Intent("com.google.android.googlequicksearchbox.GENERIC_QUERY");
-        intent.putExtra("query", "weather");
-        intent.putExtra("full_screen", false);
-        PendingIntent pendingIntent =
-                PendingIntent.getActivity(
-                        context, 1, intent,
-                        PendingIntent.FLAG_UPDATE_CURRENT);
-
-        Bitmap bitmap = Bitmap.createBitmap(200, 200, Bitmap.Config.ARGB_8888);
-        Canvas canvas = new Canvas(bitmap);
-        canvas.drawColor(color);
-        Icon icon = Icon.createWithAdaptiveBitmap(bitmap);
-
-        Bundle b = new Bundle();
-        b.putBoolean(BUNDLE_EXTRA_SHOULD_START_FOR_RESULT, true);
-        b.putBoolean(BUNDLE_EXTRA_BADGE_WITH_PACKAGE, true);
-        b.putBoolean(BUNDLE_EXTRA_PRIMARY_ICON_FROM_TITLE, true);
-
-        builder.setSearchAction(new SearchAction.Builder(id, id + TITLE)
-                .setSubtitle(id + SUBTITLE)
-                .setPendingIntent(pendingIntent)
-                .setIcon(icon)
-                .setExtras(b)
-                .build());
-        return builder.build();
-    }
-
-    /**
-     * Inside SearchServicePipeline, add following samples to test the search target.
-     *
-     * targets.add(SearchTargetUtil.generateThumbnail_SearchAction("blue", Color.BLUE));
-     * targets.add(SearchTargetUtil.generateThumbnail_SearchAction("red", Color.RED));
-     * targets.add(SearchTargetUtil.generateThumbnail_SearchAction("green", Color.GREEN));
-     */
-    public static SearchTarget generateThumbnail_usingSearchAction(String id, int color) {
-        SearchTarget.Builder builder =
-                new SearchTarget.Builder(SCREENSHOT, THUMBNAIL, id)
-                        .setPackageName(PACKAGE2) /* required */
-                        .setUserHandle(USERHANDLE); /* required */
-
-
-        Intent intent = new Intent(Intent.ACTION_VIEW)
-                .setData(Uri.parse("uri blah blah"))
-                .setType("image/*")
-                .addFlags(Intent.FLAG_GRANT_READ_URI_PERMISSION)
-                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
-
-        Bitmap bitmap = Bitmap.createBitmap(1000, 500, Bitmap.Config.ARGB_8888);
-        Canvas canvas = new Canvas(bitmap);
-        canvas.drawColor(color);
-        Icon icon = Icon.createWithBitmap(bitmap);
-
-        builder.setSearchAction(new SearchAction.Builder(id, TITLE)
-                .setSubtitle(SUBTITLE)
-                .setIcon(icon)
-                .setIntent(intent)
-                .build());
-        return builder.build();
-
-    }
-
-    /**
-     * Generate SearchTargetUtil for SMALL_ICON_HORIZONTAL_TEXT layout type.
-     *
-     * targets.add(SearchTargetUtil.generateIconHorizontalText_SearchAction(
-     * mContext, "red", Color.RED));
-     * targets.add(SearchTargetUtil.generateIconHorizontalText_SearchAction(
-     * mContext, "yellow", Color.YELLOW));
-     */
-    public static SearchTarget generateSmallIconHorizontalText_usingSearchAction(
-            Context context, String id, int color) {
-        String title = "Ask the assistant";
-        String fallbackQuery = "sourdough bread";
-        SearchTarget.Builder builder =
-                new SearchTarget.Builder(SUGGEST, SMALL_ICON_HORIZONTAL_TEXT, id)
-                        .setPackageName(PACKAGE2) /* required */
-                        .setUserHandle(USERHANDLE); /* required */
-
-        Intent intent3 = new Intent("com.google.android.googlequicksearchbox.GENERIC_QUERY");
-        intent3.putExtra("query", fallbackQuery);
-        intent3.putExtra("full_screen", false);
-        PendingIntent pendingIntent3 =
-                PendingIntent.getActivity(
-                        context, 1, intent3,
-                        PendingIntent.FLAG_UPDATE_CURRENT);
-
-        Bitmap bitmap = Bitmap.createBitmap(200, 200, Bitmap.Config.ARGB_8888);
-        Canvas canvas = new Canvas(bitmap);
-        canvas.drawColor(color);
-        Icon icon = Icon.createWithAdaptiveBitmap(bitmap);
-
-        Bundle extra = new Bundle();
-        extra.putBoolean(BUNDLE_EXTRA_SHOULD_START_FOR_RESULT, true);
-
-        SearchAction searchAction = new SearchAction.Builder(id, title)
-                .setSubtitle(fallbackQuery)
-                .setPendingIntent(pendingIntent3)
-                .setIcon(icon)
-                .setExtras(extra)
-                .build();
-        return builder.setSearchAction(searchAction).build();
-    }
-
-    public static SearchTarget generateDivider() {
-        SearchTarget.Builder builder =
-                new SearchTarget.Builder(SUGGEST, DIVIDER, "divider")
-                        .setPackageName("") /* required but not used*/
-                        .setUserHandle(USERHANDLE); /* required */
-        return builder.build();
-    }
-
-
-    /**
-     * Generate SearchTargetUtil for ICON_DOUBLE_HORIZONTAL_TEXT layout type.
-     */
-    public static SearchTarget generateIconDoubleHorizontalText_ShortcutInfo(Context context) {
-        String id = "23456";
-        SearchTarget.Builder builder =
-                new SearchTarget.Builder(ResultType.SHORTCUT, SMALL_ICON_HORIZONTAL_TEXT, id)
-                        .setPackageName("com.google.android.gm") /* required */
-                        .setUserHandle(UserHandle.CURRENT); /* required */
-
-        builder.setShortcutInfo(new ShortcutInfo.Builder(context, id).build());
-        return builder.build();
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 6098b52..d98e792 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -21,12 +21,11 @@
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
+import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.LauncherState.OVERVIEW;
 import static com.android.launcher3.LauncherState.OVERVIEW_MODAL_TASK;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.sendCustomAccessibilityEvent;
-import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.EXTENDED_CONTAINERS;
-import static com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers.ContainerCase.DEVICE_SEARCH_RESULT_CONTAINER;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_APP_LAUNCH_TAP;
 import static com.android.launcher3.testing.TestProtocol.HINT_STATE_ORDINAL;
 import static com.android.launcher3.testing.TestProtocol.OVERVIEW_STATE_ORDINAL;
@@ -45,8 +44,6 @@
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.Workspace;
-import com.android.launcher3.allapps.AllAppsContainerView;
-import com.android.launcher3.allapps.search.SearchAdapterProvider;
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.appprediction.PredictionRowView;
 import com.android.launcher3.hybridhotseat.HotseatPredictionController;
@@ -56,7 +53,6 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.popup.SystemShortcut;
-import com.android.launcher3.search.DeviceSearchAdapterProvider;
 import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory;
 import com.android.launcher3.uioverrides.states.QuickstepAtomicAnimationFactory;
 import com.android.launcher3.uioverrides.touchcontrollers.NavBarToHomeTouchController;
@@ -85,7 +81,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.stream.Stream;
 
 public class QuickstepLauncher extends BaseQuickstepLauncher {
@@ -108,13 +103,11 @@
 
     @Override
     protected void logAppLaunch(ItemInfo info, InstanceId instanceId) {
-        // If the app launch is from DeviceSearchResultContainer then add the InstanceId from
-        // LiveSearchManager to recreate the AllApps search session on the server side.
-        Optional<InstanceId> logInstanceId = this.getLiveSearchManager().getLogInstanceId();
-        if (info.getContainerInfo().getContainerCase() == EXTENDED_CONTAINERS
-                && info.getContainerInfo().getExtendedContainers().getContainerCase()
-                == DEVICE_SEARCH_RESULT_CONTAINER && logInstanceId.isPresent()) {
-            instanceId = logInstanceId.get();
+        // If the app launch is from any of the surfaces in AllApps then add the InstanceId from
+        // LiveSearchManager to recreate the AllApps session on the server side.
+        if (mAllAppsSessionLogId != null && ALL_APPS.equals(
+                getStateManager().getCurrentStableState())) {
+            instanceId = mAllAppsSessionLogId;
         }
 
         StatsLogger logger = getStatsLogManager()
@@ -278,11 +271,6 @@
     }
 
     @Override
-    public SearchAdapterProvider createSearchAdapterProvider(AllAppsContainerView appsView) {
-        return new DeviceSearchAdapterProvider(this, appsView);
-    }
-
-    @Override
     public TouchController[] createTouchControllers() {
         Mode mode = SysUINavigationMode.getMode(this);
 
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 36b51cd..feeee50 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -1071,9 +1071,7 @@
                     && runningTaskTarget != null
                     && runningTaskTarget.pictureInPictureParams != null
                     && TaskInfoCompat.isAutoEnterPipEnabled(
-                            runningTaskTarget.pictureInPictureParams)
-                    && TaskInfoCompat.getPipSourceRectHint(
-                            runningTaskTarget.pictureInPictureParams) != null;
+                            runningTaskTarget.pictureInPictureParams);
             if (mIsSwipingPipToHome) {
                 mSwipePipToHomeAnimator = getSwipePipToHomeAnimator(
                         homeAnimFactory, runningTaskTarget, start);
@@ -1176,8 +1174,10 @@
             swipePipToHomeAnimator.setFromRotation(mTaskViewSimulator, windowRotation);
         }
         swipePipToHomeAnimator.addListener(new AnimatorListenerAdapter() {
+            private boolean mHasAnimationEnded;
             @Override
             public void onAnimationStart(Animator animation) {
+                if (mHasAnimationEnded) return;
                 // Ensure Launcher ends in NORMAL state, we intentionally skip other callbacks
                 // since they are not relevant in this swipe-pip-to-home case.
                 homeAnimFactory.createActivityAnimationToHome().dispatchOnStart();
@@ -1185,6 +1185,8 @@
 
             @Override
             public void onAnimationEnd(Animator animation) {
+                if (mHasAnimationEnded) return;
+                mHasAnimationEnded = true;
                 if (mRecentsAnimationController == null) {
                     // If the recents animation is interrupted, we still end the running
                     // animation (not canceled) so this is still called. In that case, we can
diff --git a/quickstep/src/com/android/quickstep/ImageActionsApi.java b/quickstep/src/com/android/quickstep/ImageActionsApi.java
index cb4d53a..8cb64c2 100644
--- a/quickstep/src/com/android/quickstep/ImageActionsApi.java
+++ b/quickstep/src/com/android/quickstep/ImageActionsApi.java
@@ -64,6 +64,20 @@
      */
     @UiThread
     public void shareWithExplicitIntent(@Nullable Rect crop, Intent intent) {
+        addImageAndSendIntent(crop, intent, false);
+    }
+
+    /**
+     * Share the image this api was constructed with using the provided intent. The implementation
+     * should set the intent's data field to the URI pointing to the image.
+     */
+    @UiThread
+    public void shareAsDataWithExplicitIntent(@Nullable Rect crop, Intent intent) {
+        addImageAndSendIntent(crop, intent, true);
+    }
+
+    @UiThread
+    private void addImageAndSendIntent(@Nullable Rect crop, Intent intent, boolean setData) {
         if (mBitmapSupplier.get() == null) {
             Log.e(TAG, "No snapshot available, not starting share.");
             return;
@@ -71,12 +85,14 @@
 
         UI_HELPER_EXECUTOR.execute(() -> persistBitmapAndStartActivity(mContext,
                 mBitmapSupplier.get(), crop, intent, (uri, intentForUri) -> {
-                    intentForUri
-                            .addFlags(FLAG_GRANT_READ_URI_PERMISSION)
-                            .putExtra(EXTRA_STREAM, uri);
+                    intentForUri.addFlags(FLAG_GRANT_READ_URI_PERMISSION);
+                    if (setData) {
+                        intentForUri.setData(uri);
+                    } else {
+                        intentForUri.putExtra(EXTRA_STREAM, uri);
+                    }
                     return new Intent[]{intentForUri};
                 }, TAG));
-
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
index b258a10..6a32dc4 100644
--- a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
@@ -30,7 +30,6 @@
 import android.graphics.Point;
 import android.graphics.RectF;
 import android.util.Log;
-import android.util.SparseArray;
 import android.view.MotionEvent;
 import android.view.Surface;
 
@@ -40,6 +39,9 @@
 import com.android.launcher3.util.DisplayController.Info;
 
 import java.io.PrintWriter;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
 
 /**
  * Maintains state for supporting nav bars and tracking their gestures in multiple orientations.
@@ -51,6 +53,37 @@
  */
 class OrientationTouchTransformer {
 
+    class CurrentDisplay {
+        public Point size;
+        public int rotation;
+
+        CurrentDisplay() {
+            this.size = new Point(0, 0);
+            this.rotation = 0;
+        }
+
+        CurrentDisplay(Point size, int rotation) {
+            this.size = size;
+            this.rotation = rotation;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+
+            CurrentDisplay display = (CurrentDisplay) o;
+            if (rotation != display.rotation) return false;
+
+            return Objects.equals(size, display.size);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(size, rotation);
+        }
+    };
+
     private static final String TAG = "OrientationTouchTransformer";
     private static final boolean DEBUG = false;
     private static final int MAX_ORIENTATIONS = 4;
@@ -60,11 +93,12 @@
     private final Matrix mTmpMatrix = new Matrix();
     private final float[] mTmpPoint = new float[2];
 
-    private SparseArray<OrientationRectF> mSwipeTouchRegions = new SparseArray<>(MAX_ORIENTATIONS);
+    private Map<CurrentDisplay, OrientationRectF> mSwipeTouchRegions =
+            new HashMap<CurrentDisplay, OrientationRectF>();
     private final RectF mAssistantLeftRegion = new RectF();
     private final RectF mAssistantRightRegion = new RectF();
     private final RectF mOneHandedModeRegion = new RectF();
-    private int mCurrentDisplayRotation;
+    private CurrentDisplay mCurrentDisplay = new CurrentDisplay();
     private int mNavBarGesturalHeight;
     private int mNavBarLargerGesturalHeight;
     private boolean mEnableMultipleRegions;
@@ -147,21 +181,22 @@
      * @see #enableMultipleRegions(boolean, Info)
      */
     void createOrAddTouchRegion(Info info) {
-        mCurrentDisplayRotation = info.rotation;
+        mCurrentDisplay = new CurrentDisplay(info.realSize, info.rotation);
+
         if (mQuickStepStartingRotation > QUICKSTEP_ROTATION_UNINITIALIZED
-                && mCurrentDisplayRotation == mQuickStepStartingRotation) {
+                && mCurrentDisplay.rotation == mQuickStepStartingRotation) {
             // User already was swiping and the current screen is same rotation as the starting one
             // Remove active nav bars in other rotations except for the one we started out in
             resetSwipeRegions(info);
             return;
         }
-        OrientationRectF region = mSwipeTouchRegions.get(mCurrentDisplayRotation);
+        OrientationRectF region = mSwipeTouchRegions.get(mCurrentDisplay);
         if (region != null) {
             return;
         }
 
         if (mEnableMultipleRegions) {
-            mSwipeTouchRegions.put(mCurrentDisplayRotation, createRegionForDisplay(info));
+            mSwipeTouchRegions.put(mCurrentDisplay, createRegionForDisplay(info));
         } else {
             resetSwipeRegions(info);
         }
@@ -208,31 +243,31 @@
      */
     private void resetSwipeRegions(Info region) {
         if (DEBUG) {
-            Log.d(TAG, "clearing all regions except rotation: " + mCurrentDisplayRotation);
+            Log.d(TAG, "clearing all regions except rotation: " + mCurrentDisplay.rotation);
         }
 
-        mCurrentDisplayRotation = region.rotation;
-        OrientationRectF regionToKeep = mSwipeTouchRegions.get(mCurrentDisplayRotation);
+        mCurrentDisplay = new CurrentDisplay(region.realSize, region.rotation);
+        OrientationRectF regionToKeep = mSwipeTouchRegions.get(mCurrentDisplay);
         if (regionToKeep == null) {
             regionToKeep = createRegionForDisplay(region);
         }
         mSwipeTouchRegions.clear();
-        mSwipeTouchRegions.put(mCurrentDisplayRotation, regionToKeep);
+        mSwipeTouchRegions.put(mCurrentDisplay, regionToKeep);
         updateAssistantRegions(regionToKeep);
     }
 
     private void resetSwipeRegions() {
-        OrientationRectF regionToKeep = mSwipeTouchRegions.get(mCurrentDisplayRotation);
+        OrientationRectF regionToKeep = mSwipeTouchRegions.get(mCurrentDisplay);
         mSwipeTouchRegions.clear();
         if (regionToKeep != null) {
-            mSwipeTouchRegions.put(mCurrentDisplayRotation, regionToKeep);
+            mSwipeTouchRegions.put(mCurrentDisplay, regionToKeep);
             updateAssistantRegions(regionToKeep);
         }
     }
 
     private OrientationRectF createRegionForDisplay(Info display) {
         if (DEBUG) {
-            Log.d(TAG, "creating rotation region for: " + mCurrentDisplayRotation);
+            Log.d(TAG, "creating rotation region for: " + mCurrentDisplay.rotation);
         }
 
         Point size = display.realSize;
@@ -341,7 +376,9 @@
                 }
 
                 for (int i = 0; i < MAX_ORIENTATIONS; i++) {
-                    OrientationRectF rect = mSwipeTouchRegions.get(i);
+                    CurrentDisplay display = new CurrentDisplay(mCurrentDisplay.size, i);
+                    OrientationRectF rect = mSwipeTouchRegions.get(display);
+
                     if (TestProtocol.sDebugTracing) {
                         Log.d(TestProtocol.NO_SWIPE_TO_HOME, "transform:DOWN, rect=" + rect);
                     }
@@ -355,7 +392,7 @@
                         mLastRectTouched = rect;
                         mActiveTouchRotation = rect.mRotation;
                         if (mEnableMultipleRegions
-                                && mCurrentDisplayRotation == mActiveTouchRotation) {
+                                && mCurrentDisplay.rotation == mActiveTouchRotation) {
                             // TODO(b/154580671) might make this block unnecessary
                             // Start a touch session for the default nav region for the display
                             mQuickStepStartingRotation = mLastRectTouched.mRotation;
@@ -378,8 +415,8 @@
         pw.println("  lastTouchedRegion=" + mLastRectTouched);
         pw.println("  multipleRegionsEnabled=" + mEnableMultipleRegions);
         StringBuilder regions = new StringBuilder("  currentTouchableRotations=");
-        for(int i = 0; i < mSwipeTouchRegions.size(); i++) {
-            OrientationRectF rectF = mSwipeTouchRegions.get(mSwipeTouchRegions.keyAt(i));
+        for (CurrentDisplay key: mSwipeTouchRegions.keySet()) {
+            OrientationRectF rectF = mSwipeTouchRegions.get(key);
             regions.append(rectF).append(" ");
         }
         pw.println(regions.toString());
@@ -417,12 +454,12 @@
 
         boolean applyTransform(MotionEvent event, boolean forceTransform) {
             mTmpMatrix.reset();
-            postDisplayRotation(deltaRotation(mCurrentDisplayRotation, mRotation),
+            postDisplayRotation(deltaRotation(mCurrentDisplay.rotation, mRotation),
                     mHeight, mWidth, mTmpMatrix);
             if (forceTransform) {
                 if (DEBUG) {
                     Log.d(TAG, "Transforming rotation due to forceTransform, "
-                            + "mCurrentRotation: " + mCurrentDisplayRotation
+                            + "mCurrentRotation: " + mCurrentDisplay.rotation
                             + "mRotation: " + mRotation);
                 }
                 event.transform(mTmpMatrix);
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index 93ebd5a..0d2c42e 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -122,12 +122,11 @@
      */
     public static class TaskOverlay<T extends OverviewActionsView> {
 
-        private final Context mApplicationContext;
+        protected final Context mApplicationContext;
         protected final TaskThumbnailView mThumbnailView;
 
         private T mActionsView;
-        private ImageActionsApi mImageApi;
-        private boolean mIsAllowedByPolicy;
+        protected ImageActionsApi mImageApi;
 
         protected TaskOverlay(TaskThumbnailView taskThumbnailView) {
             mApplicationContext = taskThumbnailView.getContext().getApplicationContext();
@@ -153,24 +152,8 @@
 
             if (thumbnail != null) {
                 getActionsView().updateDisabledFlags(DISABLED_ROTATED, rotated);
-                final boolean isAllowedByPolicy = thumbnail.isRealSnapshot;
-
-                getActionsView().setCallbacks(new OverlayUICallbacks() {
-                    @Override
-                    public void onShare() {
-                        if (isAllowedByPolicy) {
-                            endLiveTileMode(() -> mImageApi.startShareActivity(null));
-                        } else {
-                            showBlockedByPolicyMessage();
-                        }
-                    }
-
-                    @SuppressLint("NewApi")
-                    @Override
-                    public void onScreenshot() {
-                        endLiveTileMode(() -> saveScreenshot(task));
-                    }
-                });
+                boolean isAllowedByPolicy = thumbnail.isRealSnapshot;
+                getActionsView().setCallbacks(new OverlayUICallbacksImpl(isAllowedByPolicy, task));
             }
         }
 
@@ -193,7 +176,7 @@
          * Called to save screenshot of the task thumbnail.
          */
         @SuppressLint("NewApi")
-        private void saveScreenshot(Task task) {
+        protected void saveScreenshot(Task task) {
             if (mThumbnailView.isRealSnapshot()) {
                 mImageApi.saveScreenshot(mThumbnailView.getThumbnail(),
                         getTaskSnapshotBounds(), getTaskSnapshotInsets(), task.key);
@@ -222,6 +205,12 @@
         }
 
         /**
+         * Sets full screen progress to the task overlay.
+         */
+        public void setFullscreenProgress(float progress) {
+        }
+
+        /**
          * Gets the system shortcut for the screenshot that will be added to the task menu.
          */
         public SystemShortcut getScreenshotShortcut(BaseDraggingActivity activity,
@@ -251,7 +240,7 @@
             return mThumbnailView.getScaledInsets();
         }
 
-        private void showBlockedByPolicyMessage() {
+        protected void showBlockedByPolicyMessage() {
             Toast.makeText(
                     mThumbnailView.getContext(),
                     R.string.blocked_by_policy,
@@ -273,6 +262,29 @@
                 dismissTaskMenuView(mActivity);
             }
         }
+
+        protected class OverlayUICallbacksImpl implements OverlayUICallbacks {
+            protected final boolean mIsAllowedByPolicy;
+            protected final Task mTask;
+
+            public OverlayUICallbacksImpl(boolean isAllowedByPolicy, Task task) {
+                mIsAllowedByPolicy = isAllowedByPolicy;
+                mTask = task;
+            }
+
+            public void onShare() {
+                if (mIsAllowedByPolicy) {
+                    endLiveTileMode(() -> mImageApi.startShareActivity(null));
+                } else {
+                    showBlockedByPolicyMessage();
+                }
+            }
+
+            @SuppressLint("NewApi")
+            public void onScreenshot() {
+                endLiveTileMode(() -> saveScreenshot(mTask));
+            }
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 8ebea33..e243715 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -24,7 +24,6 @@
 import static com.android.launcher3.config.FeatureFlags.ENABLE_QUICKSTEP_LIVE_TILE;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.quickstep.GestureState.DEFAULT_STATE;
-import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_INPUT_MONITOR;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SYSUI_PROXY;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_TRACING_ENABLED;
 
@@ -283,9 +282,7 @@
             return;
         }
 
-        Bundle bundle = SystemUiProxy.INSTANCE.get(this).monitorGestureInput("swipe-up",
-                mDeviceState.getDisplayId());
-        mInputMonitorCompat = InputMonitorCompat.fromBundle(bundle, KEY_EXTRA_INPUT_MONITOR);
+        mInputMonitorCompat = new InputMonitorCompat("swipe-up", mDeviceState.getDisplayId());
         mInputEventReceiver = mInputMonitorCompat.getInputReceiver(Looper.getMainLooper(),
                 mMainChoreographer, this::onInputEvent);
 
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index d22496d..f9283a4 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.EXTENDED_CONTAINERS;
 import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.FOLDER;
 import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.SEARCH_RESULT_CONTAINER;
+import static com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers.ContainerCase.DEVICE_SEARCH_RESULT_CONTAINER;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WORKSPACE_SNAPSHOT;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__ALLAPPS;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__BACKGROUND;
@@ -38,6 +39,8 @@
 import com.android.launcher3.logger.LauncherAtom.FolderIcon;
 import com.android.launcher3.logger.LauncherAtom.FromState;
 import com.android.launcher3.logger.LauncherAtom.ToState;
+import com.android.launcher3.logger.LauncherAtomExtensions.DeviceSearchResultContainer;
+import com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers;
 import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.model.AllAppsList;
@@ -85,7 +88,7 @@
     }
 
     @Override
-    public StatsLogger logger() {
+    protected StatsLogger createLogger() {
         return new StatsCompatLogger();
     }
 
@@ -301,6 +304,14 @@
                 return info.getContainerInfo().getPredictedHotseatContainer().getCardinality();
             case SEARCH_RESULT_CONTAINER:
                 return info.getContainerInfo().getSearchResultContainer().getQueryLength();
+            case EXTENDED_CONTAINERS:
+                ExtendedContainers extendedCont = info.getContainerInfo().getExtendedContainers();
+                if (extendedCont.getContainerCase() == DEVICE_SEARCH_RESULT_CONTAINER) {
+                    DeviceSearchResultContainer deviceSearchResultCont = extendedCont
+                            .getDeviceSearchResultContainer();
+                    return deviceSearchResultCont.hasQueryLength() ? deviceSearchResultCont
+                            .getQueryLength() : -1;
+                }
             default:
                 return info.getFolderIcon().getCardinality();
         }
diff --git a/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java b/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java
index b316423..1544f00 100644
--- a/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java
+++ b/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java
@@ -27,10 +27,8 @@
 
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.Workspace;
-import com.android.launcher3.allapps.AllAppsInsetTransitionController;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.hybridhotseat.HotseatPredictionController;
-import com.android.launcher3.search.DeviceSearchEdu;
 import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.uioverrides.QuickstepLauncher;
@@ -43,7 +41,6 @@
  */
 public class QuickstepOnboardingPrefs extends OnboardingPrefs<QuickstepLauncher> {
 
-
     public QuickstepOnboardingPrefs(QuickstepLauncher launcher, SharedPreferences sharedPrefs) {
         super(launcher, sharedPrefs);
 
@@ -133,22 +130,5 @@
                 }
             });
         }
-
-        if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && !getBoolean(SEARCH_EDU_SEEN)) {
-            stateManager.addStateListener(new StateListener<LauncherState>() {
-                @Override
-                public void onStateTransitionStart(LauncherState toState) {
-                    if (toState == ALL_APPS) {
-                        AllAppsInsetTransitionController insetTransitionController =
-                                mLauncher.getAllAppsController().getInsetController();
-                        insetTransitionController.setSearchEduRunnable(() -> {
-                            DeviceSearchEdu.show(launcher);
-                            insetTransitionController.setSearchEduRunnable(null);
-                        });
-                        stateManager.removeStateListener(this);
-                    }
-                }
-            });
-        }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/util/StaggeredWorkspaceAnim.java b/quickstep/src/com/android/quickstep/util/StaggeredWorkspaceAnim.java
index 5bae3c7..932ff27 100644
--- a/quickstep/src/com/android/quickstep/util/StaggeredWorkspaceAnim.java
+++ b/quickstep/src/com/android/quickstep/util/StaggeredWorkspaceAnim.java
@@ -139,7 +139,7 @@
 
         addDepthAnimationForState(launcher, NORMAL, ALPHA_DURATION_MS);
 
-        mAnimators.play(launcher.getDragLayer().getScrim().createSysuiMultiplierAnim(0f, 1f)
+        mAnimators.play(launcher.getDragLayer().getSysUiScrim().createSysuiMultiplierAnim(0f, 1f)
                 .setDuration(ALPHA_DURATION_MS));
         mAnimators.addListener(new AnimatorListenerAdapter() {
             @Override
diff --git a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
index 378f25b..0ce5072 100644
--- a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
+++ b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
@@ -31,6 +31,7 @@
 import android.view.View;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.systemui.shared.pip.PipSurfaceTransactionHelper;
@@ -60,7 +61,7 @@
     /** for calculating the transform in {@link #onAnimationUpdate(ValueAnimator)} */
     private final RectEvaluator mRectEvaluator = new RectEvaluator(new Rect());
     private final RectEvaluator mInsetsEvaluator = new RectEvaluator(new Rect());
-    private final Rect mSourceHintRectInsets = new Rect();
+    private final Rect mSourceHintRectInsets;
     private final Rect mSourceInsets = new Rect();
 
     /** for rotation via {@link #setFromRotation(TaskViewSimulator, int)} */
@@ -89,7 +90,7 @@
     public SwipePipToHomeAnimator(int taskId,
             @NonNull ComponentName componentName,
             @NonNull SurfaceControl leash,
-            @NonNull Rect sourceRectHint,
+            @Nullable Rect sourceRectHint,
             @NonNull Rect appBounds,
             @NonNull Rect startBounds,
             @NonNull Rect destinationBounds,
@@ -104,10 +105,14 @@
         mDestinationBoundsAnimation.set(mDestinationBounds);
         mSurfaceTransactionHelper = new PipSurfaceTransactionHelper();
 
-        mSourceHintRectInsets.set(sourceRectHint.left - appBounds.left,
-                sourceRectHint.top - appBounds.top,
-                appBounds.right - sourceRectHint.right,
-                appBounds.bottom - sourceRectHint.bottom);
+        if (sourceRectHint == null) {
+            mSourceHintRectInsets = null;
+        } else {
+            mSourceHintRectInsets = new Rect(sourceRectHint.left - appBounds.left,
+                    sourceRectHint.top - appBounds.top,
+                    appBounds.right - sourceRectHint.right,
+                    appBounds.bottom - sourceRectHint.bottom);
+        }
 
         addListener(new AnimationSuccessListener() {
             @Override
@@ -168,34 +173,44 @@
         final float fraction = animator.getAnimatedFraction();
         final Rect bounds = mRectEvaluator.evaluate(fraction, mStartBounds,
                 mDestinationBoundsAnimation);
-        final Rect insets = mInsetsEvaluator.evaluate(fraction, mSourceInsets,
-                mSourceHintRectInsets);
         final SurfaceControl.Transaction tx =
                 PipSurfaceTransactionHelper.newSurfaceControlTransaction();
-        if (mFromRotation == Surface.ROTATION_90 || mFromRotation == Surface.ROTATION_270) {
-            final float degree, positionX, positionY;
-            if (mFromRotation == Surface.ROTATION_90) {
-                degree = -90 * fraction;
-                positionX = fraction * (mDestinationBoundsTransformed.left - mAppBounds.left)
-                        + mAppBounds.left;
-                positionY = fraction * (mDestinationBoundsTransformed.bottom - mAppBounds.top)
-                        + mAppBounds.top;
-            } else {
-                degree = 90 * fraction;
-                positionX = fraction * (mDestinationBoundsTransformed.right - mAppBounds.left)
-                        + mAppBounds.left;
-                positionY = fraction * (mDestinationBoundsTransformed.top - mAppBounds.top)
-                        + mAppBounds.top;
-            }
-            mSurfaceTransactionHelper.scaleAndRotate(tx, mLeash, mAppBounds, bounds, insets,
-                    degree, positionX, positionY);
+        if (mSourceHintRectInsets == null) {
+            // no source rect hint been set, directly scale the window down
+            onAnimationScale(fraction, tx, bounds);
         } else {
-            mSurfaceTransactionHelper.scaleAndCrop(tx, mLeash, mAppBounds, bounds, insets);
+            // scale and crop according to the source rect hint
+            onAnimationScaleAndCrop(fraction, tx, bounds);
         }
         mSurfaceTransactionHelper.resetCornerRadius(tx, mLeash);
         tx.apply();
     }
 
+    /** scale the window directly with no source rect hint being set */
+    private void onAnimationScale(float fraction, SurfaceControl.Transaction tx, Rect bounds) {
+        if (mFromRotation == Surface.ROTATION_90 || mFromRotation == Surface.ROTATION_270) {
+            final RotatedPosition rotatedPosition = getRotatedPosition(fraction);
+            mSurfaceTransactionHelper.scale(tx, mLeash, mAppBounds, bounds,
+                    rotatedPosition.degree, rotatedPosition.positionX, rotatedPosition.positionY);
+        } else {
+            mSurfaceTransactionHelper.scale(tx, mLeash, mAppBounds, bounds);
+        }
+    }
+
+    /** scale and crop the window with source rect hint */
+    private void onAnimationScaleAndCrop(float fraction, SurfaceControl.Transaction tx,
+            Rect bounds) {
+        final Rect insets = mInsetsEvaluator.evaluate(fraction, mSourceInsets,
+                mSourceHintRectInsets);
+        if (mFromRotation == Surface.ROTATION_90 || mFromRotation == Surface.ROTATION_270) {
+            final RotatedPosition rotatedPosition = getRotatedPosition(fraction);
+            mSurfaceTransactionHelper.scaleAndRotate(tx, mLeash, mAppBounds, bounds, insets,
+                    rotatedPosition.degree, rotatedPosition.positionX, rotatedPosition.positionY);
+        } else {
+            mSurfaceTransactionHelper.scaleAndCrop(tx, mLeash, mAppBounds, bounds, insets);
+        }
+    }
+
     public int getTaskId() {
         return mTaskId;
     }
@@ -217,4 +232,34 @@
         tx.apply();
         mHasAnimationEnded = true;
     }
+
+    private RotatedPosition getRotatedPosition(float fraction) {
+        final float degree, positionX, positionY;
+        if (mFromRotation == Surface.ROTATION_90) {
+            degree = -90 * fraction;
+            positionX = fraction * (mDestinationBoundsTransformed.left - mAppBounds.left)
+                    + mAppBounds.left;
+            positionY = fraction * (mDestinationBoundsTransformed.bottom - mAppBounds.top)
+                    + mAppBounds.top;
+        } else {
+            degree = 90 * fraction;
+            positionX = fraction * (mDestinationBoundsTransformed.right - mAppBounds.left)
+                    + mAppBounds.left;
+            positionY = fraction * (mDestinationBoundsTransformed.top - mAppBounds.top)
+                    + mAppBounds.top;
+        }
+        return new RotatedPosition(degree, positionX, positionY);
+    }
+
+    private static class RotatedPosition {
+        private final float degree;
+        private final float positionX;
+        private final float positionY;
+
+        private RotatedPosition(float degree, float positionX, float positionY) {
+            this.degree = degree;
+            this.positionX = positionX;
+            this.positionY = positionY;
+        }
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
index 6f16781..edce194 100644
--- a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
+++ b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
@@ -41,7 +41,6 @@
 import com.android.quickstep.BaseActivityInterface;
 import com.android.quickstep.views.RecentsView.ScrollState;
 import com.android.quickstep.views.TaskThumbnailView.PreviewPositionHelper;
-import com.android.quickstep.views.TaskView;
 import com.android.quickstep.views.TaskView.FullscreenDrawParams;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
@@ -92,7 +91,6 @@
 
     // TaskView properties
     private final FullscreenDrawParams mCurrentFullscreenParams;
-    private float mCurveScale = 1;
     public final AnimatedFloat taskPrimaryTranslation = new AnimatedFloat();
     public final AnimatedFloat taskSecondaryTranslation = new AnimatedFloat();
 
@@ -277,8 +275,6 @@
                     .getPrimaryValue(mTaskRect.left, mTaskRect.top);
             mScrollState.screenCenter = start + mScrollState.scroll + mScrollState.halfPageSize;
             mScrollState.updateInterpolation(mDp, start);
-            mCurveScale = TaskView.getCurveScaleForInterpolation(mDp,
-                    mScrollState.linearInterpolation);
         }
 
         float progress = Utilities.boundToRange(fullScreenProgress.value, 0, 1);
@@ -295,8 +291,7 @@
         mMatrix.postTranslate(insets.left, insets.top);
         mMatrix.postScale(scale, scale);
 
-        // Apply TaskView matrix: scale, translate, scroll
-        mMatrix.postScale(mCurveScale, mCurveScale, taskWidth / 2, taskHeight / 2);
+        // Apply TaskView matrix: translate, scroll
         mMatrix.postTranslate(mTaskRect.left, mTaskRect.top);
         mOrientationState.getOrientationHandler().set(mMatrix, MATRIX_POST_TRANSLATE,
                 taskPrimaryTranslation.value);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 38d488c..5d492ac 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1984,11 +1984,6 @@
             centerToOffscreenProgress = Utilities.mapRange(centerToOffscreenProgress,
                     distanceFromMidpoint / distanceToOffscreen, 1);
         }
-        // Find the task's scale based on its offscreen progress, then see how far it still needs to
-        // move to be completely offscreen.
-        Utilities.scaleRectFAboutCenter(taskPosition,
-                TaskView.getCurveScaleForInterpolation(mActivity.getDeviceProfile(),
-                        centerToOffscreenProgress));
         distanceToOffscreen = desiredLeft - taskPosition.left;
         // Finally, we need to account for RecentsView scale, because it moves tasks based on its
         // pivot. To do this, we move the task position to where it would be offscreen at scale = 1
@@ -2114,7 +2109,7 @@
             anim.play(ObjectAnimator.ofFloat(recentsView, FULLSCREEN_PROGRESS, 1));
         } else {
             // We are launching an adjacent task, so parallax the center and other adjacent task.
-            float displacementX = tv.getWidth() * (toScale - tv.getCurveScale());
+            float displacementX = tv.getWidth() * (toScale - 1f);
             float primaryTranslation = mIsRtl ? -displacementX : displacementX;
             anim.play(ObjectAnimator.ofFloat(getPageAt(centerTaskIndex),
                     mOrientationHandler.getPrimaryViewTranslate(), primaryTranslation));
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index 4aff7e3..e21bf76 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -24,7 +24,6 @@
 import android.animation.ObjectAnimator;
 import android.content.Context;
 import android.graphics.Rect;
-import android.graphics.drawable.Drawable;
 import android.util.AttributeSet;
 import android.view.Gravity;
 import android.view.MotionEvent;
@@ -35,7 +34,6 @@
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.BaseDraggingActivity;
-import com.android.launcher3.FastBitmapDrawable;
 import com.android.launcher3.R;
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.Interpolators;
@@ -46,7 +44,6 @@
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.quickstep.TaskOverlayFactory;
 import com.android.quickstep.TaskUtils;
-import com.android.quickstep.views.IconView.OnScaleUpdateListener;
 
 /**
  * Contains options for a recent task when long-pressing its icon.
@@ -55,42 +52,15 @@
 
     private static final Rect sTempRect = new Rect();
 
-    private final OnScaleUpdateListener mTaskViewIconScaleListener = new OnScaleUpdateListener() {
-        @Override
-        public void onScaleUpdate(float scale) {
-            final Drawable drawable = mTaskIcon.getDrawable();
-            if (drawable instanceof FastBitmapDrawable) {
-                if (scale != ((FastBitmapDrawable) drawable).getScale()) {
-                    mMenuIconDrawable.setScale(scale);
-                }
-            }
-        }
-    };
-
-    private final OnScaleUpdateListener mMenuIconScaleListener = new OnScaleUpdateListener() {
-        @Override
-        public void onScaleUpdate(float scale) {
-            final Drawable taskViewDrawable = mTaskView.getIconView().getDrawable();
-            if (taskViewDrawable instanceof FastBitmapDrawable) {
-                final float currentScale = ((FastBitmapDrawable) taskViewDrawable).getScale();
-                if (currentScale != scale) {
-                    ((FastBitmapDrawable) taskViewDrawable).setScale(scale);
-                }
-            }
-        }
-    };
-
     private static final int REVEAL_OPEN_DURATION = 150;
     private static final int REVEAL_CLOSE_DURATION = 100;
 
     private final float mThumbnailTopMargin;
     private BaseDraggingActivity mActivity;
     private TextView mTaskName;
-    private IconView mTaskIcon;
     private AnimatorSet mOpenCloseAnimator;
     private TaskView mTaskView;
     private LinearLayout mOptionLayout;
-    private FastBitmapDrawable mMenuIconDrawable;
 
     public TaskMenuView(Context context, AttributeSet attrs) {
         this(context, attrs, 0);
@@ -107,7 +77,6 @@
     protected void onFinishInflate() {
         super.onFinishInflate();
         mTaskName = findViewById(R.id.task_name);
-        mTaskIcon = findViewById(R.id.task_icon);
         mOptionLayout = findViewById(R.id.menu_option_layout);
     }
 
@@ -134,15 +103,6 @@
     }
 
     @Override
-    protected void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-
-        // Remove all scale listeners when menu is removed
-        mTaskView.getIconView().removeUpdateScaleListener(mTaskViewIconScaleListener);
-        mTaskIcon.removeUpdateScaleListener(mMenuIconScaleListener);
-    }
-
-    @Override
     protected boolean isOfType(int type) {
         return (type & TYPE_TASK_MENU) != 0;
     }
@@ -203,23 +163,9 @@
     }
 
     private void addMenuOptions(TaskView taskView) {
-        Drawable icon = taskView.getTask().icon.getConstantState().newDrawable();
-        mTaskIcon.setDrawable(icon);
-        mTaskIcon.setOnClickListener(v -> close(true));
         mTaskName.setText(TaskUtils.getTitle(getContext(), taskView.getTask()));
         mTaskName.setOnClickListener(v -> close(true));
 
-        // Set the icons to match scale by listening to each other's changes
-        mMenuIconDrawable = icon instanceof FastBitmapDrawable ? (FastBitmapDrawable) icon : null;
-        taskView.getIconView().addUpdateScaleListener(mTaskViewIconScaleListener);
-        mTaskIcon.addUpdateScaleListener(mMenuIconScaleListener);
-
-        // Move the icon and text up half an icon size to lay over the TaskView
-        LinearLayout.LayoutParams params =
-                (LinearLayout.LayoutParams) mTaskIcon.getLayoutParams();
-        params.topMargin = (int) -mThumbnailTopMargin;
-        mTaskIcon.setLayoutParams(params);
-
         TaskOverlayFactory.getEnabledShortcuts(taskView).forEach(this::addMenuOption);
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 0654ab1..a1b5533 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -44,7 +44,6 @@
 import android.animation.AnimatorListenerAdapter;
 import android.animation.AnimatorSet;
 import android.animation.ObjectAnimator;
-import android.animation.TimeInterpolator;
 import android.app.ActivityOptions;
 import android.content.Context;
 import android.content.Intent;
@@ -112,10 +111,6 @@
 
     private static final String TAG = TaskView.class.getSimpleName();
 
-    /** A curve of x from 0 to 1, where 0 is the center of the screen and 1 is the edge. */
-    private static final TimeInterpolator CURVE_INTERPOLATOR
-            = x -> (float) -Math.cos(x * Math.PI) / 2f + .5f;
-
     /**
      * The alpha of a black scrim on a page in the carousel as it leaves the screen.
      * In the resting position of the carousel, the adjacent pages have about half this scrim.
@@ -252,7 +247,6 @@
     private TaskMenuView mMenuView;
     private IconView mIconView;
     private final DigitalWellBeingToast mDigitalWellBeingToast;
-    private float mCurveScale;
     private float mFullscreenProgress;
     private float mScaleAtFullscreen = 1;
     private float mFullscreenScale = 1;
@@ -612,6 +606,9 @@
         boolean isRtl = getLayoutDirection() == LAYOUT_DIRECTION_RTL;
         LayoutParams snapshotParams = (LayoutParams) mSnapshotView.getLayoutParams();
         int thumbnailPadding = (int) getResources().getDimension(R.dimen.task_thumbnail_top_margin);
+        int taskIconMargin = (int) getResources().getDimension(R.dimen.task_icon_top_margin);
+        int taskIconHeight = (int) getResources().getDimension(R.dimen.task_thumbnail_icon_size);
+        int iconTopMargin = taskIconMargin - taskIconHeight + thumbnailPadding;
         LayoutParams iconParams = (LayoutParams) mIconView.getLayoutParams();
         switch (orientationHandler.getRotation()) {
             case ROTATION_90:
@@ -623,7 +620,8 @@
             case ROTATION_180:
                 iconParams.gravity = BOTTOM | CENTER_HORIZONTAL;
                 iconParams.bottomMargin = -thumbnailPadding;
-                iconParams.leftMargin = iconParams.topMargin = iconParams.rightMargin = 0;
+                iconParams.leftMargin = iconParams.rightMargin = 0;
+                iconParams.topMargin = iconTopMargin;
                 break;
             case ROTATION_270:
                 iconParams.gravity = (isRtl ? END : START) | CENTER_VERTICAL;
@@ -634,7 +632,8 @@
             case Surface.ROTATION_0:
             default:
                 iconParams.gravity = TOP | CENTER_HORIZONTAL;
-                iconParams.leftMargin = iconParams.topMargin = iconParams.rightMargin = 0;
+                iconParams.leftMargin = iconParams.rightMargin = 0;
+                iconParams.topMargin = iconTopMargin;
                 break;
         }
         mIconView.setLayoutParams(iconParams);
@@ -650,7 +649,7 @@
             progress = 1 - progress;
         }
         mFocusTransitionProgress = progress;
-        mSnapshotView.setDimAlphaMultipler(progress);
+        mSnapshotView.setDimAlphaMultipler(0);
         float iconScalePercentage = (float) SCALE_ICON_DURATION / DIM_ANIM_DURATION;
         float lowerClamp = invert ? 1f - iconScalePercentage : 0;
         float upperClamp = invert ? 1 : iconScalePercentage;
@@ -698,7 +697,6 @@
     }
 
     protected void resetViewTransforms() {
-        setCurveScale(1);
         // fullscreenTranslation and accumulatedTranslation should not be reset, as
         // resetViewTransforms is called during Quickswitch scrolling.
         mDismissTranslationX = mTaskOffsetTranslationX = mTaskResistanceTranslationX = 0f;
@@ -733,13 +731,6 @@
             return;
         }
 
-        float curveInterpolation =
-                CURVE_INTERPOLATOR.getInterpolation(scrollState.linearInterpolation);
-        float curveScaleForCurveInterpolation = getCurveScaleForCurveInterpolation(
-                mActivity.getDeviceProfile(), curveInterpolation);
-        mSnapshotView.setDimAlpha(curveInterpolation * MAX_PAGE_SCRIM_ALPHA);
-        setCurveScale(curveScaleForCurveInterpolation);
-
         float dwbBannerAlpha = Utilities.boundToRange(1.0f - 2 * scrollState.linearInterpolation,
                 0f, 1f);
         mDigitalWellBeingToast.updateBannerAlpha(dwbBannerAlpha);
@@ -821,20 +812,6 @@
     }
 
     /**
-     * How much to scale down pages near the edge of the screen according to linearInterpolation.
-     */
-    public static float getCurveScaleForInterpolation(DeviceProfile deviceProfile,
-            float linearInterpolation) {
-        float curveInterpolation = CURVE_INTERPOLATOR.getInterpolation(linearInterpolation);
-        return getCurveScaleForCurveInterpolation(deviceProfile, curveInterpolation);
-    }
-
-    private static float getCurveScaleForCurveInterpolation(DeviceProfile deviceProfile,
-            float curveInterpolation) {
-        return 1 - curveInterpolation * getEdgeScaleDownFactor(deviceProfile);
-    }
-
-    /**
      * How much to scale down pages near the edge of the screen.
      */
     public static float getEdgeScaleDownFactor(DeviceProfile deviceProfile) {
@@ -845,23 +822,14 @@
         }
     }
 
-    private void setCurveScale(float curveScale) {
-        mCurveScale = curveScale;
-        applyScale();
-    }
-
     private void setFullscreenScale(float fullscreenScale) {
         mFullscreenScale = fullscreenScale;
         applyScale();
     }
 
     private void applyScale() {
-        setScaleX(mCurveScale * mFullscreenScale);
-        setScaleY(mCurveScale * mFullscreenScale);
-    }
-
-    public float getCurveScale() {
-        return mCurveScale;
+        setScaleX(mFullscreenScale);
+        setScaleY(mFullscreenScale);
     }
 
     private void setDismissTranslationX(float x) {
@@ -1080,6 +1048,7 @@
         progress = Utilities.boundToRange(progress, 0, 1);
         mFullscreenProgress = progress;
         mIconView.setVisibility(progress < 1 ? VISIBLE : INVISIBLE);
+        getThumbnail().getTaskOverlay().setFullscreenProgress(progress);
 
         updateTaskScaling();
 
diff --git a/res/drawable/ic_expand_less.xml b/res/drawable/ic_expand_less.xml
new file mode 100644
index 0000000..8360cee
--- /dev/null
+++ b/res/drawable/ic_expand_less.xml
@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?android:attr/textColorHint">
+    <path
+        android:fillColor="#FF000000"
+        android:pathData="M18.59,16.41L20,15l-8,-8 -8,8 1.41,1.41L12,9.83"/>
+</vector>
diff --git a/res/drawable/ic_expand_more.xml b/res/drawable/ic_expand_more.xml
new file mode 100644
index 0000000..49e24f6
--- /dev/null
+++ b/res/drawable/ic_expand_more.xml
@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?android:attr/textColorHint">
+    <path
+        android:fillColor="#FF000000"
+        android:pathData="M5.41,7.59L4,9l8,8 8,-8 -1.41,-1.41L12,14.17"/>
+</vector>
diff --git a/quickstep/res/layout/search_result_thumbnail.xml b/res/drawable/widgets_tray_expand_button.xml
similarity index 64%
copy from quickstep/res/layout/search_result_thumbnail.xml
copy to res/drawable/widgets_tray_expand_button.xml
index 5062b76..8316e0f 100644
--- a/quickstep/res/layout/search_result_thumbnail.xml
+++ b/res/drawable/widgets_tray_expand_button.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="utf-8"?>
-<!-- Copyright (C) 2020 The Android Open Source Project
+<!-- 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.
@@ -13,7 +13,9 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<com.android.launcher3.search.SearchResultThumbnailView
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="125dp"
-    android:layout_height="125dp"/>
\ No newline at end of file
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:state_checked="true"
+        android:drawable="@drawable/ic_expand_less" />
+    <item android:state_checked="false"
+        android:drawable="@drawable/ic_expand_more" />
+</selector>
diff --git a/res/layout/widgets_full_sheet.xml b/res/layout/widgets_full_sheet.xml
index f507a88..6e7cf0f 100644
--- a/res/layout/widgets_full_sheet.xml
+++ b/res/layout/widgets_full_sheet.xml
@@ -13,7 +13,7 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<com.android.launcher3.widget.WidgetsFullSheet
+<com.android.launcher3.widget.picker.WidgetsFullSheet
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
@@ -27,7 +27,7 @@
         android:background="?android:attr/colorPrimary"
         android:elevation="4dp">
 
-        <com.android.launcher3.widget.WidgetsRecyclerView
+        <com.android.launcher3.widget.picker.WidgetsRecyclerView
             android:id="@+id/widgets_list_view"
             android:layout_width="match_parent"
             android:layout_height="match_parent"
@@ -49,4 +49,4 @@
             android:layout_alignParentTop="true"
             android:layout_marginEnd="@dimen/fastscroll_end_margin" />
     </com.android.launcher3.views.TopRoundedCornerView>
-</com.android.launcher3.widget.WidgetsFullSheet>
\ No newline at end of file
+</com.android.launcher3.widget.picker.WidgetsFullSheet>
\ No newline at end of file
diff --git a/res/layout/widgets_list_row_header.xml b/res/layout/widgets_list_row_header.xml
new file mode 100644
index 0000000..faff10c
--- /dev/null
+++ b/res/layout/widgets_list_row_header.xml
@@ -0,0 +1,71 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<com.android.launcher3.widget.picker.WidgetsListHeader xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:id="@+id/widgets_list_header"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:background="?android:attr/selectableItemBackground"
+    android:paddingVertical="20dp"
+    android:orientation="horizontal">
+
+    <ImageView
+        android:id="@+id/app_icon"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_marginHorizontal="16dp"
+        android:importantForAccessibility="no"
+        tools:src="@drawable/ic_corp"/>
+
+    <LinearLayout
+        android:layout_width="0dp"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_vertical"
+        android:layout_weight="1"
+        android:orientation="vertical"
+        android:focusable="true"
+        android:descendantFocusability="afterDescendants">
+
+        <TextView
+            android:id="@+id/app_title"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_gravity="center_vertical"
+            android:textColor="?android:attr/textColorPrimary"
+            android:textSize="16sp"
+            tools:text="App name" />
+
+        <TextView
+            android:id="@+id/app_subtitle"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            tools:text="n widgets" />
+
+    </LinearLayout>
+
+    <!-- This checkbox is not clickable. The outermost LinearLayout is responsible to handle all
+         click event and update the checkbox state. -->
+    <CheckBox
+        android:id="@+id/toggle"
+        android:layout_marginHorizontal="16dp"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_vertical"
+        android:layout_alignParentEnd="true"
+        android:clickable="false"
+        android:button="@drawable/widgets_tray_expand_button"/>
+
+</com.android.launcher3.widget.picker.WidgetsListHeader>
\ No newline at end of file
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 0adde81..587df6d 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -195,4 +195,8 @@
         <attr name="android:name" />
         <attr name="android:id" />
     </declare-styleable>
+
+    <declare-styleable name="WidgetsListRowHeader">
+        <attr name="appIconSize" format="dimension" />
+    </declare-styleable>
 </resources>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 447c9ac..c30019b 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -54,6 +54,11 @@
     <string name="add_item_request_drag_hint">Touch &amp; hold to place manually</string>
     <!-- Button label to automatically add icon on home screen [CHAR_LIMIT=50] -->
     <string name="place_automatically">Add automatically</string>
+    <!-- Label for showing the number of widgets an app has in the full widgets picker. [CHAR_LIMIT=25] -->
+    <plurals name="widgets_tray_subtitle">
+        <item quantity="one"><xliff:g id="widget_count" example="1">%1$d</xliff:g> widget</item>
+        <item quantity="other"><xliff:g id="widget_count" example="2">%1$d</xliff:g> widgets</item>
+    </plurals>
 
     <!-- All Apps -->
     <!-- Search bar text in the apps view. [CHAR_LIMIT=50] -->
diff --git a/robolectric_tests/config/robolectric.properties b/robolectric_tests/config/robolectric.properties
index 0ac997f..4e811f3 100644
--- a/robolectric_tests/config/robolectric.properties
+++ b/robolectric_tests/config/robolectric.properties
@@ -1,4 +1,4 @@
-sdk=30
+sdk=29
 shadows= \
     com.android.launcher3.shadows.LShadowAppPredictionManager \
     com.android.launcher3.shadows.LShadowAppWidgetManager \
diff --git a/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
new file mode 100644
index 0000000..dbf4b3e
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
@@ -0,0 +1,41 @@
+/*
+ * 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.testing;
+
+import com.android.launcher3.BaseActivity;
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.views.BaseDragLayer;
+
+/** An empty activity for {@link android.app.Fragment}s, {@link android.view.View}s testing. */
+public class TestActivity extends BaseActivity implements ActivityContext {
+
+    private DeviceProfile mDeviceProfile;
+
+    @Override
+    public BaseDragLayer getDragLayer() {
+        return null;
+    }
+
+    @Override
+    public DeviceProfile getDeviceProfile() {
+        return mDeviceProfile;
+    }
+
+    public void setDeviceProfile(DeviceProfile deviceProfile) {
+        mDeviceProfile = deviceProfile;
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/ui/LauncherUIScrollTest.java b/robolectric_tests/src/com/android/launcher3/ui/LauncherUIScrollTest.java
index 957ae77..34cb2ad 100644
--- a/robolectric_tests/src/com/android/launcher3/ui/LauncherUIScrollTest.java
+++ b/robolectric_tests/src/com/android/launcher3/ui/LauncherUIScrollTest.java
@@ -38,7 +38,7 @@
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.LauncherLayoutBuilder.FolderBuilder;
 import com.android.launcher3.util.LauncherModelHelper;
-import com.android.launcher3.widget.WidgetsFullSheet;
+import com.android.launcher3.widget.picker.WidgetsFullSheet;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/robolectric_tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java b/robolectric_tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
deleted file mode 100644
index 5ab3106..0000000
--- a/robolectric_tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
+++ /dev/null
@@ -1,151 +0,0 @@
-/*
- * 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.
- */
-package com.android.launcher3.widget;
-
-import static org.mockito.Matchers.eq;
-import static org.mockito.Matchers.isNull;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.robolectric.Shadows.shadowOf;
-
-import android.appwidget.AppWidgetProviderInfo;
-import android.content.ComponentName;
-import android.content.Context;
-import android.graphics.Bitmap;
-import android.view.LayoutInflater;
-
-import androidx.recyclerview.widget.RecyclerView;
-
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.WidgetPreviewLoader;
-import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.model.data.PackageItemInfo;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
-
-import java.util.ArrayList;
-import java.util.Collections;
-
-@RunWith(RobolectricTestRunner.class)
-public class WidgetsListAdapterTest {
-
-    @Mock private LayoutInflater mMockLayoutInflater;
-    @Mock private WidgetPreviewLoader mMockWidgetCache;
-    @Mock private RecyclerView.AdapterDataObserver mListener;
-    @Mock private IconCache mIconCache;
-
-    private WidgetsListAdapter mAdapter;
-    private InvariantDeviceProfile mTestProfile;
-    private Context mContext;
-
-    @Before
-    public void setup() {
-        MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
-        mTestProfile = new InvariantDeviceProfile();
-        mTestProfile.numRows = 5;
-        mTestProfile.numColumns = 5;
-        mAdapter = new WidgetsListAdapter(mContext, mMockLayoutInflater, mMockWidgetCache,
-                mIconCache, null, null);
-        mAdapter.registerAdapterDataObserver(mListener);
-    }
-
-    @Test
-    public void test_notifyDataSetChanged() throws Exception {
-        mAdapter.setWidgets(generateSampleMap(1));
-        verify(mListener, times(1)).onChanged();
-    }
-
-    @Test
-    public void test_notifyItemInserted() throws Exception {
-        mAdapter.setWidgets(generateSampleMap(1));
-        mAdapter.setWidgets(generateSampleMap(2));
-        verify(mListener, times(1)).onChanged();
-        verify(mListener, times(1)).onItemRangeInserted(eq(1), eq(1));
-    }
-
-    @Test
-    public void test_notifyItemRemoved() throws Exception {
-        mAdapter.setWidgets(generateSampleMap(2));
-        mAdapter.setWidgets(generateSampleMap(1));
-        verify(mListener, times(1)).onChanged();
-        verify(mListener, times(1)).onItemRangeRemoved(eq(1), eq(1));
-    }
-
-    @Test
-    public void testNotifyItemChanged_PackageIconDiff() throws Exception {
-        mAdapter.setWidgets(generateSampleMap(1));
-        mAdapter.setWidgets(generateSampleMap(1));
-        verify(mListener, times(1)).onChanged();
-        verify(mListener, times(1)).onItemRangeChanged(eq(0), eq(1), isNull());
-    }
-
-    @Test
-    public void testNotifyItemChanged_widgetItemInfoDiff() throws Exception {
-        // TODO: same package name but item number changed
-    }
-
-    @Test
-    public void testNotifyItemInsertedRemoved_hodgepodge() throws Exception {
-        // TODO: insert and remove combined.          curMap
-        // newMap [A, C, D]                           [A, B, E]
-        // B - C < 0, removed B from index 1          [A, E]
-        // E - C > 0, C inserted to index 1           [A, C, E]
-        // E - D > 0, D inserted to index 2           [A, C, D, E]
-        // E - null = -1, E deleted from index 3      [A, C, D]
-    }
-
-    /**
-     * Helper method to generate the sample widget model map that can be used for the tests
-     * @param num the number of WidgetItem the map should contain
-     */
-    private ArrayList<WidgetListRowEntry> generateSampleMap(int num) {
-        ArrayList<WidgetListRowEntry> result = new ArrayList<>();
-        if (num <= 0) return result;
-        ShadowPackageManager spm = shadowOf(mContext.getPackageManager());
-
-        for (int i = 0; i < num; i++) {
-            ComponentName cn = new ComponentName("com.placeholder.apk" + i, "PlaceholderWidet");
-
-            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-            widgetInfo.provider = cn;
-            ReflectionHelpers.setField(widgetInfo, "providerInfo", spm.addReceiverIfNotPresent(cn));
-
-            WidgetItem wi = new WidgetItem(LauncherAppWidgetProviderInfo
-                    .fromProviderInfo(mContext, widgetInfo), mTestProfile, mIconCache);
-
-            PackageItemInfo pInfo = new PackageItemInfo(wi.componentName.getPackageName());
-            pInfo.title = pInfo.packageName;
-            pInfo.user = wi.user;
-            pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
-
-            result.add(new WidgetListRowEntry(pInfo, new ArrayList<>(Collections.singleton(wi))));
-        }
-
-        return result;
-    }
-}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
new file mode 100644
index 0000000..04797a6
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -0,0 +1,269 @@
+/*
+ * 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.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.os.UserHandle;
+
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListAdapter.WidgetListBaseRowEntryComparator;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsDiffReporterTest {
+    private static final String TEST_PACKAGE_PREFIX = "com.google.test";
+    private static final WidgetListBaseRowEntryComparator COMPARATOR =
+            new WidgetListBaseRowEntryComparator();
+
+    @Mock private IconCache mIconCache;
+    @Mock private RecyclerView.Adapter mAdapter;
+
+    private InvariantDeviceProfile mTestProfile;
+    private WidgetsDiffReporter mWidgetsDiffReporter;
+    private Context mContext;
+    private WidgetsListHeaderEntry mHeaderA;
+    private WidgetsListHeaderEntry mHeaderB;
+    private WidgetsListHeaderEntry mHeaderC;
+    private WidgetsListHeaderEntry mHeaderD;
+    private WidgetsListHeaderEntry mHeaderE;
+    private WidgetsListContentEntry mContentC;
+    private WidgetsListContentEntry mContentE;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
+                .getComponent().getPackageName())
+                .when(mIconCache).getTitleNoCache(any());
+
+        mContext = RuntimeEnvironment.application;
+        mWidgetsDiffReporter = new WidgetsDiffReporter(mIconCache, mAdapter);
+        mHeaderA = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "A",
+                /* appName= */ "A", /* numOfWidgets= */ 3);
+        mHeaderB = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "B",
+                /* appName= */ "B", /* numOfWidgets= */ 3);
+        mHeaderC = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "C",
+                /* appName= */ "C", /* numOfWidgets= */ 3);
+        mContentC = createWidgetsContentEntry(TEST_PACKAGE_PREFIX + "C",
+                /* appName= */ "C", /* numOfWidgets= */ 3);
+        mHeaderD = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "D",
+                /* appName= */ "D", /* numOfWidgets= */ 3);
+        mHeaderE = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "E",
+                /* appName= */ "E", /* numOfWidgets= */ 3);
+        mContentE = createWidgetsContentEntry(TEST_PACKAGE_PREFIX + "E",
+                /* appName= */ "E", /* numOfWidgets= */ 3);
+    }
+
+    @Test
+    public void listNotChanged_shouldNotInvokeAnyCallbacks() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderC));
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, currentList, COMPARATOR);
+
+        // THEN there is no adaptor callback.
+        verifyZeroInteractions(mAdapter);
+        // THEN the current list contains the same entries.
+        assertThat(currentList).containsExactly(mHeaderA, mHeaderB, mHeaderC);
+    }
+
+    @Test
+    public void headersOnly_emptyListToNonEmpty_shouldInvokeNotifyDataSetChanged() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>();
+
+        List<WidgetsListBaseEntry> newList = List.of(
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "A", "A", 3),
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "B", "B", 3),
+                createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "C", "C", 3));
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notifyDataSetChanged is called
+        verify(mAdapter).notifyDataSetChanged();
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersOnly_nonEmptyToEmptyList_shouldInvokeNotifyDataSetChanged() {
+        // GIVEN the current list has app headers [A, B, C].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderC));
+        // GIVEN the new list is empty.
+        List<WidgetsListBaseEntry> newList = List.of();
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notifyDataSetChanged is called.
+        verify(mAdapter).notifyDataSetChanged();
+        // THEN the current list isEmpty.
+        assertThat(currentList).isEmpty();
+    }
+
+    @Test
+    public void headersOnly_itemAddedAndRemovedInTheNewList_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, D].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mHeaderD));
+        // GIVEN the new list has app headers [A, C, E].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mHeaderC, mHeaderE);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN "B" is removed from position 1.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 1);
+        // THEN "D" is removed from position 2.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 2);
+        // THEN "C" is inserted at position 1.
+        verify(mAdapter).notifyItemInserted(/* position= */ 1);
+        // THEN "E" is inserted at position 2.
+        verify(mAdapter).notifyItemInserted(/* position= */ 2);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersContentsMix_itemAddedAndRemovedInTheNewList_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, E content].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mContentE));
+        // GIVEN the new list has app headers [A, C content, D].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mContentC, mHeaderD);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN "B" is removed from position 1.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 1);
+        // THEN "C content" is inserted at position 1.
+        verify(mAdapter).notifyItemInserted(/* position= */ 1);
+        // THEN "D" is inserted at position 2.
+        verify(mAdapter).notifyItemInserted(/* position= */ 2);
+        // THEN "E content" is removed from position 3.
+        verify(mAdapter).notifyItemRemoved(/* position= */ 3);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+    @Test
+    public void headersContentsMix_userInteractWithHeader_shouldInvokeCorrectCallbacks() {
+        // GIVEN the current list has app headers [A, B, E content].
+        ArrayList<WidgetsListBaseEntry> currentList = new ArrayList<>(
+                List.of(mHeaderA, mHeaderB, mContentE));
+        // GIVEN the new list has app headers [A, B, E content].
+        List<WidgetsListBaseEntry> newList = List.of(mHeaderA, mHeaderB, mContentE);
+        // GIVEN the user has interacted with B.
+        mHeaderB.setIsWidgetListShown(true);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notify "B" has been changed.
+        verify(mAdapter).notifyItemChanged(/* position= */ 1);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
+
+    private WidgetsListHeaderEntry createWidgetsHeaderEntry(String packageName, String appName,
+            int numOfWidgets) {
+        List<WidgetItem> widgetItems = generateWidgetItems(packageName, numOfWidgets);
+        PackageItemInfo pInfo = createPackageItemInfo(packageName, appName,
+                widgetItems.get(0).user);
+
+        return new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems);
+    }
+
+    private WidgetsListContentEntry createWidgetsContentEntry(String packageName, String appName,
+            int numOfWidgets) {
+        List<WidgetItem> widgetItems = generateWidgetItems(packageName, numOfWidgets);
+        PackageItemInfo pInfo = createPackageItemInfo(packageName, appName,
+                widgetItems.get(0).user);
+
+        return new WidgetsListContentEntry(pInfo, /* titleSectionName= */ "", widgetItems);
+    }
+
+    private PackageItemInfo createPackageItemInfo(String packageName, String appName,
+            UserHandle userHandle) {
+        PackageItemInfo pInfo = new PackageItemInfo(packageName);
+        pInfo.title = appName;
+        pInfo.user = userHandle;
+        pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+        return pInfo;
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            WidgetItem widgetItem = new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache);
+            widgetItems.add(widgetItem);
+        }
+        return widgetItems;
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
new file mode 100644
index 0000000..e94b253
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
@@ -0,0 +1,252 @@
+/*
+ * 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.
+ */
+package com.android.launcher3.widget.picker;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Matchers.isNull;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.view.LayoutInflater;
+
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.WidgetPreviewLoader;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsListAdapterTest {
+    private static final String TEST_PACKAGE_PLACEHOLDER = "com.google.test";
+
+    @Mock private LayoutInflater mMockLayoutInflater;
+    @Mock private WidgetPreviewLoader mMockWidgetCache;
+    @Mock private RecyclerView.AdapterDataObserver mListener;
+    @Mock private IconCache mIconCache;
+
+    private WidgetsListAdapter mAdapter;
+    private InvariantDeviceProfile mTestProfile;
+    private Context mContext;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+        mAdapter = new WidgetsListAdapter(mContext, mMockLayoutInflater, mMockWidgetCache,
+                mIconCache, null, null);
+        mAdapter.registerAdapterDataObserver(mListener);
+
+        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
+                        .getComponent().getPackageName())
+                .when(mIconCache).getTitleNoCache(any());
+    }
+
+    @Test
+    public void setWidgets_shouldNotifyDataSetChanged() {
+        mAdapter.setWidgets(generateSampleMap(1));
+
+        verify(mListener).onChanged();
+    }
+
+    @Test
+    public void setWidgets_withItemInserted_shouldNotifyItemInserted() {
+        mAdapter.setWidgets(generateSampleMap(1));
+        mAdapter.setWidgets(generateSampleMap(2));
+
+        verify(mListener).onItemRangeInserted(eq(1), eq(1));
+    }
+
+    @Test
+    public void setWidgets_withItemRemoved_shouldNotifyItemRemoved() {
+        mAdapter.setWidgets(generateSampleMap(2));
+        mAdapter.setWidgets(generateSampleMap(1));
+
+        verify(mListener).onItemRangeRemoved(eq(1), eq(1));
+    }
+
+    @Test
+    public void setWidgets_appIconChanged_shouldNotifyItemChanged() {
+        mAdapter.setWidgets(generateSampleMap(1));
+        mAdapter.setWidgets(generateSampleMap(1));
+
+        verify(mListener).onItemRangeChanged(eq(0), eq(1), isNull());
+    }
+
+    @Test
+    public void headerClick_expanded_shouldNotifyItemChange() {
+        // GIVEN a list of widgets entries:
+        // [com.google.test0, com.google.test0 content,
+        //  com.google.test1, com.google.test1 content,
+        //  com.google.test2, com.google.test2 content]
+        // The visible widgets entries: [com.google.test0, com.google.test1, com.google.test2].
+        mAdapter.setWidgets(generateSampleMap(3));
+
+        // WHEN com.google.test.1 header is expanded.
+        mAdapter.onHeaderClicked(/* isExpanded= */ true, TEST_PACKAGE_PLACEHOLDER + 1);
+
+        // THEN the visible entries list becomes:
+        // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
+        // com.google.test.1 content is inserted into position 2.
+        verify(mListener).onItemRangeInserted(eq(2), eq(1));
+    }
+
+    @Test
+    public void setWidgets_expandedApp_moreWidgets_shouldNotifyItemChangedWithWidgetItemInfoDiff() {
+        // GIVEN the adapter was first populated with com.google.test0 & com.google.test1. Each app
+        // has one widget.
+        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(2);
+        mAdapter.setWidgets(allEntries);
+        // GIVEN test com.google.test1 is expanded.
+        // Visible entries in the adapter are:
+        // [com.google.test0, com.google.test1, com.google.test1 content]
+        mAdapter.onHeaderClicked(/* isExpanded= */ true, TEST_PACKAGE_PLACEHOLDER + 1);
+        Mockito.reset(mListener);
+
+        // WHEN the adapter is updated with the same list of apps but com.google.test1 has 2 widgets
+        // now.
+        WidgetsListContentEntry testPackage1ContentEntry =
+                (WidgetsListContentEntry) allEntries.get(3);
+        WidgetItem widgetItem = testPackage1ContentEntry.mWidgets.get(0);
+        WidgetsListContentEntry newTestPackage1ContentEntry = new WidgetsListContentEntry(
+                testPackage1ContentEntry.mPkgItem,
+                testPackage1ContentEntry.mTitleSectionName, List.of(widgetItem, widgetItem));
+        allEntries.set(3, newTestPackage1ContentEntry);
+        mAdapter.setWidgets(allEntries);
+
+        // THEN the onItemRangeChanged is invoked for "com.google.test1 content" at index 2.
+        verify(mListener).onItemRangeChanged(eq(2), eq(1), isNull());
+    }
+
+    @Test
+    public void setWidgets_hodgepodge_shouldInvokeExpectedDataObserverCallbacks() {
+        // GIVEN a widgets entry list:
+        // Index:  0|   1      | 2|      3   | 4|     5    | 6|     7    | 8|     9    |
+        //        [A, A content, B, B content, C, C content, D, D content, E, E content]
+        List<WidgetsListBaseEntry> allAppsWithWidgets = generateSampleMap(5);
+        // GIVEN the current widgets list consist of [A, A content, B, B content, E, E content].
+        // GIVEN the visible widgets list consist of [A, B, E]
+        List<WidgetsListBaseEntry> currentList = List.of(
+                // A & A content
+                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1),
+                // B & B content
+                allAppsWithWidgets.get(2), allAppsWithWidgets.get(3),
+                // E & E content
+                allAppsWithWidgets.get(8), allAppsWithWidgets.get(9));
+        mAdapter.setWidgets(currentList);
+
+        // WHEN the widgets list is updated to [A, A content, C, C content, D, D content].
+        // WHEN the visible widgets list is updated to [A, C, D].
+        List<WidgetsListBaseEntry> newList = List.of(
+                // A & A content
+                allAppsWithWidgets.get(0), allAppsWithWidgets.get(1),
+                // C & C content
+                allAppsWithWidgets.get(4), allAppsWithWidgets.get(5),
+                // D & D content
+                allAppsWithWidgets.get(6), allAppsWithWidgets.get(7));
+        mAdapter.setWidgets(newList);
+
+        // Computation logic                           | [Intermediate list during computation]
+        // THEN B <> C < 0, removed B from index 1     | [A, E]
+        verify(mListener).onItemRangeRemoved(/* positionStart= */ 1, /* itemCount= */ 1);
+        // THEN E <> C > 0, C inserted to index 1      | [A, C, E]
+        verify(mListener).onItemRangeInserted(/* positionStart= */ 1, /* itemCount= */ 1);
+        // THEN E <> D > 0, D inserted to index 2      | [A, C, D, E]
+        verify(mListener).onItemRangeInserted(/* positionStart= */ 2, /* itemCount= */ 1);
+        // THEN E <> null = -1, E deleted from index 3 | [A, C, D]
+        verify(mListener).onItemRangeRemoved(/* positionStart= */ 3, /* itemCount= */ 1);
+    }
+
+    /**
+     * Generates a list of sample widget entries.
+     *
+     * <p>Each sample app has 1 widget only. An app is represented by 2 entries,
+     * {@link WidgetsListHeaderEntry} & {@link WidgetsListContentEntry}. Only
+     * {@link WidgetsListHeaderEntry} is always visible in the {@link WidgetsListAdapter}.
+     * {@link WidgetsListContentEntry} is only shown upon clicking the corresponding app's
+     * {@link WidgetsListHeaderEntry}. Only at most one {@link WidgetsListContentEntry} is shown at
+     * a time.
+     *
+     * @param num the number of apps that have widgets.
+     */
+    private ArrayList<WidgetsListBaseEntry> generateSampleMap(int num) {
+        ArrayList<WidgetsListBaseEntry> result = new ArrayList<>();
+        if (num <= 0) return result;
+
+        for (int i = 0; i < num; i++) {
+            String packageName = TEST_PACKAGE_PLACEHOLDER + i;
+
+            List<WidgetItem> widgetItems = generateWidgetItems(packageName, /* numOfWidgets= */ 1);
+
+            PackageItemInfo pInfo = new PackageItemInfo(packageName);
+            pInfo.title = pInfo.packageName;
+            pInfo.user = widgetItems.get(0).user;
+            pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+
+            result.add(new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems));
+            result.add(new WidgetsListContentEntry(pInfo, /* titleSectionName= */ "", widgetItems));
+        }
+
+        return result;
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            widgetItems.add(new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache));
+        }
+        return widgetItems;
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
new file mode 100644
index 0000000..ae5b9a5
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -0,0 +1,173 @@
+/*
+ * 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.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.widget.FrameLayout;
+import android.widget.TextView;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.R;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.testing.TestActivity;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListHeaderViewHolderBinder.OnHeaderClickListener;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.android.controller.ActivityController;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsListHeaderViewHolderBinderTest {
+    private static final String TEST_PACKAGE = "com.google.test";
+    private static final String APP_NAME = "Test app";
+
+    private Context mContext;
+    private WidgetsListHeaderViewHolderBinder mViewHolderBinder;
+    private InvariantDeviceProfile mTestProfile;
+    // Replace ActivityController with ActivityScenario, which is the recommended way for activity
+    // testing.
+    private ActivityController<TestActivity> mActivityController;
+    private TestActivity mTestActivity;
+    private FakeOnHeaderClickListener mFakeOnHeaderClickListener = new FakeOnHeaderClickListener();
+
+    @Mock
+    private IconCache mIconCache;
+    @Mock
+    private DeviceProfile mDeviceProfile;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        mActivityController = Robolectric.buildActivity(TestActivity.class);
+        mTestActivity = mActivityController.setup().get();
+        mTestActivity.setDeviceProfile(mDeviceProfile);
+
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            return componentWithLabel.getComponent().getShortClassName();
+        }).when(mIconCache).getTitleNoCache(any());
+
+        mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
+                LayoutInflater.from(mTestActivity),
+                mFakeOnHeaderClickListener);
+    }
+
+    @After
+    public void tearDown() {
+        mActivityController.destroy();
+    }
+
+    @Test
+    public void bindViewHolder_appWith3Widgets_shouldShowTheCorrectAppNameAndSubtitle() {
+        WidgetsListHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
+                new FrameLayout(mTestActivity));
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        WidgetsListHeaderEntry entry = generateSampleAppHeader(
+                APP_NAME,
+                TEST_PACKAGE,
+                /* numOfWidgets= */ 3);
+        mViewHolderBinder.bindViewHolder(viewHolder, entry);
+
+        TextView appTitle = widgetsListHeader.findViewById(R.id.app_title);
+        TextView appSubtitle = widgetsListHeader.findViewById(R.id.app_subtitle);
+        assertThat(appTitle.getText()).isEqualTo(APP_NAME);
+        assertThat(appSubtitle.getText()).isEqualTo("3 widgets");
+    }
+
+    private WidgetsListHeaderEntry generateSampleAppHeader(String appName, String packageName,
+            int numOfWidgets) {
+        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        appInfo.title = appName;
+        appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+
+        return new WidgetsListHeaderEntry(appInfo,
+                /* titleSectionName= */ "",
+                generateWidgetItems(packageName, numOfWidgets));
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            widgetItems.add(new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache));
+        }
+        return widgetItems;
+    }
+
+    private void assertWidgetCellWithLabel(View view, String label) {
+        assertThat(view).isInstanceOf(WidgetCell.class);
+        TextView widgetLabel = (TextView) view.findViewById(R.id.widget_name);
+        assertThat(widgetLabel.getText()).isEqualTo(label);
+    }
+
+    private final class FakeOnHeaderClickListener implements OnHeaderClickListener {
+
+        boolean mShowWidgets = false;
+        @Nullable  String mHeaderClickedPackage = null;
+
+        @Override
+        public void onHeaderClicked(boolean showWidgets, String packageName) {
+            mShowWidgets = showWidgets;
+            mHeaderClickedPackage = packageName;
+        }
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
new file mode 100644
index 0000000..ec9fde3
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinderTest.java
@@ -0,0 +1,177 @@
+/*
+ * 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 android.os.Looper.getMainLooper;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.View.OnClickListener;
+import android.view.View.OnLongClickListener;
+import android.widget.FrameLayout;
+import android.widget.TextView;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.R;
+import com.android.launcher3.WidgetPreviewLoader;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.testing.TestActivity;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.android.controller.ActivityController;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsListRowViewHolderBinderTest {
+    private static final String TEST_PACKAGE = "com.google.test";
+    private static final String APP_NAME = "Test app";
+
+    private Context mContext;
+    private WidgetsListRowViewHolderBinder mViewHolderBinder;
+    private InvariantDeviceProfile mTestProfile;
+    // Replace ActivityController with ActivityScenario, which is the recommended way for activity
+    // testing.
+    private ActivityController<TestActivity> mActivityController;
+    private TestActivity mTestActivity;
+
+    @Mock
+    private OnLongClickListener mOnLongClickListener;
+    @Mock
+    private OnClickListener mOnIconClickListener;
+    @Mock
+    private IconCache mIconCache;
+    @Mock
+    private WidgetPreviewLoader mWidgetPreviewLoader;
+    @Mock
+    private DeviceProfile mDeviceProfile;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        mActivityController = Robolectric.buildActivity(TestActivity.class);
+        mTestActivity = mActivityController.setup().get();
+        mTestActivity.setDeviceProfile(mDeviceProfile);
+
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            return componentWithLabel.getComponent().getShortClassName();
+        }).when(mIconCache).getTitleNoCache(any());
+
+        mViewHolderBinder = new WidgetsListRowViewHolderBinder(
+                mContext,
+                LayoutInflater.from(mTestActivity),
+                mOnIconClickListener,
+                mOnLongClickListener,
+                mWidgetPreviewLoader);
+    }
+
+    @After
+    public void tearDown() {
+        mActivityController.destroy();
+    }
+
+    @Test
+    public void bindViewHolder_appWith3Widgets_shouldHave3Widgets() {
+        WidgetsRowViewHolder viewHolder = mViewHolderBinder.newViewHolder(
+                new FrameLayout(mTestActivity));
+        WidgetsListContentEntry entry = generateSampleAppWithWidgets(
+                APP_NAME,
+                TEST_PACKAGE,
+                /* numOfWidgets= */ 3);
+        mViewHolderBinder.bindViewHolder(viewHolder, entry);
+        shadowOf(getMainLooper()).idle();
+
+        // THEN the cell container has 5 children: 3 widgets + 2 separators
+        // Index:        0       1        2       3        4
+        // View:  .SampleWidget0 | .SampleWidget1 | .SampleWidget2
+        assertThat(viewHolder.cellContainer.getChildCount()).isEqualTo(5);
+        // Widget 0 label is .SampleWidget0.
+        assertWidgetCellWithLabel(viewHolder.cellContainer.getChildAt(0), ".SampleWidget0");
+        // Widget 1 label is .SampleWidget1.
+        assertWidgetCellWithLabel(viewHolder.cellContainer.getChildAt(2), ".SampleWidget1");
+        // Widget 2 label is .SampleWidget2.
+        assertWidgetCellWithLabel(viewHolder.cellContainer.getChildAt(4), ".SampleWidget2");
+    }
+
+    private WidgetsListContentEntry generateSampleAppWithWidgets(String appName, String packageName,
+            int numOfWidgets) {
+        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        appInfo.title = appName;
+        appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
+
+        return new WidgetsListContentEntry(appInfo,
+                /* titleSectionName= */ "",
+                generateWidgetItems(packageName, numOfWidgets));
+    }
+
+    private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        ArrayList<WidgetItem> widgetItems = new ArrayList<>();
+        for (int i = 0; i < numOfWidgets; i++) {
+            ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                    packageManager.addReceiverIfNotPresent(cn));
+
+            widgetItems.add(new WidgetItem(
+                    LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
+                    mTestProfile, mIconCache));
+        }
+        return widgetItems;
+    }
+
+    private void assertWidgetCellWithLabel(View view, String label) {
+        assertThat(view).isInstanceOf(WidgetCell.class);
+        TextView widgetLabel = (TextView) view.findViewById(R.id.widget_name);
+        assertThat(widgetLabel.getText()).isEqualTo(label);
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
new file mode 100644
index 0000000..86df3f8
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -0,0 +1,173 @@
+/*
+ * 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.model;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsListContentEntryTest {
+    private static final String PACKAGE_NAME = "com.google.test";
+    private final PackageItemInfo mPackageItemInfo = new PackageItemInfo(PACKAGE_NAME);
+    private final ComponentName mWidget1 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget1");
+    private final ComponentName mWidget2 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget2");
+    private final ComponentName mWidget3 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget3");
+    private final Map<ComponentName, String> mWidgetsToLabels = new HashMap();
+
+    @Mock private IconCache mIconCache;
+
+    private Context mContext;
+    private InvariantDeviceProfile mTestProfile;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mWidgetsToLabels.put(mWidget1, "Cat");
+        mWidgetsToLabels.put(mWidget2, "Dog");
+        mWidgetsToLabels.put(mWidget3, "Bird");
+
+        mContext = RuntimeEnvironment.application;
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            return mWidgetsToLabels.get(componentWithLabel.getComponent());
+        }).when(mIconCache).getTitleNoCache(any());
+    }
+
+    @Test
+    public void unsortedWidgets_diffLabels_shouldSortWidgetItems() {
+        // GIVEN a list of widgets in unsorted order.
+        // Cat 2x3
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 3);
+        // Dog 2x3
+        WidgetItem widgetItem2 = createWidgetItem(mWidget2, /* spanX= */ 2, /* spanY= */ 3);
+        // Bird 2x3
+        WidgetItem widgetItem3 = createWidgetItem(mWidget3, /* spanX= */ 2, /* spanY= */ 3);
+
+        // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1, widgetItem2, widgetItem3));
+
+        // THEN the widgets list is sorted by their labels alphabetically: [Bird, Cat, Dog].
+        assertThat(widgetsListRowEntry.mWidgets)
+                .containsExactly(widgetItem3, widgetItem1, widgetItem2)
+                .inOrder();
+        assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
+    }
+
+    @Test
+    public void unsortedWidgets_sameLabels_differentSize_shouldSortWidgetItems() {
+        // GIVEN a list of widgets in unsorted order.
+        // Cat 3x3
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 3, /* spanY= */ 3);
+        // Cat 1x2
+        WidgetItem widgetItem2 = createWidgetItem(mWidget1, /* spanX= */ 1, /* spanY= */ 2);
+        // Cat 2x2
+        WidgetItem widgetItem3 = createWidgetItem(mWidget1, /* spanX= */ 2, /* spanY= */ 2);
+
+        // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1, widgetItem2, widgetItem3));
+
+        // THEN the widgets list is sorted by their gird sizes in an ascending order:
+        // [1x2, 2x2, 3x3].
+        assertThat(widgetsListRowEntry.mWidgets)
+                .containsExactly(widgetItem2, widgetItem3, widgetItem1)
+                .inOrder();
+        assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
+    }
+
+    @Test
+    public void unsortedWidgets_hodgepodge_shouldSortWidgetItems() {
+        // GIVEN a list of widgets in unsorted order.
+        // Cat 3x3
+        WidgetItem widgetItem1 = createWidgetItem(mWidget1, /* spanX= */ 3, /* spanY= */ 3);
+        // Cat 1x2
+        WidgetItem widgetItem2 = createWidgetItem(mWidget1, /* spanX= */ 1, /* spanY= */ 2);
+        // Dog 2x2
+        WidgetItem widgetItem3 = createWidgetItem(mWidget2, /* spanX= */ 2, /* spanY= */ 2);
+        // Bird 2x2
+        WidgetItem widgetItem4 = createWidgetItem(mWidget3, /* spanX= */ 2, /* spanY= */ 2);
+
+        // WHEN creates a WidgetsListRowEntry with the unsorted widgets.
+        WidgetsListContentEntry widgetsListRowEntry = new WidgetsListContentEntry(mPackageItemInfo,
+                /* titleSectionName= */ "T",
+                List.of(widgetItem1, widgetItem2, widgetItem3, widgetItem4));
+
+        // THEN the widgets list is first sorted by labels alphabetically. Then, for widgets with
+        // same labels, they are sorted by their gird sizes in an ascending order:
+        // [Bird 2x2, Cat 1x2, Cat 3x3, Dog 2x2]
+        assertThat(widgetsListRowEntry.mWidgets)
+                .containsExactly(widgetItem4, widgetItem2, widgetItem1, widgetItem3)
+                .inOrder();
+        assertThat(widgetsListRowEntry.mTitleSectionName).isEqualTo("T");
+        assertThat(widgetsListRowEntry.mPkgItem).isEqualTo(mPackageItemInfo);
+    }
+
+    private WidgetItem createWidgetItem(ComponentName componentName, int spanX, int spanY) {
+        String label = mWidgetsToLabels.get(componentName);
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+        widgetInfo.provider = componentName;
+        ReflectionHelpers.setField(widgetInfo, "providerInfo",
+                packageManager.addReceiverIfNotPresent(componentName));
+
+        LauncherAppWidgetProviderInfo launcherAppWidgetProviderInfo =
+                LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo);
+        launcherAppWidgetProviderInfo.spanX = spanX;
+        launcherAppWidgetProviderInfo.spanY = spanY;
+        launcherAppWidgetProviderInfo.label = label;
+
+        return new WidgetItem(launcherAppWidgetProviderInfo, mTestProfile, mIconCache);
+    }
+}
diff --git a/src/com/android/launcher3/BaseActivity.java b/src/com/android/launcher3/BaseActivity.java
index 5e50e27..8a03fac 100644
--- a/src/com/android/launcher3/BaseActivity.java
+++ b/src/com/android/launcher3/BaseActivity.java
@@ -141,7 +141,10 @@
         return mDeviceProfile;
     }
 
-    public final StatsLogManager getStatsLogManager() {
+    /**
+     * Returns {@link StatsLogManager} for user event logging.
+     */
+    public StatsLogManager getStatsLogManager() {
         if (mStatsLogManager == null) {
             mStatsLogManager = StatsLogManager.newInstance(this);
         }
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index b05a13a..edc7e9b 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -25,7 +25,6 @@
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ObjectAnimator;
-import android.animation.ValueAnimator;
 import android.content.Context;
 import android.content.res.ColorStateList;
 import android.content.res.TypedArray;
@@ -35,8 +34,6 @@
 import android.graphics.Paint;
 import android.graphics.Path;
 import android.graphics.PointF;
-import android.graphics.PorterDuff.Mode;
-import android.graphics.PorterDuffColorFilter;
 import android.graphics.Rect;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
@@ -53,7 +50,6 @@
 
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
-import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.Launcher.OnResumeCallback;
 import com.android.launcher3.accessibility.LauncherAccessibilityDelegate;
@@ -807,7 +803,7 @@
         if (mIcon != null
                 && mIcon instanceof PlaceHolderIconDrawable
                 && iconUpdateAnimationEnabled()) {
-            animateIconUpdate((PlaceHolderIconDrawable) mIcon, icon);
+            ((PlaceHolderIconDrawable) mIcon).animateIconUpdate(icon);
         }
 
         mDisableRelayout = false;
@@ -959,28 +955,6 @@
         }
     }
 
-    private static void animateIconUpdate(PlaceHolderIconDrawable oldIcon, Drawable newIcon) {
-        int placeholderColor = oldIcon.mPaint.getColor();
-        int originalAlpha = Color.alpha(placeholderColor);
-
-        ValueAnimator iconUpdateAnimation = ValueAnimator.ofInt(originalAlpha, 0);
-        iconUpdateAnimation.setDuration(ICON_UPDATE_ANIMATION_DURATION);
-        iconUpdateAnimation.addUpdateListener(valueAnimator -> {
-            int newAlpha = (int) valueAnimator.getAnimatedValue();
-            int newColor = ColorUtils.setAlphaComponent(placeholderColor, newAlpha);
-
-            newIcon.setColorFilter(new PorterDuffColorFilter(newColor, Mode.SRC_ATOP));
-        });
-        iconUpdateAnimation.addListener(new AnimatorListenerAdapter() {
-            @Override
-            public void onAnimationEnd(Animator animation) {
-                newIcon.setColorFilter(null);
-            }
-        });
-        iconUpdateAnimation.start();
-    }
-
-
     @Override
     public void decorate(int color) {
         mHighlightColor = color;
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 344ae0a..51e7c7d 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -40,6 +40,8 @@
 import static com.android.launcher3.dragndrop.DragLayer.ALPHA_INDEX_LAUNCHER_LOAD;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_BACKGROUND;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_HOME;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_ENTRY;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_EXIT;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ONRESUME;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ONSTOP;
 import static com.android.launcher3.model.ItemInstallQueue.FLAG_ACTIVITY_PAUSED;
@@ -124,6 +126,8 @@
 import com.android.launcher3.keyboard.ViewGroupFocusHelper;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logging.FileLog;
+import com.android.launcher3.logging.InstanceId;
+import com.android.launcher3.logging.InstanceIdSequence;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.ItemInstallQueue;
@@ -181,10 +185,10 @@
 import com.android.launcher3.widget.PendingAppWidgetHostView;
 import com.android.launcher3.widget.WidgetAddFlowHandler;
 import com.android.launcher3.widget.WidgetHostViewLoader;
-import com.android.launcher3.widget.WidgetListRowEntry;
 import com.android.launcher3.widget.WidgetManagerHelper;
-import com.android.launcher3.widget.WidgetsFullSheet;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.picker.WidgetsFullSheet;
 import com.android.systemui.plugins.OverlayPlugin;
 import com.android.systemui.plugins.PluginListener;
 import com.android.systemui.plugins.shared.LauncherExterns;
@@ -352,6 +356,13 @@
 
     private SafeCloseable mUserChangedCallbackCloseable;
 
+    // New InstanceId is assigned to mAllAppsSessionLogId for each AllApps sessions.
+    // When Launcher is not in AllApps state mAllAppsSessionLogId will be null.
+    // User actions within AllApps state are logged with this InstanceId, to recreate AllApps
+    // session on the server side.
+    protected InstanceId mAllAppsSessionLogId;
+    private LauncherState mPrevLauncherState;
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         Object traceToken = TraceHelper.INSTANCE.beginSection(ON_CREATE_EVT,
@@ -1027,6 +1038,7 @@
         }
         // When multiple pages are visible, show persistent page indicator
         mWorkspace.getPageIndicator().setShouldAutoHide(!state.hasFlag(FLAG_MULTI_PAGE));
+        mPrevLauncherState = mStateManager.getCurrentStableState();
     }
 
     @Override
@@ -1050,6 +1062,17 @@
             // Clear any rotation locks when going to normal state
             getRotationHelper().setCurrentStateRequest(REQUEST_NONE);
         }
+
+        if (ALL_APPS.equals(state)) {
+            // creates new instance ID since new all apps session is started.
+            mAllAppsSessionLogId = new InstanceIdSequence().newInstanceId();
+            getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_ENTRY);
+        } else if (ALL_APPS.equals(mPrevLauncherState)
+                // Check if mLogInstanceId is not null to make sure exit event is logged only once.
+                && mAllAppsSessionLogId != null) {
+            getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_EXIT);
+            mAllAppsSessionLogId = null;
+        }
     }
 
     @Override
@@ -2552,7 +2575,7 @@
     }
 
     @Override
-    public void bindAllWidgets(final ArrayList<WidgetListRowEntry> allWidgets) {
+    public void bindAllWidgets(final List<WidgetsListBaseEntry> allWidgets) {
         mPopupDataProvider.setAllWidgets(allWidgets);
     }
 
@@ -2818,4 +2841,9 @@
         public Configuration config;
         public Bitmap snapshot;
     }
+
+    @Override
+    public StatsLogManager getStatsLogManager() {
+        return super.getStatsLogManager().withDefaultInstanceId(mAllAppsSessionLogId);
+    }
 }
diff --git a/src/com/android/launcher3/LauncherRootView.java b/src/com/android/launcher3/LauncherRootView.java
index 51504ce..76c4518 100644
--- a/src/com/android/launcher3/LauncherRootView.java
+++ b/src/com/android/launcher3/LauncherRootView.java
@@ -4,12 +4,14 @@
 
 import android.annotation.TargetApi;
 import android.content.Context;
+import android.graphics.Canvas;
 import android.graphics.Rect;
 import android.os.Build;
 import android.util.AttributeSet;
 import android.view.ViewDebug;
 import android.view.WindowInsets;
 
+import com.android.launcher3.graphics.SysUiScrim;
 import com.android.launcher3.statemanager.StatefulActivity;
 
 import java.util.Collections;
@@ -31,6 +33,8 @@
     @ViewDebug.ExportedProperty(category = "launcher")
     private boolean mForceHideBackArrow;
 
+    private SysUiScrim mSysUiScrim;
+
     public LauncherRootView(Context context, AttributeSet attrs) {
         super(context, attrs);
         mActivity = StatefulActivity.fromContext(context);
@@ -89,6 +93,18 @@
         }
     }
 
+    public void setSysUiScrim(SysUiScrim scrim) {
+        mSysUiScrim = scrim;
+    }
+
+    @Override
+    protected void dispatchDraw(Canvas canvas) {
+        if (mSysUiScrim != null) {
+            mSysUiScrim.draw(canvas);
+        }
+        super.dispatchDraw(canvas);
+    }
+
     @Override
     protected void onLayout(boolean changed, int l, int t, int r, int b) {
         super.onLayout(changed, l, t, r, b);
@@ -119,4 +135,4 @@
 
         void onWindowVisibilityChanged(int visibility);
     }
-}
\ No newline at end of file
+}
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index be27039..dfdc53c 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -75,6 +75,37 @@
         public static final int ITEM_TYPE_SHORTCUT = 1;
 
         /**
+         * The favorite is a user created folder
+         */
+        public static final int ITEM_TYPE_FOLDER = 2;
+
+        /**
+         * The favorite is a widget
+         */
+        public static final int ITEM_TYPE_APPWIDGET = 4;
+
+        /**
+         * The favorite is a custom widget provided by the launcher
+         */
+        public static final int ITEM_TYPE_CUSTOM_APPWIDGET = 5;
+
+        /**
+         * The gesture is an application created deep shortcut
+         */
+        public static final int ITEM_TYPE_DEEP_SHORTCUT = 6;
+
+        /**
+         * Type of the item is recents task.
+         * TODO(hyunyoungs): move constants not related to Favorites DB to a better location.
+         */
+        public static final int ITEM_TYPE_TASK = 7;
+
+        /**
+         * The item is QSB
+         */
+        public static final int ITEM_TYPE_QSB = 8;
+
+        /**
          * The icon package name in Intent.ShortcutIconResource
          * <P>Type: TEXT</P>
          */
@@ -170,6 +201,7 @@
         public static final int CONTAINER_SHORTCUTS = -107;
         public static final int CONTAINER_SETTINGS = -108;
         public static final int CONTAINER_TASKSWITCHER = -109;
+        public static final int CONTAINER_QSB = -110;
 
         // Represents any of the extended containers implemented in non-AOSP variants.
         public static final int EXTENDED_CONTAINERS = -200;
@@ -195,6 +227,8 @@
                 case ITEM_TYPE_APPWIDGET: return "WIDGET";
                 case ITEM_TYPE_CUSTOM_APPWIDGET: return "CUSTOMWIDGET";
                 case ITEM_TYPE_DEEP_SHORTCUT: return "DEEPSHORTCUT";
+                case ITEM_TYPE_TASK: return "TASK";
+                case ITEM_TYPE_QSB: return "QSB";
                 default: return String.valueOf(type);
             }
         }
@@ -240,32 +274,6 @@
         public static final String PROFILE_ID = "profileId";
 
         /**
-         * The favorite is a user created folder
-         */
-        public static final int ITEM_TYPE_FOLDER = 2;
-
-        /**
-         * The favorite is a widget
-         */
-        public static final int ITEM_TYPE_APPWIDGET = 4;
-
-        /**
-         * The favorite is a custom widget provided by the launcher
-         */
-        public static final int ITEM_TYPE_CUSTOM_APPWIDGET = 5;
-
-        /**
-         * The gesture is an application created deep shortcut
-         */
-        public static final int ITEM_TYPE_DEEP_SHORTCUT = 6;
-
-        /**
-         * Type of the item is recents task.
-         * TODO(hyunyoungs): move constants not related to Favorites DB to a better location.
-         */
-        public static final int ITEM_TYPE_TASK = 7;
-
-        /**
          * The appWidgetId of the widget
          *
          * <P>Type: INTEGER</P>
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 8a45c81..2528c03 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -81,6 +81,7 @@
 import com.android.launcher3.folder.PreviewBackground;
 import com.android.launcher3.graphics.DragPreviewProvider;
 import com.android.launcher3.graphics.PreloadIconDrawable;
+import com.android.launcher3.graphics.WorkspaceDragScrim;
 import com.android.launcher3.icons.BitmapRenderer;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logging.StatsLogManager;
@@ -199,6 +200,9 @@
     private boolean mStripScreensOnPageStopMoving = false;
 
     private DragPreviewProvider mOutlineProvider = null;
+
+    private WorkspaceDragScrim mWorkspaceDragScrim;
+
     private boolean mWorkspaceFadeInAdjacentScreens;
 
     final WallpaperOffsetInterpolator mWallpaperOffset;
@@ -1161,6 +1165,19 @@
         }
     }
 
+    public void setWorkspaceDragScrim(WorkspaceDragScrim workspaceDragScrim) {
+        mWorkspaceDragScrim = workspaceDragScrim;
+    }
+
+    @Override
+    public void invalidate() {
+        // The workspace scrim may need to be re-rendered based on the workspace scroll
+        if (mWorkspaceDragScrim != null) {
+            mWorkspaceDragScrim.invalidate();
+        }
+        super.invalidate();
+    }
+
     @Override
     public void computeScroll() {
         super.computeScroll();
@@ -1706,10 +1723,9 @@
         if (dropOverView instanceof FolderIcon) {
             FolderIcon fi = (FolderIcon) dropOverView;
             if (fi.acceptDrop(d.dragInfo)) {
-                mStatsLogManager.logger().withItemInfo(d.dragInfo).withInstanceId(d.logInstanceId)
-                        .log(LauncherEvent.LAUNCHER_ITEM_DROP_COMPLETED);
+                mStatsLogManager.logger().withItemInfo(fi.mInfo).withInstanceId(d.logInstanceId)
+                        .log(LauncherEvent.LAUNCHER_ITEM_DROP_COMPLETED_ON_FOLDER_ICON);
                 fi.onDrop(d, false /* itemReturnedOnFailedDrop */);
-
                 // if the drag started here, we need to remove it from the workspace
                 if (!external) {
                     getParentCellLayoutForView(mDragInfo.cell).removeView(mDragInfo.cell);
@@ -2028,7 +2044,7 @@
         }
         // Invalidating the scrim will also force this CellLayout
         // to be invalidated so that it is highlighted if necessary.
-        mLauncher.getDragLayer().getScrim().invalidate();
+        mLauncher.getDragLayer().getWorkspaceDragScrim().invalidate();
     }
 
     public CellLayout getCurrentDragOverlappingLayout() {
diff --git a/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java b/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java
index cd938e1..0e0ddfb 100644
--- a/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java
+++ b/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java
@@ -32,7 +32,7 @@
 import static com.android.launcher3.anim.Interpolators.ZOOM_OUT;
 import static com.android.launcher3.anim.PropertySetter.NO_ANIM_PROPERTY_SETTER;
 import static com.android.launcher3.graphics.Scrim.SCRIM_PROGRESS;
-import static com.android.launcher3.graphics.WorkspaceAndHotseatScrim.SYSUI_PROGRESS;
+import static com.android.launcher3.graphics.SysUiScrim.SYSUI_PROGRESS;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_HOTSEAT_SCALE;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_HOTSEAT_TRANSLATE;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_WORKSPACE_FADE;
@@ -49,7 +49,8 @@
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.anim.PropertySetter;
 import com.android.launcher3.anim.SpringAnimationBuilder;
-import com.android.launcher3.graphics.WorkspaceAndHotseatScrim;
+import com.android.launcher3.graphics.SysUiScrim;
+import com.android.launcher3.graphics.WorkspaceDragScrim;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.util.DynamicResource;
 import com.android.systemui.plugins.ResourceProvider;
@@ -183,10 +184,12 @@
     }
 
     public void setScrim(PropertySetter propertySetter, LauncherState state) {
-        WorkspaceAndHotseatScrim scrim = mLauncher.getDragLayer().getScrim();
-        propertySetter.setFloat(scrim, SCRIM_PROGRESS, state.getWorkspaceScrimAlpha(mLauncher),
-                LINEAR);
-        propertySetter.setFloat(scrim, SYSUI_PROGRESS,
+        WorkspaceDragScrim workspaceDragScrim = mLauncher.getDragLayer().getWorkspaceDragScrim();
+        propertySetter.setFloat(workspaceDragScrim, SCRIM_PROGRESS,
+                state.getWorkspaceScrimAlpha(mLauncher), LINEAR);
+
+        SysUiScrim sysUiScrim = mLauncher.getDragLayer().getSysUiScrim();
+        propertySetter.setFloat(sysUiScrim, SYSUI_PROGRESS,
                 state.hasFlag(FLAG_HAS_SYS_UI_SCRIM) ? 1 : 0, LINEAR);
     }
 
diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index 20f7c23..505e6d8 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -213,10 +213,14 @@
         if (insets == null) return;
 
         if (insets.isVisible(WindowInsets.Type.ime())) {
-            getWindowInsetsController().hide(WindowInsets.Type.ime());
+            hideIme();
         }
     }
 
+    protected void hideIme() {
+        getWindowInsetsController().hide(WindowInsets.Type.ime());
+    }
+
     /**
      * Returns whether the view itself will handle the touch event or not.
      */
diff --git a/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java b/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
index 1d31975..6c95992 100644
--- a/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
+++ b/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
@@ -110,7 +110,6 @@
      */
     public static class SectionDecorationHandler {
         private static final int FILL_ALPHA = 0;
-        private static final int FOCUS_ALPHA = (int) (.9f * 255);
 
         protected RectF mBounds = new RectF();
         private final boolean mIsFullWidth;
@@ -123,9 +122,9 @@
 
         public SectionDecorationHandler(Context context, boolean isFullWidth) {
             mIsFullWidth = isFullWidth;
-            int endScrim = Themes.getAttrColor(context, R.attr.allAppsScrimColor);
+            int endScrim = Themes.getColorBackground(context);
             mFillcolor = ColorUtils.setAlphaComponent(endScrim, FILL_ALPHA);
-            mFocusColor = ColorUtils.setAlphaComponent(endScrim, FOCUS_ALPHA);
+            mFocusColor = endScrim;
             mRadius = Themes.getDialogCornerRadius(context);
         }
 
diff --git a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
index a6bc6cf..4876298 100644
--- a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
@@ -15,6 +15,10 @@
  */
 package com.android.launcher3.allapps;
 
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_KEYBOARD_CLOSED;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_SWITCHED_TO_MAIN_TAB;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_SWITCHED_TO_WORK_TAB;
+
 import android.content.Context;
 import android.graphics.Rect;
 import android.util.AttributeSet;
@@ -86,6 +90,13 @@
     public void onTabChanged(int pos) {
         super.onTabChanged(pos);
         if (mUsingTabs) {
+            // Log tab switches only when the launcher is in AllApps state
+            if (mLauncher.getStateManager().getCurrentStableState() == LauncherState.ALL_APPS) {
+                mLauncher.getStatsLogManager().logger()
+                        .log(pos == AdapterHolder.WORK ? LAUNCHER_ALLAPPS_SWITCHED_TO_WORK_TAB
+                                : LAUNCHER_ALLAPPS_SWITCHED_TO_MAIN_TAB);
+            }
+
             if (pos == AdapterHolder.WORK) {
                 WorkEduView.showWorkEduIfNeeded(mLauncher);
             } else {
@@ -93,4 +104,10 @@
             }
         }
     }
+
+    @Override
+    protected void hideIme() {
+        super.hideIme();
+        mLauncher.getStatsLogManager().logger().log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED);
+    }
 }
diff --git a/src/com/android/launcher3/allapps/search/LiveSearchManager.java b/src/com/android/launcher3/allapps/search/LiveSearchManager.java
index e52c790..4ef154e 100644
--- a/src/com/android/launcher3/allapps/search/LiveSearchManager.java
+++ b/src/com/android/launcher3/allapps/search/LiveSearchManager.java
@@ -16,8 +16,6 @@
 package com.android.launcher3.allapps.search;
 
 import static com.android.launcher3.LauncherState.ALL_APPS;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_ENTRY;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_EXIT;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.THREAD_POOL_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
@@ -47,9 +45,6 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.logging.InstanceId;
-import com.android.launcher3.logging.InstanceIdSequence;
-import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.SafeCloseable;
@@ -57,7 +52,6 @@
 
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.Optional;
 
 /**
  * Manages Lifecycle for Live search results
@@ -74,8 +68,6 @@
     private final HashMap<ComponentKey, SearchWidgetInfoContainer> mWidgetPlaceholders =
             new HashMap<>();
     private SearchWidgetHost mSearchWidgetHost;
-    private InstanceId mLogInstanceId;
-    private LauncherState mPrevLauncherState;
 
     public LiveSearchManager(Launcher launcher) {
         mLauncher = launcher;
@@ -139,11 +131,6 @@
     }
 
     @Override
-    public void onStateTransitionStart(LauncherState toState) {
-        mPrevLauncherState = mLauncher.getStateManager().getCurrentStableState();
-    }
-
-    @Override
     public void onStateTransitionComplete(LauncherState finalState) {
         if (finalState != ALL_APPS) {
             // Clear all search session related objects
@@ -152,15 +139,6 @@
 
             clearWidgetHost();
         }
-
-        StatsLogger logger = mLauncher.getStatsLogManager().logger();
-        if (finalState.equals(ALL_APPS)) {
-            mLogInstanceId = new InstanceIdSequence().newInstanceId();
-            logger.withInstanceId(mLogInstanceId).log(LAUNCHER_ALLAPPS_ENTRY);
-        } else if (mPrevLauncherState.equals(ALL_APPS)) {
-            logger.withInstanceId(mLogInstanceId).log(LAUNCHER_ALLAPPS_EXIT);
-            mLogInstanceId = null;
-        }
     }
 
     /**
@@ -178,14 +156,6 @@
         return () -> sliceLifeCycle.removeListener(listener);
     }
 
-    /**
-     * Returns {@link InstanceId} that should be used for logging events within search session, if
-     * available.
-     */
-    public Optional<InstanceId> getLogInstanceId() {
-        return Optional.ofNullable(mLogInstanceId);
-    }
-
     static class SearchWidgetHost extends AppWidgetHost {
         SearchWidgetHost(Context context) {
             super(context, SEARCH_APPWIDGET_HOST_ID);
diff --git a/src/com/android/launcher3/dragndrop/DragController.java b/src/com/android/launcher3/dragndrop/DragController.java
index 1cfe6ac..93df599 100644
--- a/src/com/android/launcher3/dragndrop/DragController.java
+++ b/src/com/android/launcher3/dragndrop/DragController.java
@@ -40,12 +40,14 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
 import com.android.launcher3.accessibility.DragViewStateAnnouncer;
+import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.TouchController;
 
 import java.util.ArrayList;
+import java.util.Optional;
 
 /**
  * Class for initiating a drag within a view or across multiple views.
@@ -230,6 +232,11 @@
         }
     }
 
+    public Optional<InstanceId> getLogInstanceId() {
+        return Optional.ofNullable(mDragObject)
+                .map(dragObject -> dragObject.logInstanceId);
+    }
+
     /**
      * Call this from a drag source view like this:
      *
diff --git a/src/com/android/launcher3/dragndrop/DragLayer.java b/src/com/android/launcher3/dragndrop/DragLayer.java
index e71c12d..7a6b4f9 100644
--- a/src/com/android/launcher3/dragndrop/DragLayer.java
+++ b/src/com/android/launcher3/dragndrop/DragLayer.java
@@ -41,12 +41,14 @@
 import com.android.launcher3.CellLayout;
 import com.android.launcher3.DropTargetBar;
 import com.android.launcher3.Launcher;
+import com.android.launcher3.LauncherRootView;
 import com.android.launcher3.R;
 import com.android.launcher3.ShortcutAndWidgetContainer;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.folder.Folder;
 import com.android.launcher3.graphics.OverviewScrim;
-import com.android.launcher3.graphics.WorkspaceAndHotseatScrim;
+import com.android.launcher3.graphics.SysUiScrim;
+import com.android.launcher3.graphics.WorkspaceDragScrim;
 import com.android.launcher3.keyboard.ViewGroupFocusHelper;
 import com.android.launcher3.util.Thunk;
 import com.android.launcher3.views.BaseDragLayer;
@@ -82,8 +84,10 @@
 
     // Related to adjacent page hints
     private final ViewGroupFocusHelper mFocusIndicatorHelper;
-    private final WorkspaceAndHotseatScrim mWorkspaceScrim;
     private final OverviewScrim mOverviewScrim;
+    private WorkspaceDragScrim mWorkspaceDragScrim;
+    private SysUiScrim mSysUiScrim;
+    private LauncherRootView mRootView;
 
     /**
      * Used to create a new DragLayer from XML.
@@ -99,14 +103,23 @@
         setChildrenDrawingOrderEnabled(true);
 
         mFocusIndicatorHelper = new ViewGroupFocusHelper(this);
-        mWorkspaceScrim = new WorkspaceAndHotseatScrim(this);
         mOverviewScrim = new OverviewScrim(this);
     }
 
     public void setup(DragController dragController, Workspace workspace) {
         mDragController = dragController;
-        mWorkspaceScrim.setWorkspace(workspace);
         recreateControllers();
+
+        mWorkspaceDragScrim = new WorkspaceDragScrim((this));
+        mWorkspaceDragScrim.setWorkspace(workspace);
+
+        // We delegate drawing of the workspace scrim to LauncherRootView (one level up), so as
+        // to avoid artifacts when translating the entire drag layer in the -1 transition.
+        mRootView = (LauncherRootView) getParent();
+        mSysUiScrim = new SysUiScrim(mRootView);
+        mRootView.setSysUiScrim(mSysUiScrim);
+
+
     }
 
     @Override
@@ -515,7 +528,7 @@
     @Override
     protected void dispatchDraw(Canvas canvas) {
         // Draw the background below children.
-        mWorkspaceScrim.draw(canvas);
+        mWorkspaceDragScrim.draw(canvas);
         mOverviewScrim.updateCurrentScrimmedView(this);
         mFocusIndicatorHelper.draw(canvas);
         super.dispatchDraw(canvas);
@@ -535,18 +548,22 @@
     @Override
     protected void onSizeChanged(int w, int h, int oldw, int oldh) {
         super.onSizeChanged(w, h, oldw, oldh);
-        mWorkspaceScrim.setSize(w, h);
+        mSysUiScrim.setSize(w, h);
     }
 
     @Override
     public void setInsets(Rect insets) {
         super.setInsets(insets);
-        mWorkspaceScrim.onInsetsChanged(insets, mAllowSysuiScrims);
+        mSysUiScrim.onInsetsChanged(insets, mAllowSysuiScrims);
         mOverviewScrim.onInsetsChanged(insets);
     }
 
-    public WorkspaceAndHotseatScrim getScrim() {
-        return mWorkspaceScrim;
+    public WorkspaceDragScrim getWorkspaceDragScrim() {
+        return mWorkspaceDragScrim;
+    }
+
+    public SysUiScrim getSysUiScrim() {
+        return mSysUiScrim;
     }
 
     public OverviewScrim getOverviewScrim() {
diff --git a/src/com/android/launcher3/folder/Folder.java b/src/com/android/launcher3/folder/Folder.java
index 61938d1..74d8dca 100644
--- a/src/com/android/launcher3/folder/Folder.java
+++ b/src/com/android/launcher3/folder/Folder.java
@@ -1208,7 +1208,9 @@
                         newIcon.requestFocus();
                     }
                     if (finalItem != null) {
-                        mStatsLogManager.logger().withItemInfo(finalItem)
+                        StatsLogger logger = mStatsLogManager.logger().withItemInfo(finalItem);
+                        mDragController.getLogInstanceId().map(logger::withInstanceId)
+                                .orElse(logger)
                                 .log(LAUNCHER_FOLDER_CONVERTED_TO_ICON);
                     }
                 }
diff --git a/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java b/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
index d347e8f..b6d25c4 100644
--- a/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
+++ b/src/com/android/launcher3/graphics/PlaceHolderIconDrawable.java
@@ -19,10 +19,19 @@
 
 import static com.android.launcher3.graphics.IconShape.getShapePath;
 
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
+import android.animation.ValueAnimator;
 import android.content.Context;
 import android.graphics.Canvas;
+import android.graphics.Color;
 import android.graphics.Path;
+import android.graphics.PorterDuff;
+import android.graphics.PorterDuffColorFilter;
 import android.graphics.Rect;
+import android.graphics.drawable.Drawable;
+
+import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.FastBitmapDrawable;
 import com.android.launcher3.R;
@@ -53,4 +62,27 @@
         canvas.drawPath(mProgressPath, mPaint);
         canvas.restoreToCount(saveCount);
     }
+
+    /** Updates this placeholder to {@code newIcon} with animation. */
+    public void animateIconUpdate(Drawable newIcon) {
+        int placeholderColor = mPaint.getColor();
+        int originalAlpha = Color.alpha(placeholderColor);
+
+        ValueAnimator iconUpdateAnimation = ValueAnimator.ofInt(originalAlpha, 0);
+        iconUpdateAnimation.setDuration(375);
+        iconUpdateAnimation.addUpdateListener(valueAnimator -> {
+            int newAlpha = (int) valueAnimator.getAnimatedValue();
+            int newColor = ColorUtils.setAlphaComponent(placeholderColor, newAlpha);
+
+            newIcon.setColorFilter(new PorterDuffColorFilter(newColor, PorterDuff.Mode.SRC_ATOP));
+        });
+        iconUpdateAnimation.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                newIcon.setColorFilter(null);
+            }
+        });
+        iconUpdateAnimation.start();
+    }
+
 }
diff --git a/src/com/android/launcher3/graphics/WorkspaceAndHotseatScrim.java b/src/com/android/launcher3/graphics/SysUiScrim.java
similarity index 83%
rename from src/com/android/launcher3/graphics/WorkspaceAndHotseatScrim.java
rename to src/com/android/launcher3/graphics/SysUiScrim.java
index 7b7ab5e..d9c648b 100644
--- a/src/com/android/launcher3/graphics/WorkspaceAndHotseatScrim.java
+++ b/src/com/android/launcher3/graphics/SysUiScrim.java
@@ -34,7 +34,6 @@
 import android.graphics.Paint;
 import android.graphics.Rect;
 import android.graphics.RectF;
-import android.graphics.Region;
 import android.graphics.Shader;
 import android.graphics.drawable.Drawable;
 import android.util.DisplayMetrics;
@@ -44,41 +43,39 @@
 
 import androidx.core.graphics.ColorUtils;
 
-import com.android.launcher3.CellLayout;
 import com.android.launcher3.R;
 import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.Utilities;
-import com.android.launcher3.Workspace;
 import com.android.launcher3.uioverrides.WallpaperColorInfo;
 import com.android.launcher3.util.Themes;
 
 /**
  * View scrim which draws behind hotseat and workspace
  */
-public class WorkspaceAndHotseatScrim extends Scrim {
+public class SysUiScrim extends Scrim {
 
-    public static final FloatProperty<WorkspaceAndHotseatScrim> SYSUI_PROGRESS =
-            new FloatProperty<WorkspaceAndHotseatScrim>("sysUiProgress") {
+    public static final FloatProperty<SysUiScrim> SYSUI_PROGRESS =
+            new FloatProperty<SysUiScrim>("sysUiProgress") {
                 @Override
-                public Float get(WorkspaceAndHotseatScrim scrim) {
+                public Float get(SysUiScrim scrim) {
                     return scrim.mSysUiProgress;
                 }
 
                 @Override
-                public void setValue(WorkspaceAndHotseatScrim scrim, float value) {
+                public void setValue(SysUiScrim scrim, float value) {
                     scrim.setSysUiProgress(value);
                 }
             };
 
-    private static final FloatProperty<WorkspaceAndHotseatScrim> SYSUI_ANIM_MULTIPLIER =
-            new FloatProperty<WorkspaceAndHotseatScrim>("sysUiAnimMultiplier") {
+    private static final FloatProperty<SysUiScrim> SYSUI_ANIM_MULTIPLIER =
+            new FloatProperty<SysUiScrim>("sysUiAnimMultiplier") {
                 @Override
-                public Float get(WorkspaceAndHotseatScrim scrim) {
+                public Float get(SysUiScrim scrim) {
                     return scrim.mSysUiAnimMultiplier;
                 }
 
                 @Override
-                public void setValue(WorkspaceAndHotseatScrim scrim, float value) {
+                public void setValue(SysUiScrim scrim, float value) {
                     scrim.mSysUiAnimMultiplier = value;
                     scrim.reapplySysUiAlpha();
                 }
@@ -108,10 +105,6 @@
     private static final int ALPHA_MASK_BITMAP_DP = 200;
     private static final int ALPHA_MASK_WIDTH_DP = 2;
 
-    private final Rect mHighlightRect = new Rect();
-
-    private Workspace mWorkspace;
-
     private boolean mDrawTopScrim, mDrawBottomScrim;
 
     private final RectF mFinalMaskRect = new RectF();
@@ -127,9 +120,8 @@
     private boolean mAnimateScrimOnNextDraw = false;
     private float mSysUiAnimMultiplier = 1;
 
-    public WorkspaceAndHotseatScrim(View view) {
+    public SysUiScrim(View view) {
         super(view);
-
         mMaskHeight = ResourceUtils.pxFromDp(ALPHA_MASK_BITMAP_DP,
                 view.getResources().getDisplayMetrics());
         mTopScrim = Themes.getAttrDrawable(view.getContext(), R.attr.workspaceStatusBarScrim);
@@ -139,28 +131,10 @@
         onExtractedColorsChanged(mWallpaperColorInfo);
     }
 
-    public void setWorkspace(Workspace workspace)  {
-        mWorkspace = workspace;
-    }
-
+    /**
+     * Draw the top and bottom scrims
+     */
     public void draw(Canvas canvas) {
-        // Draw the background below children.
-        if (mScrimAlpha > 0) {
-            // Update the scroll position first to ensure scrim cutout is in the right place.
-            mWorkspace.computeScrollWithoutInvalidation();
-            CellLayout currCellLayout = mWorkspace.getCurrentDragOverlappingLayout();
-            canvas.save();
-            if (currCellLayout != null && currCellLayout != mLauncher.getHotseat()) {
-                // Cut a hole in the darkening scrim on the page that should be highlighted, if any.
-                mLauncher.getDragLayer()
-                        .getDescendantRectRelativeToSelf(currCellLayout, mHighlightRect);
-                canvas.clipRect(mHighlightRect, Region.Op.DIFFERENCE);
-            }
-
-            super.draw(canvas);
-            canvas.restore();
-        }
-
         if (!mHideSysUiScrim) {
             if (mSysUiProgress <= 0) {
                 mAnimateScrimOnNextDraw = false;
@@ -247,6 +221,11 @@
         super.onExtractedColorsChanged(wallpaperColorInfo);
     }
 
+    /**
+     * Set the width and height of the view being scrimmed
+     * @param w
+     * @param h
+     */
     public void setSize(int w, int h) {
         if (mTopScrim != null) {
             mTopScrim.setBounds(0, 0, w, h);
diff --git a/src/com/android/launcher3/graphics/WorkspaceDragScrim.java b/src/com/android/launcher3/graphics/WorkspaceDragScrim.java
new file mode 100644
index 0000000..d8dc563
--- /dev/null
+++ b/src/com/android/launcher3/graphics/WorkspaceDragScrim.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright (C) 2018 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.graphics;
+
+import android.graphics.Canvas;
+import android.graphics.Rect;
+import android.graphics.Region;
+import android.view.View;
+
+import com.android.launcher3.CellLayout;
+import com.android.launcher3.Workspace;
+
+/**
+ * Scrim drawn during SpringLoaded State (ie. Drag and Drop). Darkens the workspace except for
+ * the focused CellLayout.
+ */
+public class WorkspaceDragScrim extends Scrim {
+
+    private final Rect mHighlightRect = new Rect();
+
+    private Workspace mWorkspace;
+
+    public WorkspaceDragScrim(View view) {
+        super(view);
+        onExtractedColorsChanged(mWallpaperColorInfo);
+    }
+
+    /**
+     * Set the workspace that this scrim is acting on
+     * @param workspace
+     */
+    public void setWorkspace(Workspace workspace)  {
+        mWorkspace = workspace;
+        mWorkspace.setWorkspaceDragScrim(this);
+    }
+
+    /**
+     * Cut out the focused paged of the Workspace and then draw the scrim
+     * @param canvas
+     */
+    public void draw(Canvas canvas) {
+        // Draw the background below children.
+        if (mScrimAlpha > 0) {
+            // Update the scroll position first to ensure scrim cutout is in the right place.
+            mWorkspace.computeScrollWithoutInvalidation();
+            CellLayout currCellLayout = mWorkspace.getCurrentDragOverlappingLayout();
+            canvas.save();
+            if (currCellLayout != null && currCellLayout != mLauncher.getHotseat()) {
+                // Cut a hole in the darkening scrim on the page that should be highlighted, if any.
+                mLauncher.getDragLayer()
+                        .getDescendantRectRelativeToSelf(currCellLayout, mHighlightRect);
+                canvas.clipRect(mHighlightRect, Region.Op.DIFFERENCE);
+            }
+
+            super.draw(canvas);
+            canvas.restore();
+        }
+    }
+
+}
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index 0292d20..1266bb4 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -22,6 +22,8 @@
 
 import android.content.Context;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.R;
 import com.android.launcher3.logger.LauncherAtom.ContainerInfo;
 import com.android.launcher3.logger.LauncherAtom.FromState;
@@ -29,6 +31,8 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.ResourceBasedOverride;
 
+import java.util.Optional;
+
 /**
  * Handles the user event logging in R+.
  *
@@ -47,6 +51,7 @@
     public static final int LAUNCHER_STATE_ALLAPPS = 4;
     public static final int LAUNCHER_STATE_UNCHANGED = 5;
 
+    private InstanceId mInstanceId;
     /**
      * Returns event enum based on the two state transition information when swipe
      * gesture happens(to be removed during UserEventDispatcher cleanup).
@@ -95,9 +100,13 @@
         @UiEvent(doc = "User dragged a launcher item")
         LAUNCHER_ITEM_DRAG_STARTED(383),
 
-        @UiEvent(doc = "A dragged launcher item is successfully dropped")
+        @UiEvent(doc = "A dragged launcher item is successfully dropped onto workspace, hotseat "
+                + "open folder etc")
         LAUNCHER_ITEM_DROP_COMPLETED(385),
 
+        @UiEvent(doc = "A dragged launcher item is successfully dropped onto a folder icon.")
+        LAUNCHER_ITEM_DROP_COMPLETED_ON_FOLDER_ICON(697),
+
         @UiEvent(doc = "A dragged launcher item is successfully dropped on another item "
                 + "resulting in a new folder creation")
         LAUNCHER_ITEM_DROP_FOLDER_CREATED(386),
@@ -350,6 +359,15 @@
 
         @UiEvent(doc = "Launcher exited from AllApps state.")
         LAUNCHER_ALLAPPS_EXIT(693),
+
+        @UiEvent(doc = "User closed the AllApps keyboard.")
+        LAUNCHER_ALLAPPS_KEYBOARD_CLOSED(694),
+
+        @UiEvent(doc = "User switched to Main tab in AllApps screen.")
+        LAUNCHER_ALLAPPS_SWITCHED_TO_MAIN_TAB(695),
+
+        @UiEvent(doc = "User switched to Work tab in AllApps screen.")
+        LAUNCHER_ALLAPPS_SWITCHED_TO_WORK_TAB(696),
         ;
 
         // ADD MORE
@@ -467,16 +485,30 @@
      * Returns new logger object.
      */
     public StatsLogger logger() {
+        StatsLogger logger = createLogger();
+        Optional.ofNullable(mInstanceId).ifPresent(logger::withInstanceId);
+        return logger;
+    }
+
+    protected StatsLogger createLogger() {
         return new StatsLogger() {
         };
     }
 
     /**
+     * Sets InstanceId to every new {@link StatsLogger} object returned by {@link #logger()} when
+     * not-null.
+     */
+    public StatsLogManager withDefaultInstanceId(@Nullable InstanceId instanceId) {
+        this.mInstanceId = instanceId;
+        return this;
+    }
+
+    /**
      * Creates a new instance of {@link StatsLogManager} based on provided context.
      */
     public static StatsLogManager newInstance(Context context) {
-        StatsLogManager mgr = Overrides.getObject(StatsLogManager.class,
+        return Overrides.getObject(StatsLogManager.class,
                 context.getApplicationContext(), R.string.stats_log_manager_class);
-        return mgr;
     }
 }
diff --git a/src/com/android/launcher3/model/BaseModelUpdateTask.java b/src/com/android/launcher3/model/BaseModelUpdateTask.java
index d1e5017..543ee2a 100644
--- a/src/com/android/launcher3/model/BaseModelUpdateTask.java
+++ b/src/com/android/launcher3/model/BaseModelUpdateTask.java
@@ -28,7 +28,7 @@
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.ItemInfoMatcher;
-import com.android.launcher3.widget.WidgetListRowEntry;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -123,7 +123,7 @@
     }
 
     public void bindUpdatedWidgets(BgDataModel dataModel) {
-        final ArrayList<WidgetListRowEntry> widgets =
+        final ArrayList<WidgetsListBaseEntry> widgets =
                 dataModel.widgetsModel.getWidgetsList(mApp.getContext());
         scheduleCallbackTask(c -> c.bindAllWidgets(widgets));
     }
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index 2d860a4..1d7d1a2 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -50,7 +50,7 @@
 import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.ViewOnDrawExecutor;
-import com.android.launcher3.widget.WidgetListRowEntry;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
@@ -467,7 +467,7 @@
         void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets);
         void bindRestoreItemsChange(HashSet<ItemInfo> updates);
         void bindWorkspaceComponentsRemoved(ItemInfoMatcher matcher);
-        void bindAllWidgets(ArrayList<WidgetListRowEntry> widgets);
+        void bindAllWidgets(List<WidgetsListBaseEntry> widgets);
         void onPageBoundSynchronously(int page);
         void executeOnNextDraw(ViewOnDrawExecutor executor);
         void bindDeepShortcutMap(HashMap<ComponentKey, Integer> deepShortcutMap);
diff --git a/src/com/android/launcher3/model/data/FolderInfo.java b/src/com/android/launcher3/model/data/FolderInfo.java
index cc783f7..cd2ef35 100644
--- a/src/com/android/launcher3/model/data/FolderInfo.java
+++ b/src/com/android/launcher3/model/data/FolderInfo.java
@@ -34,6 +34,7 @@
 import com.android.launcher3.folder.FolderNameInfos;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logger.LauncherAtom.Attribute;
+import com.android.launcher3.logger.LauncherAtom.FolderIcon;
 import com.android.launcher3.logger.LauncherAtom.FromState;
 import com.android.launcher3.logger.LauncherAtom.ToState;
 import com.android.launcher3.model.ModelWriter;
@@ -208,8 +209,13 @@
 
     @Override
     public LauncherAtom.ItemInfo buildProto(FolderInfo fInfo) {
+        FolderIcon.Builder folderIcon = FolderIcon.newBuilder()
+                .setCardinality(contents.size());
+        if (LabelState.SUGGESTED.equals(getLabelState())) {
+            folderIcon.setLabelInfo(title.toString());
+        }
         return getDefaultItemInfoBuilder()
-                .setFolderIcon(LauncherAtom.FolderIcon.newBuilder().setCardinality(contents.size()))
+                .setFolderIcon(folderIcon)
                 .setRank(rank)
                 .setAttribute(getLabelState().mLogAttribute)
                 .setContainerInfo(getContainerInfo())
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index 76048ba..1853632 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -31,10 +31,10 @@
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.ShortcutUtil;
-import com.android.launcher3.widget.WidgetListRowEntry;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
 
 import java.io.PrintWriter;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -59,7 +59,7 @@
     /** Maps packages to their DotInfo's . */
     private Map<PackageUserKey, DotInfo> mPackageUserToDotInfos = new HashMap<>();
     /** Maps packages to their Widgets */
-    private ArrayList<WidgetListRowEntry> mAllWidgets = new ArrayList<>();
+    private List<WidgetsListBaseEntry> mAllWidgets = List.of();
 
     private PopupDataChangeListener mChangeListener = PopupDataChangeListener.INSTANCE;
 
@@ -187,7 +187,7 @@
         notificationListener.cancelNotificationFromLauncher(notificationKey);
     }
 
-    public void setAllWidgets(ArrayList<WidgetListRowEntry> allWidgets) {
+    public void setAllWidgets(List<WidgetsListBaseEntry> allWidgets) {
         mAllWidgets = allWidgets;
         mChangeListener.onWidgetsBound();
     }
@@ -196,14 +196,15 @@
         mChangeListener = listener == null ? PopupDataChangeListener.INSTANCE : listener;
     }
 
-    public ArrayList<WidgetListRowEntry> getAllWidgets() {
+    public List<WidgetsListBaseEntry> getAllWidgets() {
         return mAllWidgets;
     }
 
     public List<WidgetItem> getWidgetsForPackageUser(PackageUserKey packageUserKey) {
         return mAllWidgets.stream()
-                .filter(row -> row.pkgItem.packageName.equals(packageUserKey.mPackageName))
-                .flatMap(row -> row.widgets.stream())
+                .filter(row -> row instanceof WidgetsListContentEntry
+                        && row.mPkgItem.packageName.equals(packageUserKey.mPackageName))
+                .flatMap(row -> ((WidgetsListContentEntry) row).mWidgets.stream())
                 .filter(widget -> packageUserKey.mUser.equals(widget.user))
                 .collect(Collectors.toList());
     }
diff --git a/src/com/android/launcher3/recyclerview/ViewHolderBinder.java b/src/com/android/launcher3/recyclerview/ViewHolderBinder.java
new file mode 100644
index 0000000..4653774
--- /dev/null
+++ b/src/com/android/launcher3/recyclerview/ViewHolderBinder.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.recyclerview;
+
+import android.view.ViewGroup;
+
+import androidx.recyclerview.widget.RecyclerView.ViewHolder;
+
+/**
+ * Creates and populates views with data
+ *
+ * @param <T> A data model which is used to populate the {@link ViewHolder}.
+ * @param <V> A subclass of {@link ViewHolder} which holds references to views.
+ */
+public interface ViewHolderBinder<T, V extends ViewHolder> {
+    /**
+     * Creates a new view, and attach it to the parent {@link ViewGroup}. Then, populates UI
+     * references in a {@link ViewHolder}.
+     */
+    V newViewHolder(ViewGroup parent);
+
+    /** Populate UI references in {@link ViewHolder} with data. */
+    void bindViewHolder(V viewHolder, T data);
+
+    /**
+     * Called when the view is recycled. Views are recycled in batches once they are sufficiently
+     * far off screen that it is unlikely the user will scroll back to them soon. Optionally
+     * override this to free expensive resources.
+     */
+    default void unbindViewHolder(V viewHolder) {}
+}
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
index 0266345..0e8d4ae 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
@@ -46,7 +46,7 @@
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.util.ViewOnDrawExecutor;
 import com.android.launcher3.views.BaseDragLayer;
-import com.android.launcher3.widget.WidgetListRowEntry;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -232,7 +232,7 @@
     public void bindWorkspaceComponentsRemoved(ItemInfoMatcher matcher) { }
 
     @Override
-    public void bindAllWidgets(ArrayList<WidgetListRowEntry> widgets) { }
+    public void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
 
     @Override
     public void onPageBoundSynchronously(int page) { }
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index d4a132e..4261d08 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -33,7 +33,7 @@
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.R;
 import com.android.launcher3.util.ResourceBasedOverride;
-import com.android.launcher3.widget.WidgetsFullSheet;
+import com.android.launcher3.widget.picker.WidgetsFullSheet;
 
 import java.util.concurrent.ExecutionException;
 import java.util.function.Function;
diff --git a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
index 10cd04c..25ecea5 100644
--- a/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
+++ b/src/com/android/launcher3/touch/AbstractStateChangeTouchController.java
@@ -543,7 +543,8 @@
             // case the user started interacting with it before the animation finished.
             mLauncher.getStateManager().goToState(targetState, false /* animated */);
         }
-        mLauncher.getDragLayer().getScrim().createSysuiMultiplierAnim(1f).setDuration(0).start();
+        mLauncher.getDragLayer().getSysUiScrim().createSysuiMultiplierAnim(
+                1f).setDuration(0).start();
     }
 
     private void logReachedState(LauncherState targetState) {
diff --git a/src/com/android/launcher3/util/Themes.java b/src/com/android/launcher3/util/Themes.java
index 55d17fc..512a286 100644
--- a/src/com/android/launcher3/util/Themes.java
+++ b/src/com/android/launcher3/util/Themes.java
@@ -82,6 +82,11 @@
     }
 
     /** Returns the floating background color attribute. */
+    public static int getColorBackground(Context context) {
+        return getAttrColor(context, android.R.attr.colorBackground);
+    }
+
+    /** Returns the floating background color attribute. */
     public static int getColorBackgroundFloating(Context context) {
         return getAttrColor(context, android.R.attr.colorBackgroundFloating);
     }
diff --git a/src/com/android/launcher3/views/OptionsPopupView.java b/src/com/android/launcher3/views/OptionsPopupView.java
index 80f0981..addaf9c 100644
--- a/src/com/android/launcher3/views/OptionsPopupView.java
+++ b/src/com/android/launcher3/views/OptionsPopupView.java
@@ -49,7 +49,7 @@
 import com.android.launcher3.shortcuts.DeepShortcutView;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.TestProtocol;
-import com.android.launcher3.widget.WidgetsFullSheet;
+import com.android.launcher3.widget.picker.WidgetsFullSheet;
 
 import java.util.ArrayList;
 import java.util.List;
diff --git a/src/com/android/launcher3/views/ScrimView.java b/src/com/android/launcher3/views/ScrimView.java
index 77cec80..c9bd284 100644
--- a/src/com/android/launcher3/views/ScrimView.java
+++ b/src/com/android/launcher3/views/ScrimView.java
@@ -42,7 +42,7 @@
  */
 public class ScrimView<T extends Launcher> extends View implements Insettable, OnChangeListener {
 
-    private static final float SCRIM_ALPHA = .75f;
+    private static final float SCRIM_ALPHA = .95f;
     protected final T mLauncher;
     private final WallpaperColorInfo mWallpaperColorInfo;
     protected final int mEndScrim;
@@ -63,6 +63,7 @@
         mWallpaperColorInfo = WallpaperColorInfo.INSTANCE.get(context);
         int endScrim = Themes.getAttrColor(context, R.attr.allAppsScrimColor);
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
+            endScrim = Themes.getColorBackgroundFloating(context);
             endScrim = ColorUtils.setAlphaComponent(endScrim, (int) (255  * SCRIM_ALPHA));
         }
         mEndScrim = endScrim;
diff --git a/src/com/android/launcher3/widget/WidgetListRowEntry.java b/src/com/android/launcher3/widget/WidgetListRowEntry.java
deleted file mode 100644
index 17e4673..0000000
--- a/src/com/android/launcher3/widget/WidgetListRowEntry.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * Copyright (C) 2016 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.launcher3.widget;
-
-import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.PackageItemInfo;
-
-import java.util.ArrayList;
-
-/**
- * Holder class to store all the information related to a single row in the widget list
- */
-public class WidgetListRowEntry {
-
-    public final PackageItemInfo pkgItem;
-
-    public final ArrayList<WidgetItem> widgets;
-
-    /**
-     * Character that is used as a section name for the {@link ItemInfo#title}.
-     * (e.g., "G" will be stored if title is "Google")
-     */
-    public String titleSectionName;
-
-    public WidgetListRowEntry(PackageItemInfo pkgItem, ArrayList<WidgetItem> items) {
-        this.pkgItem = pkgItem;
-        this.widgets = items;
-    }
-
-    @Override
-    public String toString() {
-        return pkgItem.packageName + ":" + widgets.size();
-    }
-}
diff --git a/src/com/android/launcher3/widget/WidgetsDiffReporter.java b/src/com/android/launcher3/widget/WidgetsDiffReporter.java
deleted file mode 100644
index df6e2c3..0000000
--- a/src/com/android/launcher3/widget/WidgetsDiffReporter.java
+++ /dev/null
@@ -1,143 +0,0 @@
-/*
- * 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.
- */
-
-package com.android.launcher3.widget;
-
-import android.util.Log;
-
-import androidx.recyclerview.widget.RecyclerView;
-
-import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.widget.WidgetsListAdapter.WidgetListRowEntryComparator;
-
-import java.util.ArrayList;
-import java.util.Iterator;
-
-/**
- * Do diff on widget's tray list items and call the {@link RecyclerView.Adapter}
- * methods accordingly.
- */
-public class WidgetsDiffReporter {
-    private static final boolean DEBUG = false;
-    private static final String TAG = "WidgetsDiffReporter";
-
-    private final IconCache mIconCache;
-    private final RecyclerView.Adapter mListener;
-
-    public WidgetsDiffReporter(IconCache iconCache, RecyclerView.Adapter listener) {
-        mIconCache = iconCache;
-        mListener = listener;
-    }
-
-    public void process(ArrayList<WidgetListRowEntry> currentEntries,
-            ArrayList<WidgetListRowEntry> newEntries, WidgetListRowEntryComparator comparator) {
-        if (DEBUG) {
-            Log.d(TAG, "process oldEntries#=" + currentEntries.size()
-                    + " newEntries#=" + newEntries.size());
-        }
-        // Early exit if either of the list is empty
-        if (currentEntries.isEmpty() || newEntries.isEmpty()) {
-            // Skip if both list are empty.
-            // On rotation, we open the widget tray with empty. Then try to fetch the list again
-            // when the animation completes (which still gives empty). And we get the final result
-            // when the bind actually completes.
-            if (currentEntries.size() != newEntries.size()) {
-                currentEntries.clear();
-                currentEntries.addAll(newEntries);
-                mListener.notifyDataSetChanged();
-            }
-            return;
-        }
-        ArrayList<WidgetListRowEntry> orgEntries =
-                (ArrayList<WidgetListRowEntry>) currentEntries.clone();
-        Iterator<WidgetListRowEntry> orgIter = orgEntries.iterator();
-        Iterator<WidgetListRowEntry> newIter = newEntries.iterator();
-
-        WidgetListRowEntry orgRowEntry = orgIter.next();
-        WidgetListRowEntry newRowEntry = newIter.next();
-
-        do {
-            int diff = comparePackageName(orgRowEntry, newRowEntry, comparator);
-            if (DEBUG) {
-                Log.d(TAG, String.format("diff=%d orgRowEntry (%s) newRowEntry (%s)",
-                        diff, orgRowEntry != null? orgRowEntry.toString() : null,
-                        newRowEntry != null? newRowEntry.toString() : null));
-            }
-            int index = -1;
-            if (diff < 0) {
-                index = currentEntries.indexOf(orgRowEntry);
-                mListener.notifyItemRemoved(index);
-                if (DEBUG) {
-                    Log.d(TAG, String.format("notifyItemRemoved called (%d)%s", index,
-                            orgRowEntry.titleSectionName));
-                }
-                currentEntries.remove(index);
-                orgRowEntry = orgIter.hasNext() ? orgIter.next() : null;
-            } else if (diff > 0) {
-                index = orgRowEntry != null? currentEntries.indexOf(orgRowEntry):
-                        currentEntries.size();
-                currentEntries.add(index, newRowEntry);
-                if (DEBUG) {
-                    Log.d(TAG, String.format("notifyItemInserted called (%d)%s", index,
-                            newRowEntry.titleSectionName));
-                }
-                newRowEntry = newIter.hasNext() ? newIter.next() : null;
-                mListener.notifyItemInserted(index);
-
-            } else {
-                // same package name but,
-                // did the icon, title, etc, change?
-                // or did the widget size and desc, span, etc change?
-                if (!isSamePackageItemInfo(orgRowEntry.pkgItem, newRowEntry.pkgItem) ||
-                        !orgRowEntry.widgets.equals(newRowEntry.widgets)) {
-                    index = currentEntries.indexOf(orgRowEntry);
-                    currentEntries.set(index, newRowEntry);
-                    mListener.notifyItemChanged(index);
-                    if (DEBUG) {
-                        Log.d(TAG, String.format("notifyItemChanged called (%d)%s", index,
-                                newRowEntry.titleSectionName));
-                    }
-                }
-                orgRowEntry = orgIter.hasNext() ? orgIter.next() : null;
-                newRowEntry = newIter.hasNext() ? newIter.next() : null;
-            }
-        } while(orgRowEntry != null || newRowEntry != null);
-    }
-
-    /**
-     * Compare package name using the same comparator as in {@link WidgetsListAdapter}.
-     * Also handle null row pointers.
-     */
-    private int comparePackageName(WidgetListRowEntry curRow, WidgetListRowEntry newRow,
-            WidgetListRowEntryComparator comparator) {
-        if (curRow == null && newRow == null) {
-            throw new IllegalStateException("Cannot compare PackageItemInfo if both rows are null.");
-        }
-
-        if (curRow == null && newRow != null) {
-            return 1; // new row needs to be inserted
-        } else if (curRow != null && newRow == null) {
-            return -1; // old row needs to be deleted
-        }
-        return comparator.compare(curRow, newRow);
-    }
-
-    private boolean isSamePackageItemInfo(PackageItemInfo curInfo, PackageItemInfo newInfo) {
-        return curInfo.bitmap.icon.equals(newInfo.bitmap.icon)
-                && !mIconCache.isDefaultIcon(curInfo.bitmap, curInfo.user);
-    }
-}
diff --git a/src/com/android/launcher3/widget/WidgetsListAdapter.java b/src/com/android/launcher3/widget/WidgetsListAdapter.java
deleted file mode 100644
index 5bf9690..0000000
--- a/src/com/android/launcher3/widget/WidgetsListAdapter.java
+++ /dev/null
@@ -1,219 +0,0 @@
-/*
- * Copyright (C) 2015 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.launcher3.widget;
-
-import android.content.Context;
-import android.util.Log;
-import android.view.LayoutInflater;
-import android.view.View;
-import android.view.View.OnClickListener;
-import android.view.View.OnLongClickListener;
-import android.view.ViewGroup;
-
-import androidx.recyclerview.widget.RecyclerView;
-import androidx.recyclerview.widget.RecyclerView.Adapter;
-
-import com.android.launcher3.R;
-import com.android.launcher3.WidgetPreviewLoader;
-import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.util.LabelComparator;
-
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Comparator;
-import java.util.List;
-
-/**
- * List view adapter for the widget tray.
- *
- * <p>Memory vs. Performance:
- * The less number of types of views are inserted into a {@link RecyclerView}, the more recycling
- * happens and less memory is consumed. {@link #getItemViewType} was not overridden as there is
- * only a single type of view.
- */
-public class WidgetsListAdapter extends Adapter<WidgetsRowViewHolder> {
-
-    private static final String TAG = "WidgetsListAdapter";
-    private static final boolean DEBUG = false;
-
-    private final WidgetPreviewLoader mWidgetPreviewLoader;
-    private final LayoutInflater mLayoutInflater;
-
-    private final OnClickListener mIconClickListener;
-    private final OnLongClickListener mIconLongClickListener;
-    private final int mIndent;
-    private ArrayList<WidgetListRowEntry> mEntries = new ArrayList<>();
-    private final WidgetsDiffReporter mDiffReporter;
-
-    private boolean mApplyBitmapDeferred;
-
-    public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
-            WidgetPreviewLoader widgetPreviewLoader, IconCache iconCache,
-            OnClickListener iconClickListener, OnLongClickListener iconLongClickListener) {
-        mLayoutInflater = layoutInflater;
-        mWidgetPreviewLoader = widgetPreviewLoader;
-        mIconClickListener = iconClickListener;
-        mIconLongClickListener = iconLongClickListener;
-        mIndent = context.getResources().getDimensionPixelSize(R.dimen.widget_section_indent);
-        mDiffReporter = new WidgetsDiffReporter(iconCache, this);
-    }
-
-    /**
-     * Defers applying bitmap on all the {@link WidgetCell} in the {@param rv}
-     *
-     * @see WidgetCell#setApplyBitmapDeferred(boolean)
-     */
-    public void setApplyBitmapDeferred(boolean isDeferred, RecyclerView rv) {
-        mApplyBitmapDeferred = isDeferred;
-
-        for (int i = rv.getChildCount() - 1; i >= 0; i--) {
-            WidgetsRowViewHolder holder = (WidgetsRowViewHolder)
-                    rv.getChildViewHolder(rv.getChildAt(i));
-            for (int j = holder.cellContainer.getChildCount() - 1; j >= 0; j--) {
-                View v = holder.cellContainer.getChildAt(j);
-                if (v instanceof WidgetCell) {
-                    ((WidgetCell) v).setApplyBitmapDeferred(mApplyBitmapDeferred);
-                }
-            }
-        }
-    }
-
-    /**
-     * Update the widget list.
-     */
-    public void setWidgets(ArrayList<WidgetListRowEntry> tempEntries) {
-        WidgetListRowEntryComparator rowComparator = new WidgetListRowEntryComparator();
-        Collections.sort(tempEntries, rowComparator);
-        mDiffReporter.process(mEntries, tempEntries, rowComparator);
-    }
-
-    @Override
-    public int getItemCount() {
-        return mEntries.size();
-    }
-
-    public String getSectionName(int pos) {
-        return mEntries.get(pos).titleSectionName;
-    }
-
-    @Override
-    public void onBindViewHolder(WidgetsRowViewHolder holder, int pos) {
-        WidgetListRowEntry entry = mEntries.get(pos);
-        List<WidgetItem> infoList = entry.widgets;
-
-        ViewGroup row = holder.cellContainer;
-        if (DEBUG) {
-            Log.d(TAG, String.format(
-                    "onBindViewHolder [pos=%d, widget#=%d, row.getChildCount=%d]",
-                    pos, infoList.size(), row.getChildCount()));
-        }
-
-        // Add more views.
-        // if there are too many, hide them.
-        int expectedChildCount = infoList.size() + Math.max(0, infoList.size() - 1);
-        int childCount = row.getChildCount();
-
-        if (expectedChildCount > childCount) {
-            for (int i = childCount; i < expectedChildCount; i++) {
-                if ((i & 1) == 1) {
-                    // Add a divider for odd index
-                    mLayoutInflater.inflate(R.layout.widget_list_divider, row);
-                } else {
-                    // Add cell for even index
-                    WidgetCell widget = (WidgetCell) mLayoutInflater.inflate(
-                            R.layout.widget_cell, row, false);
-
-                    // set up touch.
-                    widget.setOnClickListener(mIconClickListener);
-                    widget.setOnLongClickListener(mIconLongClickListener);
-                    row.addView(widget);
-                }
-            }
-        } else if (expectedChildCount < childCount) {
-            for (int i = expectedChildCount; i < childCount; i++) {
-                row.getChildAt(i).setVisibility(View.GONE);
-            }
-        }
-
-        // Bind the views in the application info section.
-        holder.title.applyFromItemInfoWithIcon(entry.pkgItem);
-
-        // Bind the view in the widget horizontal tray region.
-        for (int i = 0; i < infoList.size(); i++) {
-            WidgetCell widget = (WidgetCell) row.getChildAt(2 * i);
-            widget.applyFromCellItem(infoList.get(i), mWidgetPreviewLoader);
-            widget.setApplyBitmapDeferred(mApplyBitmapDeferred);
-            widget.ensurePreview();
-            widget.setVisibility(View.VISIBLE);
-
-            if (i > 0) {
-                row.getChildAt(2 * i - 1).setVisibility(View.VISIBLE);
-            }
-        }
-    }
-
-    @Override
-    public WidgetsRowViewHolder onCreateViewHolder(ViewGroup parent, int viewType) {
-        if (DEBUG) {
-            Log.v(TAG, "\nonCreateViewHolder");
-        }
-
-        ViewGroup container = (ViewGroup) mLayoutInflater.inflate(
-                R.layout.widgets_list_row_view, parent, false);
-
-        // if the end padding is 0, then container view (horizontal scroll view) doesn't respect
-        // the end of the linear layout width + the start padding and doesn't allow scrolling.
-        container.findViewById(R.id.widgets_cell_list).setPaddingRelative(mIndent, 0, 1, 0);
-
-        return new WidgetsRowViewHolder(container);
-    }
-
-    @Override
-    public void onViewRecycled(WidgetsRowViewHolder holder) {
-        int total = holder.cellContainer.getChildCount();
-        for (int i = 0; i < total; i += 2) {
-            WidgetCell widget = (WidgetCell) holder.cellContainer.getChildAt(i);
-            widget.clear();
-        }
-    }
-
-    public boolean onFailedToRecycleView(WidgetsRowViewHolder holder) {
-        // If child views are animating, then the RecyclerView may choose not to recycle the view,
-        // causing extraneous onCreateViewHolder() calls.  It is safe in this case to continue
-        // recycling this view, and take care in onViewRecycled() to cancel any existing
-        // animations.
-        return true;
-    }
-
-    @Override
-    public long getItemId(int pos) {
-        return pos;
-    }
-
-    /**
-     * Comparator for sorting WidgetListRowEntry based on package title
-     */
-    public static class WidgetListRowEntryComparator implements Comparator<WidgetListRowEntry> {
-
-        private final LabelComparator mComparator = new LabelComparator();
-
-        @Override
-        public int compare(WidgetListRowEntry a, WidgetListRowEntry b) {
-            return mComparator.compare(a.pkgItem.title.toString(), b.pkgItem.title.toString());
-        }
-    }
-}
diff --git a/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
new file mode 100644
index 0000000..09517e1
--- /dev/null
+++ b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
@@ -0,0 +1,60 @@
+/*
+ * 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.model;
+
+import static java.lang.annotation.RetentionPolicy.SOURCE;
+
+import androidx.annotation.IntDef;
+
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.PackageItemInfo;
+
+import java.lang.annotation.Retention;
+
+/** Holder class to store the package information of an entry shown in the widgets list. */
+public abstract class WidgetsListBaseEntry {
+    public final PackageItemInfo mPkgItem;
+
+    /**
+     * Character that is used as a section name for the {@link ItemInfo#title}.
+     * (e.g., "G" will be stored if title is "Google")
+     */
+    public final String mTitleSectionName;
+
+    public WidgetsListBaseEntry(PackageItemInfo pkgItem, String titleSectionName) {
+        mPkgItem = pkgItem;
+        mTitleSectionName = titleSectionName;
+    }
+
+    /**
+     * Returns the ranking of this entry in the
+     * {@link com.android.launcher3.widget.picker.WidgetsListAdapter}.
+     *
+     * <p>Entries with smaller value should be shown first. See
+     * {@link com.android.launcher3.widget.picker.WidgetsDiffReporter} for more details.
+     */
+    @Rank
+    public abstract int getRank();
+
+    @Retention(SOURCE)
+    @IntDef({RANK_WIDGETS_LIST_HEADER, RANK_WIDGETS_LIST_CONTENT})
+    public @interface Rank {
+    }
+
+    public static final int RANK_WIDGETS_LIST_HEADER = 1;
+    public static final int RANK_WIDGETS_LIST_CONTENT = 2;
+}
diff --git a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
new file mode 100644
index 0000000..b0cb8c7
--- /dev/null
+++ b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2016 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.model;
+
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.widget.WidgetItemComparator;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Holder class to store all the information related to a list of widgets from the same app which is
+ * shown in the {@link com.android.launcher3.widget.picker.WidgetsFullSheet}.
+ */
+public final class WidgetsListContentEntry extends WidgetsListBaseEntry {
+
+    public final List<WidgetItem> mWidgets;
+
+    public WidgetsListContentEntry(PackageItemInfo pkgItem, String titleSectionName,
+            List<WidgetItem> items) {
+        super(pkgItem, titleSectionName);
+        this.mWidgets =
+                items.stream().sorted(new WidgetItemComparator()).collect(Collectors.toList());
+    }
+
+    @Override
+    public String toString() {
+        return mPkgItem.packageName + ":" + mWidgets.size();
+    }
+
+    @Override
+    @Rank
+    public int getRank() {
+        return RANK_WIDGETS_LIST_CONTENT;
+    }
+}
diff --git a/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
new file mode 100644
index 0000000..6899647
--- /dev/null
+++ b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
@@ -0,0 +1,64 @@
+/*
+ * 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.model;
+
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.data.PackageItemInfo;
+
+import java.util.Collection;
+
+/** An information holder for an app which has widgets or/and shortcuts. */
+public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
+
+    public final int widgetsCount;
+    public final int shortcutsCount;
+
+    private boolean mIsWidgetListShown = false;
+    private boolean mHasEntryUpdated = false;
+
+    public WidgetsListHeaderEntry(PackageItemInfo pkgItem, String titleSectionName,
+            Collection<WidgetItem> items) {
+        super(pkgItem, titleSectionName);
+        widgetsCount = (int) items.stream().filter(item -> item.widgetInfo != null).count();
+        shortcutsCount = Math.max(0, items.size() - widgetsCount);
+    }
+
+    /** Sets if the widgets list associated with this header is shown. */
+    public void setIsWidgetListShown(boolean isWidgetListShown) {
+        if (mIsWidgetListShown != isWidgetListShown) {
+            this.mIsWidgetListShown = isWidgetListShown;
+            mHasEntryUpdated = true;
+        } else {
+            mHasEntryUpdated = false;
+        }
+    }
+
+    /** Returns {@code true} if the widgets list associated with this header is shown. */
+    public boolean isWidgetListShown() {
+        return mIsWidgetListShown;
+    }
+
+    /** Returns {@code true} if this entry has been updated due to user interactions. */
+    public boolean hasEntryUpdated() {
+        return mHasEntryUpdated;
+    }
+
+    @Override
+    @Rank
+    public int getRank() {
+        return RANK_WIDGETS_LIST_HEADER;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
new file mode 100644
index 0000000..dbd1bdf
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
@@ -0,0 +1,189 @@
+/*
+ * 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.
+ */
+
+package com.android.launcher3.widget.picker;
+
+import android.util.Log;
+
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListAdapter.WidgetListBaseRowEntryComparator;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * Do diff on widget's tray list items and call the {@link RecyclerView.Adapter}
+ * methods accordingly.
+ */
+public class WidgetsDiffReporter {
+    private static final boolean DEBUG = false;
+    private static final String TAG = "WidgetsDiffReporter";
+
+    private final IconCache mIconCache;
+    private final RecyclerView.Adapter mListener;
+
+    public WidgetsDiffReporter(IconCache iconCache, RecyclerView.Adapter listener) {
+        mIconCache = iconCache;
+        mListener = listener;
+    }
+
+    /**
+     * Notifies the difference between {@code currentEntries} & {@code newEntries} by calling the
+     * relevant {@link androidx.recyclerview.widget.RecyclerView.RecyclerViewDataObserver} methods.
+     */
+    public void process(ArrayList<WidgetsListBaseEntry> currentEntries,
+            List<WidgetsListBaseEntry> newEntries,
+            WidgetListBaseRowEntryComparator comparator) {
+        if (DEBUG) {
+            Log.d(TAG, "process oldEntries#=" + currentEntries.size()
+                    + " newEntries#=" + newEntries.size());
+        }
+        // Early exit if either of the list is empty
+        if (currentEntries.isEmpty() || newEntries.isEmpty()) {
+            // Skip if both list are empty.
+            // On rotation, we open the widget tray with empty. Then try to fetch the list again
+            // when the animation completes (which still gives empty). And we get the final result
+            // when the bind actually completes.
+            if (currentEntries.size() != newEntries.size()) {
+                currentEntries.clear();
+                currentEntries.addAll(newEntries);
+                mListener.notifyDataSetChanged();
+            }
+            return;
+        }
+        ArrayList<WidgetsListBaseEntry> orgEntries =
+                (ArrayList<WidgetsListBaseEntry>) currentEntries.clone();
+        Iterator<WidgetsListBaseEntry> orgIter = orgEntries.iterator();
+        Iterator<WidgetsListBaseEntry> newIter = newEntries.iterator();
+
+        WidgetsListBaseEntry orgRowEntry = orgIter.next();
+        WidgetsListBaseEntry newRowEntry = newIter.next();
+
+        do {
+            int diff = compareAppNameAndType(orgRowEntry, newRowEntry, comparator);
+            if (DEBUG) {
+                Log.d(TAG, String.format("diff=%d orgRowEntry (%s) newRowEntry (%s)",
+                        diff, orgRowEntry != null ? orgRowEntry.toString() : null,
+                        newRowEntry != null ? newRowEntry.toString() : null));
+            }
+            int index = -1;
+            if (diff < 0) {
+                index = currentEntries.indexOf(orgRowEntry);
+                mListener.notifyItemRemoved(index);
+                if (DEBUG) {
+                    Log.d(TAG, String.format("notifyItemRemoved called (%d)%s", index,
+                            orgRowEntry.mTitleSectionName));
+                }
+                currentEntries.remove(index);
+                orgRowEntry = orgIter.hasNext() ? orgIter.next() : null;
+            } else if (diff > 0) {
+                index = orgRowEntry != null ? currentEntries.indexOf(orgRowEntry)
+                        : currentEntries.size();
+                currentEntries.add(index, newRowEntry);
+                if (DEBUG) {
+                    Log.d(TAG, String.format("notifyItemInserted called (%d)%s", index,
+                            newRowEntry.mTitleSectionName));
+                }
+                newRowEntry = newIter.hasNext() ? newIter.next() : null;
+                mListener.notifyItemInserted(index);
+
+            } else {
+                // same app name & type but,
+                // did the icon, title, etc, change?
+                // or did the header view changed due to user interactions?
+                // or did the widget size and desc, span, etc change?
+                if (!isSamePackageItemInfo(orgRowEntry.mPkgItem, newRowEntry.mPkgItem)
+                        || hasHeaderUpdated(newRowEntry)
+                        || hasWidgetsListChanged(orgRowEntry, newRowEntry)) {
+                    index = currentEntries.indexOf(orgRowEntry);
+                    currentEntries.set(index, newRowEntry);
+                    mListener.notifyItemChanged(index);
+                    if (DEBUG) {
+                        Log.d(TAG, String.format("notifyItemChanged called (%d)%s", index,
+                                newRowEntry.mTitleSectionName));
+                    }
+                }
+                orgRowEntry = orgIter.hasNext() ? orgIter.next() : null;
+                newRowEntry = newIter.hasNext() ? newIter.next() : null;
+            }
+        } while(orgRowEntry != null || newRowEntry != null);
+    }
+
+    /**
+     * Compares the app name and then entry type for the given {@link WidgetsListBaseEntry}s.
+     *
+     * @Return 0 if both entries' order is the same. Negative integer if {@code newRowEntry} should
+     *         order before {@code orgRowEntry}. Positive integer if {@code orgRowEntry} should
+     *         order before {@code newRowEntry}.
+     */
+    private int compareAppNameAndType(WidgetsListBaseEntry curRow, WidgetsListBaseEntry newRow,
+            WidgetListBaseRowEntryComparator comparator) {
+        if (curRow == null && newRow == null) {
+            throw new IllegalStateException(
+                    "Cannot compare PackageItemInfo if both rows are null.");
+        }
+
+        if (curRow == null && newRow != null) {
+            return 1; // new row needs to be inserted
+        } else if (curRow != null && newRow == null) {
+            return -1; // old row needs to be deleted
+        }
+        int diff = comparator.compare(curRow, newRow);
+        if (diff == 0) {
+            return newRow.getRank() - curRow.getRank();
+        }
+        return diff;
+    }
+
+    /**
+     * Returns {@code true} if both {@code curRow} & {@code newRow} are
+     * {@link WidgetsListContentEntry}s with a different list of widgets.
+     */
+    private boolean hasWidgetsListChanged(WidgetsListBaseEntry curRow,
+            WidgetsListBaseEntry newRow) {
+        if (!(curRow instanceof WidgetsListContentEntry)
+                || !(newRow instanceof WidgetsListContentEntry)) {
+            return false;
+        }
+        WidgetsListContentEntry orgRowEntry = (WidgetsListContentEntry) curRow;
+        WidgetsListContentEntry newRowEntry = (WidgetsListContentEntry) newRow;
+        return !orgRowEntry.mWidgets.equals(newRowEntry.mWidgets);
+    }
+
+    /**
+     * Returns {@code true} if {@code newRow} is {@link WidgetsListHeaderEntry} and its content has
+     * been changed due to user interactions.
+     */
+    private boolean hasHeaderUpdated(WidgetsListBaseEntry newRow) {
+        if (!(newRow instanceof WidgetsListHeaderEntry)) {
+            return false;
+        }
+        WidgetsListHeaderEntry newRowEntry = (WidgetsListHeaderEntry) newRow;
+        return newRowEntry.hasEntryUpdated();
+    }
+
+    private boolean isSamePackageItemInfo(PackageItemInfo curInfo, PackageItemInfo newInfo) {
+        return curInfo.bitmap.icon.equals(newInfo.bitmap.icon)
+                && !mIconCache.isDefaultIcon(curInfo.bitmap, curInfo.user);
+    }
+}
diff --git a/src/com/android/launcher3/widget/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
similarity index 95%
rename from src/com/android/launcher3/widget/WidgetsFullSheet.java
rename to src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 4c8c339..03623d5 100644
--- a/src/com/android/launcher3/widget/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.launcher3.widget;
+package com.android.launcher3.widget.picker;
 
 import static com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_Y;
 import static com.android.launcher3.testing.TestProtocol.NORMAL_STATE_ORDINAL;
@@ -42,6 +42,7 @@
 import com.android.launcher3.compat.AccessibilityManagerCompat;
 import com.android.launcher3.views.RecyclerViewFastScroller;
 import com.android.launcher3.views.TopRoundedCornerView;
+import com.android.launcher3.widget.BaseWidgetSheet;
 
 /**
  * Popup for showing the full list of available widgets
@@ -218,8 +219,8 @@
         if (ev.getAction() == MotionEvent.ACTION_DOWN) {
             mNoIntercept = false;
             RecyclerViewFastScroller scroller = mRecyclerView.getScrollbar();
-            if (scroller.getThumbOffsetY() >= 0 &&
-                    getPopupContainer().isEventOverView(scroller, ev)) {
+            if (scroller.getThumbOffsetY() >= 0
+                    && getPopupContainer().isEventOverView(scroller, ev)) {
                 mNoIntercept = true;
             } else if (getPopupContainer().isEventOverView(mContent, ev)) {
                 mNoIntercept = !mRecyclerView.shouldContainerScroll(ev, getPopupContainer());
@@ -228,6 +229,7 @@
         return super.onControllerInterceptTouchEvent(ev);
     }
 
+    /** Shows the {@link WidgetsFullSheet} on the launcher. */
     public static WidgetsFullSheet show(Launcher launcher, boolean animate) {
         WidgetsFullSheet sheet = (WidgetsFullSheet) launcher.getLayoutInflater()
                 .inflate(R.layout.widgets_full_sheet, launcher.getDragLayer(), false);
@@ -237,6 +239,7 @@
         return sheet;
     }
 
+    /** Gets the {@link WidgetsRecyclerView} which shows all widgets in {@link WidgetsFullSheet}. */
     @VisibleForTesting
     public static WidgetsRecyclerView getWidgetsView(Launcher launcher) {
         return launcher.findViewById(R.id.widgets_list_view);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
new file mode 100644
index 0000000..5ec7f3b
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -0,0 +1,208 @@
+/*
+ * Copyright (C) 2015 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.Log;
+import android.util.SparseArray;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.View.OnClickListener;
+import android.view.View.OnLongClickListener;
+import android.view.ViewGroup;
+
+import androidx.annotation.Nullable;
+import androidx.recyclerview.widget.RecyclerView;
+import androidx.recyclerview.widget.RecyclerView.Adapter;
+import androidx.recyclerview.widget.RecyclerView.ViewHolder;
+
+import com.android.launcher3.R;
+import com.android.launcher3.WidgetPreviewLoader;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.util.LabelComparator;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsListHeaderViewHolderBinder.OnHeaderClickListener;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Recycler view adapter for the widget tray.
+ *
+ * <p>This adapter supports view binding of subclasses of {@link WidgetsListBaseEntry}. There are 2
+ * subclasses: {@link WidgetsListHeader} & {@link WidgetsListContentEntry}.
+ * {@link WidgetsListHeader} entries are always visible in the recycler view. At most one
+ * {@link WidgetsListContentEntry} is shown in the recycler view at any time. Clicking a
+ * {@link WidgetsListHeader} will result in expanding / collapsing a corresponding
+ * {@link WidgetsListContentEntry} of the same app.
+ */
+public class WidgetsListAdapter extends Adapter<ViewHolder> implements OnHeaderClickListener {
+
+    private static final String TAG = "WidgetsListAdapter";
+    private static final boolean DEBUG = false;
+
+    /** Uniquely identifies widgets list view type within the app. */
+    private static final int VIEW_TYPE_WIDGETS_LIST = R.layout.widgets_list_row_view;
+    private static final int VIEW_TYPE_WIDGETS_HEADER = R.layout.widgets_list_row_header;
+
+    private final WidgetsDiffReporter mDiffReporter;
+    private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
+    private final WidgetsListRowViewHolderBinder mWidgetsListRowViewHolderBinder;
+    private final WidgetListBaseRowEntryComparator mRowComparator =
+            new WidgetListBaseRowEntryComparator();
+
+    private List<WidgetsListBaseEntry> mAllEntries = new ArrayList<>();
+    private ArrayList<WidgetsListBaseEntry> mVisibleEntries = new ArrayList<>();
+    @Nullable private String mWidgetsContentVisiblePackage = null;
+
+    public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
+            WidgetPreviewLoader widgetPreviewLoader, IconCache iconCache,
+            OnClickListener iconClickListener, OnLongClickListener iconLongClickListener) {
+        mDiffReporter = new WidgetsDiffReporter(iconCache, this);
+        mWidgetsListRowViewHolderBinder = new WidgetsListRowViewHolderBinder(context,
+                layoutInflater, iconClickListener, iconLongClickListener, widgetPreviewLoader);
+        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_LIST, mWidgetsListRowViewHolderBinder);
+        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_HEADER,
+                new WidgetsListHeaderViewHolderBinder(layoutInflater, this::onHeaderClicked));
+    }
+
+    /**
+     * Defers applying bitmap on all the {@link WidgetCell} in the {@param rv}.
+     *
+     * @see WidgetCell#setApplyBitmapDeferred(boolean)
+     */
+    public void setApplyBitmapDeferred(boolean isDeferred, RecyclerView rv) {
+        mWidgetsListRowViewHolderBinder.setApplyBitmapDeferred(isDeferred);
+
+        for (int i = rv.getChildCount() - 1; i >= 0; i--) {
+            ViewHolder viewHolder = rv.getChildViewHolder(rv.getChildAt(i));
+            if (viewHolder.getItemViewType() == VIEW_TYPE_WIDGETS_LIST) {
+                WidgetsRowViewHolder holder = (WidgetsRowViewHolder) viewHolder;
+                for (int j = holder.cellContainer.getChildCount() - 1; j >= 0; j--) {
+                    View v = holder.cellContainer.getChildAt(j);
+                    if (v instanceof WidgetCell) {
+                        ((WidgetCell) v).setApplyBitmapDeferred(isDeferred);
+                    }
+                }
+            }
+        }
+    }
+
+    @Override
+    public int getItemCount() {
+        return mVisibleEntries.size();
+    }
+
+    /** Gets the section name for {@link com.android.launcher3.views.RecyclerViewFastScroller}. */
+    public String getSectionName(int pos) {
+        return mVisibleEntries.get(pos).mTitleSectionName;
+    }
+
+    /** Updates the widget list. */
+    public void setWidgets(List<WidgetsListBaseEntry> tempEntries) {
+        mAllEntries = tempEntries.stream().sorted(mRowComparator)
+                .collect(Collectors.toList());
+        updateVisibleEntries();
+    }
+
+    private void updateVisibleEntries() {
+        mAllEntries.forEach(entry -> {
+            if (entry instanceof WidgetsListHeaderEntry) {
+                ((WidgetsListHeaderEntry) entry).setIsWidgetListShown(
+                        entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage));
+            }
+        });
+        List<WidgetsListBaseEntry> newVisibleEntries = mAllEntries.stream()
+                .filter(entry -> entry instanceof WidgetsListHeaderEntry
+                        || entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage))
+                .collect(Collectors.toList());
+        mDiffReporter.process(mVisibleEntries, newVisibleEntries, mRowComparator);
+    }
+
+    @Override
+    public void onBindViewHolder(ViewHolder holder, int pos) {
+        ViewHolderBinder viewHolderBinder = mViewHolderBinders.get(getItemViewType(pos));
+        viewHolderBinder.bindViewHolder(holder, mVisibleEntries.get(pos));
+    }
+
+    @Override
+    public ViewHolder onCreateViewHolder(ViewGroup parent, int viewType) {
+        if (DEBUG) {
+            Log.v(TAG, "\nonCreateViewHolder");
+        }
+
+        return mViewHolderBinders.get(viewType).newViewHolder(parent);
+    }
+
+    @Override
+    public void onViewRecycled(ViewHolder holder) {
+        mViewHolderBinders.get(holder.getItemViewType()).unbindViewHolder(holder);
+    }
+
+    @Override
+    public boolean onFailedToRecycleView(ViewHolder holder) {
+        // If child views are animating, then the RecyclerView may choose not to recycle the view,
+        // causing extraneous onCreateViewHolder() calls.  It is safe in this case to continue
+        // recycling this view, and take care in onViewRecycled() to cancel any existing
+        // animations.
+        return true;
+    }
+
+    @Override
+    public long getItemId(int pos) {
+        return pos;
+    }
+
+    @Override
+    public int getItemViewType(int pos) {
+        WidgetsListBaseEntry entry = mVisibleEntries.get(pos);
+        if (entry instanceof WidgetsListContentEntry) {
+            return VIEW_TYPE_WIDGETS_LIST;
+        } else if (entry instanceof WidgetsListHeaderEntry) {
+            return VIEW_TYPE_WIDGETS_HEADER;
+        }
+        throw new UnsupportedOperationException("ViewHolderBinder not found for " + entry);
+    }
+
+    @Override
+    public void onHeaderClicked(boolean showWidgets, String expandedPackage) {
+        if (showWidgets) {
+            mWidgetsContentVisiblePackage = expandedPackage;
+            updateVisibleEntries();
+        } else if (expandedPackage.equals(mWidgetsContentVisiblePackage)) {
+            mWidgetsContentVisiblePackage = null;
+            updateVisibleEntries();
+        }
+    }
+
+    /** Comparator for sorting WidgetListRowEntry based on package title. */
+    public static class WidgetListBaseRowEntryComparator implements
+            Comparator<WidgetsListBaseEntry> {
+
+        private final LabelComparator mComparator = new LabelComparator();
+
+        @Override
+        public int compare(WidgetsListBaseEntry a, WidgetsListBaseEntry b) {
+            return mComparator.compare(a.mPkgItem.title.toString(), b.mPkgItem.title.toString());
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
new file mode 100644
index 0000000..823fb7b
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -0,0 +1,205 @@
+/*
+ * 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.FastBitmapDrawable.newIcon;
+
+import android.content.Context;
+import android.content.res.Resources;
+import android.content.res.TypedArray;
+import android.graphics.drawable.Drawable;
+import android.util.AttributeSet;
+import android.widget.CheckBox;
+import android.widget.ImageView;
+import android.widget.LinearLayout;
+import android.widget.TextView;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.UiThread;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.FastBitmapDrawable;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.R;
+import com.android.launcher3.graphics.PlaceHolderIconDrawable;
+import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver;
+import com.android.launcher3.icons.cache.HandlerRunnable;
+import com.android.launcher3.model.data.ItemInfoWithIcon;
+import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+
+/**
+ * A UI represents a header of an app shown in the full widgets tray.
+ *
+ * It is a {@link LinearLayout} which contains an app icon, an app name, a subtitle and a checkbox
+ * which indicates if the widgets content view underneath this header should be shown.
+ */
+public final class WidgetsListHeader extends LinearLayout implements ItemInfoUpdateReceiver {
+
+    private boolean mEnableIconUpdateAnimation = false;
+
+    @Nullable private HandlerRunnable mIconLoadRequest;
+    @Nullable private Drawable mIconDrawable;
+    private final int mIconSize;
+
+    private ImageView mAppIcon;
+    private TextView mTitle;
+    private TextView mSubtitle;
+
+    private CheckBox mExpandToggle;
+    private boolean mIsExpanded = false;
+
+    public WidgetsListHeader(Context context) {
+        this(context, /* attrs= */ null);
+    }
+
+    public WidgetsListHeader(Context context, @Nullable AttributeSet attrs) {
+        this(context, attrs, /* defStyle= */ 0);
+    }
+
+    public WidgetsListHeader(Context context, @Nullable AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+
+        ActivityContext activity = ActivityContext.lookupContext(context);
+        DeviceProfile grid = activity.getDeviceProfile();
+        TypedArray a = context.obtainStyledAttributes(attrs,
+                R.styleable.WidgetsListRowHeader, defStyleAttr, /* defStyleRes= */ 0);
+        mIconSize = a.getDimensionPixelSize(R.styleable.WidgetsListRowHeader_appIconSize,
+                grid.iconSizePx);
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        mAppIcon = findViewById(R.id.app_icon);
+        mTitle = findViewById(R.id.app_title);
+        mSubtitle = findViewById(R.id.app_subtitle);
+        mExpandToggle = findViewById(R.id.toggle);
+    }
+
+    /**
+     * Sets a {@link OnExpansionChangeListener} to get a callback when this app widgets section
+     * expands / collapses.
+     */
+    @UiThread
+    public void setOnExpandChangeListener(
+            @Nullable OnExpansionChangeListener onExpandChangeListener) {
+        // Use the entire touch area of this view to expand / collapse an app widgets section.
+        setOnClickListener(view -> {
+            setExpanded(!mIsExpanded);
+            onExpandChangeListener.onExpansionChange(mIsExpanded);
+        });
+    }
+
+    /** Sets the expand toggle to expand / collapse. */
+    @UiThread
+    public void setExpanded(boolean isExpanded) {
+        this.mIsExpanded = isExpanded;
+        mExpandToggle.setChecked(isExpanded);
+    }
+
+    /** Apply app icon, labels and tag using a generic {@link WidgetsListHeaderEntry}. */
+    @UiThread
+    public void applyFromItemInfoWithIcon(WidgetsListHeaderEntry entry) {
+        applyIconAndLabel(entry);
+    }
+
+    @UiThread
+    private void applyIconAndLabel(WidgetsListHeaderEntry entry) {
+        PackageItemInfo info = entry.mPkgItem;
+        setIcon(info);
+        setTitles(entry);
+        setExpanded(entry.isWidgetListShown());
+
+        super.setTag(info);
+
+        verifyHighRes();
+    }
+
+    private void setIcon(PackageItemInfo info) {
+        FastBitmapDrawable icon = newIcon(getContext(), info);
+        applyDrawables(icon);
+        mIconDrawable = icon;
+        if (mIconDrawable != null) {
+            mIconDrawable.setVisible(
+                    /* visible= */ getWindowVisibility() == VISIBLE && isShown(),
+                    /* restart= */ false);
+        }
+    }
+
+    private void applyDrawables(Drawable icon) {
+        icon.setBounds(0, 0, mIconSize, mIconSize);
+
+        mAppIcon.setImageDrawable(icon);
+
+        // If the current icon is a placeholder color, animate its update.
+        if (mIconDrawable != null
+                && mIconDrawable instanceof PlaceHolderIconDrawable
+                && mEnableIconUpdateAnimation) {
+            ((PlaceHolderIconDrawable) mIconDrawable).animateIconUpdate(icon);
+        }
+    }
+
+    private void setTitles(WidgetsListHeaderEntry entry) {
+        mTitle.setText(entry.mPkgItem.title);
+
+        if (entry.widgetsCount > 0) {
+            Resources resources = getContext().getResources();
+            mSubtitle.setText(resources.getQuantityString(R.plurals.widgets_tray_subtitle,
+                    entry.widgetsCount, entry.widgetsCount));
+            mSubtitle.setVisibility(VISIBLE);
+        } else {
+            mSubtitle.setVisibility(GONE);
+        }
+    }
+
+    @Override
+    public void reapplyItemInfo(ItemInfoWithIcon info) {
+        if (getTag() == info) {
+            mIconLoadRequest = null;
+            mEnableIconUpdateAnimation = true;
+
+            // Optimization: Starting in N, pre-uploads the bitmap to RenderThread.
+            info.bitmap.icon.prepareToDraw();
+
+            setIcon((PackageItemInfo) info);
+
+            mEnableIconUpdateAnimation = false;
+        }
+    }
+
+    /** Verifies that the current icon is high-res otherwise posts a request to load the icon. */
+    public void verifyHighRes() {
+        if (mIconLoadRequest != null) {
+            mIconLoadRequest.cancel();
+            mIconLoadRequest = null;
+        }
+        if (getTag() instanceof ItemInfoWithIcon) {
+            ItemInfoWithIcon info = (ItemInfoWithIcon) getTag();
+            if (info.usingLowResIcon()) {
+                mIconLoadRequest = LauncherAppState.getInstance(getContext()).getIconCache()
+                        .updateIconInBackground(this, info);
+            }
+        }
+    }
+
+    /** A listener for the widget section expansion / collapse events. */
+    public interface OnExpansionChangeListener {
+        /** Notifies that the widget section is expanded or collapsed. */
+        void onExpansionChange(boolean isExpanded);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java
new file mode 100644
index 0000000..d4e1b1c
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderHolder.java
@@ -0,0 +1,32 @@
+/*
+ * 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 androidx.recyclerview.widget.RecyclerView.ViewHolder;
+
+/**
+ * A {@link ViewHolder} for {@link WidgetsListHeader} of an app, which renders the app icon, the app
+ * name, label and a button for showing / hiding widgets.
+ */
+public final class WidgetsListHeaderHolder extends ViewHolder {
+    final WidgetsListHeader mWidgetsListHeader;
+
+    public WidgetsListHeaderHolder(WidgetsListHeader view) {
+        super(view);
+
+        mWidgetsListHeader = view;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
new file mode 100644
index 0000000..ed53e6f
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
@@ -0,0 +1,61 @@
+/*
+ * 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.view.LayoutInflater;
+import android.view.ViewGroup;
+
+import com.android.launcher3.R;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+
+/**
+ * Binds data from {@link WidgetsListHeaderEntry} to UI elements in {@link WidgetsListHeaderHolder}.
+ */
+public final class WidgetsListHeaderViewHolderBinder implements
+        ViewHolderBinder<WidgetsListHeaderEntry, WidgetsListHeaderHolder> {
+    private final LayoutInflater mLayoutInflater;
+    private final OnHeaderClickListener mOnHeaderClickListener;
+
+    public WidgetsListHeaderViewHolderBinder(LayoutInflater layoutInflater,
+            OnHeaderClickListener onHeaderClickListener) {
+        mLayoutInflater = layoutInflater;
+        mOnHeaderClickListener = onHeaderClickListener;
+    }
+
+    @Override
+    public WidgetsListHeaderHolder newViewHolder(ViewGroup parent) {
+        WidgetsListHeader header = (WidgetsListHeader) mLayoutInflater.inflate(
+                R.layout.widgets_list_row_header, parent, false);
+
+        return new WidgetsListHeaderHolder(header);
+    }
+
+    @Override
+    public void bindViewHolder(WidgetsListHeaderHolder viewHolder, WidgetsListHeaderEntry data) {
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        widgetsListHeader.applyFromItemInfoWithIcon(data);
+        widgetsListHeader.setExpanded(data.isWidgetListShown());
+        widgetsListHeader.setOnExpandChangeListener(isExpanded ->
+                mOnHeaderClickListener.onHeaderClicked(isExpanded, data.mPkgItem.packageName));
+    }
+
+    /** A listener to be invoked when {@link WidgetsListHeader} is clicked. */
+    public interface OnHeaderClickListener {
+        /** Calls when {@link WidgetsListHeader} is clicked to show / hide widgets for a package. */
+        void onHeaderClicked(boolean showWidgets, String packageName);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
new file mode 100644
index 0000000..cec6b80
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListRowViewHolderBinder.java
@@ -0,0 +1,147 @@
+/*
+ * 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.Log;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.View.OnClickListener;
+import android.view.View.OnLongClickListener;
+import android.view.ViewGroup;
+
+import com.android.launcher3.R;
+import com.android.launcher3.WidgetPreviewLoader;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+
+import java.util.List;
+
+/**
+ * Binds data from {@link WidgetsListContentEntry} to UI elements in {@link WidgetsRowViewHolder}.
+ */
+public class WidgetsListRowViewHolderBinder
+        implements ViewHolderBinder<WidgetsListContentEntry, WidgetsRowViewHolder> {
+    private static final boolean DEBUG = false;
+    private static final String TAG = "WidgetsListRowViewHolderBinder";
+
+    private final LayoutInflater mLayoutInflater;
+    private final int mIndent;
+    private final OnClickListener mIconClickListener;
+    private final OnLongClickListener mIconLongClickListener;
+    private final WidgetPreviewLoader mWidgetPreviewLoader;
+    private boolean mApplyBitmapDeferred = false;
+
+    public WidgetsListRowViewHolderBinder(
+            Context context,
+            LayoutInflater layoutInflater,
+            OnClickListener iconClickListener,
+            OnLongClickListener iconLongClickListener,
+            WidgetPreviewLoader widgetPreviewLoader) {
+        mLayoutInflater = layoutInflater;
+        mIndent = context.getResources().getDimensionPixelSize(R.dimen.widget_section_indent);
+        mIconClickListener = iconClickListener;
+        mIconLongClickListener = iconLongClickListener;
+        mWidgetPreviewLoader = widgetPreviewLoader;
+    }
+
+    /**
+     * Defers applying bitmap on all the {@link WidgetCell} at
+     * {@link #bindViewHolder(WidgetsRowViewHolder, WidgetsListContentEntry)} if
+     * {@code applyBitmapDeferred} is {@code true}.
+     */
+    public void setApplyBitmapDeferred(boolean applyBitmapDeferred) {
+        mApplyBitmapDeferred = applyBitmapDeferred;
+    }
+
+    @Override
+    public WidgetsRowViewHolder newViewHolder(ViewGroup parent) {
+        if (DEBUG) {
+            Log.v(TAG, "\nonCreateViewHolder");
+        }
+
+        ViewGroup container = (ViewGroup) mLayoutInflater.inflate(
+                R.layout.widgets_scroll_container, parent, false);
+
+        // if the end padding is 0, then container view (horizontal scroll view) doesn't respect
+        // the end of the linear layout width + the start padding and doesn't allow scrolling.
+        container.findViewById(R.id.widgets_cell_list).setPaddingRelative(mIndent, 0, 1, 0);
+
+        return new WidgetsRowViewHolder(container);
+    }
+
+    @Override
+    public void bindViewHolder(WidgetsRowViewHolder holder, WidgetsListContentEntry entry) {
+        List<WidgetItem> infoList = entry.mWidgets;
+
+        ViewGroup row = holder.cellContainer;
+        if (DEBUG) {
+            Log.d(TAG, String.format("onBindViewHolder [widget#=%d, row.getChildCount=%d]",
+                    infoList.size(), row.getChildCount()));
+        }
+
+        // Add more views.
+        // if there are too many, hide them.
+        int expectedChildCount = infoList.size() + Math.max(0, infoList.size() - 1);
+        int childCount = row.getChildCount();
+
+        if (expectedChildCount > childCount) {
+            for (int i = childCount; i < expectedChildCount; i++) {
+                if ((i & 1) == 1) {
+                    // Add a divider for odd index
+                    mLayoutInflater.inflate(R.layout.widget_list_divider, row);
+                } else {
+                    // Add cell for even index
+                    WidgetCell widget = (WidgetCell) mLayoutInflater.inflate(
+                            R.layout.widget_cell, row, false);
+
+                    // set up touch.
+                    widget.setOnClickListener(mIconClickListener);
+                    widget.setOnLongClickListener(mIconLongClickListener);
+                    row.addView(widget);
+                }
+            }
+        } else if (expectedChildCount < childCount) {
+            for (int i = expectedChildCount; i < childCount; i++) {
+                row.getChildAt(i).setVisibility(View.GONE);
+            }
+        }
+
+        // Bind the view in the widget horizontal tray region.
+        for (int i = 0; i < infoList.size(); i++) {
+            WidgetCell widget = (WidgetCell) row.getChildAt(2 * i);
+            widget.applyFromCellItem(infoList.get(i), mWidgetPreviewLoader);
+            widget.setApplyBitmapDeferred(mApplyBitmapDeferred);
+            widget.ensurePreview();
+            widget.setVisibility(View.VISIBLE);
+
+            if (i > 0) {
+                row.getChildAt(2 * i - 1).setVisibility(View.VISIBLE);
+            }
+        }
+    }
+
+    @Override
+    public void unbindViewHolder(WidgetsRowViewHolder holder) {
+        int total = holder.cellContainer.getChildCount();
+        for (int i = 0; i < total; i += 2) {
+            WidgetCell widget = (WidgetCell) holder.cellContainer.getChildAt(i);
+            widget.clear();
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/WidgetsRecyclerView.java b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
similarity index 97%
rename from src/com/android/launcher3/widget/WidgetsRecyclerView.java
rename to src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
index 69de12b..52e9496 100644
--- a/src/com/android/launcher3/widget/WidgetsRecyclerView.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.launcher3.widget;
+package com.android.launcher3.widget.picker;
 
 import android.content.Context;
 import android.graphics.Point;
@@ -22,13 +22,13 @@
 import android.view.MotionEvent;
 import android.view.View;
 
-import com.android.launcher3.BaseRecyclerView;
-import com.android.launcher3.R;
-
 import androidx.recyclerview.widget.LinearLayoutManager;
 import androidx.recyclerview.widget.RecyclerView;
 import androidx.recyclerview.widget.RecyclerView.OnItemTouchListener;
 
+import com.android.launcher3.BaseRecyclerView;
+import com.android.launcher3.R;
+
 /**
  * The widgets recycler view.
  */
@@ -89,7 +89,7 @@
         LinearLayoutManager layoutManager = ((LinearLayoutManager) getLayoutManager());
         layoutManager.scrollToPositionWithOffset(0, (int) -(availableScrollHeight * touchFraction));
 
-        int posInt = (int) ((touchFraction == 1)? pos -1 : pos);
+        int posInt = (int) ((touchFraction == 1) ? pos - 1 : pos);
         return mAdapter.getSectionName(posInt);
     }
 
@@ -171,4 +171,4 @@
     @Override
     public void onRequestDisallowInterceptTouchEvent(boolean disallowIntercept) {
     }
-}
\ No newline at end of file
+}
diff --git a/src/com/android/launcher3/widget/WidgetsRowViewHolder.java b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
similarity index 77%
rename from src/com/android/launcher3/widget/WidgetsRowViewHolder.java
rename to src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
index d26edb6..ae94584 100644
--- a/src/com/android/launcher3/widget/WidgetsRowViewHolder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRowViewHolder.java
@@ -13,25 +13,22 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.launcher3.widget;
+package com.android.launcher3.widget.picker;
 
 import android.view.ViewGroup;
 
-import com.android.launcher3.BubbleTextView;
-import com.android.launcher3.R;
-
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
 
-public class WidgetsRowViewHolder extends ViewHolder {
+import com.android.launcher3.R;
+
+/** A {@link ViewHolder} for showing widgets of an app in the full widget picker. */
+public final class WidgetsRowViewHolder extends ViewHolder {
 
     public final ViewGroup cellContainer;
-    public final BubbleTextView title;
 
     public WidgetsRowViewHolder(ViewGroup v) {
         super(v);
 
         cellContainer = v.findViewById(R.id.widgets_cell_list);
-        title = v.findViewById(R.id.section);
-        title.setAccessibilityDelegate(null);
     }
 }
diff --git a/src_plugins/com/android/systemui/plugins/BcSmartspaceDataPlugin.java b/src_plugins/com/android/systemui/plugins/BcSmartspaceDataPlugin.java
index b90e43b..f8a9a04 100644
--- a/src_plugins/com/android/systemui/plugins/BcSmartspaceDataPlugin.java
+++ b/src_plugins/com/android/systemui/plugins/BcSmartspaceDataPlugin.java
@@ -39,6 +39,6 @@
     /** Provides Smartspace data to registered listeners. */
     interface SmartspaceTargetListener {
         /** Each Parcelable is a SmartspaceTarget that represents a card. */
-        void onSmartspaceTargetsUpdated(List<Parcelable> targets);
+        void onSmartspaceTargetsUpdated(List<? extends Parcelable> targets);
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java b/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
index 269af7b..73b1601 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
@@ -21,10 +21,10 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.widget.WidgetListRowEntry;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
-import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 
 /**
  * Helper class to handle results of {@link com.android.launcher3.model.LoaderTask}.
@@ -47,7 +47,7 @@
 
     @Override
     public void bindWidgets() {
-        final ArrayList<WidgetListRowEntry> widgets =
+        final List<WidgetsListBaseEntry> widgets =
                 mBgDataModel.widgetsModel.getWidgetsList(mApp.getContext());
         executeCallbacksTask(c -> c.bindAllWidgets(widgets), mUiExecutor);
     }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index b4e45f8..30c9b5f 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -28,12 +28,13 @@
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
-import com.android.launcher3.widget.WidgetItemComparator;
-import com.android.launcher3.widget.WidgetListRowEntry;
 import com.android.launcher3.widget.WidgetManagerHelper;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.WidgetsDiffReporter;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -60,25 +61,24 @@
     private final Map<PackageItemInfo, List<WidgetItem>> mWidgetsList = new HashMap<>();
 
     /**
-     * Returns a list of {@link WidgetListRowEntry}. All {@link WidgetItem} in a single row
-     * are sorted (based on label and user), but the overall list of {@link WidgetListRowEntry}s
-     * is not sorted. This list is sorted at the UI when using
-     * {@link com.android.launcher3.widget.WidgetsDiffReporter}
+     * Returns a list of {@link WidgetsListBaseEntry}. All {@link WidgetItem} in a single row
+     * are sorted (based on label and user), but the overall list of
+     * {@link WidgetsListBaseEntry}s is not sorted. This list is sorted at the UI when using
+     * {@link WidgetsDiffReporter}
      *
-     * @see com.android.launcher3.widget.WidgetsListAdapter#setWidgets(ArrayList)
+     * @see com.android.launcher3.widget.picker.WidgetsListAdapter#setWidgets(List)
      */
-    public synchronized ArrayList<WidgetListRowEntry> getWidgetsList(Context context) {
-        ArrayList<WidgetListRowEntry> result = new ArrayList<>();
+    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsList(Context context) {
+        ArrayList<WidgetsListBaseEntry> result = new ArrayList<>();
         AlphabeticIndexCompat indexer = new AlphabeticIndexCompat(context);
 
-        WidgetItemComparator widgetComparator = new WidgetItemComparator();
         for (Map.Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsList.entrySet()) {
-            WidgetListRowEntry row = new WidgetListRowEntry(
-                    entry.getKey(), new ArrayList<>(entry.getValue()));
-            row.titleSectionName = (row.pkgItem.title == null) ? "" :
-                    indexer.computeSectionName(row.pkgItem.title);
-            Collections.sort(row.widgets, widgetComparator);
-            result.add(row);
+            PackageItemInfo pkgItem = entry.getKey();
+            List<WidgetItem> widgetItems = entry.getValue();
+            String sectionName = (pkgItem.title == null) ? "" :
+                    indexer.computeSectionName(pkgItem.title);
+            result.add(new WidgetsListHeaderEntry(pkgItem, sectionName, widgetItems));
+            result.add(new WidgetsListContentEntry(pkgItem, sectionName, widgetItems));
         }
         return result;
     }
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 1648bdd..4c9a8e7 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -37,8 +37,8 @@
 import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.tapl.Workspace;
 import com.android.launcher3.views.OptionsPopupView;
-import com.android.launcher3.widget.WidgetsFullSheet;
-import com.android.launcher3.widget.WidgetsRecyclerView;
+import com.android.launcher3.widget.picker.WidgetsFullSheet;
+import com.android.launcher3.widget.picker.WidgetsRecyclerView;
 
 import org.junit.Before;
 import org.junit.Ignore;
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 9d4ccff..737f891 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -92,9 +92,8 @@
 
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        widgets.
-                getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace(true, false);
+        widgets.getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .dragToWorkspace(true, false);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
 
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 49af616..f95abdb 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -31,6 +31,7 @@
 import com.android.launcher3.testing.TestProtocol;
 
 import java.util.Collection;
+import java.util.List;
 
 /**
  * All widgets container.
@@ -101,22 +102,28 @@
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
              LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                      "getting widget " + labelText + " in widgets list")) {
-            final UiObject2 widgetsContainer = verifyActiveContainer();
+            final UiObject2 fullWidgetsPicker = verifyActiveContainer();
             mLauncher.assertTrue("Widgets container didn't become scrollable",
-                    widgetsContainer.wait(Until.scrollable(true), WAIT_TIME_MS));
+                    fullWidgetsPicker.wait(Until.scrollable(true), WAIT_TIME_MS));
             final Point displaySize = mLauncher.getRealDisplaySize();
-            final BySelector labelSelector = By.clazz("android.widget.TextView").text(labelText);
 
+            final UiObject2 widgetsContainer = findTestAppWidgetsScrollContainer();
+            mLauncher.assertTrue("Can't locate widgets list for the test app: "
+                                    + mLauncher.getLauncherPackageName(),
+                    widgetsContainer != null);
+            final BySelector labelSelector = By.clazz("android.widget.TextView").text(labelText);
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> cells = mLauncher.getObjectsInContainer(
-                        widgetsContainer, "widgets_scroll_container");
-                mLauncher.assertTrue("Widgets doesn't have 2 rows", cells.size() >= 2);
+                final Collection<UiObject2> cells = widgetsContainer.getChildren();
+                mLauncher.assertTrue("Widgets doesn't have 2 rows: ", cells.size() >= 2);
                 for (UiObject2 cell : cells) {
                     final UiObject2 label = cell.findObject(labelSelector);
+                    // The logic below doesn't handle the case which a widget cell of the given
+                    // label is not yet visible on the horizontal scrolling container. This won't be
+                    // an issue once we get rid of the horizontal scrolling container.
                     if (label == null) continue;
 
-                    final UiObject2 widget = label.getParent().getParent();
+                    final UiObject2 widget = cell;
                     mLauncher.assertEquals(
                             "View is not WidgetCell",
                             "com.android.launcher3.widget.WidgetCell",
@@ -131,7 +138,7 @@
                             <= displaySize.y - mLauncher.getBottomGestureSize()) {
                         int visibleDelta = maxWidth - mLauncher.getVisibleBounds(widget).width();
                         if (visibleDelta > 0) {
-                            Rect parentBounds = mLauncher.getVisibleBounds(cell);
+                            Rect parentBounds = mLauncher.getVisibleBounds(cell.getParent());
                             mLauncher.linearGesture(parentBounds.centerX() + visibleDelta
                                             + mLauncher.getTouchSlop(),
                                     parentBounds.centerY(), parentBounds.centerX(),
@@ -153,4 +160,53 @@
             }
         }
     }
+
+    /** Finds the widgets list of this test app from the collapsed full widgets picker. */
+    private UiObject2 findTestAppWidgetsScrollContainer() {
+        final BySelector headerSelector = By.res(mLauncher.getLauncherPackageName(),
+                "widgets_list_header");
+        final BySelector targetAppSelector = By.clazz("android.widget.TextView").text(
+                mLauncher.getContext().getPackageName());
+        final BySelector widgetsContainerSelector = By.res(mLauncher.getLauncherPackageName(),
+                "widgets_cell_list");
+
+        boolean hasHeaderExpanded = false;
+        for (int i = 0; i < 40; i++) {
+            UiObject2 fullWidgetsPicker = verifyActiveContainer();
+
+            UiObject2 header = fullWidgetsPicker.findObject(headerSelector);
+            mLauncher.assertTrue("Can't find a widget header", header != null);
+
+            // Look for a header that has the test app name.
+            UiObject2 headerTitle = fullWidgetsPicker.findObject(targetAppSelector);
+            if (headerTitle != null) {
+                // If we find the header and it has not been expanded, let's click it to see the
+                // widgets list.
+                if (!hasHeaderExpanded) {
+                    hasHeaderExpanded = true;
+                    mLauncher.clickLauncherObject(headerTitle);
+                    // After clicking the header, the recyclerview has been updated. Let's refresh
+                    // the container UIObject2.
+                    fullWidgetsPicker = verifyActiveContainer();
+                    // Refresh headerTitle because the first instance is stale after
+                    // verifyActiveContainer call.
+                    headerTitle = fullWidgetsPicker.findObject(targetAppSelector);
+                }
+
+                // Look for a widgets list.
+                UiObject2 widgetsContainer = fullWidgetsPicker.findObject(widgetsContainerSelector);
+                if (widgetsContainer != null) {
+                    // Make sure the widgets list is fully visible on the screen.
+                    mLauncher.scrollToLastVisibleRow(fullWidgetsPicker,
+                            widgetsContainer.getChildren(), 0);
+                    return widgetsContainer;
+                }
+                mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, List.of(headerTitle), 0);
+            } else {
+                mLauncher.scrollToLastVisibleRow(fullWidgetsPicker, header.getChildren(), 0);
+            }
+        }
+
+        return null;
+    }
 }