Merge "Fix issue where hotseat apps don't migrate to the new grid" into main
diff --git a/src/com/android/launcher3/model/GridSizeMigrationDBController.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
index 47f13bd..5d0a7bd 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationDBController.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -238,16 +238,19 @@
         Collections.sort(hotseatToBeAdded);
         Collections.sort(workspaceToBeAdded);
 
-        List<Integer> idsInUse = dstWorkspaceItems.stream()
+        List<DbEntry> remainingDstHotseatItems = destReader.loadHotseatEntries();
+        List<DbEntry> remainingDstWorkspaceItems = destReader.loadAllWorkspaceEntries();
+        List<Integer> idsInUse = remainingDstHotseatItems.stream()
                 .map(entry -> entry.id)
                 .collect(Collectors.toList());
-        idsInUse.addAll(dstHotseatItems.stream()
+        idsInUse.addAll(remainingDstWorkspaceItems.stream()
                 .map(entry -> entry.id)
                 .collect(Collectors.toList()));
 
+
         // Migrate hotseat
         solveHotseatPlacement(helper, destHotseatSize,
-                srcReader, destReader, dstHotseatItems, hotseatToBeAdded, idsInUse);
+                srcReader, destReader, remainingDstHotseatItems, hotseatToBeAdded, idsInUse);
 
         // Migrate workspace.
         // First we create a collection of the screens
@@ -467,7 +470,7 @@
         final Context mContext;
         int mLastScreenId = -1;
 
-        final Map<Integer, List<DbEntry>> mWorkspaceEntriesByScreenId =
+        Map<Integer, List<DbEntry>> mWorkspaceEntriesByScreenId =
                 new ArrayMap<>();
 
         public DbReader(SQLiteDatabase db, String tableName, Context context) {
@@ -539,6 +542,7 @@
         }
 
         protected List<DbEntry> loadAllWorkspaceEntries() {
+            mWorkspaceEntriesByScreenId.clear();
             final List<DbEntry> workspaceEntries = new ArrayList<>();
             Cursor c = queryWorkspace(
                     new String[]{
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
index 9586bf3..5df135a 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
@@ -183,9 +183,11 @@
             )
         }
 
+        val remainingDstHotseatItems = destReader.loadHotseatEntries()
+
         placeHotseatItems(
             itemsToBeAdded,
-            dstHotseatItems,
+            remainingDstHotseatItems,
             destHotseatSize,
             helper,
             srcReader,
@@ -265,9 +267,10 @@
             )
         }
 
+        val remainingDstWorkspaceItems = destReader.loadAllWorkspaceEntries()
         placeWorkspaceItems(
             workspaceToBeAdded,
-            dstWorkspaceItems,
+            remainingDstWorkspaceItems,
             targetSize.x,
             targetSize.y,
             helper,
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationTest.kt
index adf38fe..5607bb4 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationTest.kt
@@ -275,36 +275,15 @@
         val readerGridA = DbReader(db, TMP_TABLE, context)
         val readerGridB = DbReader(db, TABLE_NAME, context)
         // migrate from A -> B
-        if (Flags.gridMigrationRefactor()) {
-            var gridSizeMigrationLogic = GridSizeMigrationLogic()
-            val idsInUse = mutableListOf<Int>()
-            gridSizeMigrationLogic.migrateHotseat(
-                5,
-                idp.numDatabaseHotseatIcons,
-                readerGridA,
-                readerGridB,
-                dbHelper,
-                idsInUse,
-            )
-            gridSizeMigrationLogic.migrateWorkspace(
-                readerGridA,
-                readerGridB,
-                dbHelper,
-                Point(idp.numColumns, idp.numRows),
-                idsInUse,
-            )
-        } else {
-            GridSizeMigrationDBController.migrate(
-                dbHelper,
-                readerGridA,
-                readerGridB,
-                5,
-                idp.numDatabaseHotseatIcons,
-                Point(idp.numColumns, idp.numRows),
-                DeviceGridState(context),
-                DeviceGridState(idp),
-            )
-        }
+        migrateGrid(
+            dbHelper,
+            readerGridA,
+            readerGridB,
+            5,
+            idp.numDatabaseHotseatIcons,
+            idp.numColumns,
+            idp.numRows,
+        )
 
         // Check hotseat items in grid B
         var c =
@@ -388,8 +367,8 @@
         locMap = parseLocMap(c)
         // Expected workspace items in grid A
         // _ _ _ _ _
-        // _ _ _ _ 5
-        // 9 _ 6 _ 7
+        // 9 _ _ _ 5
+        // _ _ 6 _ 7
         // _ _ 8 _ _
         // _ _ _ _ _
         assertThat(locMap.size.toLong()).isEqualTo(5)
@@ -399,7 +378,7 @@
         assertThat(locMap[testPackage7]).isEqualTo(Triple(0, 4, 2))
         assertThat(locMap[testPackage8]).isEqualTo(Triple(0, 2, 3))
         // Verify items that didn't exist in grid A are added in new screen
-        assertThat(locMap[testPackage9]).isEqualTo(Triple(0, 0, 2))
+        assertThat(locMap[testPackage9]).isEqualTo(Triple(0, 0, 1))
 
         // remove item from B
         db.delete(TMP_TABLE, "$_ID=7", null)
@@ -493,37 +472,15 @@
         val readerGridA = DbReader(db, TMP_TABLE, context)
         val readerGridB = DbReader(db, TABLE_NAME, context)
         // migrate from A -> B
-        if (Flags.gridMigrationRefactor()) {
-            var gridSizeMigrationLogic = GridSizeMigrationLogic()
-            val idsInUse = mutableListOf<Int>()
-            gridSizeMigrationLogic.migrateHotseat(
-                5,
-                idp.numDatabaseHotseatIcons,
-                readerGridA,
-                readerGridB,
-                dbHelper,
-                idsInUse,
-            )
-            gridSizeMigrationLogic.migrateWorkspace(
-                readerGridA,
-                readerGridB,
-                dbHelper,
-                Point(idp.numColumns, idp.numRows),
-                idsInUse,
-            )
-        } else {
-            GridSizeMigrationDBController.migrate(
-                dbHelper,
-                readerGridA,
-                readerGridB,
-                5,
-                idp.numDatabaseHotseatIcons,
-                Point(idp.numColumns, idp.numRows),
-                DeviceGridState(context),
-                DeviceGridState(idp),
-            )
-        }
-
+        migrateGrid(
+            dbHelper,
+            readerGridA,
+            readerGridB,
+            5,
+            idp.numDatabaseHotseatIcons,
+            idp.numColumns,
+            idp.numRows,
+        )
         // Check hotseat items in grid B
         var c =
             db.query(
@@ -597,6 +554,112 @@
         )
     }
 
+    @Test
+    @Throws(Exception::class)
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun testMigrationToFullGridFlagOn() {
+        testMigrationToFullGrid()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun testHotseatMigrationToFullGridFlagOff() {
+        testMigrationToFullGrid()
+    }
+
+    @Throws(Exception::class)
+    fun testMigrationToFullGrid() {
+        // Hotseat items in grid A
+        // 1 2 3 4 5
+        addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_HOTSEAT, 0, 0, testPackage1, 1, TMP_TABLE)
+        addItem(ITEM_TYPE_DEEP_SHORTCUT, 1, CONTAINER_HOTSEAT, 0, 0, testPackage2, 2, TMP_TABLE)
+        addItem(ITEM_TYPE_DEEP_SHORTCUT, 2, CONTAINER_HOTSEAT, 0, 0, testPackage3, 3, TMP_TABLE)
+        addItem(ITEM_TYPE_APPLICATION, 3, CONTAINER_HOTSEAT, 0, 0, testPackage4, 4, TMP_TABLE)
+        addItem(ITEM_TYPE_APPLICATION, 4, CONTAINER_HOTSEAT, 0, 0, testPackage5, 5, TMP_TABLE)
+
+        // Hotseat items in grid B
+        // 6 7 8 9
+        addItem(ITEM_TYPE_DEEP_SHORTCUT, 0, CONTAINER_HOTSEAT, 0, 0, testPackage6)
+        addItem(ITEM_TYPE_DEEP_SHORTCUT, 1, CONTAINER_HOTSEAT, 0, 0, testPackage7)
+        addItem(ITEM_TYPE_APPLICATION, 2, CONTAINER_HOTSEAT, 0, 0, testPackage8)
+        addItem(ITEM_TYPE_APPLICATION, 3, CONTAINER_HOTSEAT, 0, 0, testPackage9)
+
+        // Workspace items in grid A
+        // _ _ _ _ _
+        // 6 _ _ _ _
+        // _ _ _ _ _
+        // _ _ _ _ _
+        // _ _ _ _ _
+        addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_DESKTOP, 0, 1, testPackage6, 6, TMP_TABLE)
+
+        // Workspace items in grid B
+        // _ _ _ _
+        // 1 2 3 4
+        // _ _ _ _
+        // _ _ _ _
+        addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_DESKTOP, 0, 1, testPackage1)
+        addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_DESKTOP, 1, 1, testPackage2)
+        addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_DESKTOP, 2, 1, testPackage3)
+        addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_DESKTOP, 3, 1, testPackage4)
+
+        idp.numDatabaseHotseatIcons = 4
+        idp.numColumns = 4
+        idp.numRows = 4
+        val readerGridA = DbReader(db, TMP_TABLE, context)
+        val readerGridB = DbReader(db, TABLE_NAME, context)
+
+        // migrate from A -> B
+        migrateGrid(
+            dbHelper,
+            readerGridA,
+            readerGridB,
+            5,
+            idp.numDatabaseHotseatIcons,
+            idp.numColumns,
+            idp.numRows,
+        )
+
+        // Check hotseat items in grid B
+        var c =
+            db.query(
+                TABLE_NAME,
+                arrayOf(SCREEN, INTENT),
+                "container=$CONTAINER_HOTSEAT",
+                null,
+                SCREEN,
+                null,
+                null,
+            ) ?: throw IllegalStateException()
+        // Expected hotseat items in grid B
+        // 1 2 3 4
+        verifyHotseat(
+            c,
+            mutableListOf(testPackage1, testPackage2, testPackage3, testPackage4).toList(),
+            4,
+        )
+
+        // Check workspace items in grid B
+        c =
+            db.query(
+                TABLE_NAME,
+                arrayOf(SCREEN, CELLX, CELLY, INTENT),
+                "container=$CONTAINER_DESKTOP",
+                null,
+                null,
+                null,
+                null,
+            ) ?: throw IllegalStateException()
+        val locMap = parseLocMap(c)
+        // Expected workspace items in grid B
+        // _ _ _ _
+        // 6 _ _ _
+        // _ _ _ _
+        // _ _ _ _
+        assertThat(locMap.size.toLong()).isEqualTo(1)
+        assertThat(locMap[testPackage6]).isEqualTo(Triple(0, 0, 1))
+    }
+
     private fun migrateGrid(
         dbHelper: DatabaseHelper,
         srcReader: DbReader,
@@ -621,7 +684,7 @@
                 srcReader,
                 destReader,
                 dbHelper,
-                Point(idp.numColumns, idp.numRows),
+                Point(pointX, pointY),
                 idsInUse,
             )
         } else {