Merge "Resolve issue where incorrect grid migration logic was used when migrating to strictly taller new grid" into main
diff --git a/src/com/android/launcher3/model/GridSizeMigrationDBController.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
index 2d6be7e..617cac7 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationDBController.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -17,7 +17,6 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
-import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
@@ -118,13 +117,14 @@
             @NonNull DeviceGridState srcDeviceState,
             @NonNull DeviceGridState destDeviceState,
             @NonNull DatabaseHelper target,
-            @NonNull SQLiteDatabase source) {
+            @NonNull SQLiteDatabase source,
+            boolean isDestNewDb) {
 
         if (!needsToMigrate(srcDeviceState, destDeviceState)) {
             return true;
         }
 
-        if (LauncherPrefs.get(context).get(IS_FIRST_LOAD_AFTER_RESTORE)
+        if (isDestNewDb
                 && srcDeviceState.getColumns().equals(destDeviceState.getColumns())
                 && srcDeviceState.getRows() < destDeviceState.getRows()) {
             // Only use this strategy when comparing the previous grid to the new grid and the
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
index 07316ef..c856d4b 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
@@ -46,6 +46,7 @@
         destDeviceState: DeviceGridState,
         target: DatabaseHelper,
         source: SQLiteDatabase,
+        isDestNewDb: Boolean,
     ) {
         if (!GridSizeMigrationDBController.needsToMigrate(srcDeviceState, destDeviceState)) {
             return
@@ -57,7 +58,7 @@
         // This is a special case where if the grid is the same amount of columns but a larger
         // amount of rows we simply copy over the source grid to the destination grid, rather
         // than undergoing the general grid migration.
-        if (shouldMigrateToStrictlyTallerGrid(isFirstLoad, srcDeviceState, destDeviceState)) {
+        if (shouldMigrateToStrictlyTallerGrid(isDestNewDb, srcDeviceState, destDeviceState)) {
             GridSizeMigrationDBController.copyCurrentGridToNewGrid(
                 context,
                 destDeviceState,
@@ -335,11 +336,11 @@
 
     /** Only migrate the grid in this manner if the target grid is taller and not wider. */
     private fun shouldMigrateToStrictlyTallerGrid(
-        isFirstLoad: Boolean,
+        isDestNewDb: Boolean,
         srcDeviceState: DeviceGridState,
         destDeviceState: DeviceGridState,
     ): Boolean {
-        return isFirstLoad &&
+        return isDestNewDb &&
             srcDeviceState.columns == destDeviceState.columns &&
             srcDeviceState.rows < destDeviceState.rows
     }
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 094798b..6ff8547 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -89,6 +89,7 @@
 import java.io.File;
 import java.io.InputStream;
 import java.io.StringReader;
+import java.util.List;
 
 /**
  * Utility class which maintains an instance of Launcher database and provides utility methods
@@ -368,6 +369,13 @@
 
         InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
         DatabaseHelper oldHelper = mOpenHelper;
+
+        // We save the existing db's before creating the destination db helper so we know what logic
+        // to run in grid migration based on if that grid already existed before migration or not.
+        List<String> existingDBs = LauncherFiles.GRID_DB_FILES.stream()
+                .filter(dbName -> mContext.getDatabasePath(dbName).exists())
+                .toList();
+
         mOpenHelper = (mContext instanceof SandboxContext) ? oldHelper
                 : createDatabaseHelper(true, new DeviceGridState(idp).getDbFile());
         try {
@@ -376,9 +384,11 @@
             // This is the state we want to migrate to that is given by the idp
             DeviceGridState destDeviceState = new DeviceGridState(idp);
 
+            boolean isDestNewDb = !existingDBs.contains(destDeviceState.getDbFile());
+
             GridSizeMigrationLogic gridSizeMigrationLogic = new GridSizeMigrationLogic();
             gridSizeMigrationLogic.migrateGrid(mContext, srcDeviceState, destDeviceState,
-                    mOpenHelper, oldHelper.getWritableDatabase());
+                    mOpenHelper, oldHelper.getWritableDatabase(), isDestNewDb);
         } catch (Exception e) {
             resetLauncherDb(restoreEventLogger);
             throw new Exception("Failed to migrate grid", e);
@@ -441,6 +451,13 @@
             return false;
         }
         DatabaseHelper oldHelper = mOpenHelper;
+
+        // We save the existing db's before creating the destination db helper so we know what logic
+        // to run in grid migration based on if that grid already existed before migration or not.
+        List<String> existingDBs = LauncherFiles.GRID_DB_FILES.stream()
+                .filter(dbName -> mContext.getDatabasePath(dbName).exists())
+                .toList();
+
         mOpenHelper = (mContext instanceof SandboxContext) ? oldHelper
                 : createDatabaseHelper(true /* forMigration */, targetDbName);
         try {
@@ -448,8 +465,11 @@
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
             // This is the state we want to migrate to that is given by the idp
             DeviceGridState destDeviceState = new DeviceGridState(idp);
+
+            boolean isDestNewDb = !existingDBs.contains(destDeviceState.getDbFile());
+
             return GridSizeMigrationDBController.migrateGridIfNeeded(mContext, srcDeviceState,
-                    destDeviceState, mOpenHelper, oldHelper.getWritableDatabase());
+                    destDeviceState, mOpenHelper, oldHelper.getWritableDatabase(), isDestNewDb);
         } catch (Exception e) {
             FileLog.e(TAG, "Failed to migrate grid", e);
             return false;
diff --git a/tests/src/com/android/launcher3/model/GridMigrationTest.kt b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
index 666ec16..379e98d 100644
--- a/tests/src/com/android/launcher3/model/GridMigrationTest.kt
+++ b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
@@ -101,6 +101,7 @@
                 dst.gridState,
                 dst.dbHelper,
                 src.dbHelper.readableDatabase,
+                false,
             )
         } else {
             GridSizeMigrationDBController.migrateGridIfNeeded(
@@ -109,6 +110,7 @@
                 dst.gridState,
                 dst.dbHelper,
                 src.dbHelper.readableDatabase,
+                false,
             )
         }
     }