diff --git a/quickstep/res/values/override.xml b/quickstep/res/values/override.xml
index 8f4ce43..605774d 100644
--- a/quickstep/res/values/override.xml
+++ b/quickstep/res/values/override.xml
@@ -25,8 +25,6 @@
 
   <string name="main_process_initializer_class" translatable="false">com.android.quickstep.QuickstepProcessInitializer</string>
 
-  <string name="user_event_dispatcher_class" translatable="false">com.android.quickstep.logging.UserEventDispatcherExtension</string>
-
   <string name="model_delegate_class" translatable="false">com.android.launcher3.model.QuickstepModelDelegate</string>
 
 </resources>
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
index e9fc423..d8064a2 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
@@ -569,6 +569,7 @@
         pw.println("  isUserUnlocked=" + mIsUserUnlocked);
         pw.println("  isOneHandedModeEnabled=" + mIsOneHandedModeEnabled);
         pw.println("  isSwipeToNotificationEnabled=" + mIsSwipeToNotificationEnabled);
+        pw.println("  deferredGestureRegion=" + mDeferredGestureRegion);
         mRotationTouchHelper.dump(pw);
     }
 }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 5b4a239..ebdc867 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -64,8 +64,8 @@
 import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.TestProtocol;
-import com.android.launcher3.tracing.nano.LauncherTraceProto;
-import com.android.launcher3.tracing.nano.TouchInteractionServiceProto;
+import com.android.launcher3.tracing.LauncherTraceProto;
+import com.android.launcher3.tracing.TouchInteractionServiceProto;
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.OnboardingPrefs;
 import com.android.launcher3.util.TraceHelper;
@@ -105,7 +105,7 @@
  */
 @TargetApi(Build.VERSION_CODES.R)
 public class TouchInteractionService extends Service implements PluginListener<OverscrollPlugin>,
-        ProtoTraceable<LauncherTraceProto> {
+        ProtoTraceable<LauncherTraceProto.Builder> {
 
     private static final String TAG = "TouchInteractionService";
 
@@ -368,9 +368,12 @@
 
             // Update the tracing state
             if ((mDeviceState.getSystemUiStateFlags() & SYSUI_STATE_TRACING_ENABLED) != 0) {
-                ProtoTracer.INSTANCE.get(TouchInteractionService.this).start();
+                Log.d(TAG, "Starting tracing.");
+                ProtoTracer.INSTANCE.get(this).start();
             } else {
-                ProtoTracer.INSTANCE.get(TouchInteractionService.this).stop();
+                Log.d(TAG, "Stopping tracing. Dumping to file="
+                    + ProtoTracer.INSTANCE.get(this).getTraceFile());
+                ProtoTracer.INSTANCE.get(this).stop();
             }
         }
     }
@@ -394,7 +397,7 @@
         disposeEventHandlers();
         mDeviceState.destroy();
         SystemUiProxy.INSTANCE.get(this).setProxy(null);
-        ProtoTracer.INSTANCE.get(TouchInteractionService.this).stop();
+        ProtoTracer.INSTANCE.get(this).stop();
         ProtoTracer.INSTANCE.get(this).remove(this);
 
         getSystemService(AccessibilityManager.class)
@@ -503,6 +506,7 @@
             reset();
         }
         TraceHelper.INSTANCE.endFlagsOverride(traceToken);
+        ProtoTracer.INSTANCE.get(this).scheduleFrameUpdate();
     }
 
     private GestureState createGestureState(GestureState previousGestureState) {
@@ -821,8 +825,7 @@
             pw.println("  mConsumer=" + mConsumer.getName());
             ActiveGestureLog.INSTANCE.dump("", pw);
             pw.println("ProtoTrace:");
-            pw.println("  file="
-                    + ProtoTracer.INSTANCE.get(TouchInteractionService.this).getTraceFile());
+            pw.println("  file=" + ProtoTracer.INSTANCE.get(this).getTraceFile());
         }
     }
 
@@ -878,11 +881,10 @@
     }
 
     @Override
-    public void writeToProto(LauncherTraceProto proto) {
-        if (proto.touchInteractionService == null) {
-            proto.touchInteractionService = new TouchInteractionServiceProto();
-        }
-        proto.touchInteractionService.serviceConnected = true;
-        proto.touchInteractionService.serviceConnected = true;
+    public void writeToProto(LauncherTraceProto.Builder proto) {
+        TouchInteractionServiceProto.Builder serviceProto =
+            TouchInteractionServiceProto.newBuilder();
+        serviceProto.setServiceConnected(true);
+        proto.setTouchInteractionService(serviceProto);
     }
 }
diff --git a/quickstep/src/com/android/quickstep/util/ProtoTracer.java b/quickstep/src/com/android/quickstep/util/ProtoTracer.java
index 190763a..ef9586d 100644
--- a/quickstep/src/com/android/quickstep/util/ProtoTracer.java
+++ b/quickstep/src/com/android/quickstep/util/ProtoTracer.java
@@ -16,24 +16,23 @@
 
 package com.android.quickstep.util;
 
-import static com.android.launcher3.tracing.nano.LauncherTraceFileProto.MagicNumber.MAGIC_NUMBER_H;
-import static com.android.launcher3.tracing.nano.LauncherTraceFileProto.MagicNumber.MAGIC_NUMBER_L;
+import static com.android.launcher3.tracing.LauncherTraceFileProto.MagicNumber.MAGIC_NUMBER_H_VALUE;
+import static com.android.launcher3.tracing.LauncherTraceFileProto.MagicNumber.MAGIC_NUMBER_L_VALUE;
 
 import android.content.Context;
 import android.os.SystemClock;
 
-import com.android.launcher3.tracing.nano.LauncherTraceProto;
-import com.android.launcher3.tracing.nano.LauncherTraceEntryProto;
-import com.android.launcher3.tracing.nano.LauncherTraceFileProto;
+import android.os.Trace;
+import com.android.launcher3.tracing.LauncherTraceProto;
+import com.android.launcher3.tracing.LauncherTraceEntryProto;
+import com.android.launcher3.tracing.LauncherTraceFileProto;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.systemui.shared.tracing.FrameProtoTracer;
 import com.android.systemui.shared.tracing.FrameProtoTracer.ProtoTraceParams;
 import com.android.systemui.shared.tracing.ProtoTraceable;
-import com.google.protobuf.nano.MessageNano;
+import com.google.protobuf.MessageLite;
 
 import java.io.File;
-import java.io.FileDescriptor;
-import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Queue;
 
@@ -41,18 +40,20 @@
 /**
  * Controller for coordinating winscope proto tracing.
  */
-public class ProtoTracer implements ProtoTraceParams<MessageNano,
-        LauncherTraceFileProto, LauncherTraceEntryProto, LauncherTraceProto> {
+public class ProtoTracer implements ProtoTraceParams<MessageLite.Builder,
+        LauncherTraceFileProto.Builder, LauncherTraceEntryProto.Builder,
+                LauncherTraceProto.Builder> {
 
     public static final MainThreadInitializedObject<ProtoTracer> INSTANCE =
             new MainThreadInitializedObject<>(ProtoTracer::new);
 
     private static final String TAG = "ProtoTracer";
-    private static final long MAGIC_NUMBER_VALUE = ((long) MAGIC_NUMBER_H << 32) | MAGIC_NUMBER_L;
+    private static final long MAGIC_NUMBER_VALUE =
+            ((long) MAGIC_NUMBER_H_VALUE << 32) | MAGIC_NUMBER_L_VALUE;
 
     private final Context mContext;
-    private final FrameProtoTracer<MessageNano,
-            LauncherTraceFileProto, LauncherTraceEntryProto, LauncherTraceProto> mProtoTracer;
+    private final FrameProtoTracer<MessageLite.Builder, LauncherTraceFileProto.Builder,
+        LauncherTraceEntryProto.Builder, LauncherTraceProto.Builder> mProtoTracer;
 
     public ProtoTracer(Context context) {
         mContext = context;
@@ -65,40 +66,47 @@
     }
 
     @Override
-    public LauncherTraceFileProto getEncapsulatingTraceProto() {
-        return new LauncherTraceFileProto();
+    public LauncherTraceFileProto.Builder getEncapsulatingTraceProto() {
+        return LauncherTraceFileProto.newBuilder();
     }
 
     @Override
-    public LauncherTraceEntryProto updateBufferProto(LauncherTraceEntryProto reuseObj,
-            ArrayList<ProtoTraceable<LauncherTraceProto>> traceables) {
-        LauncherTraceEntryProto proto = reuseObj != null
-                ? reuseObj
-                : new LauncherTraceEntryProto();
-        proto.elapsedRealtimeNanos = SystemClock.elapsedRealtimeNanos();
-        proto.launcher = proto.launcher != null ? proto.launcher : new LauncherTraceProto();
+    public LauncherTraceEntryProto.Builder updateBufferProto(
+            LauncherTraceEntryProto.Builder reuseObj,
+            ArrayList<ProtoTraceable<LauncherTraceProto.Builder>> traceables) {
+        Trace.beginSection("ProtoTracer.updateBufferProto");
+        LauncherTraceEntryProto.Builder proto = LauncherTraceEntryProto.newBuilder();
+        proto.setElapsedRealtimeNanos(SystemClock.elapsedRealtimeNanos());
+        LauncherTraceProto.Builder launcherProto = LauncherTraceProto.newBuilder();
         for (ProtoTraceable t : traceables) {
-            t.writeToProto(proto.launcher);
+            t.writeToProto(launcherProto);
         }
+        proto.setLauncher(launcherProto);
+        Trace.endSection();
         return proto;
     }
 
     @Override
-    public byte[] serializeEncapsulatingProto(LauncherTraceFileProto encapsulatingProto,
-            Queue<LauncherTraceEntryProto> buffer) {
-        encapsulatingProto.magicNumber = MAGIC_NUMBER_VALUE;
-        encapsulatingProto.entry = buffer.toArray(new LauncherTraceEntryProto[0]);
-        return MessageNano.toByteArray(encapsulatingProto);
+    public byte[] serializeEncapsulatingProto(LauncherTraceFileProto.Builder encapsulatingProto,
+            Queue<LauncherTraceEntryProto.Builder> buffer) {
+        Trace.beginSection("ProtoTracer.serializeEncapsulatingProto");
+        encapsulatingProto.setMagicNumber(MAGIC_NUMBER_VALUE);
+        for (LauncherTraceEntryProto.Builder entry : buffer) {
+            encapsulatingProto.addEntry(entry);
+        }
+        byte[] bytes = encapsulatingProto.build().toByteArray();
+        Trace.endSection();
+        return bytes;
     }
 
     @Override
-    public byte[] getProtoBytes(MessageNano proto) {
-        return MessageNano.toByteArray(proto);
+    public byte[] getProtoBytes(MessageLite.Builder proto) {
+        return proto.build().toByteArray();
     }
 
     @Override
-    public int getProtoSize(MessageNano proto) {
-        return proto.getCachedSize();
+    public int getProtoSize(MessageLite.Builder proto) {
+        return proto.build().getSerializedSize();
     }
 
     public void start() {
@@ -109,11 +117,11 @@
         mProtoTracer.stop();
     }
 
-    public void add(ProtoTraceable<LauncherTraceProto> traceable) {
+    public void add(ProtoTraceable<LauncherTraceProto.Builder> traceable) {
         mProtoTracer.add(traceable);
     }
 
-    public void remove(ProtoTraceable<LauncherTraceProto> traceable) {
+    public void remove(ProtoTraceable<LauncherTraceProto.Builder> traceable) {
         mProtoTracer.remove(traceable);
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 982a7c4..6ae91b6 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -79,7 +79,6 @@
 import android.view.KeyEvent;
 import android.view.LayoutInflater;
 import android.view.MotionEvent;
-import android.view.TouchDelegate;
 import android.view.View;
 import android.view.ViewDebug;
 import android.view.ViewGroup;
@@ -706,12 +705,9 @@
         super.onTouchEvent(ev);
 
         TaskView taskView = getCurrentPageTaskView();
-        if (taskView != null) {
-            TouchDelegate mChildTouchDelegate = taskView.getIconTouchDelegate(ev);
-            if (mChildTouchDelegate != null && mChildTouchDelegate.onTouchEvent(ev)) {
-                // Keep consuming events to pass to delegate
-                return true;
-            }
+        if (taskView != null && taskView.offerTouchToChildren(ev)) {
+            // Keep consuming events to pass to delegate
+            return true;
         }
 
         final int x = (int) ev.getX();
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 259cd4d..321821b 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -135,6 +135,7 @@
      * delegated bounds only to be updated.
      */
     private TransformingTouchDelegate mIconTouchDelegate;
+    private TransformingTouchDelegate mChipTouchDelegate;
 
     private static final List<Rect> SYSTEM_GESTURE_EXCLUSION_RECT =
             Collections.singletonList(new Rect());
@@ -200,6 +201,7 @@
     private View mContextualChipWrapper;
     private View mContextualChip;
     private final float[] mIconCenterCoords = new float[2];
+    private final float[] mChipCenterCoords = new float[2];
 
     public TaskView(Context context) {
         this(context, null);
@@ -259,11 +261,22 @@
         mIconTouchDelegate = new TransformingTouchDelegate(mIconView);
     }
 
-    public TouchDelegate getIconTouchDelegate(MotionEvent event) {
+    /**
+     * Whether the taskview should take the touch event from parent. Events passed to children
+     * that might require special handling.
+     */
+    public boolean offerTouchToChildren(MotionEvent event) {
         if (event.getAction() == MotionEvent.ACTION_DOWN) {
             computeAndSetIconTouchDelegate();
+            computeAndSetChipTouchDelegate();
         }
-        return mIconTouchDelegate;
+        if (mIconTouchDelegate != null && mIconTouchDelegate.onTouchEvent(event)) {
+            return true;
+        }
+        if (mChipTouchDelegate != null && mChipTouchDelegate.onTouchEvent(event)) {
+            return true;
+        }
+        return false;
     }
 
     private void computeAndSetIconTouchDelegate() {
@@ -278,6 +291,23 @@
                 (int) (mIconCenterCoords[1] + iconHalfSize));
     }
 
+    private void computeAndSetChipTouchDelegate() {
+        if (mContextualChipWrapper != null) {
+            float chipHalfWidth = mContextualChipWrapper.getWidth() / 2f;
+            float chipHalfHeight = mContextualChipWrapper.getHeight() / 2f;
+            mChipCenterCoords[0] = chipHalfWidth;
+            mChipCenterCoords[1] = chipHalfHeight;
+            getDescendantCoordRelativeToAncestor(mContextualChipWrapper, mActivity.getDragLayer(),
+                    mChipCenterCoords,
+                    false);
+            mChipTouchDelegate.setBounds(
+                    (int) (mChipCenterCoords[0] - chipHalfWidth),
+                    (int) (mChipCenterCoords[1] - chipHalfHeight),
+                    (int) (mChipCenterCoords[0] + chipHalfWidth),
+                    (int) (mChipCenterCoords[1] + chipHalfHeight));
+        }
+    }
+
     /**
      * The modalness of this view is how it should be displayed when it is shown on its own in the
      * modal state of overview.
@@ -712,6 +742,7 @@
                 mContextualChip.animate().scaleX(1f).scaleY(1f).setDuration(50);
             }
             if (mContextualChipWrapper != null) {
+                mChipTouchDelegate = new TransformingTouchDelegate(mContextualChipWrapper);
                 mContextualChipWrapper.animate().alpha(1f).setDuration(50);
             }
         }
@@ -733,6 +764,7 @@
         View oldContextualChipWrapper = mContextualChipWrapper;
         mContextualChipWrapper = null;
         mContextualChip = null;
+        mChipTouchDelegate = null;
         return oldContextualChipWrapper;
     }
 
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 8da81ac..45aaa1b 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -93,6 +93,7 @@
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pageindicators.WorkspacePageIndicator;
 import com.android.launcher3.popup.PopupContainerWithArrow;
+import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.touch.WorkspaceTouchListener;
@@ -436,10 +437,6 @@
             enforceDragParity("onDragEnd", 0, 0);
         }
 
-        if (!mDeferRemoveExtraEmptyScreen) {
-            removeExtraEmptyScreen(mDragSourceInternal != null);
-        }
-
         updateChildrenLayersEnabled();
         mDragInfo = null;
         mOutlineProvider = null;
@@ -657,6 +654,7 @@
         convertFinalScreenToEmptyScreenIfNecessary();
         if (hasExtraEmptyScreen()) {
             removeView(mWorkspaceScreens.get(EXTRA_EMPTY_SCREEN_ID));
+            setCurrentPage(getNextPage());
             mWorkspaceScreens.remove(EXTRA_EMPTY_SCREEN_ID);
             mScreenOrder.removeValue(EXTRA_EMPTY_SCREEN_ID);
 
@@ -1873,6 +1871,18 @@
                             };
                         }
                     }
+                    StateManager<LauncherState> stateManager = mLauncher.getStateManager();
+                    stateManager.addStateListener(new StateManager.StateListener<LauncherState>() {
+                        @Override
+                        public void onStateTransitionComplete(LauncherState finalState) {
+                            if (finalState == NORMAL) {
+                                if (!mDeferRemoveExtraEmptyScreen) {
+                                    removeExtraEmptyScreen(true /* stripEmptyScreens */);
+                                }
+                                stateManager.removeStateListener(this);
+                            }
+                        }
+                    });
 
                     mLauncher.getModelWriter().modifyItemInDatabase(info, container, screenId,
                             lp.cellX, lp.cellY, item.spanX, item.spanY);
@@ -2453,7 +2463,7 @@
             Runnable onAnimationCompleteRunnable = new Runnable() {
                 @Override
                 public void run() {
-                    // Normally removeExtraEmptyScreen is called in Workspace#onDragEnd, but when
+                    // Normally removeExtraEmptyScreen is called in Workspace#onDrop, but when
                     // adding an item that may not be dropped right away (due to a config activity)
                     // we defer the removal until the activity returns.
                     deferRemoveExtraEmptyScreen();
diff --git a/src/com/android/launcher3/WorkspaceLayoutManager.java b/src/com/android/launcher3/WorkspaceLayoutManager.java
index c3d4aeb..ea887cc 100644
--- a/src/com/android/launcher3/WorkspaceLayoutManager.java
+++ b/src/com/android/launcher3/WorkspaceLayoutManager.java
@@ -28,7 +28,7 @@
 
     String TAG = "Launcher.Workspace";
 
-    // The screen id used for the empty screen always present to the right.
+    // The screen id used for the empty screen always present at the end.
     int EXTRA_EMPTY_SCREEN_ID = -201;
     // The is the first screen. It is always present, even if its empty.
     int FIRST_SCREEN_ID = 0;
diff --git a/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java b/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
index 2e5ed3e..29e3404 100644
--- a/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
+++ b/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
@@ -79,7 +79,7 @@
     public void beforeTextChanged(CharSequence charSequence, int i, int i1, int i2) {
         if (mSearchAlgorithm instanceof PluginWrapper) {
             ((PluginWrapper) mSearchAlgorithm).runOnPluginIfConnected(
-                    AllAppsSearchPlugin::startedTyping);
+                    AllAppsSearchPlugin::startedSearchSession);
         }
     }
 
@@ -113,7 +113,13 @@
     public boolean onEditorAction(TextView v, int actionId, KeyEvent event) {
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
             if (actionId == EditorInfo.IME_ACTION_SEARCH) {
+                // selectFocusedView should return SearchTargetEvent that is passed onto onClick
                 if (Launcher.getLauncher(mLauncher).getAppsView().selectFocusedView(v)) {
+                    if (mSearchAlgorithm instanceof PluginWrapper) {
+                        ((PluginWrapper) mSearchAlgorithm).runOnPluginIfConnected(plugin -> {
+                            plugin.onClick(false, null);
+                        });
+                    }
                     return true;
                 }
             }
diff --git a/src/com/android/launcher3/model/FirstScreenBroadcast.java b/src/com/android/launcher3/model/FirstScreenBroadcast.java
index 5112304..70d1b48 100644
--- a/src/com/android/launcher3/model/FirstScreenBroadcast.java
+++ b/src/com/android/launcher3/model/FirstScreenBroadcast.java
@@ -17,25 +17,30 @@
 
 import static android.os.Process.myUserHandle;
 
+import static com.android.launcher3.pm.InstallSessionHelper.getUserHandle;
+
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.mapping;
+
 import android.app.PendingIntent;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageInstaller.SessionInfo;
+import android.os.UserHandle;
 import android.util.Log;
 
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
-import com.android.launcher3.util.MultiHashMap;
 import com.android.launcher3.util.PackageUserKey;
 
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Helper class to send broadcasts to package installers that have:
@@ -61,26 +66,10 @@
 
     private static final String VERIFICATION_TOKEN_EXTRA = "verificationToken";
 
-    private final MultiHashMap<String, String> mPackagesForInstaller;
+    private final HashMap<PackageUserKey, SessionInfo> mSessionInfoForPackage;
 
     public FirstScreenBroadcast(HashMap<PackageUserKey, SessionInfo> sessionInfoForPackage) {
-        mPackagesForInstaller = getPackagesForInstaller(sessionInfoForPackage);
-    }
-
-    /**
-     * @return Map where the key is the package name of the installer, and the value is a list
-     *         of packages with active sessions for that installer.
-     */
-    private MultiHashMap<String, String> getPackagesForInstaller(
-            HashMap<PackageUserKey, SessionInfo> sessionInfoForPackage) {
-        MultiHashMap<String, String> packagesForInstaller = new MultiHashMap<>();
-        for (Map.Entry<PackageUserKey, SessionInfo> entry : sessionInfoForPackage.entrySet()) {
-            if (myUserHandle().equals(entry.getKey().mUser)) {
-                packagesForInstaller.addToList(entry.getValue().getInstallerPackageName(),
-                        entry.getKey().mPackageName);
-            }
-        }
-        return packagesForInstaller;
+        mSessionInfoForPackage = sessionInfoForPackage;
     }
 
     /**
@@ -88,9 +77,15 @@
      * first screen.
      */
     public void sendBroadcasts(Context context, List<ItemInfo> firstScreenItems) {
-        for (Map.Entry<String, ArrayList<String>> entry : mPackagesForInstaller.entrySet()) {
-            sendBroadcastToInstaller(context, entry.getKey(), entry.getValue(), firstScreenItems);
-        }
+        UserHandle myUser = myUserHandle();
+        mSessionInfoForPackage
+                .values()
+                .stream()
+                .filter(info -> myUser.equals(getUserHandle(info)))
+                .collect(groupingBy(SessionInfo::getInstallerPackageName,
+                        mapping(SessionInfo::getAppPackageName, Collectors.toSet())))
+                .forEach((installer, packages) ->
+                    sendBroadcastToInstaller(context, installer, packages, firstScreenItems));
     }
 
     /**
@@ -99,7 +94,7 @@
      * @param firstScreenItems List of items on the first screen.
      */
     private void sendBroadcastToInstaller(Context context, String installerPackageName,
-            List<String> packages, List<ItemInfo> firstScreenItems) {
+            Set<String> packages, List<ItemInfo> firstScreenItems) {
         Set<String> folderItems = new HashSet<>();
         Set<String> workspaceItems = new HashSet<>();
         Set<String> hotseatItems = new HashSet<>();
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 983ec0c..a9e385f 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -81,7 +81,6 @@
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IOUtils;
 import com.android.launcher3.util.LooperIdleLock;
-import com.android.launcher3.util.MultiHashMap;
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.TraceHelper;
@@ -90,8 +89,10 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CancellationException;
 
 /**
@@ -302,7 +303,7 @@
         final PackageManagerHelper pmHelper = new PackageManagerHelper(context);
         final boolean isSafeMode = pmHelper.isSafeMode();
         final boolean isSdCardReady = Utilities.isBootCompleted();
-        final MultiHashMap<UserHandle, String> pendingPackages = new MultiHashMap<>();
+        final Set<PackageUserKey> pendingPackages = new HashSet<>();
 
         boolean clearDb = false;
         try {
@@ -483,7 +484,7 @@
                                     // SdCard is not ready yet. Package might get available,
                                     // once it is ready.
                                     Log.d(TAG, "Missing pkg, will check later: " + targetPkg);
-                                    pendingPackages.addToList(c.user, targetPkg);
+                                    pendingPackages.add(new PackageUserKey(targetPkg, c.user));
                                     // Add the icon on the workspace anyway.
                                     allowMissingTarget = true;
                                 } else {
diff --git a/src/com/android/launcher3/model/ModelUtils.java b/src/com/android/launcher3/model/ModelUtils.java
index a8cc9ad..9b5fac8 100644
--- a/src/com/android/launcher3/model/ModelUtils.java
+++ b/src/com/android/launcher3/model/ModelUtils.java
@@ -36,8 +36,8 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.stream.IntStream;
 
 /**
@@ -56,13 +56,7 @@
             ArrayList<T> currentScreenItems,
             ArrayList<T> otherScreenItems) {
         // Purge any null ItemInfos
-        Iterator<T> iter = allWorkspaceItems.iterator();
-        while (iter.hasNext()) {
-            ItemInfo i = iter.next();
-            if (i == null) {
-                iter.remove();
-            }
-        }
+        allWorkspaceItems.removeIf(Objects::isNull);
         // Order the set of items by their containers first, this allows use to walk through the
         // list sequentially, build up a list of containers that are in the specified screen,
         // as well as all items in those containers.
diff --git a/src/com/android/launcher3/model/SdCardAvailableReceiver.java b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
index eb3cb52..3798575 100644
--- a/src/com/android/launcher3/model/SdCardAvailableReceiver.java
+++ b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
@@ -24,12 +24,11 @@
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
-import com.android.launcher3.util.MultiHashMap;
 import com.android.launcher3.util.PackageManagerHelper;
+import com.android.launcher3.util.PackageUserKey;
 
 import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.Map.Entry;
+import java.util.Set;
 
 /**
  * Helper class to re-query app status when SD-card becomes available.
@@ -42,10 +41,9 @@
 
     private final LauncherModel mModel;
     private final Context mContext;
-    private final MultiHashMap<UserHandle, String> mPackages;
+    private final Set<PackageUserKey> mPackages;
 
-    public SdCardAvailableReceiver(LauncherAppState app,
-            MultiHashMap<UserHandle, String> packages) {
+    public SdCardAvailableReceiver(LauncherAppState app, Set<PackageUserKey> packages) {
         mModel = app.getModel();
         mContext = app.getContext();
         mPackages = packages;
@@ -55,19 +53,17 @@
     public void onReceive(Context context, Intent intent) {
         final LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
         final PackageManagerHelper pmHelper = new PackageManagerHelper(context);
-        for (Entry<UserHandle, ArrayList<String>> entry : mPackages.entrySet()) {
-            UserHandle user = entry.getKey();
+        for (PackageUserKey puk : mPackages) {
+            UserHandle user = puk.mUser;
 
             final ArrayList<String> packagesRemoved = new ArrayList<>();
             final ArrayList<String> packagesUnavailable = new ArrayList<>();
 
-            for (String pkg : new HashSet<>(entry.getValue())) {
-                if (!launcherApps.isPackageEnabled(pkg, user)) {
-                    if (pmHelper.isAppOnSdcard(pkg, user)) {
-                        packagesUnavailable.add(pkg);
-                    } else {
-                        packagesRemoved.add(pkg);
-                    }
+            if (!launcherApps.isPackageEnabled(puk.mPackageName, user)) {
+                if (pmHelper.isAppOnSdcard(puk.mPackageName, user)) {
+                    packagesUnavailable.add(puk.mPackageName);
+                } else {
+                    packagesRemoved.add(puk.mPackageName);
                 }
             }
             if (!packagesRemoved.isEmpty()) {
diff --git a/src/com/android/launcher3/model/ShortcutsChangedTask.java b/src/com/android/launcher3/model/ShortcutsChangedTask.java
index 88006ba..6fedad1 100644
--- a/src/com/android/launcher3/model/ShortcutsChangedTask.java
+++ b/src/com/android/launcher3/model/ShortcutsChangedTask.java
@@ -25,11 +25,12 @@
 import com.android.launcher3.shortcuts.ShortcutKey;
 import com.android.launcher3.shortcuts.ShortcutRequest;
 import com.android.launcher3.util.ItemInfoMatcher;
-import com.android.launcher3.util.MultiHashMap;
 
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Handles changes due to shortcut manager updates (deep shortcut changes)
@@ -53,54 +54,53 @@
     public void execute(LauncherAppState app, BgDataModel dataModel, AllAppsList apps) {
         final Context context = app.getContext();
         // Find WorkspaceItemInfo's that have changed on the workspace.
-        HashSet<ShortcutKey> removedKeys = new HashSet<>();
-        MultiHashMap<ShortcutKey, WorkspaceItemInfo> keyToShortcutInfo = new MultiHashMap<>();
-        HashSet<String> allIds = new HashSet<>();
+        ArrayList<WorkspaceItemInfo> matchingWorkspaceItems = new ArrayList<>();
 
         synchronized (dataModel) {
             dataModel.forAllWorkspaceItemInfos(mUser, si -> {
                 if ((si.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT)
                         && mPackageName.equals(si.getIntent().getPackage())) {
-                    keyToShortcutInfo.addToList(ShortcutKey.fromItemInfo(si), si);
-                    allIds.add(si.getDeepShortcutId());
+                    matchingWorkspaceItems.add(si);
                 }
             });
         }
 
-        final ArrayList<WorkspaceItemInfo> updatedWorkspaceItemInfos = new ArrayList<>();
-        if (!keyToShortcutInfo.isEmpty()) {
+        if (!matchingWorkspaceItems.isEmpty()) {
             // Update the workspace to reflect the changes to updated shortcuts residing on it.
+            List<String> allLauncherKnownIds = matchingWorkspaceItems.stream()
+                    .map(WorkspaceItemInfo::getDeepShortcutId)
+                    .distinct()
+                    .collect(Collectors.toList());
             List<ShortcutInfo> shortcuts = new ShortcutRequest(context, mUser)
-                    .forPackage(mPackageName, new ArrayList<>(allIds))
+                    .forPackage(mPackageName, allLauncherKnownIds)
                     .query(ShortcutRequest.ALL);
+
+            Set<String> nonPinnedIds = new HashSet<>(allLauncherKnownIds);
+            ArrayList<WorkspaceItemInfo> updatedWorkspaceItemInfos = new ArrayList<>();
             for (ShortcutInfo fullDetails : shortcuts) {
-                ShortcutKey key = ShortcutKey.fromInfo(fullDetails);
-                List<WorkspaceItemInfo> workspaceItemInfos = keyToShortcutInfo.remove(key);
                 if (!fullDetails.isPinned()) {
-                    // The shortcut was previously pinned but is no longer, so remove it from
-                    // the workspace and our pinned shortcut counts.
-                    // Note that we put this check here, after querying for full details,
-                    // because there's a possible race condition between pinning and
-                    // receiving this callback.
-                    removedKeys.add(key);
                     continue;
                 }
-                for (final WorkspaceItemInfo workspaceItemInfo : workspaceItemInfos) {
-                    workspaceItemInfo.updateFromDeepShortcutInfo(fullDetails, context);
-                    app.getIconCache().getShortcutIcon(workspaceItemInfo, fullDetails);
-                    updatedWorkspaceItemInfos.add(workspaceItemInfo);
-                }
+
+                String sid = fullDetails.getId();
+                nonPinnedIds.remove(sid);
+                matchingWorkspaceItems
+                        .stream()
+                        .filter(itemInfo -> sid.equals(itemInfo.getDeepShortcutId()))
+                        .forEach(workspaceItemInfo -> {
+                            workspaceItemInfo.updateFromDeepShortcutInfo(fullDetails, context);
+                            app.getIconCache().getShortcutIcon(workspaceItemInfo, fullDetails);
+                            updatedWorkspaceItemInfos.add(workspaceItemInfo);
+                        });
             }
-        }
 
-        // If there are still entries in keyToShortcutInfo, that means that
-        // the corresponding shortcuts weren't passed in onShortcutsChanged(). This
-        // means they were cleared, so we remove and unpin them now.
-        removedKeys.addAll(keyToShortcutInfo.keySet());
-
-        bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos);
-        if (!keyToShortcutInfo.isEmpty()) {
-            deleteAndBindComponentsRemoved(ItemInfoMatcher.ofShortcutKeys(removedKeys));
+            bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos);
+            if (!nonPinnedIds.isEmpty()) {
+                deleteAndBindComponentsRemoved(ItemInfoMatcher.ofShortcutKeys(
+                        nonPinnedIds.stream()
+                                .map(id -> new ShortcutKey(mPackageName, mUser, id))
+                                .collect(Collectors.toSet())));
+            }
         }
 
         if (mUpdateIdMap) {
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index 5a5f668..76048ba 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -38,7 +38,6 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
@@ -202,20 +201,11 @@
     }
 
     public List<WidgetItem> getWidgetsForPackageUser(PackageUserKey packageUserKey) {
-        for (WidgetListRowEntry entry : mAllWidgets) {
-            if (entry.pkgItem.packageName.equals(packageUserKey.mPackageName)) {
-                ArrayList<WidgetItem> widgets = new ArrayList<>(entry.widgets);
-                // Remove widgets not associated with the correct user.
-                Iterator<WidgetItem> iterator = widgets.iterator();
-                while (iterator.hasNext()) {
-                    if (!iterator.next().user.equals(packageUserKey.mUser)) {
-                        iterator.remove();
-                    }
-                }
-                return widgets.isEmpty() ? null : widgets;
-            }
-        }
-        return null;
+        return mAllWidgets.stream()
+                .filter(row -> row.pkgItem.packageName.equals(packageUserKey.mPackageName))
+                .flatMap(row -> row.widgets.stream())
+                .filter(widget -> packageUserKey.mUser.equals(widget.user))
+                .collect(Collectors.toList());
     }
 
     /**
diff --git a/src/com/android/launcher3/popup/SystemShortcut.java b/src/com/android/launcher3/popup/SystemShortcut.java
index 7998488..81302ac 100644
--- a/src/com/android/launcher3/popup/SystemShortcut.java
+++ b/src/com/android/launcher3/popup/SystemShortcut.java
@@ -99,7 +99,7 @@
         final List<WidgetItem> widgets =
                 launcher.getPopupDataProvider().getWidgetsForPackageUser(new PackageUserKey(
                         itemInfo.getTargetComponent().getPackageName(), itemInfo.user));
-        if (widgets == null) {
+        if (widgets.isEmpty()) {
             return null;
         }
         return new Widgets(launcher, itemInfo);
diff --git a/src/com/android/launcher3/util/ItemInfoMatcher.java b/src/com/android/launcher3/util/ItemInfoMatcher.java
index e98af35..d26bb57 100644
--- a/src/com/android/launcher3/util/ItemInfoMatcher.java
+++ b/src/com/android/launcher3/util/ItemInfoMatcher.java
@@ -27,6 +27,7 @@
 import com.android.launcher3.shortcuts.ShortcutKey;
 
 import java.util.HashSet;
+import java.util.Set;
 
 /**
  * A utility class to check for {@link ItemInfo}
@@ -99,7 +100,7 @@
         return (info, cn) -> packageNames.contains(cn.getPackageName()) && info.user.equals(user);
     }
 
-    static ItemInfoMatcher ofShortcutKeys(HashSet<ShortcutKey> keys) {
+    static ItemInfoMatcher ofShortcutKeys(Set<ShortcutKey> keys) {
         return  (info, cn) -> info.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT &&
                         keys.contains(ShortcutKey.fromItemInfo(info));
     }
diff --git a/src/com/android/launcher3/util/MultiHashMap.java b/src/com/android/launcher3/util/MultiHashMap.java
deleted file mode 100644
index b7275c1..0000000
--- a/src/com/android/launcher3/util/MultiHashMap.java
+++ /dev/null
@@ -1,52 +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.util;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-
-/**
- * A utility map from keys to an ArrayList of values.
- */
-public class MultiHashMap<K, V> extends HashMap<K, ArrayList<V>> {
-
-    public MultiHashMap() { }
-
-    public MultiHashMap(int size) {
-        super(size);
-    }
-
-    public void addToList(K key, V value) {
-        ArrayList<V> list = get(key);
-        if (list == null) {
-            list = new ArrayList<>();
-            list.add(value);
-            put(key, list);
-        } else {
-            list.add(value);
-        }
-    }
-
-    @Override
-    public MultiHashMap<K, V> clone() {
-        MultiHashMap<K, V> map = new MultiHashMap<>(size());
-        for (Entry<K, ArrayList<V>> entry : entrySet()) {
-            map.put(entry.getKey(), new ArrayList<V>(entry.getValue()));
-        }
-        return map;
-    }
-}
diff --git a/src_plugins/com/android/systemui/plugins/AllAppsSearchPlugin.java b/src_plugins/com/android/systemui/plugins/AllAppsSearchPlugin.java
index be20e2d..9949678 100644
--- a/src_plugins/com/android/systemui/plugins/AllAppsSearchPlugin.java
+++ b/src_plugins/com/android/systemui/plugins/AllAppsSearchPlugin.java
@@ -16,9 +16,9 @@
 
 package com.android.systemui.plugins;
 
-import android.os.Bundle;
-
 import com.android.systemui.plugins.annotations.ProvidesInterface;
+import com.android.systemui.plugins.shared.SearchTarget;
+import com.android.systemui.plugins.shared.SearchTargetEvent;
 
 import java.util.List;
 import java.util.function.Consumer;
@@ -29,15 +29,28 @@
 @ProvidesInterface(action = AllAppsSearchPlugin.ACTION, version = AllAppsSearchPlugin.VERSION)
 public interface AllAppsSearchPlugin extends Plugin {
     String ACTION = "com.android.systemui.action.PLUGIN_ALL_APPS_SEARCH_ACTIONS";
-    int VERSION = 4;
+    int VERSION = 5;
+
+
+    /**
+     * Send signal when user enters all apps.
+     */
+    void startAllAppsSession();
 
     /**
      * Send signal when user starts typing.
      */
-    void startedTyping();
+    void startedSearchSession();
 
     /**
      * Send over the query and get the search results.
      */
-    void performSearch(String query, Consumer<List<Bundle>> results);
-}
+    void performSearch(String query, Consumer<List<SearchTarget>> results);
+
+    void onClick(boolean isTouch, SearchTargetEvent event);
+
+    /**
+     * Send signal when user exits all apps.
+     */
+    void endAllAppsSession();
+}
\ No newline at end of file
diff --git a/src_plugins/com/android/systemui/plugins/OWNERS b/src_plugins/com/android/systemui/plugins/OWNERS
deleted file mode 100644
index 0514999..0000000
--- a/src_plugins/com/android/systemui/plugins/OWNERS
+++ /dev/null
@@ -1,4 +0,0 @@
-# When changing interface for this plugin OR when increasing version code, please add Alex
-# Only add other owners if Alex is not available
-per-file AllAppsSearchPlugin.java, globs = set noparent
-per-file AllAppsSearchPlugin.java = alexmang@google.com, hyunyoungs@google.com, sunnygoyal@google.com, twickham@google.com
diff --git a/src_plugins/com/android/systemui/plugins/shared/SearchTarget.java b/src_plugins/com/android/systemui/plugins/shared/SearchTarget.java
new file mode 100644
index 0000000..913fc0e
--- /dev/null
+++ b/src_plugins/com/android/systemui/plugins/shared/SearchTarget.java
@@ -0,0 +1,99 @@
+/*
+ * 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.systemui.plugins.shared;
+
+import android.content.pm.ShortcutInfo;
+import android.os.Bundle;
+
+import java.util.List;
+
+/**
+ * Used to return all apps search targets.
+ */
+public class SearchTarget implements Comparable<SearchTarget> {
+
+    public enum ViewType {
+        TOP_HIT(0),
+        HERO(1),
+        DETAIL(2),
+        ROW(3),
+        ROW_WITH_BUTTON(4),
+        SLICE(5);
+
+        private final int mId;
+        ViewType(int id) {
+            mId = id;
+        }
+
+        public int get() {
+            return mId;
+        }
+    }
+
+    public enum ItemType {
+        PLAY_RESULTS(0, "Play Store", ViewType.DETAIL),
+        SETTINGS_ROW(1, "Settings", ViewType.ROW),
+        SETTINGS_SLICE(2, "Settings", ViewType.SLICE),
+        APP(3, "", ViewType.TOP_HIT),
+        APP_HERO(4, "", ViewType.HERO);
+
+        private final int mId;
+        private final String mTitle;
+        private final ViewType mViewType;
+
+        ItemType(int id, String title, ViewType type) {
+            mId = id;
+            mTitle = title;
+            mViewType = type;
+        }
+
+        public ViewType getViewType() {
+            return mViewType;
+        }
+
+        public String getTitle() {
+            return mTitle;
+        }
+
+        public int getId() {
+            return mId;
+        }
+    }
+
+    public ItemType type;
+    public List<ShortcutInfo> shortcuts;
+    public Bundle bundle;
+    public float score;
+
+    /**
+     * Constructor to create the search target. Bundle is currently temporary to hold
+     * search target primitives that cannot be expressed as java primitive objects
+     * or AOSP native objects.
+     *
+     */
+    public SearchTarget(ItemType itemType, List<ShortcutInfo> shortcuts,
+            Bundle bundle, float score) {
+        this.type = itemType;
+        this.shortcuts = shortcuts;
+        this.bundle = bundle;
+        this.score = score;
+    }
+
+    @Override
+    public int compareTo(SearchTarget o) {
+        return Float.compare(o.score, score);
+    }
+}
diff --git a/src_plugins/com/android/systemui/plugins/shared/SearchTargetEvent.java b/src_plugins/com/android/systemui/plugins/shared/SearchTargetEvent.java
new file mode 100644
index 0000000..00aacd0
--- /dev/null
+++ b/src_plugins/com/android/systemui/plugins/shared/SearchTargetEvent.java
@@ -0,0 +1,39 @@
+/*
+ * 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.systemui.plugins.shared;
+
+import android.content.pm.ShortcutInfo;
+import android.os.Bundle;
+
+/**
+ * Event used for the feedback loop to the plugin. (and future aiai)
+ */
+public class SearchTargetEvent {
+    public SearchTarget.ItemType type;
+    public ShortcutInfo shortcut;
+    public Bundle bundle;
+    public float score;
+
+    public SearchTargetEvent(SearchTarget.ItemType itemType,
+            ShortcutInfo shortcut,
+            Bundle bundle,
+            float score) {
+        this.type = itemType;
+        this.shortcut = shortcut;
+        this.bundle = bundle;
+        this.score = score;
+    }
+}
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index a64df62..34ebbac 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -26,7 +26,6 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
-import com.android.launcher3.util.MultiHashMap;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.widget.WidgetItemComparator;
@@ -36,11 +35,12 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 /**
  * Widgets data model that is used by the adapters of the widget views and controllers.
@@ -57,9 +57,7 @@
     private static final boolean DEBUG = false;
 
     /* Map of widgets and shortcuts that are tracked per package. */
-    private final MultiHashMap<PackageItemInfo, WidgetItem> mWidgetsList = new MultiHashMap<>();
-
-    private AppFilter mAppFilter;
+    private final Map<PackageItemInfo, List<WidgetItem>> mWidgetsList = new HashMap<>();
 
     /**
      * Returns a list of {@link WidgetListRowEntry}. All {@link WidgetItem} in a single row
@@ -74,8 +72,9 @@
         AlphabeticIndexCompat indexer = new AlphabeticIndexCompat(context);
 
         WidgetItemComparator widgetComparator = new WidgetItemComparator();
-        for (Map.Entry<PackageItemInfo, ArrayList<WidgetItem>> entry : mWidgetsList.entrySet()) {
-            WidgetListRowEntry row = new WidgetListRowEntry(entry.getKey(), entry.getValue());
+        for (Map.Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsList.entrySet()) {
+            WidgetListRowEntry row = new WidgetListRowEntry(
+                    entry.getKey(), new ArrayList<>(entry.getValue()));
             row.titleSectionName = (row.pkgItem.title == null) ? "" :
                     indexer.computeSectionName(row.pkgItem.title);
             Collections.sort(row.widgets, widgetComparator);
@@ -146,77 +145,42 @@
         if (packageUser == null) {
             mWidgetsList.clear();
         } else {
-            // Only clear the widgets for the given package/user.
-            PackageItemInfo packageItem = null;
-            for (PackageItemInfo item : mWidgetsList.keySet()) {
-                if (item.packageName.equals(packageUser.mPackageName)) {
-                    packageItem = item;
-                    break;
-                }
-            }
+            PackageItemInfo packageItem = mWidgetsList.keySet()
+                    .stream()
+                    .filter(item -> item.packageName.equals(packageUser.mPackageName))
+                    .findFirst()
+                    .orElse(null);
             if (packageItem != null) {
                 // We want to preserve the user that was on the packageItem previously,
                 // so add it to tmpPackageItemInfos here to avoid creating a new entry.
                 tmpPackageItemInfos.put(packageItem.packageName, packageItem);
 
-                Iterator<WidgetItem> widgetItemIterator = mWidgetsList.get(packageItem).iterator();
-                while (widgetItemIterator.hasNext()) {
-                    WidgetItem nextWidget = widgetItemIterator.next();
-                    if (nextWidget.componentName.getPackageName().equals(packageUser.mPackageName)
-                            && nextWidget.user.equals(packageUser.mUser)) {
-                        widgetItemIterator.remove();
-                    }
-                }
+                // Add the widgets for other users in the rawList as it only contains widgets for
+                // packageUser
+                List<WidgetItem> otherUserItems = mWidgetsList.remove(packageItem);
+                otherUserItems.removeIf(w -> w.user.equals(packageUser.mUser));
+                rawWidgetsShortcuts.addAll(otherUserItems);
             }
         }
 
-        InvariantDeviceProfile idp = app.getInvariantDeviceProfile();
         UserHandle myUser = Process.myUserHandle();
 
         // add and update.
-        for (WidgetItem item : rawWidgetsShortcuts) {
-            if (item.widgetInfo != null) {
-                if ((item.widgetInfo.getWidgetFeatures() & WIDGET_FEATURE_HIDE_FROM_PICKER) != 0) {
-                    // Widget is hidden from picker
-                    continue;
-                }
-
-                // Ensure that all widgets we show can be added on a workspace of this size
-                int minSpanX = Math.min(item.widgetInfo.spanX, item.widgetInfo.minSpanX);
-                int minSpanY = Math.min(item.widgetInfo.spanY, item.widgetInfo.minSpanY);
-                if (minSpanX > idp.numColumns || minSpanY > idp.numRows) {
-                    if (DEBUG) {
-                        Log.d(TAG, String.format(
-                                "Widget %s : (%d X %d) can't fit on this device",
-                                item.componentName, minSpanX, minSpanY));
+        mWidgetsList.putAll(rawWidgetsShortcuts.stream()
+                .filter(new WidgetValidityCheck(app))
+                .collect(Collectors.groupingBy(item -> {
+                    String packageName = item.componentName.getPackageName();
+                    PackageItemInfo pInfo = tmpPackageItemInfos.get(packageName);
+                    if (pInfo == null) {
+                        pInfo = new PackageItemInfo(packageName);
+                        pInfo.user = item.user;
+                        tmpPackageItemInfos.put(packageName,  pInfo);
+                    } else if (!myUser.equals(pInfo.user)) {
+                        // Keep updating the user, until we get the primary user.
+                        pInfo.user = item.user;
                     }
-                    continue;
-                }
-            }
-
-            if (mAppFilter == null) {
-                mAppFilter = AppFilter.newInstance(app.getContext());
-            }
-            if (!mAppFilter.shouldShowApp(item.componentName)) {
-                if (DEBUG) {
-                    Log.d(TAG, String.format("%s is filtered and not added to the widget tray.",
-                            item.componentName));
-                }
-                continue;
-            }
-
-            String packageName = item.componentName.getPackageName();
-            PackageItemInfo pInfo = tmpPackageItemInfos.get(packageName);
-            if (pInfo == null) {
-                pInfo = new PackageItemInfo(packageName);
-                pInfo.user = item.user;
-                tmpPackageItemInfos.put(packageName,  pInfo);
-            } else if (!myUser.equals(pInfo.user)) {
-                // Keep updating the user, until we get the primary user.
-                pInfo.user = item.user;
-            }
-            mWidgetsList.addToList(pInfo, item);
-        }
+                    return pInfo;
+                })));
 
         // Update each package entry
         IconCache iconCache = app.getIconCache();
@@ -227,9 +191,9 @@
 
     public void onPackageIconsUpdated(Set<String> packageNames, UserHandle user,
             LauncherAppState app) {
-        for (Entry<PackageItemInfo, ArrayList<WidgetItem>> entry : mWidgetsList.entrySet()) {
+        for (Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsList.entrySet()) {
             if (packageNames.contains(entry.getKey().packageName)) {
-                ArrayList<WidgetItem> items = entry.getValue();
+                List<WidgetItem> items = entry.getValue();
                 int count = items.size();
                 for (int i = 0; i < count; i++) {
                     WidgetItem item = items.get(i);
@@ -249,7 +213,7 @@
 
     public WidgetItem getWidgetProviderInfoByProviderName(
             ComponentName providerName) {
-        ArrayList<WidgetItem> widgetsList = mWidgetsList.get(
+        List<WidgetItem> widgetsList = mWidgetsList.get(
                 new PackageItemInfo(providerName.getPackageName()));
         if (widgetsList == null) {
             return null;
@@ -262,4 +226,46 @@
         }
         return null;
     }
+
+    private static class WidgetValidityCheck implements Predicate<WidgetItem> {
+
+        private final InvariantDeviceProfile mIdp;
+        private final AppFilter mAppFilter;
+
+        WidgetValidityCheck(LauncherAppState app) {
+            mIdp = app.getInvariantDeviceProfile();
+            mAppFilter = AppFilter.newInstance(app.getContext());
+        }
+
+        @Override
+        public boolean test(WidgetItem item) {
+            if (item.widgetInfo != null) {
+                if ((item.widgetInfo.getWidgetFeatures() & WIDGET_FEATURE_HIDE_FROM_PICKER) != 0) {
+                    // Widget is hidden from picker
+                    return false;
+                }
+
+                // Ensure that all widgets we show can be added on a workspace of this size
+                int minSpanX = Math.min(item.widgetInfo.spanX, item.widgetInfo.minSpanX);
+                int minSpanY = Math.min(item.widgetInfo.spanY, item.widgetInfo.minSpanY);
+                if (minSpanX > mIdp.numColumns || minSpanY > mIdp.numRows) {
+                    if (DEBUG) {
+                        Log.d(TAG, String.format(
+                                "Widget %s : (%d X %d) can't fit on this device",
+                                item.componentName, minSpanX, minSpanY));
+                    }
+                    return false;
+                }
+            }
+            if (!mAppFilter.shouldShowApp(item.componentName)) {
+                if (DEBUG) {
+                    Log.d(TAG, String.format("%s is filtered and not added to the widget tray.",
+                            item.componentName));
+                }
+                return false;
+            }
+
+            return true;
+        }
+    }
 }
\ No newline at end of file