diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index a949e11..6116c66 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -256,6 +256,10 @@
             "ENABLE_SPLIT_FROM_WORKSPACE", true,
             "Enable initiating split screen from workspace.");
 
+    public static final BooleanFlag ENABLE_NEW_MIGRATION_LOGIC = getDebugFlag(
+            "ENABLE_NEW_MIGRATION_LOGIC", true,
+            "Enable the new grid migration logic, keeping pages when src < dest");
+
     public static void initialize(Context context) {
         synchronized (sDebugFlags) {
             for (DebugFlag flag : sDebugFlags) {
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
index 08c3149..3e49d79 100644
--- a/src/com/android/launcher3/model/DeviceGridState.java
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -38,7 +38,7 @@
 /**
  * Utility class representing persisted grid properties.
  */
-public class DeviceGridState {
+public class DeviceGridState implements Comparable<DeviceGridState> {
 
     public static final String KEY_WORKSPACE_SIZE = "migration_src_workspace_size";
     public static final String KEY_HOTSEAT_COUNT = "migration_src_hotseat_count";
@@ -84,16 +84,16 @@
      */
     public LauncherEvent getWorkspaceSizeEvent() {
         if (!TextUtils.isEmpty(mGridSizeString)) {
-            switch (mGridSizeString.charAt(0)) {
-                case '6':
+            switch (getColumns()) {
+                case 6:
                     return LAUNCHER_GRID_SIZE_6;
-                case '5':
+                case 5:
                     return LAUNCHER_GRID_SIZE_5;
-                case '4':
+                case 4:
                     return LAUNCHER_GRID_SIZE_4;
-                case '3':
+                case 3:
                     return LAUNCHER_GRID_SIZE_3;
-                case '2':
+                case 2:
                     return LAUNCHER_GRID_SIZE_2;
             }
         }
@@ -119,4 +119,21 @@
         return mNumHotseat == other.mNumHotseat
                 && Objects.equals(mGridSizeString, other.mGridSizeString);
     }
+
+    public Integer getColumns() {
+        return Integer.parseInt(String.valueOf(mGridSizeString.charAt(0)));
+    }
+
+    public Integer getRows() {
+        return Integer.parseInt(String.valueOf(mGridSizeString.charAt(2)));
+    }
+
+    @Override
+    public int compareTo(DeviceGridState other) {
+        Integer size = getColumns() * getRows();
+        Integer otherSize = other.getColumns() * other.getRows();
+
+        return size.compareTo(otherSize);
+    }
+
 }
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
index ca680b7..74b0a6f 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
@@ -38,6 +38,7 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.graphics.LauncherPreviewRenderer;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.pm.InstallSessionHelper;
@@ -225,13 +226,21 @@
             screens.add(screenId);
         }
 
+        boolean preservePages = false;
+        if (screens.isEmpty() && FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC.get()) {
+            DeviceGridState srcDeviceState = new DeviceGridState(mContext);
+            DeviceGridState destDeviceState = new DeviceGridState(idp);
+            preservePages = destDeviceState.compareTo(srcDeviceState) >= 0
+                    && destDeviceState.getColumns() - srcDeviceState.getColumns() <= 2;
+        }
+
         // Then we place the items on the screens
         for (int screenId : screens) {
             if (DEBUG) {
                 Log.d(TAG, "Migrating " + screenId);
             }
             GridPlacementSolution workspaceSolution = new GridPlacementSolution(mDb, mSrcReader,
-                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff);
+                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff, false);
             workspaceSolution.find();
             if (mWorkspaceDiff.isEmpty()) {
                 break;
@@ -243,10 +252,12 @@
         int screenId = mDestReader.mLastScreenId + 1;
         while (!mWorkspaceDiff.isEmpty()) {
             GridPlacementSolution workspaceSolution = new GridPlacementSolution(mDb, mSrcReader,
-                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff);
+                    mDestReader, mContext, screenId, mTrgX, mTrgY, mWorkspaceDiff,
+                    preservePages);
             workspaceSolution.find();
             screenId++;
         }
+
         return true;
     }
 
@@ -363,13 +374,15 @@
         private final int mScreenId;
         private final int mTrgX;
         private final int mTrgY;
-        private final List<DbEntry> mItemsToPlace;
+        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> itemsToPlace) {
+                Context context, int screenId, int trgX, int trgY, List<DbEntry> sortedItemsToPlace,
+                boolean matchingScreenIdOnly) {
             mDb = db;
             mSrcReader = srcReader;
             mDestReader = destReader;
@@ -386,13 +399,16 @@
                     mOccupied.markCells(entry, true);
                 }
             }
-            mItemsToPlace = itemsToPlace;
+            mSortedItemsToPlace = sortedItemsToPlace;
+            mMatchingScreenIdOnly = matchingScreenIdOnly;
         }
 
         public void find() {
-            Iterator<DbEntry> iterator = mItemsToPlace.iterator();
+            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;
@@ -494,7 +510,7 @@
         private final SQLiteDatabase mDb;
         private final String mTableName;
         private final Context mContext;
-        private final HashSet<String> mValidPackages;
+        private final Set<String> mValidPackages;
         private int mLastScreenId = -1;
 
         private final ArrayList<DbEntry> mHotseatEntries = new ArrayList<>();
@@ -503,7 +519,7 @@
                 new ArrayMap<>();
 
         DbReader(SQLiteDatabase db, String tableName, Context context,
-                HashSet<String> validPackages) {
+                Set<String> validPackages) {
             mDb = db;
             mTableName = tableName;
             mContext = context;
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
deleted file mode 100644
index 005389e..0000000
--- a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
+++ /dev/null
@@ -1,278 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.model;
-
-import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
-import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
-import static com.android.launcher3.LauncherSettings.Favorites.TMP_CONTENT_URI;
-import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
-import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
-import static com.android.launcher3.util.LauncherModelHelper.DESKTOP;
-import static com.android.launcher3.util.LauncherModelHelper.HOTSEAT;
-import static com.android.launcher3.util.LauncherModelHelper.SHORTCUT;
-import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-
-import android.content.Context;
-import android.content.Intent;
-import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
-import android.graphics.Point;
-import android.os.Process;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.util.LauncherModelHelper;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-import java.util.HashMap;
-import java.util.HashSet;
-
-/** Unit tests for {@link GridSizeMigrationTaskV2} */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-public class GridSizeMigrationTaskV2Test {
-
-    private LauncherModelHelper mModelHelper;
-    private Context mContext;
-    private SQLiteDatabase mDb;
-
-    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();
-        mContext = mModelHelper.sandboxContext;
-        mDb = mModelHelper.provider.getDb();
-
-        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(
-                Process.myUserHandle());
-        dropTable(mDb, LauncherSettings.Favorites.TMP_TABLE);
-        LauncherSettings.Favorites.addTableToDb(mDb, userSerial, false,
-                LauncherSettings.Favorites.TMP_TABLE);
-    }
-
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
-    @Test
-    public void testMigration() throws Exception {
-        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),
-                -1,
-                mModelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI),
-                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,
-                mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2),
-                -1,
-        };
-        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage7);
-
-        mIdp.numDatabaseHotseatIcons = 4;
-        mIdp.numColumns = 4;
-        mIdp.numRows = 4;
-        GridSizeMigrationTaskV2.DbReader srcReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages);
-        GridSizeMigrationTaskV2.DbReader destReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages);
-        GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(mContext, mDb, srcReader,
-                destReader, mIdp.numDatabaseHotseatIcons, new Point(mIdp.numColumns, mIdp.numRows));
-        task.migrate(mIdp);
-
-        // Check hotseat items
-        Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
-                new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
-        assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
-        int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
-        int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
-        c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 0);
-        assertTrue(c.getString(intentIndex).contains(testPackage1));
-        c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 1);
-        assertTrue(c.getString(intentIndex).contains(testPackage2));
-        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);
-        intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
-        int cellXIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLX);
-        int cellYIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLY);
-
-        HashMap<String, Point> locMap = new HashMap<>();
-        while (c.moveToNext()) {
-            locMap.put(
-                    Intent.parseUri(c.getString(intentIndex), 0).getPackage(),
-                    new Point(c.getInt(cellXIndex), c.getInt(cellYIndex)));
-        }
-        c.close();
-
-        assertEquals(locMap.size(), 6);
-        assertEquals(new Point(0, 2), locMap.get(testPackage8));
-        assertEquals(new Point(0, 3), locMap.get(testPackage6));
-        assertEquals(new Point(1, 3), locMap.get(testPackage10));
-        assertEquals(new Point(2, 3), locMap.get(testPackage5));
-        assertEquals(new Point(3, 3), locMap.get(testPackage9));
-    }
-
-    @Test
-    public void migrateToLargerHotseat() {
-        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),
-                mModelHelper.addItem(APP_ICON, 2, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI),
-                mModelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI),
-        };
-
-        int numSrcDatabaseHotseatIcons = srcHotseatItems.length;
-        mIdp.numDatabaseHotseatIcons = 6;
-        mIdp.numColumns = 4;
-        mIdp.numRows = 4;
-        GridSizeMigrationTaskV2.DbReader srcReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages);
-        GridSizeMigrationTaskV2.DbReader destReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages);
-        GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(mContext, mDb, srcReader,
-                destReader, mIdp.numDatabaseHotseatIcons, new Point(mIdp.numColumns, mIdp.numRows));
-        task.migrate(mIdp);
-
-        // Check hotseat items
-        Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
-                new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
-        assertEquals(c.getCount(), numSrcDatabaseHotseatIcons);
-        int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
-        int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
-        c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 0);
-        assertTrue(c.getString(intentIndex).contains(testPackage1));
-        c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 1);
-        assertTrue(c.getString(intentIndex).contains(testPackage2));
-        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();
-    }
-
-    @Test
-    public void migrateFromLargerHotseat() {
-        int[] srcHotseatItems = {
-                mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
-                -1,
-                mModelHelper.addItem(SHORTCUT, 2, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
-                mModelHelper.addItem(APP_ICON, 3, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI),
-                mModelHelper.addItem(SHORTCUT, 4, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI),
-                mModelHelper.addItem(APP_ICON, 5, HOTSEAT, 0, 0, testPackage5, 5, TMP_CONTENT_URI),
-        };
-
-        mIdp.numDatabaseHotseatIcons = 4;
-        mIdp.numColumns = 4;
-        mIdp.numRows = 4;
-        GridSizeMigrationTaskV2.DbReader srcReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages);
-        GridSizeMigrationTaskV2.DbReader destReader = new GridSizeMigrationTaskV2.DbReader(mDb,
-                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages);
-        GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(mContext, mDb, srcReader,
-                destReader, mIdp.numDatabaseHotseatIcons, new Point(mIdp.numColumns, mIdp.numRows));
-        task.migrate(mIdp);
-
-        // Check hotseat items
-        Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
-                new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
-        assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
-        int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
-        int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
-        c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 0);
-        assertTrue(c.getString(intentIndex).contains(testPackage1));
-        c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 1);
-        assertTrue(c.getString(intentIndex).contains(testPackage2));
-        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();
-    }
-}
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.kt b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.kt
new file mode 100644
index 0000000..239e092
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.kt
@@ -0,0 +1,500 @@
+/*
+ * Copyright (C) 2022 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.content.Context
+import android.content.Intent
+import android.database.sqlite.SQLiteDatabase
+import android.graphics.Point
+import android.os.Process
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.InvariantDeviceProfile
+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.pm.UserCache
+import com.android.launcher3.provider.LauncherDbUtils
+import com.android.launcher3.util.LauncherModelHelper
+import com.android.launcher3.util.LauncherModelHelper.*
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Unit tests for [GridSizeMigrationTaskV2]  */
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class GridSizeMigrationTaskV2Test {
+    private lateinit var modelHelper: LauncherModelHelper
+    private lateinit var context: Context
+    private lateinit var db: SQLiteDatabase
+    private lateinit var validPackages: Set<String>
+    private lateinit var idp: InvariantDeviceProfile
+    private val testPackage1 = "com.android.launcher3.validpackage1"
+    private val testPackage2 = "com.android.launcher3.validpackage2"
+    private val testPackage3 = "com.android.launcher3.validpackage3"
+    private val testPackage4 = "com.android.launcher3.validpackage4"
+    private val testPackage5 = "com.android.launcher3.validpackage5"
+    private val testPackage6 = "com.android.launcher3.validpackage6"
+    private val testPackage7 = "com.android.launcher3.validpackage7"
+    private val testPackage8 = "com.android.launcher3.validpackage8"
+    private val testPackage9 = "com.android.launcher3.validpackage9"
+    private val testPackage10 = "com.android.launcher3.validpackage10"
+
+    @Before
+    fun setUp() {
+        modelHelper = LauncherModelHelper()
+        context = modelHelper.sandboxContext
+        db = modelHelper.provider.db
+
+        validPackages = setOf(
+            TEST_PACKAGE,
+            testPackage1,
+            testPackage2,
+            testPackage3,
+            testPackage4,
+            testPackage5,
+            testPackage6,
+            testPackage7,
+            testPackage8,
+            testPackage9,
+            testPackage10
+        )
+
+        idp = InvariantDeviceProfile.INSTANCE[context]
+        val userSerial = UserCache.INSTANCE[context].getSerialNumberForUser(Process.myUserHandle())
+        LauncherDbUtils.dropTable(db, TMP_TABLE)
+        addTableToDb(db, userSerial, false, TMP_TABLE)
+    }
+
+    @After
+    fun tearDown() {
+        modelHelper.destroy()
+    }
+
+    /**
+     * Old migration logic, should be modified once [FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC] is
+     * not needed anymore
+     */
+    @Test
+    @Throws(Exception::class)
+    fun testMigration() {
+        // Src Hotseat icons
+        modelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI)
+        modelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI)
+        modelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 4, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI)
+        // Src grid icons
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage5, 5, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 3, testPackage6, 6, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 1, testPackage8, 8, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 2, testPackage9, 9, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 3, testPackage10, 10, TMP_CONTENT_URI)
+
+        // Dest hotseat icons
+        modelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2)
+        // Dest grid icons
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage7)
+
+        idp.numDatabaseHotseatIcons = 4
+        idp.numColumns = 4
+        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)
+        )
+        task.migrate(idp)
+
+        // Check hotseat items
+        var c = context.contentResolver.query(
+            CONTENT_URI,
+            arrayOf(SCREEN, INTENT),
+            "container=$CONTAINER_HOTSEAT",
+            null,
+            SCREEN,
+            null
+        ) ?: throw IllegalStateException()
+
+        assertThat(c.count).isEqualTo(idp.numDatabaseHotseatIcons)
+
+        val screenIndex = c.getColumnIndex(SCREEN)
+        var intentIndex = c.getColumnIndex(INTENT)
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex).toLong()).isEqualTo(0)
+        assertThat(c.getString(intentIndex)).contains(testPackage1)
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex).toLong()).isEqualTo(1)
+        assertThat(c.getString(intentIndex)).contains(testPackage2)
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex).toLong()).isEqualTo(2)
+        assertThat(c.getString(intentIndex)).contains(testPackage3)
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex).toLong()).isEqualTo(3)
+        assertThat(c.getString(intentIndex)).contains(testPackage4)
+        c.close()
+
+        // Check workspace items
+        c = context.contentResolver.query(
+            CONTENT_URI,
+            arrayOf(CELLX, CELLY, INTENT),
+            "container=$CONTAINER_DESKTOP",
+            null,
+            null,
+            null
+        ) ?: throw IllegalStateException()
+
+        intentIndex = c.getColumnIndex(INTENT)
+        val cellXIndex = c.getColumnIndex(CELLX)
+        val cellYIndex = c.getColumnIndex(CELLY)
+        val locMap = HashMap<String, Point>()
+        while (c.moveToNext()) {
+            locMap[Intent.parseUri(c.getString(intentIndex), 0).getPackage()] =
+                Point(c.getInt(cellXIndex), c.getInt(cellYIndex))
+        }
+        c.close()
+        assertThat(locMap.size.toLong()).isEqualTo(6)
+        assertThat(locMap[testPackage8]).isEqualTo(Point(0, 2))
+        assertThat(locMap[testPackage6]).isEqualTo(Point(0, 3))
+        assertThat(locMap[testPackage10]).isEqualTo(Point(1, 3))
+        assertThat(locMap[testPackage7]).isEqualTo(Point(2, 2))
+        assertThat(locMap[testPackage5]).isEqualTo(Point(2, 3))
+        assertThat(locMap[testPackage9]).isEqualTo(Point(3, 3))
+    }
+
+    @Test
+    fun migrateToLargerHotseat() {
+        val srcHotseatItems = intArrayOf(
+            modelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
+            modelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
+            modelHelper.addItem(APP_ICON, 2, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI),
+            modelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI)
+        )
+        val numSrcDatabaseHotseatIcons = srcHotseatItems.size
+        idp.numDatabaseHotseatIcons = 6
+        idp.numColumns = 4
+        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)
+        )
+        task.migrate(idp)
+
+        // Check hotseat items
+        val c = context.contentResolver.query(
+            CONTENT_URI,
+            arrayOf(SCREEN, INTENT),
+            "container=$CONTAINER_HOTSEAT",
+            null,
+            SCREEN,
+            null
+        ) ?: throw IllegalStateException()
+
+        assertThat(c.count.toLong()).isEqualTo(numSrcDatabaseHotseatIcons.toLong())
+        val screenIndex = c.getColumnIndex(SCREEN)
+        val intentIndex = c.getColumnIndex(INTENT)
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(0)
+        assertThat(c.getString(intentIndex)).contains(testPackage1)
+
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(1)
+        assertThat(c.getString(intentIndex)).contains(testPackage2)
+
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(2)
+        assertThat(c.getString(intentIndex)).contains(testPackage3)
+
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(3)
+        assertThat(c.getString(intentIndex)).contains(testPackage4)
+
+        c.close()
+    }
+
+    @Test
+    fun migrateFromLargerHotseat() {
+        modelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI)
+        modelHelper.addItem(SHORTCUT, 2, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 3, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI)
+        modelHelper.addItem(SHORTCUT, 4, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 5, HOTSEAT, 0, 0, testPackage5, 5, TMP_CONTENT_URI)
+
+        idp.numDatabaseHotseatIcons = 4
+        idp.numColumns = 4
+        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)
+        )
+        task.migrate(idp)
+
+        // Check hotseat items
+        val c = context.contentResolver.query(
+            CONTENT_URI,
+            arrayOf(SCREEN, INTENT),
+            "container=$CONTAINER_HOTSEAT",
+            null,
+            SCREEN,
+            null
+        ) ?: throw IllegalStateException()
+
+        assertThat(c.count.toLong()).isEqualTo(idp.numDatabaseHotseatIcons.toLong())
+        val screenIndex = c.getColumnIndex(SCREEN)
+        val intentIndex = c.getColumnIndex(INTENT)
+
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(0)
+        assertThat(c.getString(intentIndex)).contains(testPackage1)
+
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(1)
+        assertThat(c.getString(intentIndex)).contains(testPackage2)
+
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(2)
+        assertThat(c.getString(intentIndex)).contains(testPackage3)
+
+        c.moveToNext()
+        assertThat(c.getInt(screenIndex)).isEqualTo(3)
+        assertThat(c.getString(intentIndex)).contains(testPackage4)
+
+        c.close()
+    }
+
+    /**
+     * Migrating from a smaller grid to a large one should keep the pages
+     * if the column difference is less than 2
+     */
+    @Test
+    @Throws(Exception::class)
+    fun migrateFromSmallerGridSmallDifference() {
+        enableNewMigrationLogic("4,4")
+
+        // Setup src grid
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage1, 5, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 3, testPackage2, 6, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 1, DESKTOP, 3, 1, testPackage3, 7, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 1, DESKTOP, 3, 2, testPackage4, 8, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 2, DESKTOP, 3, 3, testPackage5, 9, TMP_CONTENT_URI)
+
+        idp.numDatabaseHotseatIcons = 4
+        idp.numColumns = 6
+        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)
+        )
+        task.migrate(idp)
+
+        // Get workspace items
+        val c = context.contentResolver.query(
+            CONTENT_URI,
+            arrayOf(INTENT, SCREEN),
+            "container=$CONTAINER_DESKTOP",
+            null,
+            null,
+            null
+        ) ?: throw IllegalStateException()
+        val intentIndex = c.getColumnIndex(INTENT)
+        val screenIndex = c.getColumnIndex(SCREEN)
+
+        // Get in which screen the icon is
+        val locMap = HashMap<String, Int>()
+        while (c.moveToNext()) {
+            locMap[Intent.parseUri(c.getString(intentIndex), 0).getPackage()] =
+                c.getInt(screenIndex)
+        }
+        c.close()
+        assertThat(locMap.size).isEqualTo(5)
+        assertThat(locMap[testPackage1]).isEqualTo(0)
+        assertThat(locMap[testPackage2]).isEqualTo(0)
+        assertThat(locMap[testPackage3]).isEqualTo(1)
+        assertThat(locMap[testPackage4]).isEqualTo(1)
+        assertThat(locMap[testPackage5]).isEqualTo(2)
+
+        disableNewMigrationLogic()
+    }
+
+    /**
+     * Migrating from a smaller grid to a large one should reflow the pages
+     * if the column difference is more than 2
+     */
+    @Test
+    @Throws(Exception::class)
+    fun migrateFromSmallerGridBigDifference() {
+        enableNewMigrationLogic("2,2")
+
+        // Setup src grid
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 0, 1, testPackage1, 5, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 1, 1, testPackage2, 6, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 1, DESKTOP, 0, 0, testPackage3, 7, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 1, DESKTOP, 1, 0, testPackage4, 8, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 2, DESKTOP, 0, 0, testPackage5, 9, TMP_CONTENT_URI)
+
+        idp.numDatabaseHotseatIcons = 4
+        idp.numColumns = 5
+        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)
+        )
+        task.migrate(idp)
+
+        // Get workspace items
+        val c = context.contentResolver.query(
+            CONTENT_URI,
+            arrayOf(INTENT, SCREEN),
+            "container=$CONTAINER_DESKTOP",
+            null,
+            null,
+            null
+        ) ?: throw IllegalStateException()
+
+        val intentIndex = c.getColumnIndex(INTENT)
+        val screenIndex = c.getColumnIndex(SCREEN)
+
+        // Get in which screen the icon is
+        val locMap = HashMap<String, Int>()
+        while (c.moveToNext()) {
+            locMap[Intent.parseUri(c.getString(intentIndex), 0).getPackage()] =
+                c.getInt(screenIndex)
+        }
+        c.close()
+
+        // All icons fit the first screen
+        assertThat(locMap.size).isEqualTo(5)
+        assertThat(locMap[testPackage1]).isEqualTo(0)
+        assertThat(locMap[testPackage2]).isEqualTo(0)
+        assertThat(locMap[testPackage3]).isEqualTo(0)
+        assertThat(locMap[testPackage4]).isEqualTo(0)
+        assertThat(locMap[testPackage5]).isEqualTo(0)
+        disableNewMigrationLogic()
+    }
+
+    /**
+     * Migrating from a larger grid to a smaller, we reflow from page 0
+     */
+    @Test
+    @Throws(Exception::class)
+    fun migrateFromLargerGrid() {
+        enableNewMigrationLogic("5,5")
+
+        // Setup src grid
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 0, 1, testPackage1, 5, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 0, DESKTOP, 1, 1, testPackage2, 6, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 1, DESKTOP, 0, 0, testPackage3, 7, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 1, DESKTOP, 1, 0, testPackage4, 8, TMP_CONTENT_URI)
+        modelHelper.addItem(APP_ICON, 2, DESKTOP, 0, 0, testPackage5, 9, TMP_CONTENT_URI)
+
+        idp.numDatabaseHotseatIcons = 4
+        idp.numColumns = 4
+        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)
+        )
+        task.migrate(idp)
+
+        // Get workspace items
+        val c = context.contentResolver.query(
+            CONTENT_URI,
+            arrayOf(INTENT, SCREEN),
+            "container=$CONTAINER_DESKTOP",
+            null,
+            null,
+            null
+        ) ?: throw IllegalStateException()
+        val intentIndex = c.getColumnIndex(INTENT)
+        val screenIndex = c.getColumnIndex(SCREEN)
+
+        // Get in which screen the icon is
+        val locMap = HashMap<String, Int>()
+        while (c.moveToNext()) {
+            locMap[Intent.parseUri(c.getString(intentIndex), 0).getPackage()] =
+                c.getInt(screenIndex)
+        }
+        c.close()
+
+        // All icons fit the first screen
+        assertThat(locMap.size).isEqualTo(5)
+        assertThat(locMap[testPackage1]).isEqualTo(0)
+        assertThat(locMap[testPackage2]).isEqualTo(0)
+        assertThat(locMap[testPackage3]).isEqualTo(0)
+        assertThat(locMap[testPackage4]).isEqualTo(0)
+        assertThat(locMap[testPackage5]).isEqualTo(0)
+
+        disableNewMigrationLogic()
+    }
+
+    private fun enableNewMigrationLogic(srcGridSize: String) {
+        context.getSharedPreferences(FeatureFlags.FLAGS_PREF_NAME, Context.MODE_PRIVATE)
+            .edit()
+            .putBoolean(FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC.key, true)
+            .commit()
+        context.getSharedPreferences(LauncherFiles.SHARED_PREFERENCES_KEY, Context.MODE_PRIVATE)
+            .edit()
+            .putString(DeviceGridState.KEY_WORKSPACE_SIZE, srcGridSize)
+            .commit()
+        FeatureFlags.initialize(context)
+    }
+
+    private fun disableNewMigrationLogic() {
+        context.getSharedPreferences(FeatureFlags.FLAGS_PREF_NAME, Context.MODE_PRIVATE)
+            .edit()
+            .putBoolean(FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC.key, false)
+            .commit()
+    }
+}
\ No newline at end of file