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
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:2eaf15c37a051fd17c18c4443a8bd19d39057bf0)
Merged-In: I6063e3c303320096f7429a3b1d909e107d2ef2a0
Change-Id: I6063e3c303320096f7429a3b1d909e107d2ef2a0

src/com/android/launcher3/model/GridSizeMigrationUtil.java

@@ -40,6 +40,7 @@ import android.util.Log;
 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 @@ public class GridSizeMigrationUtil {
         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);

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 GridMigrationData(dbFileName: String?, val gridState: DeviceGridState) {
 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 @@ class GridMigrationTest {
      * 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 @@ class GridMigrationTest {
     }
 
     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 @@ class GridMigrationTest {
     @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 @@ class GridMigrationTest {
     @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 @@ class GridMigrationTest {
     @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 @@ class GridMigrationTest {
             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, "")
+                )
+        )
+    }
 }

tests/src/com/android/launcher3/util/rule/TestToPhoneFileCopier.kt

@@ -49,7 +49,11 @@ class TestToPhoneFileCopier(
         object : Statement() {
             override fun evaluate() {
                 before()
-                base.evaluate()
+                try {
+                    base.evaluate()
+                } finally {
+                    after()
+                }
             }
         }
 }