Merge "Removing insets controller animation as part of all-apps opening" into sc-dev
diff --git a/go/src/com/android/launcher3/model/WidgetsModel.java b/go/src/com/android/launcher3/model/WidgetsModel.java
index a202095..cc5e1cb 100644
--- a/go/src/com/android/launcher3/model/WidgetsModel.java
+++ b/go/src/com/android/launcher3/model/WidgetsModel.java
@@ -30,6 +30,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -53,10 +54,15 @@
      *
      * @see com.android.launcher3.widget.picker.WidgetsListAdapter#setWidgets(List)
      */
-    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsList(Context context) {
+    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsListForPicker(Context context) {
         return EMPTY_WIDGET_LIST;
     }
 
+    /** Returns a mapping of packages to their widgets without static shortcuts. */
+    public synchronized Map<PackageUserKey, List<WidgetItem>> getAllWidgetsWithoutShortcuts() {
+        return Map.of();
+    }
+
     /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
diff --git a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
new file mode 100644
index 0000000..70a143e
--- /dev/null
+++ b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -0,0 +1,263 @@
+/*
+ * 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.model;
+
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.app.prediction.AppTarget;
+import android.app.prediction.AppTargetId;
+import android.appwidget.AppWidgetManager;
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.os.Process;
+import android.os.UserHandle;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
+import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.ItemInfoMatcher;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.ViewOnDrawExecutor;
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.widget.PendingAddWidgetInfo;
+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 org.robolectric.shadows.ShadowAppWidgetManager;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.stream.Collectors;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsPredicationUpdateTaskTest {
+
+    private AppWidgetProviderInfo mApp1Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp1Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp2Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp5Provider1 = new AppWidgetProviderInfo();
+
+    private FakeBgDataModelCallback mCallback = new FakeBgDataModelCallback();
+    private Context mContext;
+    private LauncherModelHelper mModelHelper;
+    private UserHandle mUserHandle;
+    private InvariantDeviceProfile mTestProfile;
+
+    @Mock
+    private IconCache mIconCache;
+
+    @Before
+    public void setup() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = invocation.getArgument(0);
+            return componentWithLabel.getComponent().getShortClassName();
+        }).when(mIconCache).getTitleNoCache(any());
+
+        mContext = RuntimeEnvironment.application;
+        mModelHelper = new LauncherModelHelper();
+        mUserHandle = Process.myUserHandle();
+        mTestProfile = new InvariantDeviceProfile();
+        // 2 widgets, app4/provider1 & app5/provider1, have already been added to the workspace.
+        mModelHelper.initializeData("/widgets_predication_update_task_data.txt");
+
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        mApp1Provider1.provider = ComponentName.createRelative("app1", "provider1");
+        ReflectionHelpers.setField(mApp1Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider1.provider));
+        mApp1Provider2.provider = ComponentName.createRelative("app1", "provider2");
+        ReflectionHelpers.setField(mApp1Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider2.provider));
+        mApp2Provider1.provider = ComponentName.createRelative("app2", "provider1");
+        ReflectionHelpers.setField(mApp2Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp2Provider1.provider));
+        mApp4Provider1.provider = ComponentName.createRelative("app4", "provider1");
+        ReflectionHelpers.setField(mApp4Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider1.provider));
+        mApp4Provider2.provider = ComponentName.createRelative("app4", ".provider2");
+        ReflectionHelpers.setField(mApp4Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider2.provider));
+        mApp5Provider1.provider = ComponentName.createRelative("app5", "provider1");
+        ReflectionHelpers.setField(mApp5Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp5Provider1.provider));
+
+        ShadowAppWidgetManager shadowAppWidgetManager =
+                shadowOf(mContext.getSystemService(AppWidgetManager.class));
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp2Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp5Provider1);
+
+        mModelHelper.getModel().addCallbacks(mCallback);
+
+        MODEL_EXECUTOR.post(() -> mModelHelper.getBgDataModel().widgetsModel.update(
+                LauncherAppState.getInstance(mContext), /* packageUser= */ null));
+        waitUntilIdle();
+    }
+
+
+    @Test
+    public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder()
+            throws Exception {
+        // WHEN newPredicationTask is executed with app predication of 5 apps.
+        AppTarget app1 = new AppTarget(new AppTargetId("app1"), "app1", "className",
+                mUserHandle);
+        AppTarget app2 = new AppTarget(new AppTargetId("app2"), "app2", "className",
+                mUserHandle);
+        AppTarget app3 = new AppTarget(new AppTargetId("app3"), "app3", "className",
+                mUserHandle);
+        AppTarget app4 = new AppTarget(new AppTargetId("app4"), "app4", "className",
+                mUserHandle);
+        AppTarget app5 = new AppTarget(new AppTargetId("app5"), "app5", "className",
+                mUserHandle);
+        mModelHelper.executeTaskForTest(
+                newWidgetsPredicationTask(List.of(app5, app3, app2, app4, app1)))
+                .forEach(Runnable::run);
+
+        // THEN only 3 widgets are returned because
+        // 1. app5/provider1 & app4/provider1 have already been added to workspace. They are
+        //    excluded from the result.
+        // 2. app3 doesn't have a widget.
+        // 3. only 1 widget is picked from app1 because we only want to promote one widget per app.
+        List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
+                .stream()
+                .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
+                .collect(Collectors.toList());
+        assertThat(recommendedWidgets).hasSize(3);
+        assertWidgetInfo(recommendedWidgets.get(0).info, mApp2Provider1);
+        assertWidgetInfo(recommendedWidgets.get(1).info, mApp4Provider2);
+        assertWidgetInfo(recommendedWidgets.get(2).info, mApp1Provider1);
+    }
+
+    private void assertWidgetInfo(
+            LauncherAppWidgetProviderInfo actual, AppWidgetProviderInfo expected) {
+        assertThat(actual.provider).isEqualTo(expected.provider);
+        assertThat(actual.getUser()).isEqualTo(expected.getProfile());
+    }
+
+    private void waitUntilIdle() {
+        shadowOf(MODEL_EXECUTOR.getLooper()).idle();
+        shadowOf(MAIN_EXECUTOR.getLooper()).idle();
+    }
+
+    private WidgetsPredictionUpdateTask newWidgetsPredicationTask(List<AppTarget> appTargets) {
+        return new WidgetsPredictionUpdateTask(
+                new PredictorState(CONTAINER_WIDGETS_PREDICTION, "test_widgets_prediction"),
+                appTargets);
+    }
+
+    private final class FakeBgDataModelCallback implements BgDataModel.Callbacks {
+
+        private FixedContainerItems mRecommendedWidgets = null;
+
+        @Override
+        public void bindExtraContainerItems(FixedContainerItems item) {
+            mRecommendedWidgets = item;
+        }
+
+        @Override
+        public int getPageToBindSynchronously() {
+            return 0;
+        }
+
+        @Override
+        public void clearPendingBinds() { }
+
+        @Override
+        public void startBinding() { }
+
+        @Override
+        public void bindItems(List<ItemInfo> shortcuts, boolean forceAnimateIcons) { }
+
+        @Override
+        public void bindScreens(IntArray orderedScreenIds) { }
+
+        @Override
+        public void finishFirstPageBind(ViewOnDrawExecutor executor) { }
+
+        @Override
+        public void finishBindingItems(int pageBoundFirst) { }
+
+        @Override
+        public void preAddApps() { }
+
+        @Override
+        public void bindAppsAdded(IntArray newScreens, ArrayList<ItemInfo> addNotAnimated,
+                ArrayList<ItemInfo> addAnimated) { }
+
+        @Override
+        public void bindIncrementalDownloadProgressUpdated(AppInfo app) { }
+
+        @Override
+        public void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) { }
+
+        @Override
+        public void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets) { }
+
+        @Override
+        public void bindRestoreItemsChange(HashSet<ItemInfo> updates) { }
+
+        @Override
+        public void bindWorkspaceComponentsRemoved(ItemInfoMatcher matcher) { }
+
+        @Override
+        public void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
+
+        @Override
+        public void onPageBoundSynchronously(int page) { }
+
+        @Override
+        public void executeOnNextDraw(ViewOnDrawExecutor executor) { }
+
+        @Override
+        public void bindDeepShortcutMap(HashMap<ComponentKey, Integer> deepShortcutMap) { }
+
+        @Override
+        public void bindAllApplications(AppInfo[] apps, int flags) { }
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 2ea8bd2..b1b4d70 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -21,6 +21,7 @@
 import static com.android.launcher3.InvariantDeviceProfile.CHANGE_FLAG_GRID;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PREDICTION;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.Utilities.getDevicePrefs;
@@ -80,6 +81,8 @@
             new PredictorState(CONTAINER_PREDICTION, "all_apps_predictions");
     private final PredictorState mHotseatState =
             new PredictorState(CONTAINER_HOTSEAT_PREDICTION, "hotseat_predictions");
+    private final PredictorState mWidgetsRecommendationState =
+            new PredictorState(CONTAINER_WIDGETS_PREDICTION, "widgets_prediction");
 
     private final InvariantDeviceProfile mIDP;
     private final AppEventProducer mAppEventProducer;
@@ -111,6 +114,9 @@
         mHotseatState.items.setItems(
                 mHotseatState.storage.read(mApp.getContext(), hotseatFactory, ums.allUsers::get));
         mDataModel.extraItems.put(CONTAINER_HOTSEAT_PREDICTION, mHotseatState.items);
+
+        // Widgets prediction isn't used frequently. And thus, it is not persisted on disk.
+        mDataModel.extraItems.put(CONTAINER_WIDGETS_PREDICTION, mWidgetsRecommendationState.items);
         mActive = true;
     }
 
@@ -161,6 +167,9 @@
         if (mAllAppsState.predictor != null) {
             mAllAppsState.predictor.requestPredictionUpdate();
         }
+        if (mWidgetsRecommendationState.predictor != null) {
+            mWidgetsRecommendationState.predictor.requestPredictionUpdate();
+        }
     }
 
     @Override
@@ -176,6 +185,7 @@
     private void destroyPredictors() {
         mAllAppsState.destroyPredictor();
         mHotseatState.destroyPredictor();
+        mWidgetsRecommendationState.destroyPredictor();
     }
 
     @WorkerThread
@@ -203,6 +213,12 @@
                         .setPredictedTargetCount(mIDP.numHotseatIcons)
                         .setExtras(convertDataModelToAppTargetBundle(context, mDataModel))
                         .build()));
+
+        registerWidgetsPredictor(apm.createAppPredictionSession(
+                new AppPredictionContext.Builder(context)
+                        .setUiSurface("widgets")
+                        .setPredictedTargetCount(mIDP.numColumns)
+                        .build()));
     }
 
     private void registerPredictor(PredictorState state, AppPredictor predictor) {
@@ -220,6 +236,20 @@
         mApp.getModel().enqueueModelUpdateTask(new PredictionUpdateTask(state, targets));
     }
 
+    private void registerWidgetsPredictor(AppPredictor predictor) {
+        mWidgetsRecommendationState.predictor = predictor;
+        mWidgetsRecommendationState.predictor.registerPredictionUpdates(
+                Executors.MODEL_EXECUTOR, targets -> {
+                    if (mWidgetsRecommendationState.setTargets(targets)) {
+                        // No diff, skip
+                        return;
+                    }
+                    mApp.getModel().enqueueModelUpdateTask(
+                            new WidgetsPredictionUpdateTask(mWidgetsRecommendationState, targets));
+                });
+        mWidgetsRecommendationState.predictor.requestPredictionUpdate();
+    }
+
     @Override
     public void onIdpChanged(int changeFlags, InvariantDeviceProfile profile) {
         if ((changeFlags & CHANGE_FLAG_GRID) != 0) {
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
new file mode 100644
index 0000000..a29ac1a
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -0,0 +1,82 @@
+/*
+ * 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.model;
+
+import android.app.prediction.AppTarget;
+
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.widget.PendingAddWidgetInfo;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/** Task to update model as a result of predicted widgets update */
+public final class WidgetsPredictionUpdateTask extends BaseModelUpdateTask {
+    private final PredictorState mPredictorState;
+    private final List<AppTarget> mTargets;
+
+    WidgetsPredictionUpdateTask(PredictorState predictorState, List<AppTarget> targets) {
+        mPredictorState = predictorState;
+        mTargets = targets;
+    }
+
+    /**
+     * Uses the app predication result to infer widgets that the user may want to use.
+     *
+     * <p>The algorithm uses the app prediction ranking to create a widgets ranking which only
+     * includes one widget per app and excludes widgets that have already been added to the
+     * workspace.
+     */
+    @Override
+    public void execute(LauncherAppState appState, BgDataModel dataModel, AllAppsList apps) {
+        Set<ComponentKey> widgetsInWorkspace = dataModel.appWidgets.stream().map(
+                widget -> new ComponentKey(widget.providerName, widget.user)).collect(
+                Collectors.toSet());
+        Map<PackageUserKey, List<WidgetItem>> allWidgets =
+                dataModel.widgetsModel.getAllWidgetsWithoutShortcuts();
+
+        ArrayList<ItemInfo> recommendedWidgetsInDescendingOrder = new ArrayList<>();
+        for (AppTarget app : mTargets) {
+            PackageUserKey packageUserKey = new PackageUserKey(app.getPackageName(), app.getUser());
+            if (allWidgets.containsKey(packageUserKey)) {
+                List<WidgetItem> notAddedWidgets = allWidgets.get(packageUserKey).stream()
+                        .filter(item ->
+                                !widgetsInWorkspace.contains(
+                                        new ComponentKey(item.componentName, item.user)))
+                        .collect(Collectors.toList());
+                if (notAddedWidgets.size() > 0) {
+                    // Even an apps have more than one widgets, we only include one widget.
+                    recommendedWidgetsInDescendingOrder.add(
+                            new PendingAddWidgetInfo(notAddedWidgets.get(0).widgetInfo));
+                }
+            }
+        }
+        FixedContainerItems fixedContainerItems = mPredictorState.items;
+        fixedContainerItems.items.clear();
+        fixedContainerItems.items.addAll(recommendedWidgetsInDescendingOrder);
+        bindExtraContainerItems(fixedContainerItems);
+
+        // Don't store widgets prediction to disk because it is not used frequently.
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java
index 74a82ab..544835c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java
@@ -85,6 +85,7 @@
     // Contains all loaded Hotseat items.
     private ItemInfo[] mLatestLoadedHotseatItems;
 
+    private @Nullable Animator mAnimator;
     private boolean mIsAnimatingToLauncher;
 
     public TaskbarController(BaseQuickstepLauncher launcher,
@@ -252,6 +253,10 @@
         mTaskbarVisibilityController.cleanup();
         mHotseatController.cleanup();
         mRecentsController.cleanup();
+
+        if (mAnimator != null) {
+            mAnimator.end();
+        }
     }
 
     private void removeFromWindowManager() {
@@ -294,13 +299,21 @@
      */
     public void onLauncherResumedOrPaused(boolean isResumed) {
         long duration = QuickstepAppTransitionManagerImpl.CONTENT_ALPHA_DURATION;
-        final Animator anim;
-        if (isResumed) {
-            anim = createAnimToLauncher(null, duration);
-        } else {
-            anim = createAnimToApp(duration);
+        if (mAnimator != null) {
+            mAnimator.cancel();
         }
-        anim.start();
+        if (isResumed) {
+            mAnimator = createAnimToLauncher(null, duration);
+        } else {
+            mAnimator = createAnimToApp(duration);
+        }
+        mAnimator.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                mAnimator = null;
+            }
+        });
+        mAnimator.start();
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 55dde45..b009629 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -218,6 +218,8 @@
                     .setPredictedApps(item.items);
         } else if (item.containerId == Favorites.CONTAINER_HOTSEAT_PREDICTION) {
             mHotseatPredictionController.setPredictedItems(item);
+        } else if (item.containerId == Favorites.CONTAINER_WIDGETS_PREDICTION) {
+            getPopupDataProvider().setRecommendedWidgets(item.items);
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 7f2af6b..e0c041e 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -934,7 +934,9 @@
                     duration = Math.min(MAX_SWIPE_DURATION, 2 * baseDuration);
             }
         }
-        Interpolator interpolator = endTarget == RECENTS ? OVERSHOOT_1_2 : DEACCEL;
+        Interpolator interpolator =
+                endTarget == RECENTS ? (mDp.isTablet && FeatureFlags.ENABLE_OVERVIEW_GRID.get()
+                        ? ACCEL_DEACCEL : OVERSHOOT_1_2) : DEACCEL;
 
         if (endTarget.isLauncher) {
             mInputConsumerProxy.enable();
diff --git a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
index 6a32dc4..a749f9a 100644
--- a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
@@ -273,7 +273,6 @@
         Point size = display.realSize;
         int rotation = display.rotation;
         int touchHeight = mNavBarGesturalHeight;
-        int largerGesturalHeight = mNavBarLargerGesturalHeight;
         OrientationRectF orientationRectF =
                 new OrientationRectF(0, 0, size.x, size.y, rotation);
         if (mMode == SysUINavigationMode.Mode.NO_BUTTON) {
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java
index b10bdde..cd69cf1 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java
@@ -21,14 +21,18 @@
 import static android.view.MotionEvent.ACTION_MOVE;
 import static android.view.MotionEvent.ACTION_UP;
 
+import static com.android.launcher3.ResourceUtils.NAVBAR_BOTTOM_GESTURE_SIZE;
 import static com.android.launcher3.Utilities.squaredHypot;
 
 import android.content.Context;
+import android.graphics.Point;
 import android.graphics.PointF;
 import android.view.MotionEvent;
 
 import com.android.launcher3.R;
+import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.util.DisplayController;
 import com.android.quickstep.InputConsumer;
 import com.android.quickstep.RecentsAnimationDeviceState;
 import com.android.quickstep.SystemUiProxy;
@@ -45,11 +49,15 @@
     private static final int ANGLE_MIN = 30;
 
     private final Context mContext;
+    private final DisplayController.DisplayHolder mDisplayHolder;
+    private final Point mDisplaySize;
     private final RecentsAnimationDeviceState mDeviceState;
 
     private final float mDragDistThreshold;
     private final float mSquaredSlop;
 
+    private final int mNavBarSize;
+
     private final PointF mDownPos = new PointF();
     private final PointF mLastPos = new PointF();
 
@@ -60,10 +68,14 @@
             InputConsumer delegate, InputMonitorCompat inputMonitor) {
         super(delegate, inputMonitor);
         mContext = context;
+        mDisplayHolder = DisplayController.getDefaultDisplay(mContext);
         mDeviceState = deviceState;
         mDragDistThreshold = context.getResources().getDimensionPixelSize(
                 R.dimen.gestures_onehanded_drag_threshold);
         mSquaredSlop = Utilities.squaredTouchSlop(context);
+        mDisplaySize = mDisplayHolder.getInfo().realSize;
+        mNavBarSize = ResourceUtils.getNavbarSize(NAVBAR_BOTTOM_GESTURE_SIZE,
+                mContext.getResources());
     }
 
     @Override
@@ -96,7 +108,8 @@
                                 mDownPos.x - mLastPos.x, mDownPos.y - mLastPos.y))
                                 || (mDeviceState.isOneHandedModeActive() && isValidExitAngle(
                                 mDownPos.x - mLastPos.x, mDownPos.y - mLastPos.y))) {
-                            mPassedSlop = true;
+                            // To avoid mis-trigger when motion not touch system gesture region.
+                            mPassedSlop = isInSystemGestureRegion(mLastPos);
                             setActive(ev);
                         } else {
                             mState = STATE_DELEGATE_ACTIVE;
@@ -154,6 +167,11 @@
         SystemUiProxy.INSTANCE.get(mContext).stopOneHandedMode();
     }
 
+    private boolean isInSystemGestureRegion(PointF lastPos) {
+        final int navBarUpperBound = mDisplaySize.y - mNavBarSize;
+        return mDeviceState.isGesturalNavMode() && lastPos.y > navBarUpperBound;
+    }
+
     private boolean isValidStartAngle(float deltaX, float deltaY) {
         final float angle = (float) Math.toDegrees(Math.atan2(deltaY, deltaX));
         return angle > -(ANGLE_MAX) && angle < -(ANGLE_MIN);
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index a762cb7..66c24c8 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -31,6 +31,7 @@
 
 import android.content.Context;
 import android.util.Log;
+import android.view.View;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.WorkerThread;
@@ -55,6 +56,7 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.LogConfig;
+import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
 import com.android.systemui.shared.system.SysUiStatsLog;
 
 import java.util.Optional;
@@ -94,7 +96,7 @@
 
     @Override
     protected StatsLogger createLogger() {
-        return new StatsCompatLogger();
+        return new StatsCompatLogger(mContext);
     }
 
     /**
@@ -136,6 +138,7 @@
 
         private static final ItemInfo DEFAULT_ITEM_INFO = new ItemInfo();
 
+        private Context mContext;
         private ItemInfo mItemInfo = DEFAULT_ITEM_INFO;
         private InstanceId mInstanceId = DEFAULT_INSTANCE_ID;
         private OptionalInt mRank = OptionalInt.empty();
@@ -147,6 +150,10 @@
         private Optional<String> mEditText = Optional.empty();
         private SliceItem mSliceItem;
 
+        StatsCompatLogger(Context context) {
+            mContext = context;
+        }
+
         @Override
         public StatsLogger withItemInfo(ItemInfo itemInfo) {
             if (mContainerInfo.isPresent()) {
@@ -220,7 +227,6 @@
             if (!Utilities.ATLEAST_R) {
                 return;
             }
-
             LauncherAppState appState = LauncherAppState.getInstanceNoCreate();
 
             if (mSliceItem != null) {
@@ -256,6 +262,26 @@
             }
         }
 
+        @Override
+        public void sendToInteractionJankMonitor(EventEnum event, View view) {
+            if (!(event instanceof LauncherEvent)) {
+                return;
+            }
+            switch ((LauncherEvent) event) {
+                case LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN:
+                    InteractionJankMonitorWrapper.begin(
+                            view,
+                            InteractionJankMonitorWrapper.CUJ_ALL_APPS_SCROLL);
+                    break;
+                case LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END:
+                    InteractionJankMonitorWrapper.end(
+                            InteractionJankMonitorWrapper.CUJ_ALL_APPS_SCROLL);
+                    break;
+                default:
+                    break;
+            }
+        }
+
         private LauncherAtom.ItemInfo applyOverwrites(LauncherAtom.ItemInfo atomInfo) {
             LauncherAtom.ItemInfo.Builder itemInfoBuilder = atomInfo.toBuilder();
 
diff --git a/robolectric_tests/resources/widgets_predication_update_task_data.txt b/robolectric_tests/resources/widgets_predication_update_task_data.txt
new file mode 100644
index 0000000..941d195
--- /dev/null
+++ b/robolectric_tests/resources/widgets_predication_update_task_data.txt
@@ -0,0 +1,24 @@
+# Model data used by WidgetsPredictionUpdateTasksTest
+
+classMap s com.android.launcher3.model.data.WorkspaceItemInfo
+classMap w com.android.launcher3.model.data.LauncherAppWidgetInfo
+
+# Items for the BgDataModel
+
+# App shortcuts
+bgItem s itemType=0 title=app1-class1 intent=component=app1/class1 id=1
+bgItem s itemType=0 title=app1-class2 intent=component=app1/class2 id=2
+bgItem s itemType=0 title=app2-class1 intent=component=app2/class1 id=3
+bgItem s itemType=0 title=app2-class2 intent=component=app2/class2 id=4
+
+# Promise icons for app3
+bgItem s itemType=0 status=2 title=app3-class1 intent=component=app3/class1 id=5
+bgItem s itemType=0 status=2 title=app3-class2 intent=component=app3/class2 id=6
+bgItem s itemType=1 status=1 title=app3-shrt intent=component=app3/class3 id=7
+
+# Promise icon for app4
+bgItem s itemType=1 status=1 title=app4-shrt intent=component=app4/class1 id=8
+
+# Widget
+bgItem w providerName=app4/provider1 id=9
+bgItem w providerName=app5/provider1 id=10
\ No newline at end of file
diff --git a/src/com/android/launcher3/AppWidgetResizeFrame.java b/src/com/android/launcher3/AppWidgetResizeFrame.java
index 3b28d4d..9d6af9f 100644
--- a/src/com/android/launcher3/AppWidgetResizeFrame.java
+++ b/src/com/android/launcher3/AppWidgetResizeFrame.java
@@ -26,7 +26,6 @@
 
 import com.android.launcher3.accessibility.DragViewStateAnnouncer;
 import com.android.launcher3.dragndrop.DragLayer;
-import com.android.launcher3.util.FocusLogic;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
@@ -544,7 +543,7 @@
     @Override
     public boolean onKey(View v, int keyCode, KeyEvent event) {
         // Clear the frame and give focus to the widget host view when a directional key is pressed.
-        if (FocusLogic.shouldConsume(keyCode)) {
+        if (shouldConsume(keyCode)) {
             close(false);
             mWidgetView.requestFocus();
             return true;
@@ -667,4 +666,14 @@
             return moveEnd ? out.size() - size() : size() - out.size();
         }
     }
+
+    /**
+     * Returns true only if this utility class handles the key code.
+     */
+    public static boolean shouldConsume(int keyCode) {
+        return (keyCode == KeyEvent.KEYCODE_DPAD_LEFT || keyCode == KeyEvent.KEYCODE_DPAD_RIGHT
+                || keyCode == KeyEvent.KEYCODE_DPAD_UP || keyCode == KeyEvent.KEYCODE_DPAD_DOWN
+                || keyCode == KeyEvent.KEYCODE_MOVE_HOME || keyCode == KeyEvent.KEYCODE_MOVE_END
+                || keyCode == KeyEvent.KEYCODE_PAGE_UP || keyCode == KeyEvent.KEYCODE_PAGE_DOWN);
+    }
 }
diff --git a/src/com/android/launcher3/FocusHelper.java b/src/com/android/launcher3/FocusHelper.java
deleted file mode 100644
index e5aecf7..0000000
--- a/src/com/android/launcher3/FocusHelper.java
+++ /dev/null
@@ -1,567 +0,0 @@
-/*
- * Copyright (C) 2015 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3;
-
-import android.util.Log;
-import android.view.KeyEvent;
-import android.view.SoundEffectConstants;
-import android.view.View;
-import android.view.ViewGroup;
-
-import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.folder.Folder;
-import com.android.launcher3.folder.FolderPagedView;
-import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.util.FocusLogic;
-import com.android.launcher3.util.Thunk;
-
-/**
- * A keyboard listener we set on all the workspace icons.
- */
-class IconKeyEventListener implements View.OnKeyListener {
-    @Override
-    public boolean onKey(View v, int keyCode, KeyEvent event) {
-        return FocusHelper.handleIconKeyEvent(v, keyCode, event);
-    }
-}
-
-/**
- * A keyboard listener we set on all the hotseat buttons.
- */
-class HotseatIconKeyEventListener implements View.OnKeyListener {
-    @Override
-    public boolean onKey(View v, int keyCode, KeyEvent event) {
-        return FocusHelper.handleHotseatButtonKeyEvent(v, keyCode, event);
-    }
-}
-
-/**
- * A keyboard listener we set on full screen pages (e.g. custom content).
- */
-class FullscreenKeyEventListener implements View.OnKeyListener {
-    @Override
-    public boolean onKey(View v, int keyCode, KeyEvent event) {
-        if (keyCode == KeyEvent.KEYCODE_DPAD_LEFT || keyCode == KeyEvent.KEYCODE_DPAD_RIGHT
-                || keyCode == KeyEvent.KEYCODE_PAGE_DOWN || keyCode == KeyEvent.KEYCODE_PAGE_UP) {
-            // Handle the key event just like a workspace icon would in these cases. In this case,
-            // it will basically act as if there is a single icon in the top left (so you could
-            // think of the fullscreen page as a focusable fullscreen widget).
-            return FocusHelper.handleIconKeyEvent(v, keyCode, event);
-        }
-        return false;
-    }
-}
-
-/**
- * TODO: Reevaluate if this is still required
- */
-public class FocusHelper {
-
-    private static final String TAG = "FocusHelper";
-    private static final boolean DEBUG = false;
-
-    /**
-     * Handles key events in paged folder.
-     */
-    public static class PagedFolderKeyEventListener implements View.OnKeyListener {
-
-        private final Folder mFolder;
-
-        public PagedFolderKeyEventListener(Folder folder) {
-            mFolder = folder;
-        }
-
-        @Override
-        public boolean onKey(View v, int keyCode, KeyEvent e) {
-            boolean consume = FocusLogic.shouldConsume(keyCode);
-            if (e.getAction() == KeyEvent.ACTION_UP) {
-                return consume;
-            }
-            if (DEBUG) {
-                Log.v(TAG, String.format("Handle ALL Folders keyevent=[%s].",
-                        KeyEvent.keyCodeToString(keyCode)));
-            }
-
-            if (!(v.getParent() instanceof ShortcutAndWidgetContainer)) {
-                if (FeatureFlags.IS_STUDIO_BUILD) {
-                    throw new IllegalStateException("Parent of the focused item is not supported.");
-                } else {
-                    return false;
-                }
-            }
-
-            // Initialize variables.
-            final ShortcutAndWidgetContainer itemContainer = (ShortcutAndWidgetContainer) v.getParent();
-            final CellLayout cellLayout = (CellLayout) itemContainer.getParent();
-
-            final int iconIndex = itemContainer.indexOfChild(v);
-            final FolderPagedView pagedView = (FolderPagedView) cellLayout.getParent();
-
-            final int pageIndex = pagedView.indexOfChild(cellLayout);
-            final int pageCount = pagedView.getPageCount();
-            final boolean isLayoutRtl = Utilities.isRtl(v.getResources());
-
-            int[][] matrix = FocusLogic.createSparseMatrix(cellLayout);
-            // Process focus.
-            int newIconIndex = FocusLogic.handleKeyEvent(keyCode, matrix, iconIndex, pageIndex,
-                    pageCount, isLayoutRtl);
-            if (newIconIndex == FocusLogic.NOOP) {
-                handleNoopKey(keyCode, v);
-                return consume;
-            }
-            ShortcutAndWidgetContainer newParent = null;
-            View child = null;
-
-            switch (newIconIndex) {
-                case FocusLogic.PREVIOUS_PAGE_RIGHT_COLUMN:
-                case FocusLogic.PREVIOUS_PAGE_LEFT_COLUMN:
-                    newParent = getCellLayoutChildrenForIndex(pagedView, pageIndex - 1);
-                    if (newParent != null) {
-                        int row = ((CellLayout.LayoutParams) v.getLayoutParams()).cellY;
-                        pagedView.snapToPage(pageIndex - 1);
-                        child = newParent.getChildAt(
-                                ((newIconIndex == FocusLogic.PREVIOUS_PAGE_LEFT_COLUMN)
-                                    ^ newParent.invertLayoutHorizontally()) ? 0 : matrix.length - 1,
-                                row);
-                    }
-                    break;
-                case FocusLogic.PREVIOUS_PAGE_FIRST_ITEM:
-                    newParent = getCellLayoutChildrenForIndex(pagedView, pageIndex - 1);
-                    if (newParent != null) {
-                        pagedView.snapToPage(pageIndex - 1);
-                        child = newParent.getChildAt(0, 0);
-                    }
-                    break;
-                case FocusLogic.PREVIOUS_PAGE_LAST_ITEM:
-                    newParent = getCellLayoutChildrenForIndex(pagedView, pageIndex - 1);
-                    if (newParent != null) {
-                        pagedView.snapToPage(pageIndex - 1);
-                        child = newParent.getChildAt(matrix.length - 1, matrix[0].length - 1);
-                    }
-                    break;
-                case FocusLogic.NEXT_PAGE_FIRST_ITEM:
-                    newParent = getCellLayoutChildrenForIndex(pagedView, pageIndex + 1);
-                    if (newParent != null) {
-                        pagedView.snapToPage(pageIndex + 1);
-                        child = newParent.getChildAt(0, 0);
-                    }
-                    break;
-                case FocusLogic.NEXT_PAGE_LEFT_COLUMN:
-                case FocusLogic.NEXT_PAGE_RIGHT_COLUMN:
-                    newParent = getCellLayoutChildrenForIndex(pagedView, pageIndex + 1);
-                    if (newParent != null) {
-                        pagedView.snapToPage(pageIndex + 1);
-                        child = FocusLogic.getAdjacentChildInNextFolderPage(
-                                newParent, v, newIconIndex);
-                    }
-                    break;
-                case FocusLogic.CURRENT_PAGE_FIRST_ITEM:
-                    child = cellLayout.getChildAt(0, 0);
-                    break;
-                case FocusLogic.CURRENT_PAGE_LAST_ITEM:
-                    child = pagedView.getLastItem();
-                    break;
-                default: // Go to some item on the current page.
-                    child = itemContainer.getChildAt(newIconIndex);
-                    break;
-            }
-            if (child != null) {
-                child.requestFocus();
-                playSoundEffect(keyCode, v);
-            } else {
-                handleNoopKey(keyCode, v);
-            }
-            return consume;
-        }
-
-        public void handleNoopKey(int keyCode, View v) {
-            if (keyCode == KeyEvent.KEYCODE_DPAD_DOWN) {
-                mFolder.mFolderName.requestFocus();
-                playSoundEffect(keyCode, v);
-            }
-        }
-    }
-
-    /**
-     * Handles key events in the workspace hotseat (bottom of the screen).
-     * <p>Currently we don't special case for the phone UI in different orientations, even though
-     * the hotseat is on the side in landscape mode. This is to ensure that accessibility
-     * consistency is maintained across rotations.
-     */
-    static boolean handleHotseatButtonKeyEvent(View v, int keyCode, KeyEvent e) {
-        boolean consume = FocusLogic.shouldConsume(keyCode);
-        if (e.getAction() == KeyEvent.ACTION_UP || !consume) {
-            return consume;
-        }
-
-        final Launcher launcher = Launcher.getLauncher(v.getContext());
-        final DeviceProfile profile = launcher.getDeviceProfile();
-
-        if (DEBUG) {
-            Log.v(TAG, String.format(
-                    "Handle HOTSEAT BUTTONS keyevent=[%s] on hotseat buttons, isVertical=%s",
-                    KeyEvent.keyCodeToString(keyCode), profile.isVerticalBarLayout()));
-        }
-
-        // Initialize the variables.
-        final Workspace workspace = (Workspace) v.getRootView().findViewById(R.id.workspace);
-        final ShortcutAndWidgetContainer hotseatParent = (ShortcutAndWidgetContainer) v.getParent();
-        final CellLayout hotseatLayout = (CellLayout) hotseatParent.getParent();
-
-        final ItemInfo itemInfo = (ItemInfo) v.getTag();
-        int pageIndex = workspace.getNextPage();
-        int pageCount = workspace.getChildCount();
-        int iconIndex = hotseatParent.indexOfChild(v);
-        int iconRank = ((CellLayout.LayoutParams) hotseatLayout.getShortcutsAndWidgets()
-                .getChildAt(iconIndex).getLayoutParams()).cellX;
-
-        final CellLayout iconLayout = (CellLayout) workspace.getChildAt(pageIndex);
-        if (iconLayout == null) {
-            // This check is to guard against cases where key strokes rushes in when workspace
-            // child creation/deletion is still in flux. (e.g., during drop or fling
-            // animation.)
-            return consume;
-        }
-        final ViewGroup iconParent = iconLayout.getShortcutsAndWidgets();
-
-        ViewGroup parent = null;
-        int[][] matrix = null;
-
-        if (keyCode == KeyEvent.KEYCODE_DPAD_UP &&
-                !profile.isVerticalBarLayout()) {
-            matrix = FocusLogic.createSparseMatrixWithHotseat(iconLayout, hotseatLayout, profile);
-            iconIndex += iconParent.getChildCount();
-            parent = iconParent;
-        } else if (keyCode == KeyEvent.KEYCODE_DPAD_LEFT &&
-                profile.isVerticalBarLayout()) {
-            matrix = FocusLogic.createSparseMatrixWithHotseat(iconLayout, hotseatLayout, profile);
-            iconIndex += iconParent.getChildCount();
-            parent = iconParent;
-        } else if (keyCode == KeyEvent.KEYCODE_DPAD_RIGHT &&
-                profile.isVerticalBarLayout()) {
-            keyCode = KeyEvent.KEYCODE_PAGE_DOWN;
-        } else {
-            // For other KEYCODE_DPAD_LEFT and KEYCODE_DPAD_RIGHT navigation, do not use the
-            // matrix extended with hotseat.
-            matrix = FocusLogic.createSparseMatrix(hotseatLayout);
-            parent = hotseatParent;
-        }
-
-        // Process the focus.
-        int newIconIndex = FocusLogic.handleKeyEvent(keyCode, matrix, iconIndex, pageIndex,
-                pageCount, Utilities.isRtl(v.getResources()));
-
-        View newIcon = null;
-        switch (newIconIndex) {
-            case FocusLogic.NEXT_PAGE_FIRST_ITEM:
-                parent = getCellLayoutChildrenForIndex(workspace, pageIndex + 1);
-                newIcon = parent.getChildAt(0);
-                // TODO(hyunyoungs): handle cases where the child is not an icon but
-                // a folder or a widget.
-                workspace.snapToPage(pageIndex + 1);
-                break;
-            case FocusLogic.PREVIOUS_PAGE_FIRST_ITEM:
-                parent = getCellLayoutChildrenForIndex(workspace, pageIndex - 1);
-                newIcon = parent.getChildAt(0);
-                // TODO(hyunyoungs): handle cases where the child is not an icon but
-                // a folder or a widget.
-                workspace.snapToPage(pageIndex - 1);
-                break;
-            case FocusLogic.PREVIOUS_PAGE_LAST_ITEM:
-                parent = getCellLayoutChildrenForIndex(workspace, pageIndex - 1);
-                newIcon = parent.getChildAt(parent.getChildCount() - 1);
-                // TODO(hyunyoungs): handle cases where the child is not an icon but
-                // a folder or a widget.
-                workspace.snapToPage(pageIndex - 1);
-                break;
-            case FocusLogic.PREVIOUS_PAGE_LEFT_COLUMN:
-            case FocusLogic.PREVIOUS_PAGE_RIGHT_COLUMN:
-                // Go to the previous page but keep the focus on the same hotseat icon.
-                workspace.snapToPage(pageIndex - 1);
-                break;
-            case FocusLogic.NEXT_PAGE_LEFT_COLUMN:
-            case FocusLogic.NEXT_PAGE_RIGHT_COLUMN:
-                // Go to the next page but keep the focus on the same hotseat icon.
-                workspace.snapToPage(pageIndex + 1);
-                break;
-        }
-        if (parent == iconParent && newIconIndex >= iconParent.getChildCount()) {
-            newIconIndex -= iconParent.getChildCount();
-        }
-        if (parent != null) {
-            if (newIcon == null && newIconIndex >= 0) {
-                newIcon = parent.getChildAt(newIconIndex);
-            }
-            if (newIcon != null) {
-                newIcon.requestFocus();
-                playSoundEffect(keyCode, v);
-            }
-        }
-        return consume;
-    }
-
-    /**
-     * Handles key events in a workspace containing icons.
-     */
-    static boolean handleIconKeyEvent(View v, int keyCode, KeyEvent e) {
-        boolean consume = FocusLogic.shouldConsume(keyCode);
-        if (e.getAction() == KeyEvent.ACTION_UP || !consume) {
-            return consume;
-        }
-
-        Launcher launcher = Launcher.getLauncher(v.getContext());
-        DeviceProfile profile = launcher.getDeviceProfile();
-
-        if (DEBUG) {
-            Log.v(TAG, String.format("Handle WORKSPACE ICONS keyevent=[%s] isVerticalBar=%s",
-                    KeyEvent.keyCodeToString(keyCode), profile.isVerticalBarLayout()));
-        }
-
-        // Initialize the variables.
-        ShortcutAndWidgetContainer parent = (ShortcutAndWidgetContainer) v.getParent();
-        CellLayout iconLayout = (CellLayout) parent.getParent();
-        final Workspace workspace = (Workspace) iconLayout.getParent();
-        final ViewGroup dragLayer = (ViewGroup) workspace.getParent();
-        final ViewGroup tabs = (ViewGroup) dragLayer.findViewById(R.id.drop_target_bar);
-        final Hotseat hotseat = (Hotseat) dragLayer.findViewById(R.id.hotseat);
-
-        final ItemInfo itemInfo = (ItemInfo) v.getTag();
-        final int iconIndex = parent.indexOfChild(v);
-        final int pageIndex = workspace.indexOfChild(iconLayout);
-        final int pageCount = workspace.getChildCount();
-
-        CellLayout hotseatLayout = (CellLayout) hotseat.getChildAt(0);
-        ShortcutAndWidgetContainer hotseatParent = hotseatLayout.getShortcutsAndWidgets();
-        int[][] matrix;
-
-        // KEYCODE_DPAD_DOWN in portrait (KEYCODE_DPAD_RIGHT in landscape) is the only key allowed
-        // to take a user to the hotseat. For other dpad navigation, do not use the matrix extended
-        // with the hotseat.
-        if (keyCode == KeyEvent.KEYCODE_DPAD_DOWN && !profile.isVerticalBarLayout()) {
-            matrix = FocusLogic.createSparseMatrixWithHotseat(iconLayout, hotseatLayout, profile);
-        } else if (keyCode == KeyEvent.KEYCODE_DPAD_RIGHT &&
-                profile.isVerticalBarLayout()) {
-            matrix = FocusLogic.createSparseMatrixWithHotseat(iconLayout, hotseatLayout, profile);
-        } else {
-            matrix = FocusLogic.createSparseMatrix(iconLayout);
-        }
-
-        // Process the focus.
-        int newIconIndex = FocusLogic.handleKeyEvent(keyCode, matrix, iconIndex, pageIndex,
-                pageCount, Utilities.isRtl(v.getResources()));
-        boolean isRtl = Utilities.isRtl(v.getResources());
-        View newIcon = null;
-        CellLayout workspaceLayout = (CellLayout) workspace.getChildAt(pageIndex);
-        switch (newIconIndex) {
-            case FocusLogic.NOOP:
-                if (keyCode == KeyEvent.KEYCODE_DPAD_UP) {
-                    newIcon = tabs;
-                }
-                break;
-            case FocusLogic.PREVIOUS_PAGE_RIGHT_COLUMN:
-            case FocusLogic.NEXT_PAGE_RIGHT_COLUMN:
-                int newPageIndex = pageIndex - 1;
-                if (newIconIndex == FocusLogic.NEXT_PAGE_RIGHT_COLUMN) {
-                    newPageIndex = pageIndex + 1;
-                }
-                int row = ((CellLayout.LayoutParams) v.getLayoutParams()).cellY;
-                parent = getCellLayoutChildrenForIndex(workspace, newPageIndex);
-                if (parent != null) {
-                    iconLayout = (CellLayout) parent.getParent();
-                    matrix = FocusLogic.createSparseMatrixWithPivotColumn(iconLayout,
-                            iconLayout.getCountX(), row);
-                    newIconIndex = FocusLogic.handleKeyEvent(keyCode, matrix, FocusLogic.PIVOT,
-                            newPageIndex, pageCount, Utilities.isRtl(v.getResources()));
-                    if (newIconIndex == FocusLogic.NEXT_PAGE_FIRST_ITEM) {
-                        newIcon = handleNextPageFirstItem(workspace, hotseatLayout, pageIndex,
-                                isRtl);
-                    } else if (newIconIndex == FocusLogic.PREVIOUS_PAGE_LAST_ITEM) {
-                        newIcon = handlePreviousPageLastItem(workspace, hotseatLayout, pageIndex,
-                                isRtl);
-                    } else {
-                        newIcon = parent.getChildAt(newIconIndex);
-                    }
-                }
-                break;
-            case FocusLogic.PREVIOUS_PAGE_FIRST_ITEM:
-                workspaceLayout = (CellLayout) workspace.getChildAt(pageIndex - 1);
-                newIcon = getFirstFocusableIconInReadingOrder(workspaceLayout, isRtl);
-                if (newIcon == null) {
-                    // Check the hotseat if no focusable item was found on the workspace.
-                    newIcon = getFirstFocusableIconInReadingOrder(hotseatLayout, isRtl);
-                    workspace.snapToPage(pageIndex - 1);
-                }
-                break;
-            case FocusLogic.PREVIOUS_PAGE_LAST_ITEM:
-                newIcon = handlePreviousPageLastItem(workspace, hotseatLayout, pageIndex, isRtl);
-                break;
-            case FocusLogic.NEXT_PAGE_FIRST_ITEM:
-                newIcon = handleNextPageFirstItem(workspace, hotseatLayout, pageIndex, isRtl);
-                break;
-            case FocusLogic.NEXT_PAGE_LEFT_COLUMN:
-            case FocusLogic.PREVIOUS_PAGE_LEFT_COLUMN:
-                newPageIndex = pageIndex + 1;
-                if (newIconIndex == FocusLogic.PREVIOUS_PAGE_LEFT_COLUMN) {
-                    newPageIndex = pageIndex - 1;
-                }
-                row = ((CellLayout.LayoutParams) v.getLayoutParams()).cellY;
-                parent = getCellLayoutChildrenForIndex(workspace, newPageIndex);
-                if (parent != null) {
-                    iconLayout = (CellLayout) parent.getParent();
-                    matrix = FocusLogic.createSparseMatrixWithPivotColumn(iconLayout, -1, row);
-                    newIconIndex = FocusLogic.handleKeyEvent(keyCode, matrix, FocusLogic.PIVOT,
-                            newPageIndex, pageCount, Utilities.isRtl(v.getResources()));
-                    if (newIconIndex == FocusLogic.NEXT_PAGE_FIRST_ITEM) {
-                        newIcon = handleNextPageFirstItem(workspace, hotseatLayout, pageIndex,
-                                isRtl);
-                    } else if (newIconIndex == FocusLogic.PREVIOUS_PAGE_LAST_ITEM) {
-                        newIcon = handlePreviousPageLastItem(workspace, hotseatLayout, pageIndex,
-                                isRtl);
-                    } else {
-                        newIcon = parent.getChildAt(newIconIndex);
-                    }
-                }
-                break;
-            case FocusLogic.CURRENT_PAGE_FIRST_ITEM:
-                newIcon = getFirstFocusableIconInReadingOrder(workspaceLayout, isRtl);
-                if (newIcon == null) {
-                    // Check the hotseat if no focusable item was found on the workspace.
-                    newIcon = getFirstFocusableIconInReadingOrder(hotseatLayout, isRtl);
-                }
-                break;
-            case FocusLogic.CURRENT_PAGE_LAST_ITEM:
-                newIcon = getFirstFocusableIconInReverseReadingOrder(workspaceLayout, isRtl);
-                if (newIcon == null) {
-                    // Check the hotseat if no focusable item was found on the workspace.
-                    newIcon = getFirstFocusableIconInReverseReadingOrder(hotseatLayout, isRtl);
-                }
-                break;
-            default:
-                // current page, some item.
-                if (0 <= newIconIndex && newIconIndex < parent.getChildCount()) {
-                    newIcon = parent.getChildAt(newIconIndex);
-                } else if (parent.getChildCount() <= newIconIndex &&
-                        newIconIndex < parent.getChildCount() + hotseatParent.getChildCount()) {
-                    newIcon = hotseatParent.getChildAt(newIconIndex - parent.getChildCount());
-                }
-                break;
-        }
-        if (newIcon != null) {
-            newIcon.requestFocus();
-            playSoundEffect(keyCode, v);
-        }
-        return consume;
-    }
-
-    //
-    // Helper methods.
-    //
-
-    /**
-     * Private helper method to get the CellLayoutChildren given a CellLayout index.
-     */
-    @Thunk static ShortcutAndWidgetContainer getCellLayoutChildrenForIndex(
-            ViewGroup container, int i) {
-        CellLayout parent = (CellLayout) container.getChildAt(i);
-        return parent.getShortcutsAndWidgets();
-    }
-
-    /**
-     * Helper method to be used for playing sound effects.
-     */
-    @Thunk static void playSoundEffect(int keyCode, View v) {
-        switch (keyCode) {
-            case KeyEvent.KEYCODE_DPAD_LEFT:
-                v.playSoundEffect(SoundEffectConstants.NAVIGATION_LEFT);
-                break;
-            case KeyEvent.KEYCODE_DPAD_RIGHT:
-                v.playSoundEffect(SoundEffectConstants.NAVIGATION_RIGHT);
-                break;
-            case KeyEvent.KEYCODE_DPAD_DOWN:
-            case KeyEvent.KEYCODE_PAGE_DOWN:
-            case KeyEvent.KEYCODE_MOVE_END:
-                v.playSoundEffect(SoundEffectConstants.NAVIGATION_DOWN);
-                break;
-            case KeyEvent.KEYCODE_DPAD_UP:
-            case KeyEvent.KEYCODE_PAGE_UP:
-            case KeyEvent.KEYCODE_MOVE_HOME:
-                v.playSoundEffect(SoundEffectConstants.NAVIGATION_UP);
-                break;
-            default:
-                break;
-        }
-    }
-
-    private static View handlePreviousPageLastItem(Workspace workspace, CellLayout hotseatLayout,
-            int pageIndex, boolean isRtl) {
-        if (pageIndex - 1 < 0) {
-            return null;
-        }
-        CellLayout workspaceLayout = (CellLayout) workspace.getChildAt(pageIndex - 1);
-        View newIcon = getFirstFocusableIconInReverseReadingOrder(workspaceLayout, isRtl);
-        if (newIcon == null) {
-            // Check the hotseat if no focusable item was found on the workspace.
-            newIcon = getFirstFocusableIconInReverseReadingOrder(hotseatLayout,isRtl);
-            workspace.snapToPage(pageIndex - 1);
-        }
-        return newIcon;
-    }
-
-    private static View handleNextPageFirstItem(Workspace workspace, CellLayout hotseatLayout,
-            int pageIndex, boolean isRtl) {
-        if (pageIndex + 1 >= workspace.getPageCount()) {
-            return null;
-        }
-        CellLayout workspaceLayout = (CellLayout) workspace.getChildAt(pageIndex + 1);
-        View newIcon = getFirstFocusableIconInReadingOrder(workspaceLayout, isRtl);
-        if (newIcon == null) {
-            // Check the hotseat if no focusable item was found on the workspace.
-            newIcon = getFirstFocusableIconInReadingOrder(hotseatLayout, isRtl);
-            workspace.snapToPage(pageIndex + 1);
-        }
-        return newIcon;
-    }
-
-    private static View getFirstFocusableIconInReadingOrder(CellLayout cellLayout, boolean isRtl) {
-        View icon;
-        int countX = cellLayout.getCountX();
-        for (int y = 0; y < cellLayout.getCountY(); y++) {
-            int increment = isRtl ? -1 : 1;
-            for (int x = isRtl ? countX - 1 : 0; 0 <= x && x < countX; x += increment) {
-                if ((icon = cellLayout.getChildAt(x, y)) != null && icon.isFocusable()) {
-                    return icon;
-                }
-            }
-        }
-        return null;
-    }
-
-    private static View getFirstFocusableIconInReverseReadingOrder(CellLayout cellLayout,
-            boolean isRtl) {
-        View icon;
-        int countX = cellLayout.getCountX();
-        for (int y = cellLayout.getCountY() - 1; y >= 0; y--) {
-            int increment = isRtl ? 1 : -1;
-            for (int x = isRtl ? 0 : countX - 1; 0 <= x && x < countX; x += increment) {
-                if ((icon = cellLayout.getChildAt(x, y)) != null && icon.isFocusable()) {
-                    return icon;
-                }
-            }
-        }
-        return null;
-    }
-}
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index dfdc53c..22c257a 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -193,6 +193,7 @@
         public static final int CONTAINER_DESKTOP = -100;
         public static final int CONTAINER_HOTSEAT = -101;
         public static final int CONTAINER_PREDICTION = -102;
+        public static final int CONTAINER_WIDGETS_PREDICTION = -111;
         public static final int CONTAINER_HOTSEAT_PREDICTION = -103;
         public static final int CONTAINER_ALL_APPS = -104;
         public static final int CONTAINER_WIDGETS_TRAY = -105;
diff --git a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
index a05e036..bb175ea 100644
--- a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
+++ b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
@@ -290,8 +290,10 @@
     public void setAppsPerRow(int appsPerRow) {
         mAppsPerRow = appsPerRow;
         int totalSpans = mAppsPerRow;
-        for (int itemPerRow : mSearchAdapterProvider.getSupportedItemsPerRow()) {
-            totalSpans *= itemPerRow;
+        for (int itemPerRow : mSearchAdapterProvider.getSupportedItemsPerRowArray()) {
+            if (totalSpans % itemPerRow != 0) {
+                totalSpans *= itemPerRow;
+            }
         }
         mGridLayoutMgr.setSpanCount(totalSpans);
     }
diff --git a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
index ace9938..179cb77 100644
--- a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
@@ -19,6 +19,9 @@
 import static android.view.View.MeasureSpec.UNSPECIFIED;
 import static android.view.View.MeasureSpec.makeMeasureSpec;
 
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END;
+
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Canvas;
@@ -36,6 +39,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
+import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.views.RecyclerViewFastScroller;
 
 import java.util.ArrayList;
@@ -176,6 +180,23 @@
     }
 
     @Override
+    public void onScrollStateChanged(int state) {
+        super.onScrollStateChanged(state);
+
+        StatsLogManager mgr = BaseDraggingActivity.fromContext(getContext()).getStatsLogManager();
+        switch (state) {
+            case SCROLL_STATE_DRAGGING:
+                mgr.logger().sendToInteractionJankMonitor(
+                        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN, this);
+                break;
+            case SCROLL_STATE_IDLE:
+                mgr.logger().sendToInteractionJankMonitor(
+                        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END, this);
+                break;
+        }
+    }
+
+    @Override
     public boolean onInterceptTouchEvent(MotionEvent e) {
         boolean result = super.onInterceptTouchEvent(e);
         if (!result && e.getAction() == MotionEvent.ACTION_DOWN
diff --git a/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java b/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java
index 6d491fd..fdacd3d 100644
--- a/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java
+++ b/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java
@@ -53,7 +53,7 @@
     /**
      * Returns supported item per row combinations supported
      */
-    public int[] getSupportedItemsPerRow() {
+    public int[] getSupportedItemsPerRowArray() {
         return new int[]{};
     }
 
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index d554bb9..c5ffff9 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -21,6 +21,7 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_OVERVIEW_GESTURE;
 
 import android.content.Context;
+import android.view.View;
 
 import androidx.annotation.Nullable;
 import androidx.slice.SliceItem;
@@ -407,6 +408,12 @@
 
         @UiEvent(doc = "User switched to AllApps Work tab by tapping on it.")
         LAUNCHER_ALLAPPS_TAP_ON_WORK_TAB(722),
+
+        @UiEvent(doc = "All apps vertical fling started.")
+        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN(724),
+
+        @UiEvent(doc = "All apps vertical fling ended.")
+        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END(725)
         ;
 
         // ADD MORE
@@ -525,6 +532,13 @@
          */
         default void log(EventEnum event) {
         }
+
+        /**
+         * Builds the final message and logs it to two different atoms, one for
+         * event tracking and the other for jank tracking.
+         */
+        default void sendToInteractionJankMonitor(EventEnum event, View v) {
+        }
     }
 
     /**
diff --git a/src/com/android/launcher3/model/BaseModelUpdateTask.java b/src/com/android/launcher3/model/BaseModelUpdateTask.java
index 543ee2a..ad553d5 100644
--- a/src/com/android/launcher3/model/BaseModelUpdateTask.java
+++ b/src/com/android/launcher3/model/BaseModelUpdateTask.java
@@ -124,7 +124,7 @@
 
     public void bindUpdatedWidgets(BgDataModel dataModel) {
         final ArrayList<WidgetsListBaseEntry> widgets =
-                dataModel.widgetsModel.getWidgetsList(mApp.getContext());
+                dataModel.widgetsModel.getWidgetsListForPicker(mApp.getContext());
         scheduleCallbackTask(c -> c.bindAllWidgets(widgets));
     }
 
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index 1853632..7780894 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -58,8 +58,11 @@
     private HashMap<ComponentKey, Integer> mDeepShortcutMap = new HashMap<>();
     /** Maps packages to their DotInfo's . */
     private Map<PackageUserKey, DotInfo> mPackageUserToDotInfos = new HashMap<>();
-    /** Maps packages to their Widgets */
+
+    /** All installed widgets. */
     private List<WidgetsListBaseEntry> mAllWidgets = List.of();
+    /** Widgets that can be recommended to the users. */
+    private List<ItemInfo> mRecommendedWidgets = List.of();
 
     private PopupDataChangeListener mChangeListener = PopupDataChangeListener.INSTANCE;
 
@@ -187,6 +190,15 @@
         notificationListener.cancelNotificationFromLauncher(notificationKey);
     }
 
+    /**
+     * Sets a list of recommended widgets ordered by their order of appearance in the widgets
+     * recommendation UI.
+     */
+    public void setRecommendedWidgets(List<ItemInfo> recommendedWidgets) {
+        mRecommendedWidgets = recommendedWidgets;
+        mChangeListener.onRecommendedWidgetsBound();
+    }
+
     public void setAllWidgets(List<WidgetsListBaseEntry> allWidgets) {
         mAllWidgets = allWidgets;
         mChangeListener.onWidgetsBound();
@@ -200,6 +212,21 @@
         return mAllWidgets;
     }
 
+    /** Returns a list of recommended widgets. */
+    public List<WidgetItem> getRecommendedWidgets() {
+        HashMap<ComponentKey, WidgetItem> allWidgetItems = new HashMap<>();
+        mAllWidgets.stream()
+                .filter(entry -> entry instanceof WidgetsListContentEntry)
+                .forEach(entry -> ((WidgetsListContentEntry) entry).mWidgets
+                        .forEach(widget -> allWidgetItems.put(
+                                new ComponentKey(widget.componentName, widget.user), widget)));
+        return mRecommendedWidgets.stream()
+                .map(recommendedWidget -> allWidgetItems.get(
+                        new ComponentKey(recommendedWidget.getTargetComponent(),
+                                recommendedWidget.user)))
+                .collect(Collectors.toList());
+    }
+
     public List<WidgetItem> getWidgetsForPackageUser(PackageUserKey packageUserKey) {
         return mAllWidgets.stream()
                 .filter(row -> row instanceof WidgetsListContentEntry
@@ -244,5 +271,8 @@
         default void trimNotifications(Map<PackageUserKey, DotInfo> updatedDots) { }
 
         default void onWidgetsBound() { }
+
+        /** A callback to get notified when recommended widgets are bound. */
+        default void onRecommendedWidgetsBound() { }
     }
 }
diff --git a/src/com/android/launcher3/util/FocusLogic.java b/src/com/android/launcher3/util/FocusLogic.java
deleted file mode 100644
index 4f4cccd..0000000
--- a/src/com/android/launcher3/util/FocusLogic.java
+++ /dev/null
@@ -1,553 +0,0 @@
-/*
- * Copyright (C) 2015 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.util;
-
-import android.util.Log;
-import android.view.KeyEvent;
-import android.view.View;
-import android.view.ViewGroup;
-
-import com.android.launcher3.CellLayout;
-import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.ShortcutAndWidgetContainer;
-import com.android.launcher3.config.FeatureFlags;
-
-import java.util.Arrays;
-
-/**
- * Calculates the next item that a {@link KeyEvent} should change the focus to.
- *<p>
- * Note, this utility class calculates everything regards to icon index and its (x,y) coordinates.
- * Currently supports:
- * <ul>
- *  <li> full matrix of cells that are 1x1
- *  <li> sparse matrix of cells that are 1x1
- *     [ 1][  ][ 2][  ]
- *     [  ][  ][ 3][  ]
- *     [  ][ 4][  ][  ]
- *     [  ][ 5][ 6][ 7]
- * </ul>
- * *<p>
- * For testing, one can use a BT keyboard, or use following adb command.
- * ex. $ adb shell input keyevent 20 // KEYCODE_DPAD_LEFT
- */
-public class FocusLogic {
-
-    private static final String TAG = "FocusLogic";
-    private static final boolean DEBUG = false;
-
-    /** Item and page index related constant used by {@link #handleKeyEvent}. */
-    public static final int NOOP = -1;
-
-    public static final int PREVIOUS_PAGE_RIGHT_COLUMN  = -2;
-    public static final int PREVIOUS_PAGE_FIRST_ITEM    = -3;
-    public static final int PREVIOUS_PAGE_LAST_ITEM     = -4;
-    public static final int PREVIOUS_PAGE_LEFT_COLUMN   = -5;
-
-    public static final int CURRENT_PAGE_FIRST_ITEM     = -6;
-    public static final int CURRENT_PAGE_LAST_ITEM      = -7;
-
-    public static final int NEXT_PAGE_FIRST_ITEM        = -8;
-    public static final int NEXT_PAGE_LEFT_COLUMN       = -9;
-    public static final int NEXT_PAGE_RIGHT_COLUMN      = -10;
-
-    public static final int ALL_APPS_COLUMN = -11;
-
-    // Matrix related constant.
-    public static final int EMPTY = -1;
-    public static final int PIVOT = 100;
-
-    /**
-     * Returns true only if this utility class handles the key code.
-     */
-    public static boolean shouldConsume(int keyCode) {
-        return (keyCode == KeyEvent.KEYCODE_DPAD_LEFT || keyCode == KeyEvent.KEYCODE_DPAD_RIGHT ||
-                keyCode == KeyEvent.KEYCODE_DPAD_UP || keyCode == KeyEvent.KEYCODE_DPAD_DOWN ||
-                keyCode == KeyEvent.KEYCODE_MOVE_HOME || keyCode == KeyEvent.KEYCODE_MOVE_END ||
-                keyCode == KeyEvent.KEYCODE_PAGE_UP || keyCode == KeyEvent.KEYCODE_PAGE_DOWN);
-    }
-
-    public static int handleKeyEvent(int keyCode, int [][] map, int iconIdx, int pageIndex,
-            int pageCount, boolean isRtl) {
-
-        int cntX = map == null ? -1 : map.length;
-        int cntY = map == null ? -1 : map[0].length;
-
-        if (DEBUG) {
-            Log.v(TAG, String.format(
-                    "handleKeyEvent START: cntX=%d, cntY=%d, iconIdx=%d, pageIdx=%d, pageCnt=%d",
-                    cntX, cntY, iconIdx, pageIndex, pageCount));
-        }
-
-        int newIndex = NOOP;
-        switch (keyCode) {
-            case KeyEvent.KEYCODE_DPAD_LEFT:
-                newIndex = handleDpadHorizontal(iconIdx, cntX, cntY, map, -1 /*increment*/, isRtl);
-                if (!isRtl && newIndex == NOOP && pageIndex > 0) {
-                    newIndex = PREVIOUS_PAGE_RIGHT_COLUMN;
-                } else if (isRtl && newIndex == NOOP && pageIndex < pageCount - 1) {
-                    newIndex = NEXT_PAGE_RIGHT_COLUMN;
-                }
-                break;
-            case KeyEvent.KEYCODE_DPAD_RIGHT:
-                newIndex = handleDpadHorizontal(iconIdx, cntX, cntY, map, 1 /*increment*/, isRtl);
-                if (!isRtl && newIndex == NOOP && pageIndex < pageCount - 1) {
-                    newIndex = NEXT_PAGE_LEFT_COLUMN;
-                } else if (isRtl && newIndex == NOOP && pageIndex > 0) {
-                    newIndex = PREVIOUS_PAGE_LEFT_COLUMN;
-                }
-                break;
-            case KeyEvent.KEYCODE_DPAD_DOWN:
-                newIndex = handleDpadVertical(iconIdx, cntX, cntY, map, 1  /*increment*/);
-                break;
-            case KeyEvent.KEYCODE_DPAD_UP:
-                newIndex = handleDpadVertical(iconIdx, cntX, cntY, map, -1  /*increment*/);
-                break;
-            case KeyEvent.KEYCODE_MOVE_HOME:
-                newIndex = handleMoveHome();
-                break;
-            case KeyEvent.KEYCODE_MOVE_END:
-                newIndex = handleMoveEnd();
-                break;
-            case KeyEvent.KEYCODE_PAGE_DOWN:
-                newIndex = handlePageDown(pageIndex, pageCount);
-                break;
-            case KeyEvent.KEYCODE_PAGE_UP:
-                newIndex = handlePageUp(pageIndex);
-                break;
-            default:
-                break;
-        }
-
-        if (DEBUG) {
-            Log.v(TAG, String.format("handleKeyEvent FINISH: index [%d -> %s]",
-                    iconIdx, getStringIndex(newIndex)));
-        }
-        return newIndex;
-    }
-
-    /**
-     * Returns a matrix of size (m x n) that has been initialized with {@link #EMPTY}.
-     *
-     * @param m                 number of columns in the matrix
-     * @param n                 number of rows in the matrix
-     */
-    // TODO: get rid of dynamic matrix creation.
-    private static int[][] createFullMatrix(int m, int n) {
-        int[][] matrix = new int [m][n];
-
-        for (int i=0; i < m;i++) {
-            Arrays.fill(matrix[i], EMPTY);
-        }
-        return matrix;
-    }
-
-    /**
-     * Returns a matrix of size same as the {@link CellLayout} dimension that is initialized with the
-     * index of the child view.
-     */
-    // TODO: get rid of the dynamic matrix creation
-    public static int[][] createSparseMatrix(CellLayout layout) {
-        ShortcutAndWidgetContainer parent = layout.getShortcutsAndWidgets();
-        final int m = layout.getCountX();
-        final int n = layout.getCountY();
-        final boolean invert = parent.invertLayoutHorizontally();
-
-        int[][] matrix = createFullMatrix(m, n);
-
-        // Iterate thru the children.
-        for (int i = 0; i < parent.getChildCount(); i++ ) {
-            View cell = parent.getChildAt(i);
-            if (!cell.isFocusable()) {
-                continue;
-            }
-            int cx = ((CellLayout.LayoutParams) cell.getLayoutParams()).cellX;
-            int cy = ((CellLayout.LayoutParams) cell.getLayoutParams()).cellY;
-            int x = invert ? (m - cx - 1) : cx;
-            if (x < m && cy < n) { // check if view fits into matrix, else skip
-                matrix[x][cy] = i;
-            }
-        }
-        if (DEBUG) {
-            printMatrix(matrix);
-        }
-        return matrix;
-    }
-
-    /**
-     * Creates a sparse matrix that merges the icon and hotseat view group using the cell layout.
-     * The size of the returning matrix is [icon column count x (icon + hotseat row count)]
-     * in portrait orientation. In landscape, [(icon + hotseat) column count x (icon row count)]
-     */
-    // TODO: get rid of the dynamic matrix creation
-    public static int[][] createSparseMatrixWithHotseat(
-            CellLayout iconLayout, CellLayout hotseatLayout, DeviceProfile dp) {
-
-        ViewGroup iconParent = iconLayout.getShortcutsAndWidgets();
-        ViewGroup hotseatParent = hotseatLayout.getShortcutsAndWidgets();
-
-        boolean isHotseatHorizontal = !dp.isVerticalBarLayout();
-
-        int m, n;
-        if (isHotseatHorizontal) {
-            m = hotseatLayout.getCountX();
-            n = iconLayout.getCountY() + hotseatLayout.getCountY();
-        } else {
-            m = iconLayout.getCountX() + hotseatLayout.getCountX();
-            n = hotseatLayout.getCountY();
-        }
-        int[][] matrix = createFullMatrix(m, n);
-        // Iterate through the children of the workspace.
-        for (int i = 0; i < iconParent.getChildCount(); i++) {
-            View cell = iconParent.getChildAt(i);
-            if (!cell.isFocusable()) {
-                continue;
-            }
-            int cx = ((CellLayout.LayoutParams) cell.getLayoutParams()).cellX;
-            int cy = ((CellLayout.LayoutParams) cell.getLayoutParams()).cellY;
-            matrix[cx][cy] = i;
-        }
-
-        // Iterate thru the children of the hotseat.
-        for (int i = hotseatParent.getChildCount() - 1; i >= 0; i--) {
-            if (isHotseatHorizontal) {
-                int cx = ((CellLayout.LayoutParams)
-                        hotseatParent.getChildAt(i).getLayoutParams()).cellX;
-                matrix[cx][iconLayout.getCountY()] = iconParent.getChildCount() + i;
-            } else {
-                int cy = ((CellLayout.LayoutParams)
-                        hotseatParent.getChildAt(i).getLayoutParams()).cellY;
-                matrix[iconLayout.getCountX()][cy] = iconParent.getChildCount() + i;
-            }
-        }
-        if (DEBUG) {
-            printMatrix(matrix);
-        }
-        return matrix;
-    }
-
-    /**
-     * Creates a sparse matrix that merges the icon of previous/next page and last column of
-     * current page. When left key is triggered on the leftmost column, sparse matrix is created
-     * that combines previous page matrix and an extra column on the right. Likewise, when right
-     * key is triggered on the rightmost column, sparse matrix is created that combines this column
-     * on the 0th column and the next page matrix.
-     *
-     * @param pivotX    x coordinate of the focused item in the current page
-     * @param pivotY    y coordinate of the focused item in the current page
-     */
-    // TODO: get rid of the dynamic matrix creation
-    public static int[][] createSparseMatrixWithPivotColumn(CellLayout iconLayout,
-            int pivotX, int pivotY) {
-
-        ViewGroup iconParent = iconLayout.getShortcutsAndWidgets();
-
-        int[][] matrix = createFullMatrix(iconLayout.getCountX() + 1, iconLayout.getCountY());
-
-        // Iterate thru the children of the top parent.
-        for (int i = 0; i < iconParent.getChildCount(); i++) {
-            View cell = iconParent.getChildAt(i);
-            if (!cell.isFocusable()) {
-                continue;
-            }
-            int cx = ((CellLayout.LayoutParams) cell.getLayoutParams()).cellX;
-            int cy = ((CellLayout.LayoutParams) cell.getLayoutParams()).cellY;
-            if (pivotX < 0) {
-                matrix[cx - pivotX][cy] = i;
-            } else {
-                matrix[cx][cy] = i;
-            }
-        }
-
-        if (pivotX < 0) {
-            matrix[0][pivotY] = PIVOT;
-        } else {
-            matrix[pivotX][pivotY] = PIVOT;
-        }
-        if (DEBUG) {
-            printMatrix(matrix);
-        }
-        return matrix;
-    }
-
-    //
-    // key event handling methods.
-    //
-
-    /**
-     * Calculates icon that has is closest to the horizontal axis in reference to the cur icon.
-     *
-     * Example of the check order for KEYCODE_DPAD_RIGHT:
-     * [  ][  ][13][14][15]
-     * [  ][ 6][ 8][10][12]
-     * [ X][ 1][ 2][ 3][ 4]
-     * [  ][ 5][ 7][ 9][11]
-     */
-    // TODO: add unit tests to verify all permutation.
-    private static int handleDpadHorizontal(int iconIdx, int cntX, int cntY,
-            int[][] matrix, int increment, boolean isRtl) {
-        if(matrix == null) {
-            throw new IllegalStateException("Dpad navigation requires a matrix.");
-        }
-        int newIconIndex = NOOP;
-
-        int xPos = -1;
-        int yPos = -1;
-        // Figure out the location of the icon.
-        for (int i = 0; i < cntX; i++) {
-            for (int j = 0; j < cntY; j++) {
-                if (matrix[i][j] == iconIdx) {
-                    xPos = i;
-                    yPos = j;
-                }
-            }
-        }
-        if (DEBUG) {
-            Log.v(TAG, String.format("\thandleDpadHorizontal: \t[x, y]=[%d, %d] iconIndex=%d",
-                    xPos, yPos, iconIdx));
-        }
-
-        // Rule1: check first in the horizontal direction
-        for (int x = xPos + increment; 0 <= x && x < cntX; x += increment) {
-            if ((newIconIndex = inspectMatrix(x, yPos, cntX, cntY, matrix)) != NOOP
-                    && newIconIndex != ALL_APPS_COLUMN) {
-                return newIconIndex;
-            }
-        }
-
-        // Rule2: check (x1-n, yPos + increment),   (x1-n, yPos - increment)
-        //              (x2-n, yPos + 2*increment), (x2-n, yPos - 2*increment)
-        int nextYPos1;
-        int nextYPos2;
-        boolean haveCrossedAllAppsColumn1 = false;
-        boolean haveCrossedAllAppsColumn2 = false;
-        int x = -1;
-        for (int coeff = 1; coeff < cntY; coeff++) {
-            nextYPos1 = yPos + coeff * increment;
-            nextYPos2 = yPos - coeff * increment;
-            x = xPos + increment * coeff;
-            if (inspectMatrix(x, nextYPos1, cntX, cntY, matrix) == ALL_APPS_COLUMN) {
-                haveCrossedAllAppsColumn1 = true;
-            }
-            if (inspectMatrix(x, nextYPos2, cntX, cntY, matrix) == ALL_APPS_COLUMN) {
-                haveCrossedAllAppsColumn2 = true;
-            }
-            for (; 0 <= x && x < cntX; x += increment) {
-                int offset1 = haveCrossedAllAppsColumn1 && x < cntX - 1 ? increment : 0;
-                newIconIndex = inspectMatrix(x, nextYPos1 + offset1, cntX, cntY, matrix);
-                if (newIconIndex != NOOP) {
-                    return newIconIndex;
-                }
-                int offset2 = haveCrossedAllAppsColumn2 && x < cntX - 1 ? -increment : 0;
-                newIconIndex = inspectMatrix(x, nextYPos2 + offset2, cntX, cntY, matrix);
-                if (newIconIndex != NOOP) {
-                    return newIconIndex;
-                }
-            }
-        }
-
-        // Rule3: if switching between pages, do a brute-force search to find an item that was
-        //        missed by rules 1 and 2 (such as when going from a bottom right icon to top left)
-        if (iconIdx == PIVOT) {
-            if (isRtl) {
-                return increment < 0 ? NEXT_PAGE_FIRST_ITEM : PREVIOUS_PAGE_LAST_ITEM;
-            }
-            return increment < 0 ? PREVIOUS_PAGE_LAST_ITEM : NEXT_PAGE_FIRST_ITEM;
-        }
-        return newIconIndex;
-    }
-
-    /**
-     * Calculates icon that is closest to the vertical axis in reference to the current icon.
-     *
-     * Example of the check order for KEYCODE_DPAD_DOWN:
-     * [  ][  ][  ][ X][  ][  ][  ]
-     * [  ][  ][ 5][ 1][ 4][  ][  ]
-     * [  ][10][ 7][ 2][ 6][ 9][  ]
-     * [14][12][ 9][ 3][ 8][11][13]
-     */
-    // TODO: add unit tests to verify all permutation.
-    private static int handleDpadVertical(int iconIndex, int cntX, int cntY,
-            int [][] matrix, int increment) {
-        int newIconIndex = NOOP;
-        if(matrix == null) {
-            throw new IllegalStateException("Dpad navigation requires a matrix.");
-        }
-
-        int xPos = -1;
-        int yPos = -1;
-        // Figure out the location of the icon.
-        for (int i = 0; i< cntX; i++) {
-            for (int j = 0; j < cntY; j++) {
-                if (matrix[i][j] == iconIndex) {
-                    xPos = i;
-                    yPos = j;
-                }
-            }
-        }
-
-        if (DEBUG) {
-            Log.v(TAG, String.format("\thandleDpadVertical: \t[x, y]=[%d, %d] iconIndex=%d",
-                    xPos, yPos, iconIndex));
-        }
-
-        // Rule1: check first in the dpad direction
-        for (int y = yPos + increment; 0 <= y && y <cntY && 0 <= y; y += increment) {
-            if ((newIconIndex = inspectMatrix(xPos, y, cntX, cntY, matrix)) != NOOP
-                    && newIconIndex != ALL_APPS_COLUMN) {
-                return newIconIndex;
-            }
-        }
-
-        // Rule2: check (xPos + increment, y_(1-n)),   (xPos - increment, y_(1-n))
-        //              (xPos + 2*increment, y_(2-n))), (xPos - 2*increment, y_(2-n))
-        int nextXPos1;
-        int nextXPos2;
-        boolean haveCrossedAllAppsColumn1 = false;
-        boolean haveCrossedAllAppsColumn2 = false;
-        int y = -1;
-        for (int coeff = 1; coeff < cntX; coeff++) {
-            nextXPos1 = xPos + coeff * increment;
-            nextXPos2 = xPos - coeff * increment;
-            y = yPos + increment * coeff;
-            if (inspectMatrix(nextXPos1, y, cntX, cntY, matrix) == ALL_APPS_COLUMN) {
-                haveCrossedAllAppsColumn1 = true;
-            }
-            if (inspectMatrix(nextXPos2, y, cntX, cntY, matrix) == ALL_APPS_COLUMN) {
-                haveCrossedAllAppsColumn2 = true;
-            }
-            for (; 0 <= y && y < cntY; y = y + increment) {
-                int offset1 = haveCrossedAllAppsColumn1 && y < cntY - 1 ? increment : 0;
-                newIconIndex = inspectMatrix(nextXPos1 + offset1, y, cntX, cntY, matrix);
-                if (newIconIndex != NOOP) {
-                    return newIconIndex;
-                }
-                int offset2 = haveCrossedAllAppsColumn2 && y < cntY - 1 ? -increment : 0;
-                newIconIndex = inspectMatrix(nextXPos2 + offset2, y, cntX, cntY, matrix);
-                if (newIconIndex != NOOP) {
-                    return newIconIndex;
-                }
-            }
-        }
-        return newIconIndex;
-    }
-
-    private static int handleMoveHome() {
-        return CURRENT_PAGE_FIRST_ITEM;
-    }
-
-    private static int handleMoveEnd() {
-        return CURRENT_PAGE_LAST_ITEM;
-    }
-
-    private static int handlePageDown(int pageIndex, int pageCount) {
-        if (pageIndex < pageCount -1) {
-            return NEXT_PAGE_FIRST_ITEM;
-        }
-        return CURRENT_PAGE_LAST_ITEM;
-    }
-
-    private static int handlePageUp(int pageIndex) {
-        if (pageIndex > 0) {
-            return PREVIOUS_PAGE_FIRST_ITEM;
-        } else {
-            return CURRENT_PAGE_FIRST_ITEM;
-        }
-    }
-
-    //
-    // Helper methods.
-    //
-
-    private static boolean isValid(int xPos, int yPos, int countX, int countY) {
-        return (0 <= xPos && xPos < countX && 0 <= yPos && yPos < countY);
-    }
-
-    private static int inspectMatrix(int x, int y, int cntX, int cntY, int[][] matrix) {
-        int newIconIndex = NOOP;
-        if (isValid(x, y, cntX, cntY)) {
-            if (matrix[x][y] != -1) {
-                newIconIndex = matrix[x][y];
-                if (DEBUG) {
-                    Log.v(TAG, String.format("\t\tinspect: \t[x, y]=[%d, %d] %d",
-                            x, y, matrix[x][y]));
-                }
-                return newIconIndex;
-            }
-        }
-        return newIconIndex;
-    }
-
-    /**
-     * Only used for debugging.
-     */
-    private static String getStringIndex(int index) {
-        switch(index) {
-            case NOOP: return "NOOP";
-            case PREVIOUS_PAGE_FIRST_ITEM:  return "PREVIOUS_PAGE_FIRST";
-            case PREVIOUS_PAGE_LAST_ITEM:   return "PREVIOUS_PAGE_LAST";
-            case PREVIOUS_PAGE_RIGHT_COLUMN:return "PREVIOUS_PAGE_RIGHT_COLUMN";
-            case CURRENT_PAGE_FIRST_ITEM:   return "CURRENT_PAGE_FIRST";
-            case CURRENT_PAGE_LAST_ITEM:    return "CURRENT_PAGE_LAST";
-            case NEXT_PAGE_FIRST_ITEM:      return "NEXT_PAGE_FIRST";
-            case NEXT_PAGE_LEFT_COLUMN:     return "NEXT_PAGE_LEFT_COLUMN";
-            case ALL_APPS_COLUMN:           return "ALL_APPS_COLUMN";
-            default:
-                return Integer.toString(index);
-        }
-    }
-
-    /**
-     * Only used for debugging.
-     */
-    private static void printMatrix(int[][] matrix) {
-        Log.v(TAG, "\tprintMap:");
-        int m = matrix.length;
-        int n = matrix[0].length;
-
-        for (int j=0; j < n; j++) {
-            String colY = "\t\t";
-            for (int i=0; i < m; i++) {
-                colY +=  String.format("%3d",matrix[i][j]);
-            }
-            Log.v(TAG, colY);
-        }
-    }
-
-    /**
-     * @param edgeColumn the column of the new icon. either {@link #NEXT_PAGE_LEFT_COLUMN} or
-     * {@link #NEXT_PAGE_RIGHT_COLUMN}
-     * @return the view adjacent to {@param oldView} in the {@param nextPage} of the folder.
-     */
-    public static View getAdjacentChildInNextFolderPage(
-            ShortcutAndWidgetContainer nextPage, View oldView, int edgeColumn) {
-        final int newRow = ((CellLayout.LayoutParams) oldView.getLayoutParams()).cellY;
-
-        int column = (edgeColumn == NEXT_PAGE_LEFT_COLUMN) ^ nextPage.invertLayoutHorizontally()
-                ? 0 : (((CellLayout) nextPage.getParent()).getCountX() - 1);
-
-        for (; column >= 0; column--) {
-            for (int row = newRow; row >= 0; row--) {
-                View newView = nextPage.getChildAt(column, row);
-                if (newView != null) {
-                    return newView;
-                }
-            }
-        }
-        return null;
-    }
-}
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java b/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
index 73b1601..abce2a2 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
@@ -48,7 +48,7 @@
     @Override
     public void bindWidgets() {
         final List<WidgetsListBaseEntry> widgets =
-                mBgDataModel.widgetsModel.getWidgetsList(mApp.getContext());
+                mBgDataModel.widgetsModel.getWidgetsListForPicker(mApp.getContext());
         executeCallbacksTask(c -> c.bindAllWidgets(widgets), mUiExecutor);
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index a8c7c48..3ea4766 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -5,6 +5,8 @@
 
 import static com.android.launcher3.pm.ShortcutConfigActivityInfo.queryList;
 
+import static java.util.stream.Collectors.toList;
+
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Context;
@@ -67,7 +69,7 @@
      *
      * @see com.android.launcher3.widget.picker.WidgetsListAdapter#setWidgets(List)
      */
-    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsList(Context context) {
+    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsListForPicker(Context context) {
         ArrayList<WidgetsListBaseEntry> result = new ArrayList<>();
         AlphabeticIndexCompat indexer = new AlphabeticIndexCompat(context);
 
@@ -82,6 +84,22 @@
         return result;
     }
 
+    /** Returns a mapping of packages to their widgets without static shortcuts. */
+    public synchronized Map<PackageUserKey, List<WidgetItem>> getAllWidgetsWithoutShortcuts() {
+        Map<PackageUserKey, List<WidgetItem>> packagesToWidgets = new HashMap<>();
+        mWidgetsList.forEach((packageItemInfo, widgetsAndShortcuts) -> {
+            List<WidgetItem> widgets = widgetsAndShortcuts.stream()
+                        .filter(item -> item.widgetInfo != null)
+                        .collect(toList());
+            if (widgets.size() > 0) {
+                packagesToWidgets.put(
+                        new PackageUserKey(packageItemInfo.packageName, packageItemInfo.user),
+                        widgets);
+            }
+        });
+        return packagesToWidgets;
+    }
+
     /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
diff --git a/tests/src/com/android/launcher3/ui/WorkTabTest.java b/tests/src/com/android/launcher3/ui/WorkTabTest.java
index 63220ad..a6c7c03 100644
--- a/tests/src/com/android/launcher3/ui/WorkTabTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkTabTest.java
@@ -18,6 +18,7 @@
 import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST;
+import static com.android.launcher3.tapl.LauncherInstrumentation.LONG_WAIT_TIME_MS;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -140,9 +141,13 @@
                 WorkEduView.KEY_WORK_EDU_STEP).remove(
                 WorkEduView.KEY_LEGACY_WORK_EDU_SEEN).commit());
 
-        waitForLauncherCondition("Work tab not setup",
-                launcher -> launcher.getAppsView().getContentView() instanceof AllAppsPagedView,
-                60000);
+        waitForLauncherCondition("Work tab not setup", launcher -> {
+            if (launcher.getAppsView().getContentView() instanceof AllAppsPagedView) {
+                launcher.getAppsView().getAppsStore().enableDeferUpdates(DEFER_UPDATES_TEST);
+                return true;
+            }
+            return false;
+        }, LONG_WAIT_TIME_MS);
 
         executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
         WorkEduView workEduView = getEduView();
@@ -153,16 +158,6 @@
             workEduView.findViewById(R.id.proceed).callOnClick();
         });
 
-        executeOnLauncher(launcher -> Log.d(TestProtocol.WORK_PROFILE_REMOVED,
-                "work profile status (" + mProfileUserId + ") :"
-                        + launcher.getAppsView().isWorkTabVisible()));
-
-
-        executeOnLauncher(launcher -> {
-            launcher.getAppsView().getAppsStore().enableDeferUpdates(DEFER_UPDATES_TEST);
-            Log.d(TestProtocol.WORK_PROFILE_REMOVED, "Defer all apps update");
-        });
-
         AtomicInteger attempt = new AtomicInteger(0);
         // verify work edu is seen next
         waitForLauncherCondition("Launcher did not show the next edu screen", l -> {
@@ -178,7 +173,7 @@
             }
             return ((TextView) workEduView.findViewById(R.id.content_text)).getText().equals(
                     l.getResources().getString(R.string.work_profile_edu_work_apps));
-        }, 60000);
+        });
     }
 
     @Test