In the case the gird migration goes to a taller grid keep the grid configuration

For example going from 4x5 to 4x6 it will keep the 4x5 configuration
which fits in the 4x6 grid.

Also, it fixes GridMigrationTest which was flaky because of
how it sorted arrays.

Bug: 325286145
Flag: ACONFIG grid_migration_fix disabled
Test: NA
Change-Id: I6063e3c303320096f7429a3b1d909e107d2ef2a0
diff --git a/src/com/android/launcher3/model/GridSizeMigrationUtil.java b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
index 30d2cfb..c40484d 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationUtil.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
@@ -40,6 +40,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
@@ -122,6 +123,16 @@
         if (!needsToMigrate(srcDeviceState, destDeviceState)) {
             return true;
         }
+
+        if (Flags.gridMigrationFix()
+                && srcDeviceState.getColumns().equals(destDeviceState.getColumns())
+                && srcDeviceState.getRows() < destDeviceState.getRows()) {
+            // Only use this strategy when comparing the previous grid to the new grid and the
+            // columns are the same and the destination has more rows
+            copyTable(source, TABLE_NAME, target.getWritableDatabase(), TABLE_NAME, context);
+            destDeviceState.writeToPrefs(context);
+            return true;
+        }
         copyTable(source, TABLE_NAME, target.getWritableDatabase(), TMP_TABLE, context);
 
         HashSet<String> validPackages = getValidPackages(context);
diff --git a/tests/assets/databases/GridMigrationTest/flagged_result5x5to5x8.db b/tests/assets/databases/GridMigrationTest/flagged_result5x5to5x8.db
new file mode 100644
index 0000000..8bea3ce
--- /dev/null
+++ b/tests/assets/databases/GridMigrationTest/flagged_result5x5to5x8.db
Binary files differ
diff --git a/tests/src/com/android/launcher3/model/GridMigrationTest.kt b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
index eb8604e..da4a208 100644
--- a/tests/src/com/android/launcher3/model/GridMigrationTest.kt
+++ b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
@@ -16,14 +16,18 @@
 
 package com.android.launcher3.model
 
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.Flags
 import com.android.launcher3.InvariantDeviceProfile.TYPE_PHONE
 import com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME
 import com.android.launcher3.celllayout.board.CellLayoutBoard
 import com.android.launcher3.pm.UserCache
 import com.android.launcher3.util.rule.TestToPhoneFileCopier
+import com.android.launcher3.util.rule.setFlags
+import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -65,11 +69,24 @@
 class GridMigrationTest {
     private val DB_FILE = "test_launcher.db"
 
+    @JvmField
+    @Rule
+    val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+
     // Copying the src db for all tests.
     @JvmField
     @Rule
     val fileCopier =
-        TestToPhoneFileCopier("databases/GridMigrationTest/$DB_FILE", "databases/$DB_FILE", true)
+        TestToPhoneFileCopier(
+            src = "databases/GridMigrationTest/$DB_FILE",
+            dest = "databases/$DB_FILE",
+            removeOnFinish = true
+        )
+
+    @Before
+    fun setup() {
+        setFlagsRule.setFlags(false, Flags.FLAG_GRID_MIGRATION_FIX)
+    }
 
     private fun migrate(src: GridMigrationData, dst: GridMigrationData) {
         GridSizeMigrationUtil.migrateGridIfNeeded(
@@ -86,8 +103,10 @@
      * same space in the db.
      */
     private fun validateDb(data: GridMigrationData) {
-        val cellLayoutBoard = CellLayoutBoard(data.gridState.columns, data.gridState.rows)
+        // The array size is just a big enough number to fit all the number of workspaces
+        val boards = Array(100) { CellLayoutBoard(data.gridState.columns, data.gridState.rows) }
         data.readEntries().forEach {
+            val cellLayoutBoard = boards[it.screenId]
             assert(cellLayoutBoard.isEmpty(it.cellX, it.cellY, it.spanX, it.spanY)) {
                 "Db has overlapping items"
             }
@@ -96,13 +115,13 @@
     }
 
     private fun compare(dst: GridMigrationData, target: GridMigrationData) {
-        val sortX = { it: GridSizeMigrationUtil.DbEntry -> it.cellX }
-        val sortY = { it: GridSizeMigrationUtil.DbEntry -> it.cellX }
+        val sort = compareBy<GridSizeMigrationUtil.DbEntry>({ it.cellX }, { it.cellY })
         val mapF = { it: GridSizeMigrationUtil.DbEntry ->
             EntryData(it.cellX, it.cellY, it.spanX, it.spanY, it.rank)
         }
-        val entriesDst = dst.readEntries().sortedBy(sortX).sortedBy(sortY).map(mapF)
-        val entriesTarget = target.readEntries().sortedBy(sortX).sortedBy(sortY).map(mapF)
+        val entriesDst = dst.readEntries().sortedWith(sort).map(mapF)
+        val entriesTarget = target.readEntries().sortedWith(sort).map(mapF)
+
         assert(entriesDst == entriesTarget) {
             "The elements on the dst database is not the same as in the target"
         }
@@ -128,9 +147,9 @@
     @Rule
     val result5x5to3x3 =
         TestToPhoneFileCopier(
-            "databases/GridMigrationTest/result5x5to3x3.db",
-            "databases/result5x5to3x3.db",
-            true
+            src = "databases/GridMigrationTest/result5x5to3x3.db",
+            dest = "databases/result5x5to3x3.db",
+            removeOnFinish = true
         )
 
     @Test
@@ -151,9 +170,9 @@
     @Rule
     val result5x5to4x7 =
         TestToPhoneFileCopier(
-            "databases/GridMigrationTest/result5x5to4x7.db",
-            "databases/result5x5to4x7.db",
-            true
+            src = "databases/GridMigrationTest/result5x5to4x7.db",
+            dest = "databases/result5x5to4x7.db",
+            removeOnFinish = true
         )
 
     @Test
@@ -174,9 +193,9 @@
     @Rule
     val result5x5to5x8 =
         TestToPhoneFileCopier(
-            "databases/GridMigrationTest/result5x5to5x8.db",
-            "databases/result5x5to5x8.db",
-            true
+            src = "databases/GridMigrationTest/result5x5to5x8.db",
+            dest = "databases/result5x5to5x8.db",
+            removeOnFinish = true
         )
 
     @Test
@@ -192,4 +211,32 @@
             target =
                 GridMigrationData("result5x5to5x8.db", DeviceGridState(5, 8, 5, TYPE_PHONE, ""))
         )
+
+    @JvmField
+    @Rule
+    val flaggedResult5x5to5x8 =
+        TestToPhoneFileCopier(
+            src = "databases/GridMigrationTest/flagged_result5x5to5x8.db",
+            dest = "databases/flagged_result5x5to5x8.db",
+            removeOnFinish = true
+        )
+
+    @Test
+    fun `flagged 5x5 to 5x8`() {
+        setFlagsRule.setFlags(true, Flags.FLAG_GRID_MIGRATION_FIX)
+        runTest(
+            src = GridMigrationData(DB_FILE, DeviceGridState(5, 5, 5, TYPE_PHONE, DB_FILE)),
+            dst =
+                GridMigrationData(
+                    null, // in memory db, to download a new db change null for the filename of the
+                    // db name to store it. Do not use existing names.
+                    DeviceGridState(5, 8, 5, TYPE_PHONE, "")
+                ),
+            target =
+                GridMigrationData(
+                    "flagged_result5x5to5x8.db",
+                    DeviceGridState(5, 8, 5, TYPE_PHONE, "")
+                )
+        )
+    }
 }
diff --git a/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt b/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt
index 72c4f16..d3516d1 100644
--- a/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt
+++ b/tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt
@@ -49,7 +49,11 @@
         object : Statement() {
             override fun evaluate() {
                 before()
-                base.evaluate()
+                try {
+                    base.evaluate()
+                } finally {
+                    after()
+                }
             }
         }
 }