Remove gap between screenIds when restoring from single display device

- In RestoreDbTask.sanitizeDb, add extra step to fill all screenIds gaps so that no empty pages are left in between
  e.g. [0, 3, 4, 6, 7] -> [0, 1, 2, 3, 4]

Fix: 201085524
Test: RestoreDbTaskTest
Change-Id: If0b577bf789347cc490ada0b22f2cd39ee1a33bb
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index df09f29..9f3d445 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -483,7 +483,7 @@
                             LauncherSettings.Favorites.CONTAINER + " FROM "
                                 + Favorites.TABLE_NAME + ")";
 
-            IntArray folderIds = LauncherDbUtils.queryIntArray(db, Favorites.TABLE_NAME,
+            IntArray folderIds = LauncherDbUtils.queryIntArray(false, db, Favorites.TABLE_NAME,
                     Favorites._ID, selection, null, null);
             if (!folderIds.isEmpty()) {
                 db.delete(Favorites.TABLE_NAME, Utilities.createDbSelectionQuery(
@@ -835,8 +835,8 @@
                 case 27: {
                     // Update the favorites table so that the screen ids are ordered based on
                     // workspace page rank.
-                    IntArray finalScreens = LauncherDbUtils.queryIntArray(db, "workspaceScreens",
-                            BaseColumns._ID, null, null, "screenRank");
+                    IntArray finalScreens = LauncherDbUtils.queryIntArray(false, db,
+                            "workspaceScreens", BaseColumns._ID, null, null, "screenRank");
                     int[] original = finalScreens.toArray();
                     Arrays.sort(original);
                     String updatemap = "";
@@ -919,7 +919,7 @@
                 Log.e(TAG, "getAppWidgetIds not supported", e);
                 return;
             }
-            final IntSet validWidgets = IntSet.wrap(LauncherDbUtils.queryIntArray(db,
+            final IntSet validWidgets = IntSet.wrap(LauncherDbUtils.queryIntArray(false, db,
                     Favorites.TABLE_NAME, Favorites.APPWIDGET_ID,
                     "itemType=" + Favorites.ITEM_TYPE_APPWIDGET, null, null));
             for (int widgetId : allWidgets) {
diff --git a/src/com/android/launcher3/provider/LauncherDbUtils.java b/src/com/android/launcher3/provider/LauncherDbUtils.java
index 6855bb1..b510378 100644
--- a/src/com/android/launcher3/provider/LauncherDbUtils.java
+++ b/src/com/android/launcher3/provider/LauncherDbUtils.java
@@ -31,11 +31,11 @@
  */
 public class LauncherDbUtils {
 
-    public static IntArray queryIntArray(SQLiteDatabase db, String tableName, String columnName,
-            String selection, String groupBy, String orderBy) {
+    public static IntArray queryIntArray(boolean distinct, SQLiteDatabase db, String tableName,
+            String columnName, String selection, String groupBy, String orderBy) {
         IntArray out = new IntArray();
-        try (Cursor c = db.query(tableName, new String[] { columnName }, selection, null,
-                groupBy, null, orderBy)) {
+        try (Cursor c = db.query(distinct, tableName, new String[] { columnName }, selection, null,
+                groupBy, null, orderBy, null)) {
             while (c.moveToNext()) {
                 out.add(c.getInt(0));
             }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 257d732..8f607a1 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.provider;
 
+import static com.android.launcher3.model.DeviceGridState.TYPE_MULTI_DISPLAY;
 import static com.android.launcher3.model.DeviceGridState.TYPE_PHONE;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 
@@ -96,7 +97,7 @@
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
             RestoreDbTask task = new RestoreDbTask();
             task.backupWorkspace(context, db);
-            task.sanitizeDB(helper, db, new BackupManager(context));
+            task.sanitizeDB(context, helper, db, new BackupManager(context));
             task.restoreAppWidgetIdsIfExists(context);
             t.commit();
             return true;
@@ -139,7 +140,7 @@
         GridBackupTable backupTable = new GridBackupTable(context, db, idp.numDatabaseHotseatIcons,
                 idp.numColumns, idp.numRows);
         if (backupTable.restoreFromRawBackupIfAvailable(getDefaultProfileId(db))) {
-            int itemsDeleted = sanitizeDB(helper, db, backupManager);
+            int itemsDeleted = sanitizeDB(context, helper, db, backupManager);
             LauncherAppState.getInstance(context).getModel().forceReload();
             restoreAppWidgetIdsIfExists(context);
             if (itemsDeleted == 0) {
@@ -156,11 +157,12 @@
      *      the restored apps get installed.
      *   3. If the user serial for any restored profile is different than that of the previous
      *      device, update the entries to the new profile id.
+     *   4. If restored from a single display backup, remove gaps between screenIds
      *
      * @return number of items deleted.
      */
-    private int sanitizeDB(DatabaseHelper helper, SQLiteDatabase db, BackupManager backupManager)
-            throws Exception {
+    private int sanitizeDB(Context context, DatabaseHelper helper, SQLiteDatabase db,
+            BackupManager backupManager) throws Exception {
         // Primary user ids
         long myProfileId = helper.getDefaultUserSerial();
         long oldProfileId = getDefaultProfileId(db);
@@ -236,10 +238,43 @@
         if (myProfileId != oldProfileId) {
             changeDefaultColumn(db, myProfileId);
         }
+
+        // If restored from a single display backup, remove gaps between screenIds
+        if (Utilities.getPrefs(context).getInt(RESTORED_DEVICE_TYPE, TYPE_PHONE)
+                != TYPE_MULTI_DISPLAY) {
+            removeScreenIdGaps(db);
+        }
+
         return itemsDeleted;
     }
 
     /**
+     * Remove gaps between screenIds to make sure no empty pages are left in between.
+     *
+     * e.g. [0, 3, 4, 6, 7] -> [0, 1, 2, 3, 4]
+     */
+    protected void removeScreenIdGaps(SQLiteDatabase db) {
+        FileLog.d(TAG, "Removing gaps between screenIds");
+        IntArray distinctScreens = LauncherDbUtils.queryIntArray(true, db, Favorites.TABLE_NAME,
+                Favorites.SCREEN, Favorites.CONTAINER + " = " + Favorites.CONTAINER_DESKTOP, null,
+                Favorites.SCREEN);
+        if (distinctScreens.isEmpty()) {
+            return;
+        }
+
+        StringBuilder sql = new StringBuilder("UPDATE ").append(Favorites.TABLE_NAME)
+                .append(" SET ").append(Favorites.SCREEN).append(" =\nCASE\n");
+        int screenId = distinctScreens.contains(0) ? 0 : 1;
+        for (int i = 0; i < distinctScreens.size(); i++) {
+            sql.append("WHEN ").append(Favorites.SCREEN).append(" == ")
+                    .append(distinctScreens.get(i)).append(" THEN ").append(screenId++).append("\n");
+        }
+        sql.append("ELSE screen\nEND WHERE ").append(Favorites.CONTAINER).append(" = ")
+                .append(Favorites.CONTAINER_DESKTOP).append(";");
+        db.execSQL(sql.toString());
+    }
+
+    /**
      * Updates profile id of all entries from {@param oldProfileId} to {@param newProfileId}.
      */
     protected void migrateProfileId(SQLiteDatabase db, long oldProfileId, long newProfileId) {