Add new heuristic for deciding whether we should add extra rows on grid migration to the bottom or top of the grid

Bug: 364711064
Flag: com.android.launcher3.one_grid_specs
Test: GridSizeMigrationTest
Change-Id: I070ed1eddc2c7ab475267268ebdbb2e559ab6dda
diff --git a/src/com/android/launcher3/model/GridSizeMigrationDBController.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
index 617cac7..bfa00bd 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationDBController.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -17,12 +17,14 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.Flags.oneGridSpecs;
 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;
 import static com.android.launcher3.model.LoaderTask.SMARTSPACE_ON_HOME_SCREEN;
 import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+import static com.android.launcher3.provider.LauncherDbUtils.shiftTableByXCells;
 
 import android.content.ComponentName;
 import android.content.ContentValues;
@@ -130,6 +132,20 @@
             // 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);
+
+            if (oneGridSpecs()) {
+                DbReader destReader = new DbReader(
+                        target.getWritableDatabase(), TABLE_NAME, context);
+                boolean shouldShiftCells = shouldShiftCells(destReader, srcDeviceState.getRows());
+                if (shouldShiftCells) {
+                    shiftTableByXCells(
+                            target.getWritableDatabase(),
+                            (destDeviceState.getRows() - srcDeviceState.getRows()),
+                            TABLE_NAME);
+                }
+            }
+
+            // Save current configuration, so that the migration does not run again.
             destDeviceState.writeToPrefs(context);
             return true;
         }
@@ -427,17 +443,22 @@
         }
     }
 
-    static void copyCurrentGridToNewGrid(
-            @NonNull Context context,
-            @NonNull DeviceGridState destDeviceState,
-            @NonNull DatabaseHelper target,
-            @NonNull SQLiteDatabase source) {
-        // 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);
+    private static boolean shouldShiftCells(final DbReader destReader, final int srcGridRowCount) {
+        List<DbEntry> workspaceItems = destReader.loadAllWorkspaceEntries();
+        int firstPageItemsRowPosSum = workspaceItems.stream()
+                .filter(entry -> entry.screenId == 0)
+                .mapToInt(entry -> entry.cellY).sum();
+        int firstPageWorkspaceItemsCount = (int) workspaceItems.stream()
+                .filter(entry -> entry.screenId == 0).count();
+        if (firstPageWorkspaceItemsCount == 0) {
+            return false;
+        }
+        float srcGridMidPoint = srcGridRowCount / 2f;
+        float firstPageItemPosAvg = (float) firstPageItemsRowPosSum / firstPageWorkspaceItemsCount;
+        return (firstPageItemPosAvg >= srcGridMidPoint);
     }
 
+
     @VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
     public static class DbReader {
 
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
index c856d4b..3f52d8a 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
@@ -21,15 +21,20 @@
 import android.util.Log
 import androidx.annotation.VisibleForTesting
 import com.android.launcher3.Flags
+import com.android.launcher3.Flags.oneGridSpecs
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.get
 import com.android.launcher3.LauncherPrefs.Companion.getPrefs
 import com.android.launcher3.LauncherSettings
+import com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME
+import com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE
 import com.android.launcher3.Utilities
 import com.android.launcher3.config.FeatureFlags
 import com.android.launcher3.model.GridSizeMigrationDBController.DbReader
-import com.android.launcher3.provider.LauncherDbUtils
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction
+import com.android.launcher3.provider.LauncherDbUtils.copyTable
+import com.android.launcher3.provider.LauncherDbUtils.dropTable
+import com.android.launcher3.provider.LauncherDbUtils.shiftTableByXCells
 import com.android.launcher3.util.CellAndSpan
 import com.android.launcher3.util.GridOccupancy
 import com.android.launcher3.util.IntArray
@@ -59,27 +64,30 @@
         // amount of rows we simply copy over the source grid to the destination grid, rather
         // than undergoing the general grid migration.
         if (shouldMigrateToStrictlyTallerGrid(isDestNewDb, srcDeviceState, destDeviceState)) {
-            GridSizeMigrationDBController.copyCurrentGridToNewGrid(
-                context,
-                destDeviceState,
-                target,
-                source,
-            )
+            copyTable(source, TABLE_NAME, target.writableDatabase, TABLE_NAME, context)
+            if (oneGridSpecs()) {
+                val destReader = DbReader(target.writableDatabase, TABLE_NAME, context)
+                val shouldShiftCells = shouldShiftCells(destReader, srcDeviceState.rows)
+                if (shouldShiftCells) {
+                    shiftTableByXCells(
+                        target.writableDatabase,
+                        (destDeviceState.rows - srcDeviceState.rows),
+                        TABLE_NAME,
+                    )
+                }
+            }
+            // Save current configuration, so that the migration does not run again.
+            destDeviceState.writeToPrefs(context)
             return
         }
-        LauncherDbUtils.copyTable(
-            source,
-            LauncherSettings.Favorites.TABLE_NAME,
-            target.writableDatabase,
-            LauncherSettings.Favorites.TMP_TABLE,
-            context,
-        )
+
+        copyTable(source, TABLE_NAME, target.writableDatabase, TMP_TABLE, context)
 
         val migrationStartTime = System.currentTimeMillis()
         try {
             SQLiteTransaction(target.writableDatabase).use { t ->
-                val srcReader = DbReader(t.db, LauncherSettings.Favorites.TMP_TABLE, context)
-                val destReader = DbReader(t.db, LauncherSettings.Favorites.TABLE_NAME, context)
+                val srcReader = DbReader(t.db, TMP_TABLE, context)
+                val destReader = DbReader(t.db, TABLE_NAME, context)
 
                 val targetSize = Point(destDeviceState.columns, destDeviceState.rows)
 
@@ -95,7 +103,7 @@
                 // Migrate workspace.
                 migrateWorkspace(srcReader, destReader, target, targetSize, idsInUse)
 
-                LauncherDbUtils.dropTable(t.db, LauncherSettings.Favorites.TMP_TABLE)
+                dropTable(t.db, TMP_TABLE)
                 t.commit()
             }
         } catch (e: Exception) {
@@ -112,6 +120,19 @@
         }
     }
 
+    private fun shouldShiftCells(destReader: DbReader, srcGridRowCount: Int): Boolean {
+        val workspaceItems = destReader.loadAllWorkspaceEntries()
+        val firstPageItemsRowPosSum =
+            workspaceItems.sumOf { entry -> if (entry.screenId == 0) entry.cellY else 0 }
+        val firstPageWorkspaceItemsCount = workspaceItems.count { entry -> entry.screenId == 0 }
+        if (firstPageWorkspaceItemsCount == 0) {
+            return false
+        }
+        val srcGridMidPoint = srcGridRowCount / 2f
+        val firstPageItemPosAvg = firstPageItemsRowPosSum / firstPageWorkspaceItemsCount.toFloat()
+        return (firstPageItemPosAvg >= srcGridMidPoint)
+    }
+
     /** Handles hotseat migration. */
     @VisibleForTesting
     fun migrateHotseat(
diff --git a/src/com/android/launcher3/provider/LauncherDbUtils.kt b/src/com/android/launcher3/provider/LauncherDbUtils.kt
index 3c68e46..6f1d0dd 100644
--- a/src/com/android/launcher3/provider/LauncherDbUtils.kt
+++ b/src/com/android/launcher3/provider/LauncherDbUtils.kt
@@ -131,6 +131,11 @@
         }
     }
 
+    @JvmStatic
+    fun shiftTableByXCells(db: SQLiteDatabase, x: Int, toTable: String) {
+        db.run { execSQL("UPDATE $toTable SET cellY = cellY + $x") }
+    }
+
     /**
      * Migrates the legacy shortcuts to deep shortcuts pinned under Launcher. Removes any invalid
      * shortcut or any shortcut which requires some permission to launch