Merge "Remove on drag listener after drag ended" into tm-qpr-dev
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index c0b6657..50f1236 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -292,10 +292,14 @@
     <!-- An additional touch slop to prevent x-axis movement during the swipe up to show taskbar -->
     <dimen name="transient_taskbar_clamped_offset_bound">16dp</dimen>
     <!-- Taskbar swipe up thresholds -->
+    <dimen name="taskbar_nav_threshold">40dp</dimen>
     <dimen name="taskbar_app_window_threshold">150dp</dimen>
     <dimen name="taskbar_home_overview_threshold">225dp</dimen>
     <dimen name="taskbar_catch_up_threshold">300dp</dimen>
-    <dimen name="taskbar_nav_threshold">40dp</dimen>
+
+    <dimen name="taskbar_nav_threshold_v2">30dp</dimen>
+    <dimen name="taskbar_app_window_threshold_v2">100dp</dimen>
+    <dimen name="taskbar_home_overview_threshold_v2">200dp</dimen>
 
     <!--  Taskbar 3 button spacing  -->
     <dimen name="taskbar_button_space_inbetween">24dp</dimen>
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 305fa93..9aedbf8 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -27,6 +27,7 @@
 import static com.android.launcher3.anim.Interpolators.ACCEL_DEACCEL;
 import static com.android.launcher3.anim.Interpolators.DEACCEL;
 import static com.android.launcher3.anim.Interpolators.OVERSHOOT_1_2;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_TASKBAR_REVISED_THRESHOLDS;
 import static com.android.launcher3.logging.StatsLogManager.LAUNCHER_STATE_BACKGROUND;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.IGNORE;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_HOME_GESTURE;
@@ -344,7 +345,9 @@
 
         Resources res = context.getResources();
         mTaskbarAppWindowThreshold = res
-                .getDimensionPixelSize(R.dimen.taskbar_app_window_threshold);
+                .getDimensionPixelSize(ENABLE_TASKBAR_REVISED_THRESHOLDS.get()
+                        ? R.dimen.taskbar_app_window_threshold_v2
+                        : R.dimen.taskbar_app_window_threshold);
         mTaskbarCatchUpThreshold = res.getDimensionPixelSize(R.dimen.taskbar_catch_up_threshold);
         mIsTransientTaskbar = DisplayController.isTransientTaskbar(mActivity);
 
@@ -822,7 +825,7 @@
             return;
         }
         mLauncherTransitionController.setProgress(
-                Math.max(getTaskbarProgress(), getScaleProgressDueToScroll()), mDragLengthFactor);
+                Math.max(mCurrentShift.value, getScaleProgressDueToScroll()), mDragLengthFactor);
     }
 
     /**
@@ -2184,7 +2187,7 @@
             AnimatorControllerWithResistance playbackController =
                     remoteHandle.getPlaybackController();
             if (playbackController != null) {
-                playbackController.setProgress(Math.max(getTaskbarProgress(),
+                playbackController.setProgress(Math.max(mCurrentShift.value,
                         getScaleProgressDueToScroll()), mDragLengthFactor);
             }
 
@@ -2236,31 +2239,32 @@
     }
 
     /**
-     * Overrides the current shift progress to keep the app window at the bottom of the screen
-     * while the transient taskbar is being swiped in.
+     * Overrides the gesture displacement to keep the app window at the bottom of the screen while
+     * the transient taskbar is being swiped in.
      *
      * There is also a catch up period so that the window can start moving 1:1 with the swipe.
      */
-    private float getTaskbarProgress() {
+    @Override
+    protected float overrideDisplacementForTransientTaskbar(float displacement) {
         if (!mIsTransientTaskbar) {
-            return mCurrentShift.value;
+            return displacement;
         }
 
         if (mTaskbarAlreadyOpen) {
-            return mCurrentShift.value;
+            return displacement;
         }
 
-        if (mCurrentDisplacement < mTaskbarAppWindowThreshold) {
+        if (displacement < mTaskbarAppWindowThreshold) {
             return 0;
         }
 
-        // "Catch up" with `mCurrentShift.value`.
-        if (mCurrentDisplacement < mTaskbarCatchUpThreshold) {
-            return Utilities.mapToRange(mCurrentDisplacement, mTaskbarAppWindowThreshold,
-                    mTaskbarCatchUpThreshold, 0, mCurrentShift.value, ACCEL_DEACCEL);
+        // "Catch up" with the displacement at mTaskbarCatchUpThreshold.
+        if (displacement < mTaskbarCatchUpThreshold) {
+            return Utilities.mapToRange(displacement, mTaskbarAppWindowThreshold,
+                    mTaskbarCatchUpThreshold, 0, mTaskbarCatchUpThreshold, ACCEL_DEACCEL);
         }
 
-        return mCurrentShift.value;
+        return displacement;
     }
 
     private void setDividerShown(boolean shown, boolean immediate) {
diff --git a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
index ddb06ce..fdde45a 100644
--- a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
+++ b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
@@ -116,7 +116,7 @@
     @UiThread
     public void updateDisplacement(float displacement) {
         // We are moving in the negative x/y direction
-        displacement = -displacement;
+        displacement = overrideDisplacementForTransientTaskbar(-displacement);
         mCurrentDisplacement = displacement;
 
         float shift;
@@ -131,6 +131,17 @@
     }
 
     /**
+     * When Transient Taskbar is enabled, subclasses can override the displacement to keep the app
+     * window at the bottom of the screen while taskbar is being swiped in.
+     * @param displacement The distance the user has swiped up from the bottom of the screen. This
+     *                     value will be positive unless the user swipe downwards.
+     * @return the overridden displacement.
+     */
+    protected float overrideDisplacementForTransientTaskbar(float displacement) {
+        return displacement;
+    }
+
+    /**
      * Called when the value of {@link #mCurrentShift} changes
      */
     @UiThread
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
index b39d13c..bf666ea 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
@@ -27,6 +27,7 @@
 import static com.android.launcher3.PagedView.DEBUG_FAILED_QUICKSWITCH;
 import static com.android.launcher3.Utilities.EDGE_NAV_BAR;
 import static com.android.launcher3.Utilities.squaredHypot;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_TASKBAR_REVISED_THRESHOLDS;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.TraceHelper.FLAG_CHECK_FOR_RACE_CONDITIONS;
 import static com.android.launcher3.util.VelocityUtils.PX_PER_MS;
@@ -164,7 +165,9 @@
         mTaskbarAlreadyOpen = controller != null && !controller.isTaskbarStashed();
         mIsTransientTaskbar = DisplayController.isTransientTaskbar(base);
         mTaskbarHomeOverviewThreshold = base.getResources()
-                .getDimensionPixelSize(R.dimen.taskbar_home_overview_threshold);
+                .getDimensionPixelSize(ENABLE_TASKBAR_REVISED_THRESHOLDS.get()
+                        ? R.dimen.taskbar_home_overview_threshold_v2
+                        : R.dimen.taskbar_home_overview_threshold);
 
         boolean continuingPreviousGesture = mTaskAnimationManager.isRecentsAnimationRunning();
         mIsDeferredDownTarget = !continuingPreviousGesture && isDeferredDownTarget;
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/TaskbarStashInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/TaskbarStashInputConsumer.java
index b7f2022..3a09490 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/TaskbarStashInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/TaskbarStashInputConsumer.java
@@ -18,6 +18,7 @@
 import static android.view.MotionEvent.INVALID_POINTER_ID;
 
 import static com.android.launcher3.Utilities.squaredHypot;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_TASKBAR_REVISED_THRESHOLDS;
 import static com.android.launcher3.taskbar.TaskbarAutohideSuspendController.FLAG_AUTOHIDE_SUSPEND_TOUCHING;
 
 import android.content.Context;
@@ -72,7 +73,9 @@
 
         Resources res = context.getResources();
         mUnstashArea = res.getDimensionPixelSize(R.dimen.taskbar_unstash_input_area);
-        int taskbarThreshold = res.getDimensionPixelSize(R.dimen.taskbar_nav_threshold);
+        int taskbarThreshold = res.getDimensionPixelSize(ENABLE_TASKBAR_REVISED_THRESHOLDS.get()
+                ? R.dimen.taskbar_nav_threshold_v2
+                : R.dimen.taskbar_nav_threshold);
         int screenHeight = taskbarActivityContext.getDeviceProfile().heightPx;
         mTaskbarThresholdY = screenHeight - taskbarThreshold;
 
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 50b7563..db604f2 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -388,6 +388,9 @@
     <dimen name="taskbar_app_window_threshold">0dp</dimen>
     <dimen name="taskbar_home_overview_threshold">0dp</dimen>
     <dimen name="taskbar_catch_up_threshold">0dp</dimen>
+    <dimen name="taskbar_nav_threshold_v2">0dp</dimen>
+    <dimen name="taskbar_app_window_threshold_v2">0dp</dimen>
+    <dimen name="taskbar_home_overview_threshold_v2">0dp</dimen>
 
     <!-- Size of the maximum radius for the enforced rounded rectangles. -->
     <dimen name="enforced_rounded_corner_max_radius">16dp</dimen>
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 191d063..4f2f093 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -1074,13 +1074,14 @@
         mXDown = ev.getX();
         mYDown = ev.getY();
         if (mFirstPagePinnedItem != null) {
-            mTempFXY[0] = mXDown + getScrollX();
-            mTempFXY[1] = mYDown + getScrollY();
-            Utilities.mapCoordInSelfToDescendant(mFirstPagePinnedItem, this, mTempFXY);
-            mIsEventOverFirstPagePinnedItem = mFirstPagePinnedItem.getLeft() <= mTempFXY[0]
-                    && mFirstPagePinnedItem.getRight() >= mTempFXY[0]
-                    && mFirstPagePinnedItem.getTop() <= mTempFXY[1]
-                    && mFirstPagePinnedItem.getBottom() >= mTempFXY[1];
+            final float[] tempFXY = new float[2];
+            tempFXY[0] = mXDown;
+            tempFXY[1] = mYDown;
+            Utilities.mapCoordInSelfToDescendant(mFirstPagePinnedItem, this, tempFXY);
+            mIsEventOverFirstPagePinnedItem = mFirstPagePinnedItem.getLeft() <= tempFXY[0]
+                    && mFirstPagePinnedItem.getRight() >= tempFXY[0]
+                    && mFirstPagePinnedItem.getTop() <= tempFXY[1]
+                    && mFirstPagePinnedItem.getBottom() >= tempFXY[1];
         } else {
             mIsEventOverFirstPagePinnedItem = false;
         }
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 982f6a9..22bc88d 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -262,6 +262,10 @@
             "ENABLE_NEW_MIGRATION_LOGIC", true,
             "Enable the new grid migration logic, keeping pages when src < dest");
 
+    public static final BooleanFlag ENABLE_WIDGET_HOST_IN_BACKGROUND = getDebugFlag(
+            "ENABLE_WIDGET_HOST_IN_BACKGROUND", false,
+            "Enable background widget updates listening for widget holder");
+
     public static final BooleanFlag ENABLE_ONE_SEARCH_MOTION = new DeviceFlag(
             "ENABLE_ONE_SEARCH_MOTION", true, "Enables animations in OneSearch.");
 
@@ -322,6 +326,10 @@
             "HOME_GARDENING_WORKSPACE_BUTTONS", false,
             "Change workspace edit buttons to reflect home gardening");
 
+    public static final BooleanFlag ENABLE_TASKBAR_REVISED_THRESHOLDS = getDebugFlag(
+            "ENABLE_TASKBAR_REVISED_THRESHOLDS", false,
+            "Uses revised thresholds for transient taskbar.");
+
     public static final BooleanFlag ENABLE_TRANSIENT_TASKBAR = getDebugFlag(
             "ENABLE_TRANSIENT_TASKBAR", false, "Enables transient taskbar.");
 
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 287b976..85c0a7a 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -50,7 +50,7 @@
 import com.android.launcher3.Workspace;
 import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
 import com.android.launcher3.model.BgDataModel;
-import com.android.launcher3.model.GridSizeMigrationTaskV2;
+import com.android.launcher3.model.GridSizeMigrationUtil;
 import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.RunnableList;
@@ -241,10 +241,10 @@
 
     @WorkerThread
     private boolean doGridMigrationIfNecessary() {
-        if (!GridSizeMigrationTaskV2.needsToMigrate(mContext, mIdp)) {
+        if (!GridSizeMigrationUtil.needsToMigrate(mContext, mIdp)) {
             return false;
         }
-        return GridSizeMigrationTaskV2.migrateGridIfNeeded(mContext, mIdp);
+        return GridSizeMigrationUtil.migrateGridIfNeeded(mContext, mIdp);
     }
 
     @UiThread
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
similarity index 75%
rename from src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
rename to src/com/android/launcher3/model/GridSizeMigrationUtil.java
index 341372e..d63408b 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
@@ -32,7 +32,6 @@
 import android.util.Log;
 
 import androidx.annotation.NonNull;
-import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
@@ -64,42 +63,13 @@
  * This class takes care of shrinking the workspace (by maximum of one row and one column), as a
  * result of restoring from a larger device or device density change.
  */
-public class GridSizeMigrationTaskV2 {
+public class GridSizeMigrationUtil {
 
-    private static final String TAG = "GridSizeMigrationTaskV2";
+    private static final String TAG = "GridSizeMigrationUtil";
     private static final boolean DEBUG = false;
 
-    private final Context mContext;
-    private final SQLiteDatabase mDb;
-    private final DbReader mSrcReader;
-    private final DbReader mDestReader;
-
-    private final List<DbEntry> mHotseatItems;
-    private final List<DbEntry> mWorkspaceItems;
-
-    private final List<DbEntry> mHotseatDiff;
-    private final List<DbEntry> mWorkspaceDiff;
-
-    private final int mDestHotseatSize;
-    private final int mTrgX, mTrgY;
-
-    @VisibleForTesting
-    protected GridSizeMigrationTaskV2(Context context, SQLiteDatabase db, DbReader srcReader,
-            DbReader destReader, int destHotseatSize, Point targetSize) {
-        mContext = context;
-        mDb = db;
-        mSrcReader = srcReader;
-        mDestReader = destReader;
-
-        mHotseatItems = destReader.loadHotseatEntries();
-        mWorkspaceItems = destReader.loadAllWorkspaceEntries();
-
-        mHotseatDiff = calcDiff(mSrcReader.loadHotseatEntries(), mHotseatItems);
-        mWorkspaceDiff = calcDiff(mSrcReader.loadAllWorkspaceEntries(), mWorkspaceItems);
-        mDestHotseatSize = destHotseatSize;
-
-        mTrgX = targetSize.x;
-        mTrgY = targetSize.y;
+    private GridSizeMigrationUtil() {
+        // Util class should not be instantiated
     }
 
     /**
@@ -187,9 +157,8 @@
                     context, validPackages);
 
             Point targetSize = new Point(destDeviceState.getColumns(), destDeviceState.getRows());
-            GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(context, t.getDb(),
-                    srcReader, destReader, destDeviceState.getNumHotseat(), targetSize);
-            task.migrate(srcDeviceState, destDeviceState);
+            migrate(context, t.getDb(), srcReader, destReader, destDeviceState.getNumHotseat(),
+                    targetSize, srcDeviceState, destDeviceState);
 
             if (!migrateForPreview) {
                 dropTable(t.getDb(), LauncherSettings.Favorites.TMP_TABLE);
@@ -212,25 +181,39 @@
         }
     }
 
-    @VisibleForTesting
-    protected boolean migrate(DeviceGridState srcDeviceState, DeviceGridState destDeviceState) {
-        if (mHotseatDiff.isEmpty() && mWorkspaceDiff.isEmpty()) {
+    public static boolean migrate(
+            @NonNull final Context context, @NonNull final SQLiteDatabase db,
+            @NonNull final DbReader srcReader, @NonNull final DbReader destReader,
+            final int destHotseatSize, @NonNull final Point targetSize,
+            @NonNull final DeviceGridState srcDeviceState,
+            @NonNull final DeviceGridState destDeviceState) {
+
+        final List<DbEntry> hotseatItems = destReader.loadHotseatEntries();
+        final List<DbEntry> workspaceItems = destReader.loadAllWorkspaceEntries();
+        final List<DbEntry> hotseatDiff =
+                calcDiff(srcReader.loadHotseatEntries(), hotseatItems);
+        final List<DbEntry> workspaceDiff =
+                calcDiff(srcReader.loadAllWorkspaceEntries(), workspaceItems);
+
+        final int trgX = targetSize.x;
+        final int trgY = targetSize.y;
+
+        if (hotseatDiff.isEmpty() && workspaceDiff.isEmpty()) {
             return false;
         }
 
         // Sort the items by the reading order.
-        Collections.sort(mHotseatDiff);
-        Collections.sort(mWorkspaceDiff);
+        Collections.sort(hotseatDiff);
+        Collections.sort(workspaceDiff);
 
         // Migrate hotseat
-        HotseatPlacementSolution hotseatSolution = new HotseatPlacementSolution(mDb, mSrcReader,
-                mDestReader, mContext, mDestHotseatSize, mHotseatItems, mHotseatDiff);
-        hotseatSolution.find();
+        solveHotseatPlacement(db, srcReader,
+                destReader, context, destHotseatSize, hotseatItems, hotseatDiff);
 
         // Migrate workspace.
         // First we create a collection of the screens
         List<Integer> screens = new ArrayList<>();
-        for (int screenId = 0; screenId <= mDestReader.mLastScreenId; screenId++) {
+        for (int screenId = 0; screenId <= destReader.mLastScreenId; screenId++) {
             screens.add(screenId);
         }
 
@@ -245,22 +228,19 @@
             if (DEBUG) {
                 Log.d(TAG, "Migrating " + screenId);
             }
-            GridPlacementSolution workspaceSolution = new GridPlacementSolution(mDb, mSrcReader,
-                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff, false);
-            workspaceSolution.find();
-            if (mWorkspaceDiff.isEmpty()) {
+            solveGridPlacement(db, srcReader,
+                    destReader, context, screenId, trgX, trgY, workspaceDiff, false);
+            if (workspaceDiff.isEmpty()) {
                 break;
             }
         }
 
         // In case the new grid is smaller, there might be some leftover items that don't fit on
         // any of the screens, in this case we add them to new screens until all of them are placed.
-        int screenId = mDestReader.mLastScreenId + 1;
-        while (!mWorkspaceDiff.isEmpty()) {
-            GridPlacementSolution workspaceSolution = new GridPlacementSolution(mDb, mSrcReader,
-                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff,
-                    preservePages);
-            workspaceSolution.find();
+        int screenId = destReader.mLastScreenId + 1;
+        while (!workspaceDiff.isEmpty()) {
+            solveGridPlacement(db, srcReader,
+                    destReader, context, screenId, trgX, trgY, workspaceDiff, preservePages);
             screenId++;
         }
 
@@ -365,144 +345,88 @@
         return validPackages;
     }
 
-    protected static class GridPlacementSolution {
-
-        private final SQLiteDatabase mDb;
-        private final DbReader mSrcReader;
-        private final DbReader mDestReader;
-        private final Context mContext;
-        private final GridOccupancy mOccupied;
-        private final int mScreenId;
-        private final int mTrgX;
-        private final int mTrgY;
-        private final List<DbEntry> mSortedItemsToPlace;
-        private final boolean mMatchingScreenIdOnly;
-
-        private int mNextStartX;
-        private int mNextStartY;
-
-        GridPlacementSolution(SQLiteDatabase db, DbReader srcReader, DbReader destReader,
-                Context context, int screenId, int trgX, int trgY, List<DbEntry> sortedItemsToPlace,
-                boolean matchingScreenIdOnly) {
-            mDb = db;
-            mSrcReader = srcReader;
-            mDestReader = destReader;
-            mContext = context;
-            mOccupied = new GridOccupancy(trgX, trgY);
-            mScreenId = screenId;
-            mTrgX = trgX;
-            mTrgY = trgY;
-            mNextStartX = 0;
-            mNextStartY = mScreenId == 0 && FeatureFlags.QSB_ON_FIRST_SCREEN
-                    ? 1 /* smartspace */ : 0;
-            List<DbEntry> existedEntries = mDestReader.mWorkspaceEntriesByScreenId.get(screenId);
-            if (existedEntries != null) {
-                for (DbEntry entry : existedEntries) {
-                    mOccupied.markCells(entry, true);
-                }
-            }
-            mSortedItemsToPlace = sortedItemsToPlace;
-            mMatchingScreenIdOnly = matchingScreenIdOnly;
-        }
-
-        public void find() {
-            Iterator<DbEntry> iterator = mSortedItemsToPlace.iterator();
-            while (iterator.hasNext()) {
-                final DbEntry entry = iterator.next();
-                if (mMatchingScreenIdOnly && entry.screenId < mScreenId) continue;
-                if (mMatchingScreenIdOnly && entry.screenId > mScreenId) break;
-                if (entry.minSpanX > mTrgX || entry.minSpanY > mTrgY) {
-                    iterator.remove();
-                    continue;
-                }
-                if (findPlacement(entry)) {
-                    insertEntryInDb(mDb, mContext, entry, mSrcReader.mTableName,
-                            mDestReader.mTableName);
-                    iterator.remove();
-                }
+    private static void solveGridPlacement(@NonNull final SQLiteDatabase db,
+            @NonNull final DbReader srcReader, @NonNull final DbReader destReader,
+            @NonNull final Context context, final int screenId, final int trgX, final int trgY,
+            @NonNull final List<DbEntry> sortedItemsToPlace, final boolean matchingScreenIdOnly) {
+        final GridOccupancy occupied = new GridOccupancy(trgX, trgY);
+        final Point trg = new Point(trgX, trgY);
+        final Point next = new Point(0, screenId == 0 && FeatureFlags.QSB_ON_FIRST_SCREEN
+                ? 1 /* smartspace */ : 0);
+        List<DbEntry> existedEntries = destReader.mWorkspaceEntriesByScreenId.get(screenId);
+        if (existedEntries != null) {
+            for (DbEntry entry : existedEntries) {
+                occupied.markCells(entry, true);
             }
         }
-
-        /**
-         * Search for the next possible placement of an icon. (mNextStartX, mNextStartY) serves as
-         * a memoization of last placement, we can start our search for next placement from there
-         * to speed up the search.
-         */
-        private boolean findPlacement(DbEntry entry) {
-            for (int y = mNextStartY; y <  mTrgY; y++) {
-                for (int x = mNextStartX; x < mTrgX; x++) {
-                    boolean fits = mOccupied.isRegionVacant(x, y, entry.spanX, entry.spanY);
-                    boolean minFits = mOccupied.isRegionVacant(x, y, entry.minSpanX,
-                            entry.minSpanY);
-                    if (minFits) {
-                        entry.spanX = entry.minSpanX;
-                        entry.spanY = entry.minSpanY;
-                    }
-                    if (fits || minFits) {
-                        entry.screenId = mScreenId;
-                        entry.cellX = x;
-                        entry.cellY = y;
-                        mOccupied.markCells(entry, true);
-                        mNextStartX = x + entry.spanX;
-                        mNextStartY = y;
-                        return true;
-                    }
-                }
-                mNextStartX = 0;
+        Iterator<DbEntry> iterator = sortedItemsToPlace.iterator();
+        while (iterator.hasNext()) {
+            final DbEntry entry = iterator.next();
+            if (matchingScreenIdOnly && entry.screenId < screenId) continue;
+            if (matchingScreenIdOnly && entry.screenId > screenId) break;
+            if (entry.minSpanX > trgX || entry.minSpanY > trgY) {
+                iterator.remove();
+                continue;
             }
-            return false;
+            if (findPlacementForEntry(entry, next, trg, occupied, screenId)) {
+                insertEntryInDb(db, context, entry, srcReader.mTableName, destReader.mTableName);
+                iterator.remove();
+            }
         }
     }
 
-    protected static class HotseatPlacementSolution {
-
-        private final SQLiteDatabase mDb;
-        private final DbReader mSrcReader;
-        private final DbReader mDestReader;
-        private final Context mContext;
-        private final HotseatOccupancy mOccupied;
-        private final List<DbEntry> mItemsToPlace;
-
-        HotseatPlacementSolution(SQLiteDatabase db, DbReader srcReader, DbReader destReader,
-                Context context, int hotseatSize, List<DbEntry> placedHotseatItems,
-                List<DbEntry> itemsToPlace) {
-            mDb = db;
-            mSrcReader = srcReader;
-            mDestReader = destReader;
-            mContext = context;
-            mOccupied = new HotseatOccupancy(hotseatSize);
-            for (DbEntry entry : placedHotseatItems) {
-                mOccupied.markCells(entry, true);
-            }
-            mItemsToPlace = itemsToPlace;
-        }
-
-        public void find() {
-            for (int i = 0; i < mOccupied.mCells.length; i++) {
-                if (!mOccupied.mCells[i] && !mItemsToPlace.isEmpty()) {
-                    DbEntry entry = mItemsToPlace.remove(0);
-                    entry.screenId = i;
-                    // These values does not affect the item position, but we should set them
-                    // to something other than -1.
-                    entry.cellX = i;
-                    entry.cellY = 0;
-                    insertEntryInDb(mDb, mContext, entry, mSrcReader.mTableName,
-                            mDestReader.mTableName);
-                    mOccupied.markCells(entry, true);
+    /**
+     * Search for the next possible placement of an icon. (mNextStartX, mNextStartY) serves as
+     * a memoization of last placement, we can start our search for next placement from there
+     * to speed up the search.
+     */
+    private static boolean findPlacementForEntry(@NonNull final DbEntry entry,
+            @NonNull final Point next, @NonNull final Point trg,
+            @NonNull final GridOccupancy occupied, final int screenId) {
+        for (int y = next.y; y <  trg.y; y++) {
+            for (int x = next.x; x < trg.x; x++) {
+                boolean fits = occupied.isRegionVacant(x, y, entry.spanX, entry.spanY);
+                boolean minFits = occupied.isRegionVacant(x, y, entry.minSpanX,
+                        entry.minSpanY);
+                if (minFits) {
+                    entry.spanX = entry.minSpanX;
+                    entry.spanY = entry.minSpanY;
+                }
+                if (fits || minFits) {
+                    entry.screenId = screenId;
+                    entry.cellX = x;
+                    entry.cellY = y;
+                    occupied.markCells(entry, true);
+                    next.set(x + entry.spanX, y);
+                    return true;
                 }
             }
+            next.set(0, next.y);
+        }
+        return false;
+    }
+
+    private static void solveHotseatPlacement(@NonNull final SQLiteDatabase db,
+            @NonNull final DbReader srcReader, @NonNull final DbReader destReader,
+            @NonNull final Context context, final int hotseatSize,
+            @NonNull final  List<DbEntry> placedHotseatItems,
+            @NonNull final List<DbEntry> itemsToPlace) {
+
+        final boolean[] occupied = new boolean[hotseatSize];
+        for (DbEntry entry : placedHotseatItems) {
+            occupied[entry.screenId] = true;
         }
 
-        private class HotseatOccupancy {
-
-            private final boolean[] mCells;
-
-            private HotseatOccupancy(int hotseatSize) {
-                mCells = new boolean[hotseatSize];
-            }
-
-            private void markCells(ItemInfo item, boolean value) {
-                mCells[item.screenId] = value;
+        for (int i = 0; i < occupied.length; i++) {
+            if (!occupied[i] && !itemsToPlace.isEmpty()) {
+                DbEntry entry = itemsToPlace.remove(0);
+                entry.screenId = i;
+                // These values does not affect the item position, but we should set them
+                // to something other than -1.
+                entry.cellX = i;
+                entry.cellY = 0;
+                insertEntryInDb(db, context, entry, srcReader.mTableName, destReader.mTableName);
+                occupied[entry.screenId] = true;
             }
         }
     }
@@ -515,8 +439,6 @@
         private final Set<String> mValidPackages;
         private int mLastScreenId = -1;
 
-        private final ArrayList<DbEntry> mHotseatEntries = new ArrayList<>();
-        private final ArrayList<DbEntry> mWorkspaceEntries = new ArrayList<>();
         private final Map<Integer, ArrayList<DbEntry>> mWorkspaceEntriesByScreenId =
                 new ArrayMap<>();
 
@@ -528,7 +450,8 @@
             mValidPackages = validPackages;
         }
 
-        protected ArrayList<DbEntry> loadHotseatEntries() {
+        protected List<DbEntry> loadHotseatEntries() {
+            final List<DbEntry> hotseatEntries = new ArrayList<>();
             Cursor c = queryWorkspace(
                     new String[]{
                             LauncherSettings.Favorites._ID,                  // 0
@@ -577,14 +500,15 @@
                     entriesToRemove.add(entry.id);
                     continue;
                 }
-                mHotseatEntries.add(entry);
+                hotseatEntries.add(entry);
             }
             removeEntryFromDb(mDb, mTableName, entriesToRemove);
             c.close();
-            return mHotseatEntries;
+            return hotseatEntries;
         }
 
-        protected ArrayList<DbEntry> loadAllWorkspaceEntries() {
+        protected List<DbEntry> loadAllWorkspaceEntries() {
+            final List<DbEntry> workspaceEntries = new ArrayList<>();
             Cursor c = queryWorkspace(
                     new String[]{
                             LauncherSettings.Favorites._ID,                  // 0
@@ -599,10 +523,6 @@
                             LauncherSettings.Favorites.APPWIDGET_ID},        // 9
                         LauncherSettings.Favorites.CONTAINER + " = "
                             + LauncherSettings.Favorites.CONTAINER_DESKTOP);
-            return loadWorkspaceEntries(c);
-        }
-
-        private ArrayList<DbEntry> loadWorkspaceEntries(Cursor c) {
             final int indexId = c.getColumnIndexOrThrow(LauncherSettings.Favorites._ID);
             final int indexItemType = c.getColumnIndexOrThrow(LauncherSettings.Favorites.ITEM_TYPE);
             final int indexScreen = c.getColumnIndexOrThrow(LauncherSettings.Favorites.SCREEN);
@@ -678,7 +598,7 @@
                     entriesToRemove.add(entry.id);
                     continue;
                 }
-                mWorkspaceEntries.add(entry);
+                workspaceEntries.add(entry);
                 if (!mWorkspaceEntriesByScreenId.containsKey(entry.screenId)) {
                     mWorkspaceEntriesByScreenId.put(entry.screenId, new ArrayList<>());
                 }
@@ -686,7 +606,7 @@
             }
             removeEntryFromDb(mDb, mTableName, entriesToRemove);
             c.close();
-            return mWorkspaceEntries;
+            return workspaceEntries;
         }
 
         private int getFolderItemsCount(DbEntry entry) {
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index b644b6b..1d6971e 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -349,7 +349,7 @@
         final WidgetManagerHelper widgetHelper = new WidgetManagerHelper(context);
 
         boolean clearDb = false;
-        if (!GridSizeMigrationTaskV2.migrateGridIfNeeded(context)) {
+        if (!GridSizeMigrationUtil.migrateGridIfNeeded(context)) {
             // Migration failed. Clear workspace.
             clearDb = true;
         }
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.kt b/tests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
similarity index 89%
rename from tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.kt
rename to tests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
index 90d7b43..85d7bf9 100644
--- a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.kt
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
@@ -26,7 +26,7 @@
 import com.android.launcher3.LauncherFiles
 import com.android.launcher3.LauncherSettings.Favorites.*
 import com.android.launcher3.config.FeatureFlags
-import com.android.launcher3.model.GridSizeMigrationTaskV2.DbReader
+import com.android.launcher3.model.GridSizeMigrationUtil.DbReader
 import com.android.launcher3.pm.UserCache
 import com.android.launcher3.provider.LauncherDbUtils
 import com.android.launcher3.util.LauncherModelHelper
@@ -37,10 +37,10 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 
-/** Unit tests for [GridSizeMigrationTaskV2]  */
+/** Unit tests for [GridSizeMigrationUtil]  */
 @SmallTest
 @RunWith(AndroidJUnit4::class)
-class GridSizeMigrationTaskV2Test {
+class GridSizeMigrationUtilTest {
     private lateinit var modelHelper: LauncherModelHelper
     private lateinit var context: Context
     private lateinit var db: SQLiteDatabase
@@ -122,15 +122,16 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context, validPackages)
         val destReader = DbReader(db, TABLE_NAME, context, validPackages)
-        val task = GridSizeMigrationTaskV2(
-            context,
-            db,
-            srcReader,
-            destReader,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows)
+        GridSizeMigrationUtil.migrate(
+                context,
+                db,
+                srcReader,
+                destReader,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp)
         )
-        task.migrate(DeviceGridState(context), DeviceGridState(idp))
 
         // Check hotseat items
         var c = context.contentResolver.query(
@@ -207,15 +208,16 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context, validPackages)
         val destReader = DbReader(db, TABLE_NAME, context, validPackages)
-        val task = GridSizeMigrationTaskV2(
-            context,
-            db,
-            srcReader,
-            destReader,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows)
+        GridSizeMigrationUtil.migrate(
+                context,
+                db,
+                srcReader,
+                destReader,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp)
         )
-        task.migrate(DeviceGridState(context), DeviceGridState(idp))
 
         // Check hotseat items
         val c = context.contentResolver.query(
@@ -262,15 +264,16 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context, validPackages)
         val destReader = DbReader(db, TABLE_NAME, context, validPackages)
-        val task = GridSizeMigrationTaskV2(
-            context,
-            db,
-            srcReader,
-            destReader,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows)
+        GridSizeMigrationUtil.migrate(
+                context,
+                db,
+                srcReader,
+                destReader,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp)
         )
-        task.migrate(DeviceGridState(context), DeviceGridState(idp))
 
         // Check hotseat items
         val c = context.contentResolver.query(
@@ -327,15 +330,16 @@
 
         val srcReader = DbReader(db, TMP_TABLE, context, validPackages)
         val destReader = DbReader(db, TABLE_NAME, context, validPackages)
-        val task = GridSizeMigrationTaskV2(
-            context,
-            db,
-            srcReader,
-            destReader,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows)
+        GridSizeMigrationUtil.migrate(
+                context,
+                db,
+                srcReader,
+                destReader,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp)
         )
-        task.migrate(DeviceGridState(context), DeviceGridState(idp))
 
         // Get workspace items
         val c = context.contentResolver.query(
@@ -387,15 +391,16 @@
         idp.numRows = 5
         val srcReader = DbReader(db, TMP_TABLE, context, validPackages)
         val destReader = DbReader(db, TABLE_NAME, context, validPackages)
-        val task = GridSizeMigrationTaskV2(
-            context,
-            db,
-            srcReader,
-            destReader,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows)
+        GridSizeMigrationUtil.migrate(
+                context,
+                db,
+                srcReader,
+                destReader,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp)
         )
-        task.migrate(DeviceGridState(context), DeviceGridState(idp))
 
         // Get workspace items
         val c = context.contentResolver.query(
@@ -448,15 +453,16 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context, validPackages)
         val destReader = DbReader(db, TABLE_NAME, context, validPackages)
-        val task = GridSizeMigrationTaskV2(
-            context,
-            db,
-            srcReader,
-            destReader,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows)
+        GridSizeMigrationUtil.migrate(
+                context,
+                db,
+                srcReader,
+                destReader,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp)
         )
-        task.migrate(DeviceGridState(context), DeviceGridState(idp))
 
         // Get workspace items
         val c = context.contentResolver.query(