diff --git a/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
index 59c9480..cca333c 100644
--- a/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
+++ b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
@@ -59,6 +59,17 @@
     private HashSet<String> mValidPackages;
     private InvariantDeviceProfile mIdp;
 
+    private final String testPackage1 = "com.android.launcher3.validpackage1";
+    private final String testPackage2 = "com.android.launcher3.validpackage2";
+    private final String testPackage3 = "com.android.launcher3.validpackage3";
+    private final String testPackage4 = "com.android.launcher3.validpackage4";
+    private final String testPackage5 = "com.android.launcher3.validpackage5";
+    private final String testPackage6 = "com.android.launcher3.validpackage6";
+    private final String testPackage7 = "com.android.launcher3.validpackage7";
+    private final String testPackage8 = "com.android.launcher3.validpackage8";
+    private final String testPackage9 = "com.android.launcher3.validpackage9";
+    private final String testPackage10 = "com.android.launcher3.validpackage10";
+
     @Before
     public void setUp() {
         mModelHelper = new LauncherModelHelper();
@@ -67,6 +78,17 @@
 
         mValidPackages = new HashSet<>();
         mValidPackages.add(TEST_PACKAGE);
+        mValidPackages.add(testPackage1);
+        mValidPackages.add(testPackage2);
+        mValidPackages.add(testPackage3);
+        mValidPackages.add(testPackage4);
+        mValidPackages.add(testPackage5);
+        mValidPackages.add(testPackage6);
+        mValidPackages.add(testPackage7);
+        mValidPackages.add(testPackage8);
+        mValidPackages.add(testPackage9);
+        mValidPackages.add(testPackage10);
+
         mIdp = InvariantDeviceProfile.INSTANCE.get(mContext);
 
         long userSerial = UserCache.INSTANCE.get(mContext).getSerialNumberForUser(
@@ -78,20 +100,6 @@
 
     @Test
     public void testMigration() {
-        final String testPackage1 = "com.android.launcher3.validpackage1";
-        final String testPackage2 = "com.android.launcher3.validpackage2";
-        final String testPackage3 = "com.android.launcher3.validpackage3";
-        final String testPackage4 = "com.android.launcher3.validpackage4";
-        final String testPackage5 = "com.android.launcher3.validpackage5";
-        final String testPackage7 = "com.android.launcher3.validpackage7";
-
-        mValidPackages.add(testPackage1);
-        mValidPackages.add(testPackage2);
-        mValidPackages.add(testPackage3);
-        mValidPackages.add(testPackage4);
-        mValidPackages.add(testPackage5);
-        mValidPackages.add(testPackage7);
-
         int[] srcHotseatItems = {
                 mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
                 mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
@@ -100,6 +108,10 @@
                 mModelHelper.addItem(APP_ICON, 4, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI),
         };
         mModelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage5, 5, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 3, testPackage6, 6, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 1, testPackage8, 8, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 2, testPackage9, 9, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 3, testPackage10, 10, TMP_CONTENT_URI);
 
         int[] destHotseatItems = {
                 -1,
@@ -108,21 +120,24 @@
         };
         mModelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage7);
 
-        mIdp.numHotseatIcons = 3;
-        mIdp.numColumns = 3;
-        mIdp.numRows = 3;
+        mIdp.numHotseatIcons = 4;
+        mIdp.numColumns = 4;
+        mIdp.numRows = 4;
         GridSizeMigrationTaskV2.DbReader srcReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages, 5);
+                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages,
+                srcHotseatItems.length);
         GridSizeMigrationTaskV2.DbReader destReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages, 3);
+                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages,
+                mIdp.numHotseatIcons);
         GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(mContext, mDb, srcReader,
-                destReader, 3, new Point(mIdp.numColumns, mIdp.numRows));
+                destReader, mIdp.numHotseatIcons, new Point(mIdp.numColumns, mIdp.numRows));
         task.migrate();
 
+        // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
                 "container=" + CONTAINER_HOTSEAT, null, null, null);
-        assertEquals(c.getCount(), 3);
+        assertEquals(c.getCount(), mIdp.numHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
         c.moveToNext();
@@ -134,13 +149,17 @@
         c.moveToNext();
         assertEquals(c.getInt(screenIndex), 2);
         assertTrue(c.getString(intentIndex).contains(testPackage3));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 3);
+        assertTrue(c.getString(intentIndex).contains(testPackage4));
         c.close();
 
+        // Check workspace items
         c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.CELLX, LauncherSettings.Favorites.CELLY,
                         LauncherSettings.Favorites.INTENT},
                 "container=" + CONTAINER_DESKTOP, null, null, null);
-        assertEquals(c.getCount(), 2);
+        assertEquals(c.getCount(), 6);
         intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
         int cellXIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLX);
         int cellYIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLY);
@@ -148,7 +167,23 @@
         c.moveToNext();
         assertTrue(c.getString(intentIndex).contains(testPackage7));
         c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage6));
+        assertEquals(c.getInt(cellXIndex), 0);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage10));
+        assertEquals(c.getInt(cellXIndex), 1);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
         assertTrue(c.getString(intentIndex).contains(testPackage5));
+        assertEquals(c.getInt(cellXIndex), 2);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage9));
+        assertEquals(c.getInt(cellXIndex), 3);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage8));
         assertEquals(c.getInt(cellXIndex), 0);
         assertEquals(c.getInt(cellYIndex), 2);
     }
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
index 4a28218..4bfabb0 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
@@ -33,6 +33,7 @@
 import android.database.sqlite.SQLiteDatabase;
 import android.graphics.Point;
 import android.text.TextUtils;
+import android.util.ArrayMap;
 import android.util.Log;
 
 import androidx.annotation.VisibleForTesting;
@@ -55,6 +56,7 @@
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -226,9 +228,8 @@
             if (DEBUG) {
                 Log.d(TAG, "Migrating " + screenId);
             }
-            List<DbEntry> entries = mDestReader.loadWorkspaceEntries(screenId);
             GridPlacementSolution workspaceSolution = new GridPlacementSolution(mDb, mSrcReader,
-                    mDestReader, mContext, entries, screenId, mTrgX, mTrgY, mWorkspaceDiff);
+                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff);
             workspaceSolution.find();
             if (mWorkspaceDiff.isEmpty()) {
                 break;
@@ -238,8 +239,7 @@
         int screenId = mDestReader.mLastScreenId + 1;
         while (!mWorkspaceDiff.isEmpty()) {
             GridPlacementSolution workspaceSolution = new GridPlacementSolution(mDb, mSrcReader,
-                    mDestReader, mContext, new ArrayList<>(), screenId, mTrgX, mTrgY,
-                    mWorkspaceDiff);
+                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff);
             workspaceSolution.find();
             screenId++;
         }
@@ -352,8 +352,7 @@
         private int mNextStartY;
 
         GridPlacementSolution(SQLiteDatabase db, DbReader srcReader, DbReader destReader,
-                Context context, List<DbEntry> placedWorkspaceItems, int screenId, int trgX,
-                int trgY, List<DbEntry> itemsToPlace) {
+                Context context, int screenId, int trgX, int trgY, List<DbEntry> itemsToPlace) {
             mDb = db;
             mSrcReader = srcReader;
             mDestReader = destReader;
@@ -364,8 +363,11 @@
             mTrgY = trgY;
             mNextStartX = 0;
             mNextStartY = mTrgY - 1;
-            for (DbEntry entry : placedWorkspaceItems) {
-                mOccupied.markCells(entry, true);
+            List<DbEntry> existedEntries = mDestReader.mWorkspaceEntriesByScreenId.get(screenId);
+            if (existedEntries != null) {
+                for (DbEntry entry : existedEntries) {
+                    mOccupied.markCells(entry, true);
+                }
             }
             mItemsToPlace = itemsToPlace;
         }
@@ -386,6 +388,11 @@
             }
         }
 
+        /**
+         * 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 >= 0; y--) {
                 for (int x = mNextStartX; x < mTrgX; x++) {
@@ -406,6 +413,7 @@
                         return true;
                     }
                 }
+                mNextStartX = 0;
             }
             return false;
         }
@@ -475,6 +483,8 @@
 
         private final ArrayList<DbEntry> mHotseatEntries = new ArrayList<>();
         private final ArrayList<DbEntry> mWorkspaceEntries = new ArrayList<>();
+        private final Map<Integer, ArrayList<DbEntry>> mWorkspaceEntriesByScreenId =
+                new ArrayMap<>();
 
         DbReader(SQLiteDatabase db, String tableName, Context context,
                 HashSet<String> validPackages, int hotseatSize) {
@@ -564,25 +574,6 @@
             return loadWorkspaceEntries(c);
         }
 
-        protected ArrayList<DbEntry> loadWorkspaceEntries(int screen) {
-            Cursor c = queryWorkspace(
-                    new String[]{
-                            LauncherSettings.Favorites._ID,                  // 0
-                            LauncherSettings.Favorites.ITEM_TYPE,            // 1
-                            LauncherSettings.Favorites.SCREEN,               // 2
-                            LauncherSettings.Favorites.CELLX,                // 3
-                            LauncherSettings.Favorites.CELLY,                // 4
-                            LauncherSettings.Favorites.SPANX,                // 5
-                            LauncherSettings.Favorites.SPANY,                // 6
-                            LauncherSettings.Favorites.INTENT,               // 7
-                            LauncherSettings.Favorites.APPWIDGET_PROVIDER,   // 8
-                            LauncherSettings.Favorites.APPWIDGET_ID},        // 9
-                    LauncherSettings.Favorites.CONTAINER + " = "
-                            + LauncherSettings.Favorites.CONTAINER_DESKTOP
-                            + " AND " + LauncherSettings.Favorites.SCREEN + " = " + screen);
-            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);
@@ -660,6 +651,10 @@
                     continue;
                 }
                 mWorkspaceEntries.add(entry);
+                if (!mWorkspaceEntriesByScreenId.containsKey(entry.screenId)) {
+                    mWorkspaceEntriesByScreenId.put(entry.screenId, new ArrayList<>());
+                }
+                mWorkspaceEntriesByScreenId.get(entry.screenId).add(entry);
             }
             removeEntryFromDb(mDb, mTableName, entriesToRemove);
             c.close();