diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index f04c58d..343b87e 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -86,6 +86,7 @@
 import com.android.quickstep.util.StaggeredWorkspaceAnim;
 import com.android.quickstep.util.SurfaceTransactionApplier;
 import com.android.quickstep.views.RecentsView;
+import com.android.systemui.shared.recents.IStartingWindowListener;
 import com.android.systemui.shared.system.ActivityCompat;
 import com.android.systemui.shared.system.ActivityOptionsCompat;
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
@@ -97,7 +98,6 @@
 import com.android.systemui.shared.system.RemoteTransitionCompat;
 import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
 import com.android.systemui.shared.system.WindowManagerWrapper;
-import com.android.wm.shell.startingsurface.IStartingWindowListener;
 
 import java.util.LinkedHashMap;
 
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index a1cbec7..619103c 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -35,17 +35,12 @@
 import android.view.MotionEvent;
 
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.systemui.shared.recents.IPinnedStackAnimationListener;
+import com.android.systemui.shared.recents.ISplitScreenListener;
+import com.android.systemui.shared.recents.IStartingWindowListener;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.RemoteTransitionCompat;
-import com.android.wm.shell.onehanded.IOneHanded;
-import com.android.wm.shell.pip.IPip;
-import com.android.wm.shell.pip.IPipAnimationListener;
-import com.android.wm.shell.splitscreen.ISplitScreen;
-import com.android.wm.shell.splitscreen.ISplitScreenListener;
-import com.android.wm.shell.startingsurface.IStartingWindow;
-import com.android.wm.shell.startingsurface.IStartingWindowListener;
-import com.android.wm.shell.transition.IShellTransitions;
 
 /**
  * Holds the reference to SystemUI.
@@ -58,13 +53,8 @@
             new MainThreadInitializedObject<>(SystemUiProxy::new);
 
     private ISystemUiProxy mSystemUiProxy;
-    private IPip mPip;
-    private ISplitScreen mSplitScreen;
-    private IOneHanded mOneHanded;
-    private IShellTransitions mShellTransitions;
-    private IStartingWindow mStartingWindow;
     private final DeathRecipient mSystemUiProxyDeathRecipient = () -> {
-        MAIN_EXECUTOR.execute(() -> clearProxy());
+        MAIN_EXECUTOR.execute(() -> setProxy(null));
     };
 
     // Used to dedupe calls to SystemUI
@@ -92,23 +82,12 @@
         return null;
     }
 
-    public void setProxy(ISystemUiProxy proxy, IPip pip, ISplitScreen splitScreen,
-            IOneHanded oneHanded, IShellTransitions shellTransitions,
-            IStartingWindow startingWindow) {
+    public void setProxy(ISystemUiProxy proxy) {
         unlinkToDeath();
         mSystemUiProxy = proxy;
-        mPip = pip;
-        mSplitScreen = splitScreen;
-        mOneHanded = oneHanded;
-        mShellTransitions = shellTransitions;
-        mStartingWindow = startingWindow;
         linkToDeath();
     }
 
-    public void clearProxy() {
-        setProxy(null, null, null, null, null, null);
-    }
-
     // TODO(141886704): Find a way to remove this
     public void setLastSystemUiStateFlags(int stateFlags) {
         mLastSystemUiStateFlags = stateFlags;
@@ -286,6 +265,21 @@
     }
 
     @Override
+    public void setShelfHeight(boolean visible, int shelfHeight) {
+        boolean changed = visible != mLastShelfVisible || shelfHeight != mLastShelfHeight;
+        if (mSystemUiProxy != null && changed) {
+            mLastShelfVisible = visible;
+            mLastShelfHeight = shelfHeight;
+            try {
+                mSystemUiProxy.setShelfHeight(visible, shelfHeight);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call setShelfHeight visible: " + visible
+                        + " height: " + shelfHeight, e);
+            }
+        }
+    }
+
+    @Override
     public void handleImageAsScreenshot(Bitmap bitmap, Rect rect, Insets insets, int i) {
         if (mSystemUiProxy != null) {
             try {
@@ -321,6 +315,20 @@
         }
     }
 
+    /**
+     * Sets listener to get pinned stack animation callbacks.
+     */
+    @Override
+    public void setPinnedStackAnimationListener(IPinnedStackAnimationListener listener) {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.setPinnedStackAnimationListener(listener);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call setPinnedStackAnimationListener", e);
+            }
+        }
+    }
+
     @Override
     public void onQuickSwitchToNewTask(int rotation) {
         if (mSystemUiProxy != null) {
@@ -346,6 +354,28 @@
     }
 
     @Override
+    public void startOneHandedMode() {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.startOneHandedMode();
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call startOneHandedMode", e);
+            }
+        }
+    }
+
+    @Override
+    public void stopOneHandedMode() {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.stopOneHandedMode();
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call stopOneHandedMode", e);
+            }
+        }
+    }
+
+    @Override
     public void expandNotificationPanel() {
         if (mSystemUiProxy != null) {
             try {
@@ -356,45 +386,12 @@
         }
     }
 
-    //
-    // Pip
-    //
-
-    /**
-     * Sets the shelf height.
-     */
-    public void setShelfHeight(boolean visible, int shelfHeight) {
-        boolean changed = visible != mLastShelfVisible || shelfHeight != mLastShelfHeight;
-        if (mPip != null && changed) {
-            mLastShelfVisible = visible;
-            mLastShelfHeight = shelfHeight;
-            try {
-                mPip.setShelfHeight(visible, shelfHeight);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call setShelfHeight visible: " + visible
-                        + " height: " + shelfHeight, e);
-            }
-        }
-    }
-
-    /**
-     * Sets listener to get pinned stack animation callbacks.
-     */
-    public void setPinnedStackAnimationListener(IPipAnimationListener listener) {
-        if (mPip != null) {
-            try {
-                mPip.setPinnedStackAnimationListener(listener);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call setPinnedStackAnimationListener", e);
-            }
-        }
-    }
-
+    @Override
     public Rect startSwipePipToHome(ComponentName componentName, ActivityInfo activityInfo,
             PictureInPictureParams pictureInPictureParams, int launcherRotation, int shelfHeight) {
-        if (mPip != null) {
+        if (mSystemUiProxy != null) {
             try {
-                return mPip.startSwipePipToHome(componentName, activityInfo,
+                return mSystemUiProxy.startSwipePipToHome(componentName, activityInfo,
                         pictureInPictureParams, launcherRotation, shelfHeight);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call startSwipePipToHome", e);
@@ -403,85 +400,111 @@
         return null;
     }
 
+    @Override
     public void stopSwipePipToHome(ComponentName componentName, Rect destinationBounds) {
-        if (mPip != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mPip.stopSwipePipToHome(componentName, destinationBounds);
+                mSystemUiProxy.stopSwipePipToHome(componentName, destinationBounds);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call stopSwipePipToHome");
             }
         }
     }
 
-    //
-    // Splitscreen
-    //
-
-    public void registerSplitScreenListener(ISplitScreenListener listener) {
-        if (mSplitScreen != null) {
+    @Override
+    public void registerRemoteTransition(RemoteTransitionCompat remoteTransition) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.registerSplitScreenListener(listener);
+                mSystemUiProxy.registerRemoteTransition(remoteTransition);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call registerRemoteTransition");
+            }
+        }
+    }
+
+    @Override
+    public void unregisterRemoteTransition(RemoteTransitionCompat remoteTransition) {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.unregisterRemoteTransition(remoteTransition);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call registerRemoteTransition");
+            }
+        }
+    }
+
+    @Override
+    public void registerSplitScreenListener(ISplitScreenListener listener) {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.registerSplitScreenListener(listener);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call registerSplitScreenListener");
             }
         }
     }
 
+    @Override
     public void unregisterSplitScreenListener(ISplitScreenListener listener) {
-        if (mSplitScreen != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.unregisterSplitScreenListener(listener);
+                mSystemUiProxy.unregisterSplitScreenListener(listener);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call unregisterSplitScreenListener");
             }
         }
     }
 
+    @Override
     public void setSideStageVisibility(boolean visible) {
-        if (mSplitScreen != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.setSideStageVisibility(visible);
+                mSystemUiProxy.setSideStageVisibility(visible);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call setSideStageVisibility");
             }
         }
     }
 
+    @Override
     public void exitSplitScreen() {
-        if (mSplitScreen != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.exitSplitScreen();
+                mSystemUiProxy.exitSplitScreen();
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call exitSplitScreen");
             }
         }
     }
 
+    @Override
     public void exitSplitScreenOnHide(boolean exitSplitScreenOnHide) {
-        if (mSplitScreen != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.exitSplitScreenOnHide(exitSplitScreenOnHide);
+                mSystemUiProxy.exitSplitScreenOnHide(exitSplitScreenOnHide);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call exitSplitScreen");
             }
         }
     }
 
+    @Override
     public void startTask(int taskId, int stage, int position, Bundle options) {
-        if (mSplitScreen != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.startTask(taskId, stage, position, options);
+                mSystemUiProxy.startTask(taskId, stage, position, options);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call startTask");
             }
         }
     }
 
+    @Override
     public void startShortcut(String packageName, String shortcutId, int stage, int position,
             Bundle options, UserHandle user) {
-        if (mSplitScreen != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.startShortcut(packageName, shortcutId, stage, position, options,
+                mSystemUiProxy.startShortcut(packageName, shortcutId, stage, position, options,
                         user);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call startShortcut");
@@ -489,87 +512,38 @@
         }
     }
 
-    public void startIntent(PendingIntent intent, Intent fillInIntent, int stage, int position,
-            Bundle options) {
-        if (mSplitScreen != null) {
+    @Override
+    public void startIntent(PendingIntent intent, Intent fillInIntent, int stage,
+            int position, Bundle options) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.startIntent(intent, fillInIntent, stage, position, options);
+                mSystemUiProxy.startIntent(intent, fillInIntent, stage, position,
+                        options);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call startIntent");
             }
         }
     }
 
+    @Override
     public void removeFromSideStage(int taskId) {
-        if (mSplitScreen != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mSplitScreen.removeFromSideStage(taskId);
+                mSystemUiProxy.removeFromSideStage(taskId);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call removeFromSideStage");
             }
         }
     }
 
-    //
-    // One handed
-    //
-
-    public void startOneHandedMode() {
-        if (mOneHanded != null) {
-            try {
-                mOneHanded.startOneHanded();
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call startOneHandedMode", e);
-            }
-        }
-    }
-
-    public void stopOneHandedMode() {
-        if (mOneHanded != null) {
-            try {
-                mOneHanded.stopOneHanded();
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call stopOneHandedMode", e);
-            }
-        }
-    }
-
-    //
-    // Remote transitions
-    //
-
-    public void registerRemoteTransition(RemoteTransitionCompat remoteTransition) {
-        if (mShellTransitions != null) {
-            try {
-                mShellTransitions.registerRemote(remoteTransition.getFilter(),
-                        remoteTransition.getTransition());
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call registerRemoteTransition");
-            }
-        }
-    }
-
-    public void unregisterRemoteTransition(RemoteTransitionCompat remoteTransition) {
-        if (mShellTransitions != null) {
-            try {
-                mShellTransitions.unregisterRemote(remoteTransition.getTransition());
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call registerRemoteTransition");
-            }
-        }
-    }
-
-    //
-    // Starting window
-    //
-
     /**
      * Sets listener to get callbacks when launching a task.
      */
+    @Override
     public void setStartingWindowListener(IStartingWindowListener listener) {
-        if (mStartingWindow != null) {
+        if (mSystemUiProxy != null) {
             try {
-                mStartingWindow.setStartingWindowListener(listener);
+                mSystemUiProxy.setStartingWindowListener(listener);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call setStartingWindowListener", e);
             }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 1cb5f5d..fc805d0 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -24,11 +24,6 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.quickstep.GestureState.DEFAULT_STATE;
 import static com.android.quickstep.util.NavigationModeFeatureFlag.LIVE_TILE;
-import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_ONE_HANDED;
-import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_PIP;
-import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_SHELL_TRANSITIONS;
-import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_SPLIT_SCREEN;
-import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_STARTING_WINDOW;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SYSUI_PROXY;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_TRACING_ENABLED;
 
@@ -106,11 +101,6 @@
 import com.android.systemui.shared.system.InputConsumerController;
 import com.android.systemui.shared.system.InputMonitorCompat;
 import com.android.systemui.shared.tracing.ProtoTraceable;
-import com.android.wm.shell.onehanded.IOneHanded;
-import com.android.wm.shell.pip.IPip;
-import com.android.wm.shell.splitscreen.ISplitScreen;
-import com.android.wm.shell.startingsurface.IStartingWindow;
-import com.android.wm.shell.transition.IShellTransitions;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
@@ -150,18 +140,8 @@
         public void onInitialize(Bundle bundle) {
             ISystemUiProxy proxy = ISystemUiProxy.Stub.asInterface(
                     bundle.getBinder(KEY_EXTRA_SYSUI_PROXY));
-            IPip pip = IPip.Stub.asInterface(bundle.getBinder(KEY_EXTRA_SHELL_PIP));
-            ISplitScreen splitscreen = ISplitScreen.Stub.asInterface(bundle.getBinder(
-                    KEY_EXTRA_SHELL_SPLIT_SCREEN));
-            IOneHanded onehanded = IOneHanded.Stub.asInterface(
-                    bundle.getBinder(KEY_EXTRA_SHELL_ONE_HANDED));
-            IShellTransitions shellTransitions = IShellTransitions.Stub.asInterface(
-                    bundle.getBinder(KEY_EXTRA_SHELL_SHELL_TRANSITIONS));
-            IStartingWindow startingWindow = IStartingWindow.Stub.asInterface(
-                    bundle.getBinder(KEY_EXTRA_SHELL_STARTING_WINDOW));
             MAIN_EXECUTOR.execute(() -> {
-                SystemUiProxy.INSTANCE.get(TouchInteractionService.this).setProxy(proxy, pip,
-                        splitscreen, onehanded, shellTransitions, startingWindow);
+                SystemUiProxy.INSTANCE.get(TouchInteractionService.this).setProxy(proxy);
                 TouchInteractionService.this.initInputMonitor();
                 preloadOverview(true /* fromInit */);
                 mDeviceState.runOnUserUnlocked(() -> {
@@ -441,7 +421,7 @@
         }
         disposeEventHandlers();
         mDeviceState.destroy();
-        SystemUiProxy.INSTANCE.get(this).clearProxy();
+        SystemUiProxy.INSTANCE.get(this).setProxy(null);
         ProtoTracer.INSTANCE.get(this).stop();
         ProtoTracer.INSTANCE.get(this).remove(this);
 
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 261c43f..ad98071 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -39,7 +39,6 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_DISMISS_SWIPE_UP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_LAUNCH_SWIPE_DOWN;
 import static com.android.launcher3.statehandlers.DepthController.DEPTH;
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.SystemUiController.UI_STATE_OVERVIEW;
 import static com.android.quickstep.TaskUtils.checkCurrentOrManagedUserId;
@@ -134,13 +133,13 @@
 import com.android.quickstep.util.TaskViewSimulator;
 import com.android.quickstep.util.TransformParams;
 import com.android.systemui.plugins.ResourceProvider;
+import com.android.systemui.shared.recents.IPinnedStackAnimationListener;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.Task.TaskKey;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.PackageManagerWrapper;
 import com.android.systemui.shared.system.TaskStackChangeListener;
-import com.android.wm.shell.pip.IPipAnimationListener;
 
 import java.util.ArrayList;
 import java.util.function.Consumer;
@@ -378,7 +377,7 @@
         }
     };
 
-    private final PinnedStackAnimationListener mIPipAnimationListener =
+    private final PinnedStackAnimationListener mIPinnedStackAnimationListener =
             new PinnedStackAnimationListener();
 
     // Used to keep track of the last requested task list id, so that we do not request to load the
@@ -598,9 +597,9 @@
         mLiveTileParams.setSyncTransactionApplier(mSyncTransactionApplier);
         RecentsModel.INSTANCE.get(getContext()).addThumbnailChangeListener(this);
         mIdp.addOnChangeListener(this);
-        mIPipAnimationListener.setActivity(mActivity);
+        mIPinnedStackAnimationListener.setActivity(mActivity);
         SystemUiProxy.INSTANCE.get(getContext()).setPinnedStackAnimationListener(
-                mIPipAnimationListener);
+                mIPinnedStackAnimationListener);
         mOrientationState.initListeners();
         SplitScreenBounds.INSTANCE.addOnChangeListener(this);
         mTaskOverlayFactory.initListeners();
@@ -619,7 +618,7 @@
         mIdp.removeOnChangeListener(this);
         SystemUiProxy.INSTANCE.get(getContext()).setPinnedStackAnimationListener(null);
         SplitScreenBounds.INSTANCE.removeOnChangeListener(this);
-        mIPipAnimationListener.setActivity(null);
+        mIPinnedStackAnimationListener.setActivity(null);
         mOrientationState.destroyListeners();
         mTaskOverlayFactory.removeListeners();
     }
@@ -1625,12 +1624,12 @@
         float clearAllShorterRowCompensation =
                 mIsRtl ? -shorterRowCompensation : shorterRowCompensation;
 
-        // If the total width is shorter than one task's width, move ClearAllButton further away
+        // If the total width is shorter than one grid's width, move ClearAllButton further away
         // accordingly.
         float clearAllShortTotalCompensation = 0;
         float longRowWidth = Math.max(topRowWidth, bottomRowWidth);
-        if (longRowWidth < mTaskWidth) {
-            float shortTotalCompensation = mTaskWidth - longRowWidth;
+        if (longRowWidth < mLastComputedGridSize.width()) {
+            float shortTotalCompensation = mLastComputedGridSize.width() - longRowWidth;
             clearAllShortTotalCompensation =
                     mIsRtl ? -shortTotalCompensation : shortTotalCompensation;
         }
@@ -2923,7 +2922,7 @@
     }
 
     private static class PinnedStackAnimationListener<T extends BaseActivity> extends
-            IPipAnimationListener.Stub {
+            IPinnedStackAnimationListener.Stub {
         private T mActivity;
 
         public void setActivity(T activity) {
@@ -2931,12 +2930,10 @@
         }
 
         @Override
-        public void onPipAnimationStarted() {
-            MAIN_EXECUTOR.execute(() -> {
-                // Needed for activities that auto-enter PiP, which will not trigger a remote
-                // animation to be created
-                mActivity.clearForceInvisibleFlag(STATE_HANDLER_INVISIBILITY_FLAGS);
-            });
+        public void onPinnedStackAnimationStarted() {
+            // Needed for activities that auto-enter PiP, which will not trigger a remote
+            // animation to be created
+            mActivity.clearForceInvisibleFlag(STATE_HANDLER_INVISIBILITY_FLAGS);
         }
     }
 }
diff --git a/res/drawable/ic_gm_close_24.xml b/res/drawable/ic_gm_close_24.xml
new file mode 100644
index 0000000..2c9c932
--- /dev/null
+++ b/res/drawable/ic_gm_close_24.xml
@@ -0,0 +1,9 @@
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24">
+    <path
+        android:fillColor="?android:attr/textColorTertiary"
+        android:pathData="M19,6.41L17.59,5 12,10.59 6.41,5 5,6.41 10.59,12 5,17.59 6.41,19 12,13.41 17.59,19 19,17.59 13.41,12 19,6.41z"/>
+</vector>
diff --git a/res/layout/widgets_full_sheet.xml b/res/layout/widgets_full_sheet.xml
index 6c18d7a..226c4f7 100644
--- a/res/layout/widgets_full_sheet.xml
+++ b/res/layout/widgets_full_sheet.xml
@@ -51,5 +51,13 @@
             android:layout_alignParentEnd="true"
             android:layout_alignParentTop="true"
             android:layout_marginEnd="@dimen/fastscroll_end_margin" />
+
+        <com.android.launcher3.widget.picker.WidgetsRecyclerView
+            android:id="@+id/search_widgets_list_view"
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            android:visibility="gone"
+            android:clipToPadding="false" />
+
     </com.android.launcher3.views.TopRoundedCornerView>
 </com.android.launcher3.widget.picker.WidgetsFullSheet>
\ No newline at end of file
diff --git a/res/layout/widgets_full_sheet_search_and_recommendations.xml b/res/layout/widgets_full_sheet_search_and_recommendations.xml
index 9a6f922..6182255 100644
--- a/res/layout/widgets_full_sheet_search_and_recommendations.xml
+++ b/res/layout/widgets_full_sheet_search_and_recommendations.xml
@@ -34,16 +34,5 @@
         android:textSize="24sp"
         android:layout_marginTop="16dp"
         android:text="@string/widget_button_text"/>
-    <!-- Disable the search bar because it has not been implemented. -->
-    <EditText
-        android:id="@+id/widgets_search_bar"
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:visibility="gone"
-        android:layout_marginTop="16dp"
-        android:background="@drawable/bg_widgets_searchbox"
-        android:drawablePadding="8dp"
-        android:drawableStart="@drawable/ic_allapps_search"
-        android:hint="@string/widgets_full_sheet_search_bar_hint"
-        android:padding="12dp" />
+    <include layout="@layout/widgets_search_bar"/>
 </LinearLayout>
diff --git a/res/layout/widgets_list_row_header.xml b/res/layout/widgets_list_row_header.xml
index 041e007..1590286 100644
--- a/res/layout/widgets_list_row_header.xml
+++ b/res/layout/widgets_list_row_header.xml
@@ -19,7 +19,7 @@
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
     android:background="?android:attr/selectableItemBackground"
-    android:paddingVertical="20dp"
+    android:paddingVertical="@dimen/widget_list_header_view_vertical_padding"
     android:orientation="horizontal">
 
     <ImageView
@@ -52,6 +52,8 @@
             android:id="@+id/app_subtitle"
             android:layout_width="wrap_content"
             android:layout_height="wrap_content"
+            android:ellipsize="end"
+            android:maxLines="1"
             tools:text="m widgets, n shortcuts" />
 
     </LinearLayout>
diff --git a/res/layout/widgets_search_bar.xml b/res/layout/widgets_search_bar.xml
new file mode 100644
index 0000000..252637d
--- /dev/null
+++ b/res/layout/widgets_search_bar.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<com.android.launcher3.widget.picker.search.WidgetsSearchBar
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/widgets_search_bar"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:orientation="horizontal"
+    android:layout_marginTop="16dp"
+    android:background="@drawable/bg_widgets_searchbox"
+    android:padding="12dp"
+    android:visibility="gone">
+
+    <EditText
+        android:id="@+id/widgets_search_bar_edit_text"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:drawablePadding="8dp"
+        android:drawableStart="@drawable/ic_allapps_search"
+        android:background="@null"
+        android:hint="@string/widgets_full_sheet_search_bar_hint"
+        android:maxLines="1"
+        android:layout_weight="1"
+        android:inputType="text"/>
+
+    <ImageButton
+        android:id="@+id/widgets_search_cancel_button"
+        android:layout_height="wrap_content"
+        android:layout_width="wrap_content"
+        android:src="@drawable/ic_gm_close_24"
+        android:background="?android:selectableItemBackground"
+        android:layout_gravity="center"
+        android:visibility="gone"/>
+</com.android.launcher3.widget.picker.search.WidgetsSearchBar>
\ No newline at end of file
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index da43758..d135b43 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -108,6 +108,8 @@
     <dimen name="widget_cell_vertical_padding">8dp</dimen>
     <dimen name="widget_cell_horizontal_padding">16dp</dimen>
 
+    <dimen name="widget_list_header_view_vertical_padding">20dp</dimen>
+
     <dimen name="widget_preview_shadow_blur">0.5dp</dimen>
     <dimen name="widget_preview_key_shadow_distance">1dp</dimen>
     <dimen name="widget_preview_corner_radius">2dp</dimen>
diff --git a/res/values/id.xml b/res/values/id.xml
new file mode 100644
index 0000000..39c49bd
--- /dev/null
+++ b/res/values/id.xml
@@ -0,0 +1,21 @@
+<?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.
+-->
+<resources>
+    <item type="id" name="view_type_widgets_list" />
+    <item type="id" name="view_type_widgets_header" />
+    <item type="id" name="view_type_widgets_search_header" />
+</resources>
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
index b972c6f..cc36f63 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -221,6 +221,27 @@
         assertThat(currentList).containsExactlyElementsIn(newList);
     }
 
+    @Test
+    public void headersContentsMix_headerWidgetsModified_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 one of the headers widgets list modified.
+        List<WidgetsListBaseEntry> newList = List.of(
+                new WidgetsListHeaderEntry(
+                        mHeaderA.mPkgItem, mHeaderA.mTitleSectionName,
+                        mHeaderA.mWidgets.subList(0, 1)),
+                mHeaderB, mContentE);
+
+        // WHEN computing the list difference.
+        mWidgetsDiffReporter.process(currentList, newList, COMPARATOR);
+
+        // THEN notify "A" has been changed.
+        verify(mAdapter).notifyItemChanged(/* position= */ 0);
+        // THEN the current list contains all elements from the new list.
+        assertThat(currentList).containsExactlyElementsIn(newList);
+    }
+
 
     private WidgetsListHeaderEntry createWidgetsHeaderEntry(String packageName, String appName,
             int numOfWidgets) {
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
index a7c8d92..e1214ff 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
@@ -26,6 +26,8 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.os.Process;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 
 import androidx.recyclerview.widget.RecyclerView;
@@ -37,6 +39,7 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
@@ -67,6 +70,7 @@
 
     private WidgetsListAdapter mAdapter;
     private InvariantDeviceProfile mTestProfile;
+    private UserHandle mUserHandle;
     private Context mContext;
 
     @Before
@@ -76,6 +80,7 @@
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
+        mUserHandle = Process.myUserHandle();
         mAdapter = new WidgetsListAdapter(mContext, mMockLayoutInflater, mMockWidgetCache,
                 mIconCache, null, null);
         mAdapter.registerAdapterDataObserver(mListener);
@@ -126,7 +131,8 @@
         mAdapter.setWidgets(generateSampleMap(3));
 
         // WHEN com.google.test.1 header is expanded.
-        mAdapter.onHeaderClicked(/* isExpanded= */ true, TEST_PACKAGE_PLACEHOLDER + 1);
+        mAdapter.onHeaderClicked(/* showWidgets= */ true,
+                new PackageUserKey(TEST_PACKAGE_PLACEHOLDER + 1, mUserHandle));
 
         // THEN the visible entries list becomes:
         // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
@@ -143,7 +149,8 @@
         // 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);
+        mAdapter.onHeaderClicked(/* showWidgets= */ true,
+                new PackageUserKey(TEST_PACKAGE_PLACEHOLDER + 1, mUserHandle));
         Mockito.reset(mListener);
 
         // WHEN the adapter is updated with the same list of apps but com.google.test1 has 2 widgets
@@ -200,6 +207,30 @@
         verify(mListener).onItemRangeRemoved(/* positionStart= */ 3, /* itemCount= */ 1);
     }
 
+    @Test
+    public void setWidgetsOnSearch_expandedApp_shouldResetExpandedApp() {
+        // GIVEN a list of widgets entries:
+        // [com.google.test0, com.google.test0 content,
+        //  com.google.test1, com.google.test1 content,
+        //  com.google.test2, com.google.test2 content]
+        // The visible widgets entries: [com.google.test0, com.google.test1, com.google.test2].
+        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(2);
+        mAdapter.setWidgetsOnSearch(allEntries);
+        // GIVEN com.google.test.1 header is expanded. The visible entries list becomes:
+        // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
+        mAdapter.onHeaderClicked(/* showWidgets= */ true,
+                new PackageUserKey(TEST_PACKAGE_PLACEHOLDER + 1, mUserHandle));
+        Mockito.reset(mListener);
+
+        // WHEN same widget entries are set again.
+        mAdapter.setWidgetsOnSearch(allEntries);
+
+        // THEN expanded app is reset and the visible entries list becomes:
+        // [com.google.test0, com.google.test1, com.google.test2]
+        verify(mListener).onItemRangeChanged(eq(1), eq(1), isNull());
+        verify(mListener).onItemRangeRemoved(/* positionStart= */ 2, /* itemCount= */ 1);
+    }
+
     /**
      * Generates a list of sample widget entries.
      *
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index 848630e..e8c11da 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -18,7 +18,9 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
 import static org.robolectric.Shadows.shadowOf;
 
 import android.appwidget.AppWidgetProviderInfo;
@@ -26,12 +28,9 @@
 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.R;
@@ -41,10 +40,9 @@
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.testing.TestActivity;
+import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
-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;
@@ -74,12 +72,13 @@
     // testing.
     private ActivityController<TestActivity> mActivityController;
     private TestActivity mTestActivity;
-    private FakeOnHeaderClickListener mFakeOnHeaderClickListener = new FakeOnHeaderClickListener();
 
     @Mock
     private IconCache mIconCache;
     @Mock
     private DeviceProfile mDeviceProfile;
+    @Mock
+    private OnHeaderClickListener mOnHeaderClickListener;
 
     @Before
     public void setUp() {
@@ -99,8 +98,7 @@
         }).when(mIconCache).getTitleNoCache(any());
 
         mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
-                LayoutInflater.from(mTestActivity),
-                mFakeOnHeaderClickListener);
+                LayoutInflater.from(mTestActivity), mOnHeaderClickListener);
     }
 
     @After
@@ -125,6 +123,23 @@
         assertThat(appSubtitle.getText()).isEqualTo("3 widgets");
     }
 
+    @Test
+    public void bindViewHolder_shouldAttachOnHeaderClickListener() {
+        WidgetsListHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
+                new FrameLayout(mTestActivity));
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        WidgetsListHeaderEntry entry = generateSampleAppHeader(
+                APP_NAME,
+                TEST_PACKAGE,
+                /* numOfWidgets= */ 3);
+
+        mViewHolderBinder.bindViewHolder(viewHolder, entry);
+        widgetsListHeader.callOnClick();
+
+        verify(mOnHeaderClickListener).onHeaderClicked(eq(true),
+                eq(new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)));
+    }
+
     private WidgetsListHeaderEntry generateSampleAppHeader(String appName, String packageName,
             int numOfWidgets) {
         PackageItemInfo appInfo = new PackageItemInfo(packageName);
@@ -152,22 +167,4 @@
         }
         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/WidgetsListSearchHeaderViewHolderBinderTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
new file mode 100644
index 0000000..07fbfd2
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
@@ -0,0 +1,171 @@
+/*
+ * 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.ArgumentMatchers.eq;
+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 android.widget.FrameLayout;
+import android.widget.TextView;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+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.util.PackageUserKey;
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
+
+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 WidgetsListSearchHeaderViewHolderBinderTest {
+    private static final String TEST_PACKAGE = "com.google.test";
+    private static final String APP_NAME = "Test app";
+
+    private Context mContext;
+    private WidgetsListSearchHeaderViewHolderBinder 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 IconCache mIconCache;
+    @Mock
+    private DeviceProfile mDeviceProfile;
+    @Mock
+    private OnHeaderClickListener mOnHeaderClickListener;
+
+    @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 WidgetsListSearchHeaderViewHolderBinder(
+                LayoutInflater.from(mTestActivity), mOnHeaderClickListener);
+    }
+
+    @After
+    public void tearDown() {
+        mActivityController.destroy();
+    }
+
+    @Test
+    public void bindViewHolder_appWith3Widgets_shouldShowTheCorrectAppNameAndSubtitle() {
+        WidgetsListSearchHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
+                new FrameLayout(mTestActivity));
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        WidgetsListSearchHeaderEntry entry = generateSampleSearchHeader(
+                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(".SampleWidget0, .SampleWidget1, .SampleWidget2");
+    }
+
+    @Test
+    public void bindViewHolder_shouldAttachOnHeaderClickListener() {
+        WidgetsListSearchHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
+                new FrameLayout(mTestActivity));
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        WidgetsListSearchHeaderEntry entry = generateSampleSearchHeader(
+                APP_NAME,
+                TEST_PACKAGE,
+                /* numOfWidgets= */ 3);
+
+        mViewHolderBinder.bindViewHolder(viewHolder, entry);
+        widgetsListHeader.callOnClick();
+
+        verify(mOnHeaderClickListener).onHeaderClicked(eq(true),
+                eq(new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)));
+    }
+
+    private WidgetsListSearchHeaderEntry generateSampleSearchHeader(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 WidgetsListSearchHeaderEntry(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;
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipelineTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipelineTest.java
index 8aebf12..17ededd 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipelineTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipelineTest.java
@@ -19,8 +19,6 @@
 import static android.os.Looper.getMainLooper;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.robolectric.Shadows.shadowOf;
@@ -40,6 +38,7 @@
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -56,9 +55,6 @@
 
 @RunWith(RobolectricTestRunner.class)
 public class SimpleWidgetsSearchPipelineTest {
-    private static final SimpleWidgetsSearchPipeline.StringMatcher MATCHER =
-            SimpleWidgetsSearchPipeline.StringMatcher.getInstance();
-
     @Mock private IconCache mIconCache;
 
     private InvariantDeviceProfile mTestProfile;
@@ -73,9 +69,10 @@
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
-                .getComponent().getPackageName())
-                .when(mIconCache).getTitleNoCache(any());
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            return componentWithLabel.getComponent().getShortClassName();
+        }).when(mIconCache).getTitleNoCache(any());
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
@@ -85,54 +82,60 @@
                 createWidgetsHeaderEntry("com.example.android.Calendar", "Calendar", 2);
         mCalendarContentEntry =
                 createWidgetsContentEntry("com.example.android.Calendar", "Calendar", 2);
-        mCameraHeaderEntry = createWidgetsHeaderEntry("com.example.android.Camera", "Camera", 5);
-        mCameraContentEntry = createWidgetsContentEntry("com.example.android.Camera", "Camera", 5);
+        mCameraHeaderEntry = createWidgetsHeaderEntry("com.example.android.Camera", "Camera", 11);
+        mCameraContentEntry = createWidgetsContentEntry("com.example.android.Camera", "Camera", 11);
         mClockHeaderEntry = createWidgetsHeaderEntry("com.example.android.Clock", "Clock", 3);
         mClockContentEntry = createWidgetsContentEntry("com.example.android.Clock", "Clock", 3);
     }
 
     @Test
-    public void query_shouldInformCallbackWithResultsMatchedOnAppName() {
+    public void query_shouldMatchOnAppName() {
         SimpleWidgetsSearchPipeline pipeline = new SimpleWidgetsSearchPipeline(
                 List.of(mCalendarHeaderEntry, mCalendarContentEntry, mCameraHeaderEntry,
                         mCameraContentEntry, mClockHeaderEntry, mClockContentEntry));
 
         pipeline.query("Ca", results ->
-                assertEquals(results, List.of(mCalendarHeaderEntry, mCalendarContentEntry,
-                        mCameraHeaderEntry, mCameraContentEntry)));
+                assertEquals(results,
+                        List.of(
+                                new WidgetsListSearchHeaderEntry(
+                                        mCalendarHeaderEntry.mPkgItem,
+                                        mCalendarHeaderEntry.mTitleSectionName,
+                                        mCalendarHeaderEntry.mWidgets),
+                                mCalendarContentEntry,
+                                new WidgetsListSearchHeaderEntry(
+                                        mCameraHeaderEntry.mPkgItem,
+                                        mCameraHeaderEntry.mTitleSectionName,
+                                        mCameraHeaderEntry.mWidgets),
+                                mCameraContentEntry)));
         shadowOf(getMainLooper()).idle();
     }
 
     @Test
-    public void testMatches() {
-        assertTrue(MATCHER.matches("q", "Q"));
-        assertTrue(MATCHER.matches("q", "  Q"));
-        assertTrue(MATCHER.matches("e", "elephant"));
-        assertTrue(MATCHER.matches("eL", "Elephant"));
-        assertTrue(MATCHER.matches("elephant ", "elephant"));
-        assertTrue(MATCHER.matches("whitec", "white cow"));
-        assertTrue(MATCHER.matches("white  c", "white cow"));
-        assertTrue(MATCHER.matches("white ", "white cow"));
-        assertTrue(MATCHER.matches("white c", "white cow"));
-        assertTrue(MATCHER.matches("电", "电子邮件"));
-        assertTrue(MATCHER.matches("电子", "电子邮件"));
-        assertTrue(MATCHER.matches("다", "다운로드"));
-        assertTrue(MATCHER.matches("드", "드라이브"));
-        assertTrue(MATCHER.matches("åbç", "abc"));
-        assertTrue(MATCHER.matches("ål", "Alpha"));
+    public void query_shouldMatchOnWidgetLabel() {
+        SimpleWidgetsSearchPipeline pipeline = new SimpleWidgetsSearchPipeline(
+                List.of(mCalendarHeaderEntry, mCalendarContentEntry, mCameraHeaderEntry,
+                        mCameraContentEntry));
 
-        assertFalse(MATCHER.matches("phant", "elephant"));
-        assertFalse(MATCHER.matches("elephants", "elephant"));
-        assertFalse(MATCHER.matches("cow", "white cow"));
-        assertFalse(MATCHER.matches("cow", "whiteCow"));
-        assertFalse(MATCHER.matches("dog", "cats&Dogs"));
-        assertFalse(MATCHER.matches("ba", "Bot"));
-        assertFalse(MATCHER.matches("ba", "bot"));
-        assertFalse(MATCHER.matches("子", "电子邮件"));
-        assertFalse(MATCHER.matches("邮件", "电子邮件"));
-        assertFalse(MATCHER.matches("ㄷ", "다운로드 드라이브"));
-        assertFalse(MATCHER.matches("ㄷㄷ", "다운로드 드라이브"));
-        assertFalse(MATCHER.matches("åç", "abc"));
+        pipeline.query("Widget1", results ->
+                assertEquals(results,
+                        List.of(
+                                new WidgetsListSearchHeaderEntry(
+                                        mCalendarHeaderEntry.mPkgItem,
+                                        mCalendarHeaderEntry.mTitleSectionName,
+                                        mCalendarHeaderEntry.mWidgets.subList(1, 2)),
+                                new WidgetsListContentEntry(
+                                        mCalendarHeaderEntry.mPkgItem,
+                                        mCalendarHeaderEntry.mTitleSectionName,
+                                        mCalendarHeaderEntry.mWidgets.subList(1, 2)),
+                                new WidgetsListSearchHeaderEntry(
+                                        mCameraHeaderEntry.mPkgItem,
+                                        mCameraHeaderEntry.mTitleSectionName,
+                                        mCameraHeaderEntry.mWidgets.subList(1, 3)),
+                                new WidgetsListContentEntry(
+                                        mCameraHeaderEntry.mPkgItem,
+                                        mCameraHeaderEntry.mTitleSectionName,
+                                        mCameraHeaderEntry.mWidgets.subList(1, 3)))));
+        shadowOf(getMainLooper()).idle();
     }
 
     private WidgetsListHeaderEntry createWidgetsHeaderEntry(String packageName, String appName,
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
new file mode 100644
index 0000000..7fc9650
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
@@ -0,0 +1,140 @@
+/*
+ * 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.search;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import android.content.Context;
+import android.view.View;
+import android.widget.EditText;
+import android.widget.ImageButton;
+
+import com.android.launcher3.search.SearchAlgorithm;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+
+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 java.util.ArrayList;
+
+@RunWith(RobolectricTestRunner.class)
+public class WidgetsSearchBarControllerTest {
+
+    private WidgetsSearchBarController mController;
+    private Context mContext;
+    private EditText mEditText;
+    private ImageButton mCancelButton;
+    @Mock
+    private SearchModeListener mSearchModeListener;
+    @Mock
+    private SearchAlgorithm<WidgetsListBaseEntry> mSearchAlgorithm;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mEditText = new EditText(mContext);
+        mCancelButton = new ImageButton(mContext);
+        mController = new WidgetsSearchBarController(
+                mSearchAlgorithm, mEditText, mCancelButton, mSearchModeListener);
+    }
+
+    @Test
+    public void onSearchResult_shouldInformSearchModeListener() {
+        ArrayList<WidgetsListBaseEntry> entries = new ArrayList<>();
+        mController.onSearchResult("abc", entries);
+
+        verify(mSearchModeListener).onSearchResults(entries);
+    }
+
+    @Test
+    public void afterTextChanged_shouldInformSearchModeListenerToEnterSearch() {
+        mEditText.setText("abc");
+
+        verify(mSearchModeListener).enterSearchMode();
+        verifyNoMoreInteractions(mSearchModeListener);
+    }
+
+    @Test
+    public void afterTextChanged_shouldDoSearch() {
+        mEditText.setText("abc");
+
+        verify(mSearchAlgorithm).doSearch(eq("abc"), any());
+    }
+
+    @Test
+    public void afterTextChanged_shouldShowCancelButton() {
+        mEditText.setText("abc");
+
+        assertEquals(mCancelButton.getVisibility(), View.VISIBLE);
+    }
+
+    @Test
+    public void afterTextChanged_empty_shouldInformSearchModeListenerToExitSearch() {
+        mEditText.setText("");
+
+        verify(mSearchModeListener).exitSearchMode();
+        verifyNoMoreInteractions(mSearchModeListener);
+    }
+
+    @Test
+    public void afterTextChanged_empty_shouldCancelSearch() {
+        mEditText.setText("");
+
+        verify(mSearchAlgorithm).cancel(true);
+        verifyNoMoreInteractions(mSearchAlgorithm);
+    }
+
+    @Test
+    public void afterTextChanged_empty_shouldHideCancelButton() {
+        mEditText.setText("");
+
+        assertEquals(mCancelButton.getVisibility(), View.GONE);
+    }
+
+    @Test
+    public void cancelSearch_shouldInformSearchModeListenerToExitSearch() {
+        mCancelButton.performClick();
+
+        verify(mSearchModeListener).exitSearchMode();
+        verifyNoMoreInteractions(mSearchModeListener);
+    }
+
+    @Test
+    public void cancelSearch_shouldCancelSearch() {
+        mCancelButton.performClick();
+
+        verify(mSearchAlgorithm).cancel(true);
+        verifyNoMoreInteractions(mSearchAlgorithm);
+    }
+
+    @Test
+    public void cancelSearch_shouldClearSearchBar() {
+        mCancelButton.performClick();
+
+        assertEquals(mEditText.getText().toString(), "");
+    }
+}
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 1546ee3..c57f621 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -113,7 +113,6 @@
 import com.android.launcher3.allapps.AllAppsStore;
 import com.android.launcher3.allapps.AllAppsTransitionController;
 import com.android.launcher3.allapps.DiscoveryBounce;
-import com.android.launcher3.allapps.search.LiveSearchManager;
 import com.android.launcher3.anim.PropertyListBuilder;
 import com.android.launcher3.compat.AccessibilityManagerCompat;
 import com.android.launcher3.config.FeatureFlags;
@@ -277,8 +276,6 @@
 
     private Configuration mOldConfig;
 
-    private LiveSearchManager mLiveSearchManager;
-
     @Thunk
     Workspace mWorkspace;
     @Thunk
@@ -401,8 +398,6 @@
         mAllAppsController = new AllAppsTransitionController(this);
         mStateManager = new StateManager<>(this, NORMAL);
 
-        mLiveSearchManager = new LiveSearchManager(this);
-
         mOnboardingPrefs = createOnboardingPrefs(mSharedPrefs);
 
         mAppWidgetManager = new WidgetManagerHelper(this);
@@ -490,10 +485,6 @@
         }
     }
 
-    public LiveSearchManager getLiveSearchManager() {
-        return mLiveSearchManager;
-    }
-
     protected LauncherOverlayManager getDefaultOverlay() {
         return new LauncherOverlayManager() { };
     }
@@ -1594,7 +1585,6 @@
 
         mOverlayManager.onActivityDestroyed(this);
         mUserChangedCallbackCloseable.close();
-        mLiveSearchManager.stop();
     }
 
     public LauncherAccessibilityDelegate getAccessibilityDelegate() {
diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index edd9a9f..9ede94c 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -201,9 +201,9 @@
         }
         if (!mAH[AdapterHolder.MAIN].appsList.hasFilter()) {
             rebindAdapters(hasWorkApps);
-        }
-        if (hasWorkApps) {
-            resetWorkProfile();
+            if (hasWorkApps) {
+                resetWorkProfile();
+            }
         }
     }
 
diff --git a/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java b/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
index f4d735e..269e390 100644
--- a/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
+++ b/src/com/android/launcher3/allapps/AllAppsSectionDecorator.java
@@ -56,11 +56,10 @@
                 SectionDecorationInfo sectionInfo = adapterItem.sectionDecorationInfo;
                 SectionDecorationHandler decorationHandler = sectionInfo.getDecorationHandler();
                 if (decorationHandler != null) {
-                    decorationHandler.extendBounds(view);
                     if (sectionInfo.isFocusedView()) {
                         decorationHandler.onFocusDraw(c, view);
                     } else {
-                        decorationHandler.onGroupDraw(c);
+                        decorationHandler.onGroupDraw(c, view);
                     }
                 }
             }
@@ -131,26 +130,13 @@
         }
 
         /**
-         * Extends current bounds to include the view.
-         */
-        public void extendBounds(View view) {
-            if (mBounds.isEmpty()) {
-                mBounds.set(view.getLeft(), view.getTop(), view.getRight(), view.getBottom());
-            } else {
-                mBounds.set(
-                        Math.min(mBounds.left, view.getLeft()),
-                        Math.min(mBounds.top, view.getTop()),
-                        Math.max(mBounds.right, view.getRight()),
-                        Math.max(mBounds.bottom, view.getBottom())
-                );
-            }
-        }
-
-        /**
          * Draw bounds onto canvas.
          */
-        public void onGroupDraw(Canvas canvas) {
+        public void onGroupDraw(Canvas canvas, View view) {
+            if (view == null) return;
+
             mPaint.setColor(mFillcolor);
+            mBounds.set(view.getLeft(), view.getTop(), view.getRight(), view.getBottom());
             onDraw(canvas);
         }
 
diff --git a/src/com/android/launcher3/allapps/search/AppsSearchPipeline.java b/src/com/android/launcher3/allapps/search/AppsSearchPipeline.java
index f9fb22e..34895ed 100644
--- a/src/com/android/launcher3/allapps/search/AppsSearchPipeline.java
+++ b/src/com/android/launcher3/allapps/search/AppsSearchPipeline.java
@@ -25,6 +25,7 @@
 import com.android.launcher3.model.BaseModelUpdateTask;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.search.StringMatcherUtility;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -67,10 +68,10 @@
         // apps that don't match all of the words in the query.
         final String queryTextLower = query.toLowerCase();
         final ArrayList<AppInfo> result = new ArrayList<>();
-        DefaultAppSearchAlgorithm.StringMatcher matcher =
-                DefaultAppSearchAlgorithm.StringMatcher.getInstance();
+        StringMatcherUtility.StringMatcher matcher =
+                StringMatcherUtility.StringMatcher.getInstance();
         for (AppInfo info : apps) {
-            if (DefaultAppSearchAlgorithm.matches(info, queryTextLower, matcher)) {
+            if (StringMatcherUtility.matches(queryTextLower, info.title.toString(), matcher)) {
                 result.add(info);
             }
         }
diff --git a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
index 4e213b0..a386ef8 100644
--- a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
+++ b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
@@ -20,12 +20,9 @@
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.allapps.AllAppsGridAdapter.AdapterItem;
-import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.search.SearchAlgorithm;
 import com.android.launcher3.search.SearchCallback;
 
-import java.text.Collator;
-
 /**
  * The default search implementation.
  */
@@ -54,132 +51,4 @@
                         () -> callback.onSearchResult(query, results)),
                 null);
     }
-
-    public static boolean matches(AppInfo info, String query, StringMatcher matcher) {
-        int queryLength = query.length();
-
-        String title = info.title.toString();
-        int titleLength = title.length();
-
-        if (titleLength < queryLength || queryLength <= 0) {
-            return false;
-        }
-
-        if (requestSimpleFuzzySearch(query)) {
-            return title.toLowerCase().contains(query);
-        }
-
-        int lastType;
-        int thisType = Character.UNASSIGNED;
-        int nextType = Character.getType(title.codePointAt(0));
-
-        int end = titleLength - queryLength;
-        for (int i = 0; i <= end; i++) {
-            lastType = thisType;
-            thisType = nextType;
-            nextType = i < (titleLength - 1) ?
-                    Character.getType(title.codePointAt(i + 1)) : Character.UNASSIGNED;
-            if (isBreak(thisType, lastType, nextType) &&
-                    matcher.matches(query, title.substring(i, i + queryLength))) {
-                return true;
-            }
-        }
-        return false;
-    }
-
-    /**
-     * Returns true if the current point should be a break point. Following cases
-     * are considered as break points:
-     *      1) Any non space character after a space character
-     *      2) Any digit after a non-digit character
-     *      3) Any capital character after a digit or small character
-     *      4) Any capital character before a small character
-     */
-    private static boolean isBreak(int thisType, int prevType, int nextType) {
-        switch (prevType) {
-            case Character.UNASSIGNED:
-            case Character.SPACE_SEPARATOR:
-            case Character.LINE_SEPARATOR:
-            case Character.PARAGRAPH_SEPARATOR:
-                return true;
-        }
-        switch (thisType) {
-            case Character.UPPERCASE_LETTER:
-                if (nextType == Character.UPPERCASE_LETTER) {
-                    return true;
-                }
-                // Follow through
-            case Character.TITLECASE_LETTER:
-                // Break point if previous was not a upper case
-                return prevType != Character.UPPERCASE_LETTER;
-            case Character.LOWERCASE_LETTER:
-                // Break point if previous was not a letter.
-                return prevType > Character.OTHER_LETTER || prevType <= Character.UNASSIGNED;
-            case Character.DECIMAL_DIGIT_NUMBER:
-            case Character.LETTER_NUMBER:
-            case Character.OTHER_NUMBER:
-                // Break point if previous was not a number
-                return !(prevType == Character.DECIMAL_DIGIT_NUMBER
-                        || prevType == Character.LETTER_NUMBER
-                        || prevType == Character.OTHER_NUMBER);
-            case Character.MATH_SYMBOL:
-            case Character.CURRENCY_SYMBOL:
-            case Character.OTHER_PUNCTUATION:
-            case Character.DASH_PUNCTUATION:
-                // Always a break point for a symbol
-                return true;
-            default:
-                return  false;
-        }
-    }
-
-    public static class StringMatcher {
-
-        private static final char MAX_UNICODE = '\uFFFF';
-
-        private final Collator mCollator;
-
-        StringMatcher() {
-            // On android N and above, Collator uses ICU implementation which has a much better
-            // support for non-latin locales.
-            mCollator = Collator.getInstance();
-            mCollator.setStrength(Collator.PRIMARY);
-            mCollator.setDecomposition(Collator.CANONICAL_DECOMPOSITION);
-        }
-
-        /**
-         * Returns true if {@param query} is a prefix of {@param target}
-         */
-        public boolean matches(String query, String target) {
-            switch (mCollator.compare(query, target)) {
-                case 0:
-                    return true;
-                case -1:
-                    // The target string can contain a modifier which would make it larger than
-                    // the query string (even though the length is same). If the query becomes
-                    // larger after appending a unicode character, it was originally a prefix of
-                    // the target string and hence should match.
-                    return mCollator.compare(query + MAX_UNICODE, target) > -1;
-                default:
-                    return false;
-            }
-        }
-
-        public static StringMatcher getInstance() {
-            return new StringMatcher();
-        }
-    }
-
-    private static boolean requestSimpleFuzzySearch(String s) {
-        for (int i = 0; i < s.length(); ) {
-            int codepoint = s.codePointAt(i);
-            i += Character.charCount(codepoint);
-            switch (Character.UnicodeScript.of(codepoint)) {
-                case HAN:
-                    //Character.UnicodeScript.HAN: use String.contains to match
-                    return true;
-            }
-        }
-        return false;
-    }
 }
diff --git a/src/com/android/launcher3/allapps/search/LiveSearchManager.java b/src/com/android/launcher3/allapps/search/LiveSearchManager.java
deleted file mode 100644
index adb882a..0000000
--- a/src/com/android/launcher3/allapps/search/LiveSearchManager.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.allapps.search;
-
-import static com.android.launcher3.LauncherState.ALL_APPS;
-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;
-import static com.android.launcher3.widget.WidgetHostViewLoader.getDefaultOptionsForWidget;
-
-import android.app.Activity;
-import android.app.Application.ActivityLifecycleCallbacks;
-import android.appwidget.AppWidgetHost;
-import android.appwidget.AppWidgetHostView;
-import android.appwidget.AppWidgetManager;
-import android.appwidget.AppWidgetProviderInfo;
-import android.content.ComponentName;
-import android.content.Context;
-import android.net.Uri;
-import android.os.Bundle;
-import android.os.UserHandle;
-import android.util.Log;
-
-import androidx.annotation.Nullable;
-import androidx.annotation.UiThread;
-import androidx.annotation.WorkerThread;
-import androidx.lifecycle.Observer;
-import androidx.slice.Slice;
-import androidx.slice.SliceViewManager;
-import androidx.slice.SliceViewManager.SliceCallback;
-
-import com.android.launcher3.Alarm;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherState;
-import com.android.launcher3.statemanager.StateManager.StateListener;
-import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.util.SafeCloseable;
-import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.widget.PendingAddWidgetInfo;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.function.Consumer;
-
-/**
- * Manages Lifecycle for Live search results
- */
-public class LiveSearchManager implements StateListener<LauncherState> {
-
-    private static final String TAG = "LiveSearchManager";
-
-    private static final long SLICE_TIMEOUT_MS = 50;
-    public static final int SEARCH_APPWIDGET_HOST_ID = 2048;
-
-    private final Launcher mLauncher;
-    private final HashMap<Uri, SliceLifeCycle> mUriSliceMap = new HashMap<>();
-
-    private final HashMap<ComponentKey, SearchWidgetInfoContainer> mWidgetPlaceholders =
-            new HashMap<>();
-    private SearchWidgetHost mSearchWidgetHost;
-
-    public LiveSearchManager(Launcher launcher) {
-        mLauncher = launcher;
-        mLauncher.getStateManager().addStateListener(this);
-    }
-
-    /**
-     * Creates new {@link AppWidgetHostView} from {@link AppWidgetProviderInfo}. Caches views for
-     * quicker result within the same search session
-     */
-    public SearchWidgetInfoContainer getPlaceHolderWidget(AppWidgetProviderInfo providerInfo) {
-        if (mSearchWidgetHost == null) {
-            mSearchWidgetHost = new SearchWidgetHost(mLauncher);
-            mSearchWidgetHost.startListening();
-        }
-
-        ComponentName provider = providerInfo.provider;
-        UserHandle userHandle = providerInfo.getProfile();
-
-        ComponentKey key = new ComponentKey(provider, userHandle);
-        if (mWidgetPlaceholders.containsKey(key)) {
-            return mWidgetPlaceholders.get(key);
-        }
-
-        LauncherAppWidgetProviderInfo pinfo = LauncherAppWidgetProviderInfo.fromProviderInfo(
-                mLauncher, providerInfo);
-        PendingAddWidgetInfo pendingAddWidgetInfo = new PendingAddWidgetInfo(pinfo);
-
-        Bundle options = getDefaultOptionsForWidget(mLauncher, pendingAddWidgetInfo);
-        int appWidgetId = mSearchWidgetHost.allocateAppWidgetId();
-        boolean success = AppWidgetManager.getInstance(mLauncher)
-                .bindAppWidgetIdIfAllowed(appWidgetId, userHandle, provider, options);
-        if (!success) {
-            mSearchWidgetHost.deleteAppWidgetId(appWidgetId);
-            mWidgetPlaceholders.put(key, null);
-            return null;
-        }
-
-        SearchWidgetInfoContainer view = (SearchWidgetInfoContainer) mSearchWidgetHost.createView(
-                mLauncher, appWidgetId, providerInfo);
-        view.setTag(pendingAddWidgetInfo);
-        mWidgetPlaceholders.put(key, view);
-        return view;
-    }
-
-    /**
-     * Stop search session
-     */
-    public void stop() {
-        clearWidgetHost();
-    }
-
-    private void clearWidgetHost() {
-        if (mSearchWidgetHost != null) {
-            mSearchWidgetHost.stopListening();
-            mSearchWidgetHost.clearViews();
-            mSearchWidgetHost.deleteHost();
-            mWidgetPlaceholders.clear();
-            mSearchWidgetHost = null;
-        }
-    }
-
-    @Override
-    public void onStateTransitionComplete(LauncherState finalState) {
-        if (finalState != ALL_APPS) {
-            // Clear all search session related objects
-            mUriSliceMap.values().forEach(SliceLifeCycle::destroy);
-            mUriSliceMap.clear();
-
-            clearWidgetHost();
-        }
-    }
-
-    /**
-     * Adds a new observer for the provided uri and returns a callback to cancel this observer
-     */
-    public SafeCloseable addObserver(Uri uri, Observer<Slice> listener,
-            Consumer<Uri> timeoutConsumer) {
-        SliceLifeCycle slc = mUriSliceMap.get(uri);
-        if (slc == null) {
-            slc = new SliceLifeCycle(uri, mLauncher);
-            mUriSliceMap.put(uri, slc);
-        }
-        if (slc.mLastValue != null) {
-            listener.onChanged(slc.mLastValue);
-        }
-
-        // Use a listener wrapper to handle error timeout.
-        Observer<Slice> listenerWrapper = new Observer<Slice>() {
-            final Alarm mErrorTimeout = new Alarm();
-            {
-                mErrorTimeout.setOnAlarmListener(alarm -> {
-                    alarm.cancelAlarm();
-                    timeoutConsumer.accept(uri);
-                });
-                mErrorTimeout.setAlarm(SLICE_TIMEOUT_MS);
-            }
-
-            @Override
-            public void onChanged(Slice slice) {
-                if (slice == null) {
-                    return;
-                }
-
-                if (mErrorTimeout.alarmPending()) {
-                    mErrorTimeout.cancelAlarm();
-                }
-
-                if (mUriSliceMap.get(uri) != null) {
-                    mUriSliceMap.get(uri).mLastValue = slice;
-                }
-
-                listener.onChanged(slice);
-            }
-        };
-
-        slc.addListener(listenerWrapper);
-
-        final SliceLifeCycle sliceLifeCycle = slc;
-        return () -> sliceLifeCycle.removeListener(listenerWrapper);
-    }
-
-    static class SearchWidgetHost extends AppWidgetHost {
-        SearchWidgetHost(Context context) {
-            super(context, SEARCH_APPWIDGET_HOST_ID);
-        }
-
-        @Override
-        protected AppWidgetHostView onCreateView(Context context, int appWidgetId,
-                AppWidgetProviderInfo appWidget) {
-            return new SearchWidgetInfoContainer(context);
-        }
-
-        @Override
-        public void clearViews() {
-            super.clearViews();
-        }
-    }
-
-    private static class SliceLifeCycle
-            implements ActivityLifecycleCallbacks, SliceCallback {
-
-        private final Uri mUri;
-        private final Launcher mLauncher;
-        private final SliceViewManager mSliceViewManager;
-        private final ArrayList<Observer<Slice>> mListeners = new ArrayList<>();
-
-        private boolean mDestroyed = false;
-        private boolean mWasListening = false;
-
-        Slice mLastValue;
-
-        SliceLifeCycle(Uri uri, Launcher launcher) {
-            mUri = uri;
-            mLauncher = launcher;
-            mSliceViewManager = SliceViewManager.getInstance(launcher);
-            launcher.registerActivityLifecycleCallbacks(this);
-
-            if (launcher.isDestroyed()) {
-                onActivityDestroyed(launcher);
-            } else if (launcher.isStarted()) {
-                onActivityStarted(launcher);
-            }
-        }
-
-        @Override
-        public void onActivityDestroyed(Activity activity) {
-            destroy();
-        }
-
-        @Override
-        public void onActivityStarted(Activity activity) {
-            updateListening();
-        }
-
-        @Override
-        public void onActivityStopped(Activity activity) {
-            updateListening();
-        }
-
-        private void updateListening() {
-            boolean isListening = mDestroyed
-                    ? false
-                    : (mLauncher.isStarted() && !mListeners.isEmpty());
-            UI_HELPER_EXECUTOR.execute(() -> uploadListeningBg(isListening));
-        }
-
-        @WorkerThread
-        private void uploadListeningBg(boolean isListening) {
-            if (mWasListening != isListening) {
-                mWasListening = isListening;
-                if (isListening) {
-                    mSliceViewManager.registerSliceCallback(mUri, MAIN_EXECUTOR, this);
-                    // Update slice one-time on the different thread so that we can display
-                    // multiple slices in parallel
-                    THREAD_POOL_EXECUTOR.execute(this::updateSlice);
-                } else {
-                    mSliceViewManager.unregisterSliceCallback(mUri, this);
-                }
-            }
-        }
-
-        @UiThread
-        private void addListener(Observer<Slice> listener) {
-            mListeners.add(listener);
-            updateListening();
-        }
-
-        @UiThread
-        private void removeListener(Observer<Slice> listener) {
-            mListeners.remove(listener);
-            updateListening();
-        }
-
-        @WorkerThread
-        private void updateSlice() {
-            try {
-                Slice s = mSliceViewManager.bindSlice(mUri);
-                MAIN_EXECUTOR.execute(() -> onSliceUpdated(s));
-            } catch (Exception e) {
-                Log.d(TAG, "Error fetching slice", e);
-            }
-        }
-
-        @UiThread
-        @Override
-        public void onSliceUpdated(@Nullable Slice s) {
-            mListeners.forEach(l -> l.onChanged(s));
-        }
-
-        private void destroy() {
-            if (mDestroyed) {
-                return;
-            }
-            mDestroyed = true;
-            mLauncher.unregisterActivityLifecycleCallbacks(this);
-            mListeners.clear();
-        }
-
-        @Override
-        public void onActivityCreated(Activity activity, Bundle bundle) { }
-
-        @Override
-        public void onActivityPaused(Activity activity) { }
-
-        @Override
-        public void onActivityResumed(Activity activity) { }
-
-        @Override
-        public void onActivitySaveInstanceState(Activity activity, Bundle bundle) { }
-    }
-}
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 61f2c2a..988794c 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -131,6 +131,7 @@
 
     /**
      * Fetches high-res icon for the provided ItemInfo and updates the caller when done.
+     *
      * @return a request ID that can be used to cancel the request.
      */
     public HandlerRunnable updateIconInBackground(final ItemInfoUpdateReceiver caller,
@@ -139,7 +140,7 @@
         if (mPendingIconRequestCount <= 0) {
             MODEL_EXECUTOR.setThreadPriority(Process.THREAD_PRIORITY_FOREGROUND);
         }
-        mPendingIconRequestCount ++;
+        mPendingIconRequestCount++;
 
         HandlerRunnable<ItemInfoWithIcon> request = new HandlerRunnable<>(mWorkerHandler,
                 () -> {
@@ -158,7 +159,7 @@
     }
 
     private void onIconRequestEnd() {
-        mPendingIconRequestCount --;
+        mPendingIconRequestCount--;
         if (mPendingIconRequestCount <= 0) {
             MODEL_EXECUTOR.setThreadPriority(Process.THREAD_PRIORITY_BACKGROUND);
         }
@@ -289,7 +290,8 @@
             @NonNull Supplier<LauncherActivityInfo> activityInfoProvider,
             boolean usePkgIcon, boolean useLowResIcon) {
         CacheEntry entry = cacheLocked(infoInOut.getTargetComponent(), infoInOut.user,
-                activityInfoProvider, mLauncherActivityInfoCachingLogic, usePkgIcon, useLowResIcon);
+                activityInfoProvider, mLauncherActivityInfoCachingLogic, usePkgIcon,
+                useLowResIcon);
         applyCacheEntry(entry, infoInOut);
     }
 
@@ -315,7 +317,8 @@
     }
 
     public void updateSessionCache(PackageUserKey key, PackageInstaller.SessionInfo info) {
-        cachePackageInstallInfo(key.mPackageName, key.mUser, info.getAppIcon(), info.getAppLabel());
+        cachePackageInstallInfo(key.mPackageName, key.mUser, info.getAppIcon(),
+                info.getAppLabel());
     }
 
     @Override
diff --git a/src/com/android/launcher3/search/SearchAlgorithm.java b/src/com/android/launcher3/search/SearchAlgorithm.java
index 1665354..a1720c7 100644
--- a/src/com/android/launcher3/search/SearchAlgorithm.java
+++ b/src/com/android/launcher3/search/SearchAlgorithm.java
@@ -31,4 +31,9 @@
      * Cancels any active request.
      */
     void cancel(boolean interruptActiveRequests);
+
+    /**
+     * Cleans up after search is no longer needed.
+     */
+    default void destroy() {};
 }
diff --git a/src/com/android/launcher3/search/StringMatcherUtility.java b/src/com/android/launcher3/search/StringMatcherUtility.java
new file mode 100644
index 0000000..acab52b
--- /dev/null
+++ b/src/com/android/launcher3/search/StringMatcherUtility.java
@@ -0,0 +1,162 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.search;
+
+import java.text.Collator;
+
+/**
+ * Utilities for matching query string to target string.
+ */
+public class StringMatcherUtility {
+
+    /**
+     * Returns {@code true} is {@code query} is a prefix substring of a complete word/phrase in
+     * {@code target}.
+     */
+    public static boolean matches(String query, String target, StringMatcher matcher) {
+        int queryLength = query.length();
+
+        int targetLength = target.length();
+
+        if (targetLength < queryLength || queryLength <= 0) {
+            return false;
+        }
+
+        if (requestSimpleFuzzySearch(query)) {
+            return target.toLowerCase().contains(query);
+        }
+
+        int lastType;
+        int thisType = Character.UNASSIGNED;
+        int nextType = Character.getType(target.codePointAt(0));
+
+        int end = targetLength - queryLength;
+        for (int i = 0; i <= end; i++) {
+            lastType = thisType;
+            thisType = nextType;
+            nextType = i < (targetLength - 1)
+                    ? Character.getType(target.codePointAt(i + 1)) : Character.UNASSIGNED;
+            if (isBreak(thisType, lastType, nextType)
+                    && matcher.matches(query, target.substring(i, i + queryLength))) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    /**
+     * Returns true if the current point should be a break point. Following cases
+     * are considered as break points:
+     *      1) Any non space character after a space character
+     *      2) Any digit after a non-digit character
+     *      3) Any capital character after a digit or small character
+     *      4) Any capital character before a small character
+     */
+    private static boolean isBreak(int thisType, int prevType, int nextType) {
+        switch (prevType) {
+            case Character.UNASSIGNED:
+            case Character.SPACE_SEPARATOR:
+            case Character.LINE_SEPARATOR:
+            case Character.PARAGRAPH_SEPARATOR:
+                return true;
+        }
+        switch (thisType) {
+            case Character.UPPERCASE_LETTER:
+                if (nextType == Character.UPPERCASE_LETTER) {
+                    return true;
+                }
+                // Follow through
+            case Character.TITLECASE_LETTER:
+                // Break point if previous was not a upper case
+                return prevType != Character.UPPERCASE_LETTER;
+            case Character.LOWERCASE_LETTER:
+                // Break point if previous was not a letter.
+                return prevType > Character.OTHER_LETTER || prevType <= Character.UNASSIGNED;
+            case Character.DECIMAL_DIGIT_NUMBER:
+            case Character.LETTER_NUMBER:
+            case Character.OTHER_NUMBER:
+                // Break point if previous was not a number
+                return !(prevType == Character.DECIMAL_DIGIT_NUMBER
+                        || prevType == Character.LETTER_NUMBER
+                        || prevType == Character.OTHER_NUMBER);
+            case Character.MATH_SYMBOL:
+            case Character.CURRENCY_SYMBOL:
+            case Character.OTHER_PUNCTUATION:
+            case Character.DASH_PUNCTUATION:
+                // Always a break point for a symbol
+                return true;
+            default:
+                return  false;
+        }
+    }
+
+    /**
+     * Performs locale sensitive string comparison using {@link Collator}.
+     */
+    public static class StringMatcher {
+
+        private static final char MAX_UNICODE = '\uFFFF';
+
+        private final Collator mCollator;
+
+        StringMatcher() {
+            // On android N and above, Collator uses ICU implementation which has a much better
+            // support for non-latin locales.
+            mCollator = Collator.getInstance();
+            mCollator.setStrength(Collator.PRIMARY);
+            mCollator.setDecomposition(Collator.CANONICAL_DECOMPOSITION);
+        }
+
+        /**
+         * Returns true if {@param query} is a prefix of {@param target}
+         */
+        public boolean matches(String query, String target) {
+            switch (mCollator.compare(query, target)) {
+                case 0:
+                    return true;
+                case -1:
+                    // The target string can contain a modifier which would make it larger than
+                    // the query string (even though the length is same). If the query becomes
+                    // larger after appending a unicode character, it was originally a prefix of
+                    // the target string and hence should match.
+                    return mCollator.compare(query + MAX_UNICODE, target) > -1;
+                default:
+                    return false;
+            }
+        }
+
+        public static StringMatcher getInstance() {
+            return new StringMatcher();
+        }
+    }
+
+    /**
+     * Matching optimization to search in Chinese.
+     */
+    private static boolean requestSimpleFuzzySearch(String s) {
+        for (int i = 0; i < s.length(); ) {
+            int codepoint = s.codePointAt(i);
+            i += Character.charCount(codepoint);
+            switch (Character.UnicodeScript.of(codepoint)) {
+                case HAN:
+                    //Character.UnicodeScript.HAN: use String.contains to match
+                    return true;
+            }
+        }
+        return false;
+    }
+}
diff --git a/src/com/android/launcher3/views/ScrimView.java b/src/com/android/launcher3/views/ScrimView.java
index c9bd284..7f0765b 100644
--- a/src/com/android/launcher3/views/ScrimView.java
+++ b/src/com/android/launcher3/views/ScrimView.java
@@ -32,7 +32,6 @@
 import com.android.launcher3.Insettable;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.uioverrides.WallpaperColorInfo;
 import com.android.launcher3.uioverrides.WallpaperColorInfo.OnChangeListener;
 import com.android.launcher3.util.Themes;
@@ -42,7 +41,6 @@
  */
 public class ScrimView<T extends Launcher> extends View implements Insettable, OnChangeListener {
 
-    private static final float SCRIM_ALPHA = .95f;
     protected final T mLauncher;
     private final WallpaperColorInfo mWallpaperColorInfo;
     protected final int mEndScrim;
@@ -61,12 +59,7 @@
         super(context, attrs);
         mLauncher = Launcher.cast(Launcher.getLauncher(context));
         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;
+        mEndScrim = Themes.getAttrColor(context, R.attr.allAppsScrimColor);
         mIsScrimDark = ColorUtils.calculateLuminance(mEndScrim) < 0.5f;
 
         mMaxScrimAlpha = 0.7f;
diff --git a/src/com/android/launcher3/widget/WidgetAddFlowHandler.java b/src/com/android/launcher3/widget/WidgetAddFlowHandler.java
index 1ac5a33..9313266 100644
--- a/src/com/android/launcher3/widget/WidgetAddFlowHandler.java
+++ b/src/com/android/launcher3/widget/WidgetAddFlowHandler.java
@@ -15,6 +15,9 @@
  */
 package com.android.launcher3.widget;
 
+import static android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_CONFIGURATION_OPTIONAL;
+import static android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_RECONFIGURABLE;
+
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.os.Parcel;
@@ -78,8 +81,22 @@
         return true;
     }
 
+    /**
+     * Checks whether the widget needs configuration.
+     *
+     * A widget needs configuration if (1) it has a configuration activity and (2)
+     * it's configuration is not optional.
+     *
+     * @return true if the widget needs configuration, false otherwise.
+     */
     public boolean needsConfigure() {
-        return mProviderInfo.configure != null;
+        int featureFlags = mProviderInfo.widgetFeatures;
+        // A widget's configuration is optional only if it's configuration is marked as optional AND
+        // it can be reconfigured later.
+        boolean configurationOptional = (featureFlags & WIDGET_FEATURE_CONFIGURATION_OPTIONAL) != 0
+                && (featureFlags & WIDGET_FEATURE_RECONFIGURABLE) != 0;
+
+        return mProviderInfo.configure != null && !configurationOptional;
     }
 
     public LauncherAppWidgetProviderInfo getProviderInfo(Context context) {
diff --git a/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
index 09517e1..73bae6f 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListBaseEntry.java
@@ -20,10 +20,14 @@
 
 import androidx.annotation.IntDef;
 
+import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.widget.WidgetItemComparator;
 
 import java.lang.annotation.Retention;
+import java.util.List;
+import java.util.stream.Collectors;
 
 /** Holder class to store the package information of an entry shown in the widgets list. */
 public abstract class WidgetsListBaseEntry {
@@ -35,9 +39,14 @@
      */
     public final String mTitleSectionName;
 
-    public WidgetsListBaseEntry(PackageItemInfo pkgItem, String titleSectionName) {
+    public final List<WidgetItem> mWidgets;
+
+    public WidgetsListBaseEntry(PackageItemInfo pkgItem, String titleSectionName,
+            List<WidgetItem> items) {
         mPkgItem = pkgItem;
         mTitleSectionName = titleSectionName;
+        this.mWidgets =
+                items.stream().sorted(new WidgetItemComparator()).collect(Collectors.toList());
     }
 
     /**
@@ -51,10 +60,11 @@
     public abstract int getRank();
 
     @Retention(SOURCE)
-    @IntDef({RANK_WIDGETS_LIST_HEADER, RANK_WIDGETS_LIST_CONTENT})
+    @IntDef({RANK_WIDGETS_LIST_HEADER, RANK_WIDGETS_LIST_SEARCH_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;
+    public static final int RANK_WIDGETS_LIST_SEARCH_HEADER = 2;
+    public static final int RANK_WIDGETS_LIST_CONTENT = 3;
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
index b0cb8c7..0328cf6 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListContentEntry.java
@@ -17,10 +17,8 @@
 
 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
@@ -28,18 +26,14 @@
  */
 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());
+        super(pkgItem, titleSectionName, items);
     }
 
     @Override
     public String toString() {
-        return mPkgItem.packageName + ":" + mWidgets.size();
+        return "Content:" + mPkgItem.packageName + ":" + mWidgets.size();
     }
 
     @Override
@@ -47,4 +41,12 @@
     public int getRank() {
         return RANK_WIDGETS_LIST_CONTENT;
     }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (!(obj instanceof WidgetsListContentEntry)) return false;
+        WidgetsListContentEntry otherEntry = (WidgetsListContentEntry) obj;
+        return mWidgets.equals(otherEntry.mWidgets) && mPkgItem.equals(otherEntry.mPkgItem)
+                && mTitleSectionName.equals(otherEntry.mTitleSectionName);
+    }
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
index 6899647..1fdc399 100644
--- a/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetsListHeaderEntry.java
@@ -18,7 +18,7 @@
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 
-import java.util.Collection;
+import java.util.List;
 
 /** An information holder for an app which has widgets or/and shortcuts. */
 public final class WidgetsListHeaderEntry extends WidgetsListBaseEntry {
@@ -30,8 +30,8 @@
     private boolean mHasEntryUpdated = false;
 
     public WidgetsListHeaderEntry(PackageItemInfo pkgItem, String titleSectionName,
-            Collection<WidgetItem> items) {
-        super(pkgItem, titleSectionName);
+            List<WidgetItem> items) {
+        super(pkgItem, titleSectionName, items);
         widgetsCount = (int) items.stream().filter(item -> item.widgetInfo != null).count();
         shortcutsCount = Math.max(0, items.size() - widgetsCount);
     }
@@ -57,8 +57,21 @@
     }
 
     @Override
+    public String toString() {
+        return "Header:" + mPkgItem.packageName + ":" + mWidgets.size();
+    }
+
+    @Override
     @Rank
     public int getRank() {
         return RANK_WIDGETS_LIST_HEADER;
     }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (!(obj instanceof WidgetsListHeaderEntry)) return false;
+        WidgetsListHeaderEntry otherEntry = (WidgetsListHeaderEntry) obj;
+        return mWidgets.equals(otherEntry.mWidgets) && mPkgItem.equals(otherEntry.mPkgItem)
+                && mTitleSectionName.equals(otherEntry.mTitleSectionName);
+    }
 }
diff --git a/src/com/android/launcher3/widget/model/WidgetsListSearchHeaderEntry.java b/src/com/android/launcher3/widget/model/WidgetsListSearchHeaderEntry.java
new file mode 100644
index 0000000..2aec3f8
--- /dev/null
+++ b/src/com/android/launcher3/widget/model/WidgetsListSearchHeaderEntry.java
@@ -0,0 +1,72 @@
+/*
+ * 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.List;
+
+/** An information holder for an app which has widgets or/and shortcuts, to be shown in search. */
+public final class WidgetsListSearchHeaderEntry extends WidgetsListBaseEntry {
+
+    private boolean mIsWidgetListShown = false;
+    private boolean mHasEntryUpdated = false;
+
+    public WidgetsListSearchHeaderEntry(PackageItemInfo pkgItem, String titleSectionName,
+            List<WidgetItem> items) {
+        super(pkgItem, titleSectionName, items);
+    }
+
+    /** 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
+    public String toString() {
+        return "SearchHeader:" + mPkgItem.packageName + ":" + mWidgets.size();
+    }
+
+    @Override
+    @Rank
+    public int getRank() {
+        return RANK_WIDGETS_LIST_SEARCH_HEADER;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (!(obj instanceof WidgetsListSearchHeaderEntry)) return false;
+        WidgetsListSearchHeaderEntry otherEntry = (WidgetsListSearchHeaderEntry) obj;
+        return mWidgets.equals(otherEntry.mWidgets) && mPkgItem.equals(otherEntry.mPkgItem)
+                && mTitleSectionName.equals(otherEntry.mTitleSectionName);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java b/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
new file mode 100644
index 0000000..7372751
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
@@ -0,0 +1,28 @@
+/*
+ * 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 com.android.launcher3.util.PackageUserKey;
+
+/**
+ * A listener to be invoked when a header is clicked.
+ */
+public interface OnHeaderClickListener {
+    /**
+     * Calls when a header is clicked to show / hide widgets for a package.
+     */
+    void onHeaderClicked(boolean showWidgets, PackageUserKey packageUserKey);
+}
diff --git a/src/com/android/launcher3/widget/picker/SearchAndRecommendationsScrollController.java b/src/com/android/launcher3/widget/picker/SearchAndRecommendationsScrollController.java
index 95fa05f..7eb5b83 100644
--- a/src/com/android/launcher3/widget/picker/SearchAndRecommendationsScrollController.java
+++ b/src/com/android/launcher3/widget/picker/SearchAndRecommendationsScrollController.java
@@ -34,6 +34,7 @@
     private final boolean mHasWorkProfile;
     private final SearchAndRecommendationViewHolder mViewHolder;
     private final WidgetsRecyclerView mPrimaryRecyclerView;
+    private final WidgetsRecyclerView mSearchRecyclerView;
 
     // The following are only non null if mHasWorkProfile is true.
     @Nullable private final WidgetsRecyclerView mWorkRecyclerView;
@@ -48,12 +49,14 @@
             SearchAndRecommendationViewHolder viewHolder,
             WidgetsRecyclerView primaryRecyclerView,
             @Nullable WidgetsRecyclerView workRecyclerView,
+            WidgetsRecyclerView searchRecyclerView,
             @Nullable View personalWorkTabsView,
             @Nullable PersonalWorkPagedView primaryWorkViewPager) {
         mHasWorkProfile = hasWorkProfile;
         mViewHolder = viewHolder;
         mPrimaryRecyclerView = primaryRecyclerView;
         mWorkRecyclerView = workRecyclerView;
+        mSearchRecyclerView = searchRecyclerView;
         mPrimaryWorkTabsView = personalWorkTabsView;
         mPrimaryWorkViewPager = primaryWorkViewPager;
         mCurrentRecyclerView = mPrimaryRecyclerView;
@@ -149,6 +152,11 @@
                     mPrimaryRecyclerView.getPaddingRight(),
                     mPrimaryRecyclerView.getPaddingBottom());
         }
+        mSearchRecyclerView.setPadding(
+                mSearchRecyclerView.getPaddingLeft(),
+                topContainerHeight,
+                mSearchRecyclerView.getPaddingRight(),
+                mSearchRecyclerView.getPaddingBottom());
     }
 
     /**
diff --git a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
index dbd1bdf..2366609 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsDiffReporter.java
@@ -25,6 +25,7 @@
 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.model.WidgetsListSearchHeaderEntry;
 import com.android.launcher3.widget.picker.WidgetsListAdapter.WidgetListBaseRowEntryComparator;
 
 import java.util.ArrayList;
@@ -113,7 +114,7 @@
                 // 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)
+                        || hasHeaderUpdated(orgRowEntry, newRowEntry)
                         || hasWidgetsListChanged(orgRowEntry, newRowEntry)) {
                     index = currentEntries.indexOf(orgRowEntry);
                     currentEntries.set(index, newRowEntry);
@@ -174,12 +175,16 @@
      * 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;
+    private boolean hasHeaderUpdated(WidgetsListBaseEntry curRow, WidgetsListBaseEntry newRow) {
+        if (newRow instanceof WidgetsListHeaderEntry && curRow instanceof WidgetsListHeaderEntry) {
+            return ((WidgetsListHeaderEntry) newRow).hasEntryUpdated() || !curRow.equals(newRow);
         }
-        WidgetsListHeaderEntry newRowEntry = (WidgetsListHeaderEntry) newRow;
-        return newRowEntry.hasEntryUpdated();
+        if (newRow instanceof WidgetsListSearchHeaderEntry
+                && curRow instanceof WidgetsListSearchHeaderEntry) {
+            return ((WidgetsListSearchHeaderEntry) newRow).hasEntryUpdated()
+                    || !curRow.equals(newRow);
+        }
+        return false;
     }
 
     private boolean isSamePackageItemInfo(PackageItemInfo curInfo, PackageItemInfo newInfo) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 330175f..6b3c71a 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -34,7 +34,6 @@
 import android.view.View;
 import android.view.animation.AnimationUtils;
 import android.view.animation.Interpolator;
-import android.widget.EditText;
 import android.widget.TextView;
 
 import androidx.annotation.Nullable;
@@ -53,6 +52,8 @@
 import com.android.launcher3.widget.BaseWidgetSheet;
 import com.android.launcher3.widget.LauncherAppWidgetHost.ProviderChangedListener;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.picker.search.SearchModeListener;
+import com.android.launcher3.widget.picker.search.WidgetsSearchBar;
 import com.android.launcher3.workprofile.PersonalWorkPagedView;
 import com.android.launcher3.workprofile.PersonalWorkSlidingTabStrip.OnActivePageChangedListener;
 
@@ -64,7 +65,7 @@
  */
 public class WidgetsFullSheet extends BaseWidgetSheet
         implements Insettable, ProviderChangedListener, OnActivePageChangedListener,
-        WidgetsRecyclerView.HeaderViewDimensionsProvider {
+        WidgetsRecyclerView.HeaderViewDimensionsProvider, SearchModeListener {
 
     private static final long DEFAULT_OPEN_DURATION = 267;
     private static final long FADE_IN_DURATION = 150;
@@ -81,6 +82,7 @@
 
     @Nullable private PersonalWorkPagedView mViewPager;
     private int mInitialTabsHeight = 0;
+    private boolean mIsInSearchMode;
     private View mTabsView;
     private TextView mNoWidgetsView;
     private SearchAndRecommendationViewHolder mSearchAndRecommendationViewHolder;
@@ -91,6 +93,7 @@
         mHasWorkProfile = context.getSystemService(LauncherApps.class).getProfiles().size() > 1;
         mAdapters.put(AdapterHolder.PRIMARY, new AdapterHolder(AdapterHolder.PRIMARY));
         mAdapters.put(AdapterHolder.WORK, new AdapterHolder(AdapterHolder.WORK));
+        mAdapters.put(AdapterHolder.SEARCH, new AdapterHolder(AdapterHolder.SEARCH));
     }
 
     public WidgetsFullSheet(Context context, AttributeSet attrs) {
@@ -138,6 +141,7 @@
                 mSearchAndRecommendationViewHolder,
                 findViewById(R.id.primary_widgets_list_view),
                 mHasWorkProfile ? findViewById(R.id.work_widgets_list_view) : null,
+                findViewById(R.id.search_widgets_list_view),
                 mTabsView,
                 mViewPager);
         fastScroller.setOnFastScrollChangeListener(mSearchAndRecommendationsScrollController);
@@ -145,17 +149,25 @@
         mNoWidgetsView = findViewById(R.id.no_widgets_text);
 
         onWidgetsBound();
+
+        mSearchAndRecommendationViewHolder.mSearchBar.initialize(
+                mLauncher.getPopupDataProvider().getAllWidgets(), /* searchModeListener= */ this);
     }
 
     @Override
     public void onActivePageChanged(int currentActivePage) {
         AdapterHolder currentAdapterHolder = mAdapters.get(currentActivePage);
-        WidgetsRecyclerView currentRecyclerView = currentAdapterHolder.mWidgetsRecyclerView;
-        currentRecyclerView.bindFastScrollbar();
-        mSearchAndRecommendationsScrollController.setCurrentRecyclerView(currentRecyclerView);
+        WidgetsRecyclerView currentRecyclerView =
+                mAdapters.get(currentActivePage).mWidgetsRecyclerView;
 
         updateNoWidgetsView(currentAdapterHolder);
 
+        attachScrollbarToRecyclerView(currentRecyclerView);
+    }
+
+    private void attachScrollbarToRecyclerView(WidgetsRecyclerView recyclerView) {
+        recyclerView.bindFastScrollbar();
+        mSearchAndRecommendationsScrollController.setCurrentRecyclerView(recyclerView);
         reset();
     }
 
@@ -173,11 +185,15 @@
         if (mHasWorkProfile) {
             mAdapters.get(AdapterHolder.WORK).mWidgetsRecyclerView.scrollToTop();
         }
+        mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView.scrollToTop();
         mSearchAndRecommendationsScrollController.reset();
     }
 
     @VisibleForTesting
     public WidgetsRecyclerView getRecyclerView() {
+        if (mIsInSearchMode) {
+            return mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView;
+        }
         if (!mHasWorkProfile || mViewPager.getCurrentPage() == AdapterHolder.PRIMARY) {
             return mAdapters.get(AdapterHolder.PRIMARY).mWidgetsRecyclerView;
         }
@@ -289,6 +305,8 @@
 
         AdapterHolder primaryUserAdapterHolder = mAdapters.get(AdapterHolder.PRIMARY);
         primaryUserAdapterHolder.setup(findViewById(R.id.primary_widgets_list_view));
+        AdapterHolder searchAdapterHolder = mAdapters.get(AdapterHolder.SEARCH);
+        searchAdapterHolder.setup(findViewById(R.id.search_widgets_list_view));
         primaryUserAdapterHolder.mWidgetsListAdapter.setWidgets(allWidgets);
         updateNoWidgetsView(primaryUserAdapterHolder);
 
@@ -300,6 +318,40 @@
         }
     }
 
+    @Override
+    public void enterSearchMode() {
+        if (mIsInSearchMode) return;
+        setViewVisibilityBasedOnSearch(/*isInSearchMode= */ true);
+        attachScrollbarToRecyclerView(mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView);
+    }
+
+    @Override
+    public void exitSearchMode() {
+        setViewVisibilityBasedOnSearch(/*isInSearchMode=*/ false);
+        if (mHasWorkProfile) {
+            mViewPager.snapToPage(AdapterHolder.PRIMARY);
+        }
+        attachScrollbarToRecyclerView(mAdapters.get(AdapterHolder.PRIMARY).mWidgetsRecyclerView);
+    }
+
+    @Override
+    public void onSearchResults(List<WidgetsListBaseEntry> entries) {
+        mAdapters.get(AdapterHolder.SEARCH).mWidgetsListAdapter.setWidgetsOnSearch(entries);
+    }
+
+    private void setViewVisibilityBasedOnSearch(boolean isInSearchMode) {
+        mIsInSearchMode = isInSearchMode;
+        if (mHasWorkProfile) {
+            mViewPager.setVisibility(isInSearchMode ? GONE : VISIBLE);
+            mTabsView.setVisibility(isInSearchMode ? GONE : VISIBLE);
+        } else {
+            mAdapters.get(AdapterHolder.PRIMARY).mWidgetsRecyclerView
+                    .setVisibility(isInSearchMode ? GONE : VISIBLE);
+        }
+        mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView
+                .setVisibility(mIsInSearchMode ? VISIBLE : GONE);
+    }
+
     private void open(boolean animate) {
         if (animate) {
             if (getPopupContainer().getInsets().bottom > 0) {
@@ -385,14 +437,16 @@
 
     @Override
     public int getHeaderViewHeight() {
-        // No need to check work profile here because mInitialTabHeight is always 0 if there is no
-        // work profile.
-        return mInitialTabsHeight
-                + measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mContainer);
+        return measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mCollapseHandle)
+                + measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mHeaderTitle)
+                + measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mSearchBar);
     }
 
     /** private the height, in pixel, + the vertical margins of a given view. */
     private static int measureHeightWithVerticalMargins(View view) {
+        if (view.getVisibility() != VISIBLE) {
+            return 0;
+        }
         MarginLayoutParams marginLayoutParams = (MarginLayoutParams) view.getLayoutParams();
         return view.getMeasuredHeight() + marginLayoutParams.bottomMargin
                 + marginLayoutParams.topMargin;
@@ -402,6 +456,7 @@
     private final class AdapterHolder {
         static final int PRIMARY = 0;
         static final int WORK = 1;
+        static final int SEARCH = 2;
 
         private final int mAdapterType;
         private final WidgetsListAdapter mWidgetsListAdapter;
@@ -420,8 +475,16 @@
                     apps.getIconCache(),
                     /* iconClickListener= */ WidgetsFullSheet.this,
                     /* iconLongClickListener= */ WidgetsFullSheet.this);
-            mWidgetsListAdapter.setFilter(
-                    mAdapterType == PRIMARY ? mPrimaryWidgetsFilter : mWorkWidgetsFilter);
+            switch (mAdapterType) {
+                case PRIMARY:
+                    mWidgetsListAdapter.setFilter(mPrimaryWidgetsFilter);
+                    break;
+                case WORK:
+                    mWidgetsListAdapter.setFilter(mWorkWidgetsFilter);
+                    break;
+                default:
+                    break;
+            }
         }
 
         void setup(WidgetsRecyclerView recyclerView) {
@@ -437,7 +500,7 @@
     final class SearchAndRecommendationViewHolder {
         final View mContainer;
         final View mCollapseHandle;
-        final EditText mSearchBar;
+        final WidgetsSearchBar mSearchBar;
         final TextView mHeaderTitle;
 
         SearchAndRecommendationViewHolder(View searchAndRecommendationContainer) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 8b49d1e..9009eb1 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -34,11 +34,12 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.recyclerview.ViewHolderBinder;
 import com.android.launcher3.util.LabelComparator;
+import com.android.launcher3.util.PackageUserKey;
 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 com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -62,8 +63,9 @@
     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 static final int VIEW_TYPE_WIDGETS_LIST = R.id.view_type_widgets_list;
+    private static final int VIEW_TYPE_WIDGETS_HEADER = R.id.view_type_widgets_header;
+    private static final int VIEW_TYPE_WIDGETS_SEARCH_HEADER = R.id.view_type_widgets_search_header;
 
     private final WidgetsDiffReporter mDiffReporter;
     private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
@@ -73,11 +75,13 @@
 
     private List<WidgetsListBaseEntry> mAllEntries = new ArrayList<>();
     private ArrayList<WidgetsListBaseEntry> mVisibleEntries = new ArrayList<>();
-    @Nullable private String mWidgetsContentVisiblePackage = null;
+    @Nullable private PackageUserKey mWidgetsContentVisiblePackageUserKey = null;
 
     private Predicate<WidgetsListBaseEntry> mHeaderAndSelectedContentFilter = entry ->
             entry instanceof WidgetsListHeaderEntry
-                    || entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage);
+                    || entry instanceof WidgetsListSearchHeaderEntry
+                    || new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)
+                    .equals(mWidgetsContentVisiblePackageUserKey);
     @Nullable private Predicate<WidgetsListBaseEntry> mFilter = null;
 
     public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
@@ -87,8 +91,14 @@
         mWidgetsListTableViewHolderBinder = new WidgetsListTableViewHolderBinder(context,
                 layoutInflater, iconClickListener, iconLongClickListener, widgetPreviewLoader);
         mViewHolderBinders.put(VIEW_TYPE_WIDGETS_LIST, mWidgetsListTableViewHolderBinder);
-        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_HEADER,
-                new WidgetsListHeaderViewHolderBinder(layoutInflater, this::onHeaderClicked));
+        mViewHolderBinders.put(
+                VIEW_TYPE_WIDGETS_HEADER,
+                new WidgetsListHeaderViewHolderBinder(
+                        layoutInflater, /*onHeaderClickListener=*/this));
+        mViewHolderBinders.put(
+                VIEW_TYPE_WIDGETS_SEARCH_HEADER,
+                new WidgetsListSearchHeaderViewHolderBinder(
+                        layoutInflater, /*onHeaderClickListener=*/ this));
     }
 
     public void setFilter(Predicate<WidgetsListBaseEntry> filter) {
@@ -122,23 +132,40 @@
         return mVisibleEntries.size();
     }
 
+    /** Returns all items that will be drawn in a recycler view. */
+    public List<WidgetsListBaseEntry> getItems() {
+        return mVisibleEntries;
+    }
+
     /** 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. */
+    /** Updates the widget list based on {@code tempEntries}. */
     public void setWidgets(List<WidgetsListBaseEntry> tempEntries) {
         mAllEntries = tempEntries.stream().sorted(mRowComparator)
                 .collect(Collectors.toList());
         updateVisibleEntries();
     }
 
+    /** Updates the widget list based on {@code searchResults}. */
+    public void setWidgetsOnSearch(List<WidgetsListBaseEntry> searchResults) {
+        // Forget the expanded package every time widget list is refreshed in search mode.
+        mWidgetsContentVisiblePackageUserKey = null;
+        setWidgets(searchResults);
+    }
+
     private void updateVisibleEntries() {
         mAllEntries.forEach(entry -> {
             if (entry instanceof WidgetsListHeaderEntry) {
                 ((WidgetsListHeaderEntry) entry).setIsWidgetListShown(
-                        entry.mPkgItem.packageName.equals(mWidgetsContentVisiblePackage));
+                        new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)
+                                .equals(mWidgetsContentVisiblePackageUserKey));
+            } else if (entry instanceof WidgetsListSearchHeaderEntry) {
+                ((WidgetsListSearchHeaderEntry) entry).setIsWidgetListShown(
+                        new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)
+                                .equals(mWidgetsContentVisiblePackageUserKey));
             }
         });
         List<WidgetsListBaseEntry> newVisibleEntries = mAllEntries.stream()
@@ -189,17 +216,19 @@
             return VIEW_TYPE_WIDGETS_LIST;
         } else if (entry instanceof WidgetsListHeaderEntry) {
             return VIEW_TYPE_WIDGETS_HEADER;
+        } else if (entry instanceof WidgetsListSearchHeaderEntry) {
+            return VIEW_TYPE_WIDGETS_SEARCH_HEADER;
         }
         throw new UnsupportedOperationException("ViewHolderBinder not found for " + entry);
     }
 
     @Override
-    public void onHeaderClicked(boolean showWidgets, String expandedPackage) {
+    public void onHeaderClicked(boolean showWidgets, PackageUserKey packageUserKey) {
         if (showWidgets) {
-            mWidgetsContentVisiblePackage = expandedPackage;
+            mWidgetsContentVisiblePackageUserKey = packageUserKey;
             updateVisibleEntries();
-        } else if (expandedPackage.equals(mWidgetsContentVisiblePackage)) {
-            mWidgetsContentVisiblePackage = null;
+        } else if (packageUserKey.equals(mWidgetsContentVisiblePackageUserKey)) {
+            mWidgetsContentVisiblePackageUserKey = null;
             updateVisibleEntries();
         }
     }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
index 070a9aa..119d094 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -41,6 +41,9 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
+
+import java.util.stream.Collectors;
 
 /**
  * A UI represents a header of an app shown in the full widgets tray.
@@ -173,7 +176,7 @@
                     shortcutsCount);
         } else if (entry.widgetsCount > 0) {
             subtitle = resources.getQuantityString(R.plurals.widgets_count,
-                     entry.widgetsCount, entry.widgetsCount);
+                    entry.widgetsCount, entry.widgetsCount);
         } else {
             subtitle = resources.getQuantityString(R.plurals.shortcuts_count,
                     entry.shortcutsCount, entry.shortcutsCount);
@@ -182,6 +185,32 @@
         mSubtitle.setVisibility(VISIBLE);
     }
 
+    /** Apply app icon, labels and tag using a generic {@link WidgetsListSearchHeaderEntry}. */
+    @UiThread
+    public void applyFromItemInfoWithIcon(WidgetsListSearchHeaderEntry entry) {
+        applyIconAndLabel(entry);
+    }
+
+    @UiThread
+    private void applyIconAndLabel(WidgetsListSearchHeaderEntry entry) {
+        PackageItemInfo info = entry.mPkgItem;
+        setIcon(info);
+        setTitles(entry);
+        setExpanded(entry.isWidgetListShown());
+
+        super.setTag(info);
+
+        verifyHighRes();
+    }
+
+    private void setTitles(WidgetsListSearchHeaderEntry entry) {
+        mTitle.setText(entry.mPkgItem.title);
+
+        mSubtitle.setText(entry.mWidgets.stream()
+                .map(item -> item.label).sorted().collect(Collectors.joining(", ")));
+        mSubtitle.setVisibility(VISIBLE);
+    }
+
     @Override
     public void reapplyItemInfo(ItemInfoWithIcon info) {
         if (getTag() == info) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
index ed53e6f..fcefe3a 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
@@ -20,6 +20,7 @@
 
 import com.android.launcher3.R;
 import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 
 /**
@@ -50,12 +51,9 @@
         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);
+                mOnHeaderClickListener.onHeaderClicked(
+                        isExpanded,
+                        new PackageUserKey(data.mPkgItem.packageName, data.mPkgItem.user)
+                ));
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderHolder.java b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderHolder.java
new file mode 100644
index 0000000..9562af3
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderHolder.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 WidgetsListSearchHeaderHolder extends ViewHolder {
+    final WidgetsListHeader mWidgetsListHeader;
+
+    public WidgetsListSearchHeaderHolder(WidgetsListHeader view) {
+        super(view);
+
+        mWidgetsListHeader = view;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
new file mode 100644
index 0000000..83c7948
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.widget.picker;
+
+import android.view.LayoutInflater;
+import android.view.ViewGroup;
+
+import com.android.launcher3.R;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
+
+/**
+ * Binds data from {@link WidgetsListHeaderEntry} to UI elements in {@link WidgetsListHeaderHolder}.
+ */
+public final class WidgetsListSearchHeaderViewHolderBinder implements
+        ViewHolderBinder<WidgetsListSearchHeaderEntry, WidgetsListSearchHeaderHolder> {
+    private final LayoutInflater mLayoutInflater;
+    private final OnHeaderClickListener mOnHeaderClickListener;
+
+    public WidgetsListSearchHeaderViewHolderBinder(LayoutInflater layoutInflater,
+            OnHeaderClickListener onHeaderClickListener) {
+        mLayoutInflater = layoutInflater;
+        mOnHeaderClickListener = onHeaderClickListener;
+    }
+
+    @Override
+    public WidgetsListSearchHeaderHolder newViewHolder(ViewGroup parent) {
+        WidgetsListHeader header = (WidgetsListHeader) mLayoutInflater.inflate(
+                R.layout.widgets_list_row_header, parent, false);
+
+        return new WidgetsListSearchHeaderHolder(header);
+    }
+
+    @Override
+    public void bindViewHolder(WidgetsListSearchHeaderHolder viewHolder,
+            WidgetsListSearchHeaderEntry data) {
+        WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
+        widgetsListHeader.applyFromItemInfoWithIcon(data);
+        widgetsListHeader.setExpanded(data.isWidgetListShown());
+        widgetsListHeader.setOnExpandChangeListener(isExpanded ->
+                mOnHeaderClickListener.onHeaderClicked(isExpanded,
+                        new PackageUserKey(data.mPkgItem.packageName, data.mPkgItem.user)));
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
index d65a809..9ab6424 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecyclerView.java
@@ -21,13 +21,19 @@
 import android.util.AttributeSet;
 import android.view.MotionEvent;
 import android.view.View;
+import android.widget.TableLayout;
 
 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.DeviceProfile;
 import com.android.launcher3.R;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 
 /**
  * The widgets recycler view.
@@ -39,8 +45,10 @@
     private final int mScrollbarTop;
 
     private final Point mFastScrollerOffset = new Point();
+    private final int mEstimatedWidgetListHeaderHeight;
     private boolean mTouchDownOnScroller;
     private HeaderViewDimensionsProvider mHeaderViewDimensionsProvider;
+    private int mLastVisibleWidgetContentTableHeight = 0;
 
     public WidgetsRecyclerView(Context context) {
         this(context, null);
@@ -55,6 +63,12 @@
         super(context, attrs, defStyleAttr);
         mScrollbarTop = getResources().getDimensionPixelSize(R.dimen.dynamic_grid_edge_margin);
         addOnItemTouchListener(this);
+
+        ActivityContext activity = ActivityContext.lookupContext(getContext());
+        DeviceProfile grid = activity.getDeviceProfile();
+        mEstimatedWidgetListHeaderHeight = grid.iconSizePx
+                + 2 * context.getResources().getDimensionPixelSize(
+                        R.dimen.widget_list_header_view_vertical_padding);
     }
 
     @Override
@@ -123,21 +137,32 @@
 
         View child = getChildAt(0);
         int rowIndex = getChildPosition(child);
-        int y = (child.getMeasuredHeight() * rowIndex);
+        for (int i = 0; i < getChildCount(); i++) {
+            View view = getChildAt(i);
+            if (view instanceof TableLayout) {
+                // This assumes there is ever only one content shown in this recycler view.
+                mLastVisibleWidgetContentTableHeight = view.getMeasuredHeight();
+            }
+        }
+
+        int scrollPosition = getItemsHeight(rowIndex);
         int offset = getLayoutManager().getDecoratedTop(child);
 
-        return getPaddingTop() + y - offset;
+        return getPaddingTop() + scrollPosition - offset;
     }
 
     /**
-     * Returns the available scroll height:
-     *   AvailableScrollHeight = Total height of the all items - last page height
+     * Returns the available scroll height, in pixel.
+     *
+     * <p>If the recycler view can't be scrolled, returns 0.
      */
     @Override
     protected int getAvailableScrollHeight() {
-        View child = getChildAt(0);
-        return child.getMeasuredHeight() * mAdapter.getItemCount() + getScrollBarTop()
-                + getPaddingBottom() - mScrollbar.getHeight();
+        // AvailableScrollHeight = Total height of the all items - first page height
+        int firstPageHeight = getMeasuredHeight() - getPaddingTop() - getPaddingBottom();
+        int totalHeightOfAllItems = getItemsHeight(/* untilIndex= */ mAdapter.getItemCount());
+        int availableScrollHeight = totalHeightOfAllItems - firstPageHeight;
+        return Math.max(0, availableScrollHeight);
     }
 
     private boolean isModelNotReady() {
@@ -181,6 +206,31 @@
     }
 
     /**
+     * Returns the sum of the height, in pixels, of this list adapter's items from index 0 until
+     * {@code untilIndex}.
+     *
+     * <p>If the untilIndex is larger than the total number of items in this adapter, returns the
+     * sum of all items' height.
+     */
+    private int getItemsHeight(int untilIndex) {
+        if (untilIndex > mAdapter.getItems().size()) {
+            untilIndex = mAdapter.getItems().size();
+        }
+        int totalItemsHeight = 0;
+        for (int i = 0; i < untilIndex; i++) {
+            WidgetsListBaseEntry entry = mAdapter.getItems().get(i);
+            if (entry instanceof WidgetsListHeaderEntry) {
+                totalItemsHeight += mEstimatedWidgetListHeaderHeight;
+            } else if (entry instanceof WidgetsListContentEntry) {
+                totalItemsHeight += mLastVisibleWidgetContentTableHeight;
+            } else {
+                throw new UnsupportedOperationException("Can't estimate height for " + entry);
+            }
+        }
+        return totalItemsHeight;
+    }
+
+    /**
      * Provides dimensions of the header view that is shown at the top of a
      * {@link WidgetsRecyclerView}.
      */
diff --git a/src/com/android/launcher3/widget/picker/search/SearchModeListener.java b/src/com/android/launcher3/widget/picker/search/SearchModeListener.java
new file mode 100644
index 0000000..cee7d67
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/search/SearchModeListener.java
@@ -0,0 +1,40 @@
+/*
+ * 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.search;
+
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+
+import java.util.List;
+
+/**
+ * A listener to help with widgets picker search.
+ */
+public interface SearchModeListener {
+    /**
+     * Notifies the subscriber when user enters widget picker search mode.
+     */
+    void enterSearchMode();
+
+    /**
+     * Notifies the subscriber when user exits widget picker search mode.
+     */
+    void exitSearchMode();
+
+    /**
+     * Notifies the subscriber with search results.
+     */
+    void onSearchResults(List<WidgetsListBaseEntry> entries);
+}
diff --git a/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipeline.java b/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipeline.java
index 9911495..5222e8e 100644
--- a/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipeline.java
+++ b/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchPipeline.java
@@ -16,12 +16,19 @@
 
 package com.android.launcher3.widget.picker.search;
 
-import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import static com.android.launcher3.search.StringMatcherUtility.matches;
 
-import java.text.Collator;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.search.StringMatcherUtility.StringMatcher;
+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.model.WidgetsListSearchHeaderEntry;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 /**
  * Implementation of {@link WidgetsPickerSearchPipeline} that performs search by prefix matching on
@@ -37,52 +44,29 @@
 
     @Override
     public void query(String input, Consumer<List<WidgetsListBaseEntry>> callback) {
-        StringMatcher matcher =  StringMatcher.getInstance();
         ArrayList<WidgetsListBaseEntry> results = new ArrayList<>();
-        // TODO(b/157286785): Filter entries based on query prefix matching on widget labels also.
-        for (WidgetsListBaseEntry e : mAllEntries) {
-            if (matcher.matches(input, e.mPkgItem.title.toString())) {
-                results.add(e);
-            }
-        }
+        mAllEntries.stream().filter(entry -> entry instanceof WidgetsListHeaderEntry)
+                .forEach(headerEntry -> {
+                    List<WidgetItem> matchedWidgetItems = filterWidgetItems(
+                            input, headerEntry.mPkgItem.title.toString(), headerEntry.mWidgets);
+                    if (matchedWidgetItems.size() > 0) {
+                        results.add(new WidgetsListSearchHeaderEntry(headerEntry.mPkgItem,
+                                headerEntry.mTitleSectionName, matchedWidgetItems));
+                        results.add(new WidgetsListContentEntry(headerEntry.mPkgItem,
+                                headerEntry.mTitleSectionName, matchedWidgetItems));
+                    }
+                });
         callback.accept(results);
     }
 
-    /**
-     * Performs locale sensitive string comparison using {@link Collator}.
-     */
-    public static class StringMatcher {
-
-        private static final char MAX_UNICODE = '\uFFFF';
-
-        private final Collator mCollator;
-
-        StringMatcher() {
-            mCollator = Collator.getInstance();
-            mCollator.setStrength(Collator.PRIMARY);
-            mCollator.setDecomposition(Collator.CANONICAL_DECOMPOSITION);
+    private List<WidgetItem> filterWidgetItems(String query, String packageTitle,
+            List<WidgetItem> items) {
+        StringMatcher matcher = StringMatcher.getInstance();
+        if (matches(query, packageTitle, matcher)) {
+            return items;
         }
-
-        /**
-         * Returns true if {@param query} is a prefix of {@param target}.
-         */
-        public boolean matches(String query, String target) {
-            switch (mCollator.compare(query, target)) {
-                case 0:
-                    return true;
-                case -1:
-                    // The target string can contain a modifier which would make it larger than
-                    // the query string (even though the length is same). If the query becomes
-                    // larger after appending a unicode character, it was originally a prefix of
-                    // the target string and hence should match.
-                    return mCollator.compare(query + MAX_UNICODE, target) > -1;
-                default:
-                    return false;
-            }
-        }
-
-        public static StringMatcher getInstance() {
-            return new StringMatcher();
-        }
+        return items.stream()
+                .filter(item -> matches(query, item.label, matcher))
+                .collect(Collectors.toList());
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
new file mode 100644
index 0000000..d8e9733
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
@@ -0,0 +1,79 @@
+/*
+ * 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.search;
+
+import android.content.Context;
+import android.util.AttributeSet;
+import android.widget.EditText;
+import android.widget.ImageButton;
+import android.widget.LinearLayout;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.R;
+import com.android.launcher3.search.SearchAlgorithm;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+
+import java.util.List;
+
+/**
+ * View for a search bar with an edit text with a cancel button.
+ */
+public class WidgetsSearchBar extends LinearLayout {
+    private WidgetsSearchBarController mController;
+    private EditText mEditText;
+    private ImageButton mCancelButton;
+
+    public WidgetsSearchBar(Context context) {
+        this(context, null, 0);
+    }
+
+    public WidgetsSearchBar(@NonNull Context context,
+            @Nullable AttributeSet attrs) {
+        this(context, attrs, 0);
+    }
+
+    public WidgetsSearchBar(@NonNull Context context, @Nullable AttributeSet attrs,
+            int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+    }
+
+    /**
+     * Attaches a controller to the search bar which interacts with {@code searchModeListener}.
+     */
+    public void initialize(List<WidgetsListBaseEntry> allWidgets,
+            SearchModeListener searchModeListener) {
+        SearchAlgorithm<WidgetsListBaseEntry> algo =
+                new SimpleWidgetsSearchAlgorithm(new SimpleWidgetsSearchPipeline(allWidgets));
+        mController = new WidgetsSearchBarController(
+                algo, mEditText, mCancelButton, searchModeListener);
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        mEditText = findViewById(R.id.widgets_search_bar_edit_text);
+        mCancelButton = findViewById(R.id.widgets_search_cancel_button);
+    }
+
+    @Override
+    protected void onDetachedFromWindow() {
+        super.onDetachedFromWindow();
+        mController.onDestroy();
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarController.java b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarController.java
new file mode 100644
index 0000000..6c37484
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarController.java
@@ -0,0 +1,111 @@
+/*
+ * 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.search;
+
+import static android.view.View.GONE;
+import static android.view.View.VISIBLE;
+
+import android.text.Editable;
+import android.text.TextWatcher;
+import android.util.Log;
+import android.widget.EditText;
+import android.widget.ImageButton;
+
+import com.android.launcher3.search.SearchAlgorithm;
+import com.android.launcher3.search.SearchCallback;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+
+import java.util.ArrayList;
+
+/**
+ * Controller for a search bar with an edit text and a cancel button.
+ */
+public class WidgetsSearchBarController implements TextWatcher,
+        SearchCallback<WidgetsListBaseEntry> {
+    private static final String TAG = "WidgetsSearchBarController";
+    private static final boolean DEBUG = false;
+
+    protected SearchAlgorithm<WidgetsListBaseEntry> mSearchAlgorithm;
+    protected EditText mInput;
+    protected ImageButton mCancelButton;
+    protected SearchModeListener mSearchModeListener;
+    protected String mQuery;
+
+    public WidgetsSearchBarController(
+            SearchAlgorithm<WidgetsListBaseEntry> algo, EditText editText, ImageButton cancelButton,
+            SearchModeListener searchModeListener) {
+        mSearchAlgorithm = algo;
+        mInput = editText;
+        mInput.addTextChangedListener(this);
+        mCancelButton = cancelButton;
+        mCancelButton.setOnClickListener(v -> clearSearchResult());
+        mSearchModeListener = searchModeListener;
+    }
+
+    @Override
+    public void afterTextChanged(final Editable s) {
+        mQuery = s.toString();
+        if (mQuery.isEmpty()) {
+            mSearchAlgorithm.cancel(/* interruptActiveRequests= */ true);
+            mSearchModeListener.exitSearchMode();
+            mCancelButton.setVisibility(GONE);
+        } else {
+            mSearchAlgorithm.cancel(/* interruptActiveRequests= */ false);
+            mSearchModeListener.enterSearchMode();
+            mSearchAlgorithm.doSearch(mQuery, this);
+            mCancelButton.setVisibility(VISIBLE);
+        }
+    }
+
+    @Override
+    public void beforeTextChanged(CharSequence charSequence, int i, int i1, int i2) {
+        // Do nothing.
+    }
+
+    @Override
+    public void onTextChanged(CharSequence charSequence, int i, int i1, int i2) {
+        // Do nothing.
+    }
+
+    @Override
+    public void onSearchResult(String query, ArrayList<WidgetsListBaseEntry> items) {
+        if (DEBUG) {
+            Log.d(TAG, "onSearchResult query: " + query + " items: " + items);
+        }
+        mSearchModeListener.onSearchResults(items);
+    }
+
+    @Override
+    public void onAppendSearchResult(String query, ArrayList<WidgetsListBaseEntry> items) {
+        // Not needed.
+    }
+
+    @Override
+    public void clearSearchResult() {
+        mSearchAlgorithm.cancel(/* interruptActiveRequests= */ true);
+        mInput.getText().clear();
+        mInput.clearFocus();
+        mSearchModeListener.exitSearchMode();
+    }
+
+    /**
+     * Cleans up after search is no longer needed.
+     */
+    public void onDestroy() {
+        mSearchAlgorithm.destroy();
+    }
+}
diff --git a/src_plugins/com/android/systemui/plugins/BcSmartspaceDataPlugin.java b/src_plugins/com/android/systemui/plugins/BcSmartspaceDataPlugin.java
deleted file mode 100644
index f8a9a04..0000000
--- a/src_plugins/com/android/systemui/plugins/BcSmartspaceDataPlugin.java
+++ /dev/null
@@ -1,44 +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.systemui.plugins;
-
-import android.os.Parcelable;
-
-import com.android.systemui.plugins.annotations.ProvidesInterface;
-
-import java.util.List;
-
-/**
- * Interface to provide SmartspaceTargets to BcSmartspace.
- */
-@ProvidesInterface(action = BcSmartspaceDataPlugin.ACTION, version = BcSmartspaceDataPlugin.VERSION)
-public interface BcSmartspaceDataPlugin extends Plugin {
-    String ACTION = "com.android.systemui.action.PLUGIN_BC_SMARTSPACE_DATA";
-    int VERSION = 1;
-
-    /** Register a listener to get Smartspace data. */
-    void registerListener(SmartspaceTargetListener listener);
-
-    /** Unregister a listener. */
-    void unregisterListener(SmartspaceTargetListener listener);
-
-    /** Provides Smartspace data to registered listeners. */
-    interface SmartspaceTargetListener {
-        /** Each Parcelable is a SmartspaceTarget that represents a card. */
-        void onSmartspaceTargetsUpdated(List<? extends Parcelable> targets);
-    }
-}
diff --git a/tests/Launcher3Tests.xml b/tests/Launcher3Tests.xml
new file mode 100644
index 0000000..3fff622
--- /dev/null
+++ b/tests/Launcher3Tests.xml
@@ -0,0 +1,50 @@
+<?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.
+-->
+<!-- This test config file is auto-generated. -->
+<configuration description="Runs Launcher3 tests.">
+    <option name="test-suite-tag" value="apct" />
+    <option name="test-suite-tag" value="apct-instrumentation" />
+
+    <target_preparer class="com.android.tradefed.targetprep.DeviceSetup">
+        <option name="set-test-harness" value="true" />
+        <option name="run-command" value="am force-stop com.android.launcher3" />
+        <option name="run-command" value="pm uninstall com.google.android.apps.nexuslauncher" />
+        <option name="run-command" value="pm uninstall com.google.android.apps.nexuslauncher.out_of_proc_tests" />
+        <option name="run-command" value="pm uninstall com.google.android.apps.nexuslauncher.tests" />
+        <option name="run-command" value="pm disable com.google.android.googlequicksearchbox" />
+
+        <option name="run-command" value="input keyevent 82" />
+        <option name="run-command" value="settings delete secure assistant" />
+        <option name="run-command" value="settings put global airplane_mode_on 1" />
+        <option name="run-command" value="am broadcast -a android.intent.action.AIRPLANE_MODE" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="Launcher3Tests.apk" />
+        <option name="test-file-name" value="Launcher3.apk" />
+    </target_preparer>
+
+    <metrics_collector class="com.android.tradefed.device.metric.FilePullerLogCollector">
+        <option name="directory-keys" value="/data/user/0/com.android.launcher3/files" />
+        <option name="collect-on-run-ended-only" value="true" />
+    </metrics_collector>
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.launcher3.tests" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+    </test>
+</configuration>
diff --git a/tests/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithmTest.java b/tests/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithmTest.java
deleted file mode 100644
index 39709a9..0000000
--- a/tests/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithmTest.java
+++ /dev/null
@@ -1,98 +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.allapps.search;
-
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-
-import android.content.ComponentName;
-
-import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
-
-import com.android.launcher3.model.data.AppInfo;
-
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-/**
- * Unit tests for {@link DefaultAppSearchAlgorithm}
- */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-public class DefaultAppSearchAlgorithmTest {
-    private static final DefaultAppSearchAlgorithm.StringMatcher MATCHER =
-            DefaultAppSearchAlgorithm.StringMatcher.getInstance();
-
-    @Test
-    public void testMatches() {
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("white cow"), "cow", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whiteCow"), "cow", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whiteCOW"), "cow", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whitecowCOW"), "cow", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("white2cow"), "cow", MATCHER));
-
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whitecow"), "cow", MATCHER));
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whitEcow"), "cow", MATCHER));
-
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whitecowCow"), "cow", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whitecow cow"), "cow", MATCHER));
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whitecowcow"), "cow", MATCHER));
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whit ecowcow"), "cow", MATCHER));
-
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("cats&dogs"), "dog", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("cats&Dogs"), "dog", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("cats&Dogs"), "&", MATCHER));
-
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("2+43"), "43", MATCHER));
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("2+43"), "3", MATCHER));
-
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("Q"), "q", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("  Q"), "q", MATCHER));
-
-        // match lower case words
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("elephant"), "e", MATCHER));
-
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "电", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "电子", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "子", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "邮件", MATCHER));
-
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("Bot"), "ba", MATCHER));
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("bot"), "ba", MATCHER));
-    }
-
-    @Test
-    public void testMatchesVN() {
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("다운로드"), "다", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("드라이브"), "드", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("다운로드 드라이브"), "ㄷ", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("운로 드라이브"), "ㄷ", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("abc"), "åbç", MATCHER));
-        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("Alpha"), "ål", MATCHER));
-
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("다운로드 드라이브"), "ㄷㄷ", MATCHER));
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("로드라이브"), "ㄷ", MATCHER));
-        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("abc"), "åç", MATCHER));
-    }
-
-    private AppInfo getInfo(String title) {
-        AppInfo info = new AppInfo();
-        info.title = title;
-        info.componentName = new ComponentName("Test", title);
-        return info;
-    }
-}
diff --git a/tests/src/com/android/launcher3/search/StringMatcherUtilityTest.java b/tests/src/com/android/launcher3/search/StringMatcherUtilityTest.java
new file mode 100644
index 0000000..413f404
--- /dev/null
+++ b/tests/src/com/android/launcher3/search/StringMatcherUtilityTest.java
@@ -0,0 +1,96 @@
+/*
+ * 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.search;
+
+import static com.android.launcher3.search.StringMatcherUtility.matches;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.launcher3.search.StringMatcherUtility.StringMatcher;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Unit tests for {@link StringMatcherUtility}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class StringMatcherUtilityTest {
+    private static final StringMatcher MATCHER =
+            StringMatcher.getInstance();
+
+    @Test
+    public void testMatches() {
+        assertTrue(matches("white ", "white cow", MATCHER));
+        assertTrue(matches("white c", "white cow", MATCHER));
+        assertTrue(matches("cow", "white cow", MATCHER));
+        assertTrue(matches("cow", "whiteCow", MATCHER));
+        assertTrue(matches("cow", "whiteCOW", MATCHER));
+        assertTrue(matches("cow", "whitecowCOW", MATCHER));
+        assertTrue(matches("cow", "white2cow", MATCHER));
+
+        assertFalse(matches("cow", "whitecow", MATCHER));
+        assertFalse(matches("cow", "whitEcow", MATCHER));
+
+        assertTrue(matches("cow", "whitecowCow", MATCHER));
+        assertTrue(matches("cow", "whitecow cow", MATCHER));
+        assertFalse(matches("cow", "whitecowcow", MATCHER));
+        assertFalse(matches("cow", "whit ecowcow", MATCHER));
+
+        assertTrue(matches("dog", "cats&dogs", MATCHER));
+        assertTrue(matches("dog", "cats&Dogs", MATCHER));
+        assertTrue(matches("&", "cats&Dogs", MATCHER));
+
+        assertTrue(matches("43", "2+43", MATCHER));
+        assertFalse(matches("3", "2+43", MATCHER));
+
+        assertTrue(matches("q", "Q", MATCHER));
+        assertTrue(matches("q", "  Q", MATCHER));
+
+        // match lower case words
+        assertTrue(matches("e", "elephant", MATCHER));
+        assertTrue(matches("eL", "Elephant", MATCHER));
+
+        assertTrue(matches("电", "电子邮件", MATCHER));
+        assertTrue(matches("电子", "电子邮件", MATCHER));
+        assertTrue(matches("子", "电子邮件", MATCHER));
+        assertTrue(matches("邮件", "电子邮件", MATCHER));
+
+        assertFalse(matches("ba", "Bot", MATCHER));
+        assertFalse(matches("ba", "bot", MATCHER));
+        assertFalse(matches("phant", "elephant", MATCHER));
+        assertFalse(matches("elephants", "elephant", MATCHER));
+    }
+
+    @Test
+    public void testMatchesVN() {
+        assertTrue(matches("다", "다운로드", MATCHER));
+        assertTrue(matches("드", "드라이브", MATCHER));
+        assertTrue(matches("ㄷ", "다운로드 드라이브", MATCHER));
+        assertTrue(matches("ㄷ", "운로 드라이브", MATCHER));
+        assertTrue(matches("åbç", "abc", MATCHER));
+        assertTrue(matches("ål", "Alpha", MATCHER));
+
+        assertFalse(matches("ㄷㄷ", "다운로드 드라이브", MATCHER));
+        assertFalse(matches("ㄷ", "로드라이브", MATCHER));
+        assertFalse(matches("åç", "abc", MATCHER));
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/AllApps.java b/tests/tapl/com/android/launcher3/tapl/AllApps.java
index b6c17df..e32250e 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllApps.java
@@ -27,7 +27,6 @@
 import androidx.test.uiautomator.Direction;
 import androidx.test.uiautomator.UiObject2;
 
-import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.testing.TestProtocol;
 
 import java.util.stream.Collectors;
@@ -108,26 +107,24 @@
                     "apps_list_view");
             final UiObject2 searchBox = getSearchBox(allAppsContainer);
 
-            int bottomGestureMargin = ResourceUtils.getNavbarSize(
-                    ResourceUtils.NAVBAR_BOTTOM_GESTURE_SIZE, mLauncher.getResources()) + 1;
-            int deviceHeight = mLauncher.getDevice().getDisplayHeight();
-            int displayBottom = deviceHeight - bottomGestureMargin;
+            int deviceHeight = mLauncher.getRealDisplaySize().y;
+            int bottomGestureStartOnScreen = mLauncher.getBottomGestureStartOnScreen();
             final BySelector appIconSelector = AppIcon.getAppIconSelector(appName, mLauncher);
             if (!hasClickableIcon(allAppsContainer, appListRecycler, appIconSelector,
-                    displayBottom)) {
+                    bottomGestureStartOnScreen)) {
                 scrollBackToBeginning();
                 int attempts = 0;
                 int scroll = getAllAppsScroll();
                 try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer("scrolled")) {
                     while (!hasClickableIcon(allAppsContainer, appListRecycler, appIconSelector,
-                            displayBottom)) {
+                            bottomGestureStartOnScreen)) {
                         mLauncher.scrollToLastVisibleRow(
                                 allAppsContainer,
                                 mLauncher.getObjectsInContainer(allAppsContainer, "icon")
                                         .stream()
                                         .filter(icon ->
-                                                        mLauncher.getVisibleBounds(icon).bottom
-                                                        <= displayBottom)
+                                                mLauncher.getVisibleBounds(icon).top
+                                                        < bottomGestureStartOnScreen)
                                         .collect(Collectors.toList()),
                                 mLauncher.getVisibleBounds(searchBox).bottom
                                         - mLauncher.getVisibleBounds(allAppsContainer).top);
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index a3f37d2..5138f02 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1027,16 +1027,20 @@
                 expectedState);
     }
 
-    int getBottomGestureSize() {
+    private int getBottomGestureSize() {
         return ResourceUtils.getNavbarSize(
                 ResourceUtils.NAVBAR_BOTTOM_GESTURE_SIZE, getResources()) + 1;
     }
 
     int getBottomGestureMarginInContainer(UiObject2 container) {
-        final int bottomGestureStartOnScreen = getRealDisplaySize().y - getBottomGestureSize();
+        final int bottomGestureStartOnScreen = getBottomGestureStartOnScreen();
         return getVisibleBounds(container).bottom - bottomGestureStartOnScreen;
     }
 
+    int getBottomGestureStartOnScreen() {
+        return getRealDisplaySize().y - getBottomGestureSize();
+    }
+
     void clickLauncherObject(UiObject2 object) {
         expectEvent(TestProtocol.SEQUENCE_MAIN, LauncherInstrumentation.EVENT_TOUCH_DOWN);
         expectEvent(TestProtocol.SEQUENCE_MAIN, LauncherInstrumentation.EVENT_TOUCH_UP);