Bye bye workspace screens table

Removing a separate table for workspace screens. List of screens are
automatically parsed using the items in the favorites DB. Order of the
screen based on the screen id and rearranging screens is no longer
supported. In case the screens need to be rearranged, all the items
in the favorites db will need to be updated with new screen ids.

This makes backing up the DB (in the same database) easier as only
one table needs to be duplicates.

Change-Id: I8ba947a898f637d780e2f49925e78604263126e8

16 files changed, 179 insertions, 469 deletions

diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index b4cc1ceff..f8b44d00c 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1747,7 +1747,6 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
                 orderedScreenIds.indexOf(Workspace.FIRST_SCREEN_ID) != 0) {
             orderedScreenIds.removeValue(Workspace.FIRST_SCREEN_ID);
             orderedScreenIds.add(0, Workspace.FIRST_SCREEN_ID);
-            LauncherModel.updateWorkspaceScreenOrder(this, orderedScreenIds);
         } else if (!FeatureFlags.QSB_ON_FIRST_SCREEN.get()
                 && orderedScreenIds.isEmpty()) {
             // If there are no screens, we need to have an empty screen
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index a5f97de7a..9b4c5fd09 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -20,12 +20,8 @@ import static com.android.launcher3.LauncherAppState.ACTION_FORCE_ROLOAD;
 import static com.android.launcher3.config.FeatureFlags.IS_DOGFOOD_BUILD;
 
 import android.content.BroadcastReceiver;
-import android.content.ContentProviderOperation;
-import android.content.ContentResolver;
-import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
-import android.net.Uri;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.Looper;
@@ -38,8 +34,8 @@ import android.util.Pair;
 import com.android.launcher3.compat.LauncherAppsCompat;
 import com.android.launcher3.compat.PackageInstallerCompat.PackageInstallInfo;
 import com.android.launcher3.compat.UserManagerCompat;
-import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.model.AddWorkspaceItemsTask;
 import com.android.launcher3.model.BaseModelUpdateTask;
 import com.android.launcher3.model.BgDataModel;
@@ -51,7 +47,6 @@ import com.android.launcher3.model.PackageInstallStateChangedTask;
 import com.android.launcher3.model.PackageUpdatedTask;
 import com.android.launcher3.model.ShortcutsChangedTask;
 import com.android.launcher3.model.UserLockStateChangedTask;
-import com.android.launcher3.provider.LauncherDbUtils;
 import com.android.launcher3.shortcuts.DeepShortcutManager;
 import com.android.launcher3.shortcuts.ShortcutInfoCompat;
 import com.android.launcher3.util.ComponentKey;
@@ -70,7 +65,6 @@ import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.Executor;
@@ -269,53 +263,6 @@ public class LauncherModel extends BroadcastReceiver
     }
 
     /**
-     * Update the order of the workspace screens in the database. The array list contains
-     * a list of screen ids in the order that they should appear.
-     */
-    public static void updateWorkspaceScreenOrder(Context context, IntArray screens) {
-        final ContentResolver cr = context.getContentResolver();
-        final Uri uri = LauncherSettings.WorkspaceScreens.CONTENT_URI;
-
-        // Create a copy with only non-negative values
-        final IntArray screensCopy = new IntArray();
-        for (int i = 0; i < screens.size(); i++) {
-            int id = screens.get(i);
-            if (id >= 0) {
-                screensCopy.add(id);
-            }
-        }
-
-        Runnable r = new Runnable() {
-            @Override
-            public void run() {
-                ArrayList<ContentProviderOperation> ops = new ArrayList<ContentProviderOperation>();
-                // Clear the table
-                ops.add(ContentProviderOperation.newDelete(uri).build());
-                int count = screensCopy.size();
-                for (int i = 0; i < count; i++) {
-                    ContentValues v = new ContentValues();
-                    int screenId = screensCopy.get(i);
-                    v.put(LauncherSettings.WorkspaceScreens._ID, screenId);
-                    v.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
-                    ops.add(ContentProviderOperation.newInsert(uri).withValues(v).build());
-                }
-
-                try {
-                    cr.applyBatch(LauncherProvider.AUTHORITY, ops);
-                } catch (Exception ex) {
-                    throw new RuntimeException(ex);
-                }
-
-                synchronized (sBgDataModel) {
-                    sBgDataModel.workspaceScreens.clear();
-                    sBgDataModel.workspaceScreens.addAll(screensCopy);
-                }
-            }
-        };
-        runOnWorkerThread(r);
-    }
-
-    /**
      * Set this as the current Launcher activity object for the loader.
      */
     public void initialize(Callbacks callbacks) {
@@ -519,18 +466,6 @@ public class LauncherModel extends BroadcastReceiver
         }
     }
 
-    /**
-     * Loads the workspace screen ids in an ordered list.
-     */
-    public static IntArray loadWorkspaceScreensDb(Context context) {
-        final ContentResolver contentResolver = context.getContentResolver();
-        final Uri screensUri = LauncherSettings.WorkspaceScreens.CONTENT_URI;
-
-        // Get screens ordered by rank.
-        return LauncherDbUtils.getScreenIdsFromCursor(contentResolver.query(
-                screensUri, null, null, null, LauncherSettings.WorkspaceScreens.SCREEN_RANK));
-    }
-
     public void onInstallSessionCreated(final PackageInstallInfo sessionInfo) {
         enqueueModelUpdateTask(new BaseModelUpdateTask() {
             @Override
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 7db3d5b9d..8ed3314ef 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -32,6 +32,7 @@ import android.content.SharedPreferences;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.res.Resources;
 import android.database.Cursor;
+import android.database.DatabaseUtils;
 import android.database.SQLException;
 import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteQueryBuilder;
@@ -45,12 +46,12 @@ import android.os.Message;
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.provider.BaseColumns;
 import android.text.TextUtils;
 import android.util.Log;
 
 import com.android.launcher3.AutoInstallsLayout.LayoutParserCallback;
 import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.LauncherSettings.WorkspaceScreens;
 import com.android.launcher3.compat.UserManagerCompat;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.logging.FileLog;
@@ -70,6 +71,7 @@ import java.io.PrintWriter;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Locale;
 
 public class LauncherProvider extends ContentProvider {
     private static final String TAG = "LauncherProvider";
@@ -79,8 +81,9 @@ public class LauncherProvider extends ContentProvider {
 
     /**
      * Represents the schema of the database. Changes in scheme need not be backwards compatible.
+     * When increasing the scheme version, ensure that downgrade_schema.json is updated
      */
-    public static final int SCHEMA_VERSION = 27;
+    public static final int SCHEMA_VERSION = 28;
 
     public static final String AUTHORITY = BuildConfig.APPLICATION_ID + ".settings";
 
@@ -175,10 +178,10 @@ public class LauncherProvider extends ContentProvider {
         if (values == null) {
             throw new RuntimeException("Error: attempting to insert null values");
         }
-        if (!values.containsKey(LauncherSettings.ChangeLogColumns._ID)) {
+        if (!values.containsKey(LauncherSettings.Favorites._ID)) {
             throw new RuntimeException("Error: attempting to add item without specifying an id");
         }
-        helper.checkId(table, values);
+        helper.checkId(values);
         return (int) db.insert(table, nullColumnHack, values);
     }
 
@@ -262,24 +265,7 @@ public class LauncherProvider extends ContentProvider {
             }
         }
 
-        // Add screen id if not present
-        int screenId = values.getAsInteger(LauncherSettings.Favorites.SCREEN);
-        SQLiteStatement stmp = null;
-        try {
-            stmp = mOpenHelper.getWritableDatabase().compileStatement(
-                    "INSERT OR IGNORE INTO workspaceScreens (_id, screenRank) " +
-                            "select ?, (ifnull(MAX(screenRank), -1)+1) from workspaceScreens");
-            stmp.bindLong(1, screenId);
-
-            ContentValues valuesInserted = new ContentValues();
-            valuesInserted.put(LauncherSettings.BaseLauncherColumns._ID, stmp.executeInsert());
-            mOpenHelper.checkId(WorkspaceScreens.TABLE_NAME, valuesInserted);
-            return true;
-        } catch (Exception e) {
-            return false;
-        } finally {
-            Utilities.closeSilently(stmp);
-        }
+        return true;
     }
 
     @Override
@@ -404,7 +390,6 @@ public class LauncherProvider extends ContentProvider {
      * @return Ids of deleted folders.
      */
     private IntArray deleteEmptyFolders() {
-        IntArray folderIds = new IntArray();
         SQLiteDatabase db = mOpenHelper.getWritableDatabase();
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
             // Select folders whose id do not match any container value.
@@ -413,21 +398,19 @@ public class LauncherProvider extends ContentProvider {
                     + LauncherSettings.Favorites._ID +  " NOT IN (SELECT " +
                             LauncherSettings.Favorites.CONTAINER + " FROM "
                                 + Favorites.TABLE_NAME + ")";
-            try (Cursor c = db.query(Favorites.TABLE_NAME,
-                    new String[] {LauncherSettings.Favorites._ID},
-                    selection, null, null, null, null)) {
-                LauncherDbUtils.iterateCursor(c, 0, folderIds);
-            }
+
+            IntArray folderIds = LauncherDbUtils.queryIntArray(db, Favorites.TABLE_NAME,
+                    Favorites._ID, selection, null, null);
             if (!folderIds.isEmpty()) {
                 db.delete(Favorites.TABLE_NAME, Utilities.createDbSelectionQuery(
                         LauncherSettings.Favorites._ID, folderIds), null);
             }
             t.commit();
+            return folderIds;
         } catch (SQLException ex) {
             Log.e(TAG, ex.getMessage(), ex);
-            folderIds.clear();
+            return new IntArray();
         }
-        return folderIds;
     }
 
     /**
@@ -438,7 +421,7 @@ public class LauncherProvider extends ContentProvider {
     }
 
     @Thunk static void addModifiedTime(ContentValues values) {
-        values.put(LauncherSettings.ChangeLogColumns.MODIFIED, System.currentTimeMillis());
+        values.put(LauncherSettings.Favorites.MODIFIED, System.currentTimeMillis());
     }
 
     private void clearFlagEmptyDbCreated() {
@@ -551,11 +534,10 @@ public class LauncherProvider extends ContentProvider {
             // Table creation sometimes fails silently, which leads to a crash loop.
             // This way, we will try to create a table every time after crash, so the device
             // would eventually be able to recover.
-            if (!tableExists(Favorites.TABLE_NAME) || !tableExists(WorkspaceScreens.TABLE_NAME)) {
+            if (!tableExists(Favorites.TABLE_NAME)) {
                 Log.e(TAG, "Tables are missing after onCreate has been called. Trying to recreate");
                 // This operation is a no-op if the table already exists.
                 addFavoritesTable(getWritableDatabase(), true);
-                addWorkspacesTable(getWritableDatabase(), true);
             }
 
             initIds();
@@ -602,7 +584,6 @@ public class LauncherProvider extends ContentProvider {
             mMaxScreenId = 0;
 
             addFavoritesTable(db, false);
-            addWorkspacesTable(db, false);
 
             // Fresh and clean launcher DB.
             mMaxItemId = initializeMaxItemId(db);
@@ -633,46 +614,6 @@ public class LauncherProvider extends ContentProvider {
             Favorites.addTableToDb(db, getDefaultUserSerial(), optional);
         }
 
-        private void addWorkspacesTable(SQLiteDatabase db, boolean optional) {
-            String ifNotExists = optional ? " IF NOT EXISTS " : "";
-            db.execSQL("CREATE TABLE " + ifNotExists + WorkspaceScreens.TABLE_NAME + " (" +
-                    LauncherSettings.WorkspaceScreens._ID + " INTEGER PRIMARY KEY," +
-                    LauncherSettings.WorkspaceScreens.SCREEN_RANK + " INTEGER," +
-                    LauncherSettings.ChangeLogColumns.MODIFIED + " INTEGER NOT NULL DEFAULT 0" +
-                    ");");
-        }
-
-        private void removeOrphanedItems(SQLiteDatabase db) {
-            // Delete items directly on the workspace who's screen id doesn't exist
-            //  "DELETE FROM favorites WHERE screen NOT IN (SELECT _id FROM workspaceScreens)
-            //   AND container = -100"
-            String removeOrphanedDesktopItems = "DELETE FROM " + Favorites.TABLE_NAME +
-                    " WHERE " +
-                    LauncherSettings.Favorites.SCREEN + " NOT IN (SELECT " +
-                    LauncherSettings.WorkspaceScreens._ID + " FROM " + WorkspaceScreens.TABLE_NAME + ")" +
-                    " AND " +
-                    LauncherSettings.Favorites.CONTAINER + " = " +
-                    LauncherSettings.Favorites.CONTAINER_DESKTOP;
-            db.execSQL(removeOrphanedDesktopItems);
-
-            // Delete items contained in folders which no longer exist (after above statement)
-            //  "DELETE FROM favorites  WHERE container <> -100 AND container <> -101 AND container
-            //   NOT IN (SELECT _id FROM favorites WHERE itemType = 2)"
-            String removeOrphanedFolderItems = "DELETE FROM " + Favorites.TABLE_NAME +
-                    " WHERE " +
-                    LauncherSettings.Favorites.CONTAINER + " <> " +
-                    LauncherSettings.Favorites.CONTAINER_DESKTOP +
-                    " AND "
-                    + LauncherSettings.Favorites.CONTAINER + " <> " +
-                    LauncherSettings.Favorites.CONTAINER_HOTSEAT +
-                    " AND "
-                    + LauncherSettings.Favorites.CONTAINER + " NOT IN (SELECT " +
-                    LauncherSettings.Favorites._ID + " FROM " + Favorites.TABLE_NAME +
-                    " WHERE " + LauncherSettings.Favorites.ITEM_TYPE + " = " +
-                    LauncherSettings.Favorites.ITEM_TYPE_FOLDER + ")";
-            db.execSQL(removeOrphanedFolderItems);
-        }
-
         @Override
         public void onOpen(SQLiteDatabase db) {
             super.onOpen(db);
@@ -681,8 +622,7 @@ public class LauncherProvider extends ContentProvider {
             if (!schemaFile.exists()) {
                 handleOneTimeDataUpgrade(db);
             }
-            DbDowngradeHelper.updateSchemaFile(schemaFile, SCHEMA_VERSION, mContext,
-                    R.raw.downgrade_schema);
+            DbDowngradeHelper.updateSchemaFile(schemaFile, SCHEMA_VERSION, mContext);
         }
 
         /**
@@ -709,12 +649,8 @@ public class LauncherProvider extends ContentProvider {
             switch (oldVersion) {
                 // The version cannot be lower that 12, as Launcher3 never supported a lower
                 // version of the DB.
-                case 12: {
-                    // With the new shrink-wrapped and re-orderable workspaces, it makes sense
-                    // to persist workspace screens and their relative order.
-                    mMaxScreenId = 0;
-                    addWorkspacesTable(db, false);
-                }
+                case 12:
+                    // No-op
                 case 13: {
                     try (SQLiteTransaction t = new SQLiteTransaction(db)) {
                         // Insert new column for holding widget provider name
@@ -728,15 +664,7 @@ public class LauncherProvider extends ContentProvider {
                     }
                 }
                 case 14: {
-                    try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-                        // Insert new column for holding update timestamp
-                        db.execSQL("ALTER TABLE favorites " +
-                                "ADD COLUMN modified INTEGER NOT NULL DEFAULT 0;");
-                        db.execSQL("ALTER TABLE workspaceScreens " +
-                                "ADD COLUMN modified INTEGER NOT NULL DEFAULT 0;");
-                        t.commit();
-                    } catch (SQLException ex) {
-                        Log.e(TAG, ex.getMessage(), ex);
+                    if (!addIntegerColumn(db, Favorites.MODIFIED, 0)) {
                         // Old version remains, which means we wipe old data
                         break;
                     }
@@ -747,22 +675,15 @@ public class LauncherProvider extends ContentProvider {
                         break;
                     }
                 }
-                case 16: {
+                case 16:
                     // No-op
-                }
-                case 17: {
+                case 17:
+                    // No-op
+                case 18:
                     // No-op
-                }
-                case 18: {
-                    // Due to a data loss bug, some users may have items associated with screen ids
-                    // which no longer exist. Since this can cause other problems, and since the user
-                    // will never see these items anyway, we use database upgrade as an opportunity to
-                    // clean things up.
-                    removeOrphanedItems(db);
-                }
                 case 19: {
                     // Add userId column
-                    if (!addProfileColumn(db)) {
+                    if (!addIntegerColumn(db, Favorites.PROFILE_ID, getDefaultUserSerial())) {
                         // Old version remains, which means we wipe old data
                         break;
                     }
@@ -772,10 +693,7 @@ public class LauncherProvider extends ContentProvider {
                         break;
                     }
                 case 21:
-                    // Recreate workspace table with screen id a primary key
-                    if (!recreateWorkspaceTable(db)) {
-                        break;
-                    }
+                    // No-op
                 case 22: {
                     if (!addIntegerColumn(db, Favorites.OPTIONS, 0)) {
                         // Old version remains, which means we wipe old data
@@ -794,7 +712,30 @@ public class LauncherProvider extends ContentProvider {
                             !LauncherDbUtils.prepareScreenZeroToHostQsb(mContext, db)) {
                         break;
                     }
-                case 27:
+                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");
+                    int[] original = finalScreens.toArray();
+                    Arrays.sort(original);
+                    String updatemap = "";
+                    for (int i = 0; i < original.length; i++) {
+                        if (finalScreens.get(i) != original[i]) {
+                            updatemap += String.format(Locale.ENGLISH, " WHEN %1$s=%2$d THEN %3$d",
+                                    Favorites.SCREEN, finalScreens.get(i), original[i]);
+                        }
+                    }
+                    if (!TextUtils.isEmpty(updatemap)) {
+                        String query = String.format(Locale.ENGLISH,
+                                "UPDATE %1$s SET %2$s=CASE %3$s ELSE %2$s END WHERE %4$s = %5$d",
+                                Favorites.TABLE_NAME, Favorites.SCREEN, updatemap,
+                                Favorites.CONTAINER, Favorites.CONTAINER_DESKTOP);
+                        db.execSQL(query);
+                    }
+                    db.execSQL("DROP TABLE IF EXISTS workspaceScreens");
+                }
+                case 28:
                     // DB Upgraded successfully
                     return;
             }
@@ -822,7 +763,7 @@ public class LauncherProvider extends ContentProvider {
         public void createEmptyDB(SQLiteDatabase db) {
             try (SQLiteTransaction t = new SQLiteTransaction(db)) {
                 db.execSQL("DROP TABLE IF EXISTS " + Favorites.TABLE_NAME);
-                db.execSQL("DROP TABLE IF EXISTS " + WorkspaceScreens.TABLE_NAME);
+                db.execSQL("DROP TABLE IF EXISTS workspaceScreens");
                 onCreate(db);
                 t.commit();
             }
@@ -845,17 +786,9 @@ public class LauncherProvider extends ContentProvider {
                 Log.e(TAG, "getAppWidgetIds not supported", e);
                 return;
             }
-            final IntSet validWidgets = new IntSet();
-            try (Cursor c = db.query(Favorites.TABLE_NAME,
-                    new String[] {Favorites.APPWIDGET_ID },
-                    "itemType=" + Favorites.ITEM_TYPE_APPWIDGET, null, null, null, null)) {
-                while (c.moveToNext()) {
-                    validWidgets.add(c.getInt(0));
-                }
-            } catch (SQLException ex) {
-                Log.w(TAG, "Error getting widgets list", ex);
-                return;
-            }
+            final IntSet validWidgets = IntSet.wrap(LauncherDbUtils.queryIntArray(db,
+                    Favorites.TABLE_NAME, Favorites.APPWIDGET_ID,
+                    "itemType=" + Favorites.ITEM_TYPE_APPWIDGET, null, null));
             for (int widgetId : allWidgets) {
                 if (!validWidgets.contains(widgetId)) {
                     try {
@@ -910,46 +843,6 @@ public class LauncherProvider extends ContentProvider {
             }
         }
 
-        /**
-         * Recreates workspace table and migrates data to the new table.
-         */
-        public boolean recreateWorkspaceTable(SQLiteDatabase db) {
-            try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-                final IntArray sortedIDs;
-
-                try (Cursor c = db.query(WorkspaceScreens.TABLE_NAME,
-                        new String[] {LauncherSettings.WorkspaceScreens._ID},
-                        null, null, null, null,
-                        LauncherSettings.WorkspaceScreens.SCREEN_RANK)) {
-                    // Use LinkedHashSet so that ordering is preserved
-                    sortedIDs = LauncherDbUtils.getScreenIdsFromCursor(c);
-                }
-                db.execSQL("DROP TABLE IF EXISTS " + WorkspaceScreens.TABLE_NAME);
-                addWorkspacesTable(db, false);
-
-                // Add all screen ids back
-                int total = sortedIDs.size();
-                for (int i = 0; i < total; i++) {
-                    ContentValues values = new ContentValues();
-                    values.put(LauncherSettings.WorkspaceScreens._ID, sortedIDs.get(i));
-                    values.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
-                    addModifiedTime(values);
-                    db.insertOrThrow(WorkspaceScreens.TABLE_NAME, null, values);
-                }
-                t.commit();
-
-                mMaxScreenId = 0;
-                for (int i = 0; i < sortedIDs.size(); i++) {
-                    mMaxScreenId = Math.max(mMaxScreenId, sortedIDs.get(i));
-                }
-            } catch (SQLException ex) {
-                // Old version remains, which means we wipe old data
-                Log.e(TAG, ex.getMessage(), ex);
-                return false;
-            }
-            return true;
-        }
-
         @Thunk boolean updateFolderItemsRank(SQLiteDatabase db, boolean addRankColumn) {
             try (SQLiteTransaction t = new SQLiteTransaction(db)) {
                 if (addRankColumn) {
@@ -979,10 +872,6 @@ public class LauncherProvider extends ContentProvider {
             return true;
         }
 
-        private boolean addProfileColumn(SQLiteDatabase db) {
-            return addIntegerColumn(db, Favorites.PROFILE_ID, getDefaultUserSerial());
-        }
-
         private boolean addIntegerColumn(SQLiteDatabase db, String columnName, long defaultValue) {
             try (SQLiteTransaction t = new SQLiteTransaction(db)) {
                 db.execSQL("ALTER TABLE favorites ADD COLUMN "
@@ -1018,17 +907,20 @@ public class LauncherProvider extends ContentProvider {
             return dbInsertAndCheck(this, db, Favorites.TABLE_NAME, null, values);
         }
 
-        public void checkId(String table, ContentValues values) {
-            int id = values.getAsInteger(LauncherSettings.BaseLauncherColumns._ID);
-            if (WorkspaceScreens.TABLE_NAME.equals(table)) {
-                mMaxScreenId = Math.max(id, mMaxScreenId);
-            }  else {
-                mMaxItemId = Math.max(id, mMaxItemId);
+        public void checkId(ContentValues values) {
+            int id = values.getAsInteger(Favorites._ID);
+            mMaxItemId = Math.max(id, mMaxItemId);
+
+            Integer screen = values.getAsInteger(Favorites.SCREEN);
+            Integer container = values.getAsInteger(Favorites.CONTAINER);
+            if (screen != null && container != null
+                    && container.intValue() == Favorites.CONTAINER_DESKTOP) {
+                mMaxScreenId = Math.max(screen, mMaxScreenId);
             }
         }
 
         private int initializeMaxItemId(SQLiteDatabase db) {
-            return getMaxId(db, Favorites.TABLE_NAME);
+            return getMaxId(db, "SELECT MAX(%1$s) FROM %2$s", Favorites._ID, Favorites.TABLE_NAME);
         }
 
         // Generates a new ID to use for an workspace screen in your database. This method
@@ -1045,34 +937,18 @@ public class LauncherProvider extends ContentProvider {
         }
 
         private int initializeMaxScreenId(SQLiteDatabase db) {
-            return getMaxId(db, WorkspaceScreens.TABLE_NAME);
+            return getMaxId(db, "SELECT MAX(%1$s) FROM %2$s WHERE %3$s = %4$d",
+                    Favorites.SCREEN, Favorites.TABLE_NAME, Favorites.CONTAINER,
+                    Favorites.CONTAINER_DESKTOP);
         }
 
         @Thunk int loadFavorites(SQLiteDatabase db, AutoInstallsLayout loader) {
-            IntArray screenIds = new IntArray();
             // TODO: Use multiple loaders with fall-back and transaction.
-            int count = loader.loadLayout(db, screenIds);
-
-            // Add the screens specified by the items above
-            int[] sortedScreenIds = screenIds.toArray();
-            Arrays.sort(sortedScreenIds);
-            int rank = 0;
-            ContentValues values = new ContentValues();
-            for (int id : sortedScreenIds) {
-                values.clear();
-                values.put(LauncherSettings.WorkspaceScreens._ID, id);
-                values.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, rank);
-                if (dbInsertAndCheck(this, db, WorkspaceScreens.TABLE_NAME, null, values) < 0) {
-                    throw new RuntimeException("Failed initialize screen table"
-                            + "from default layout");
-                }
-                rank++;
-            }
+            int count = loader.loadLayout(db, new IntArray());
 
             // Ensure that the max ids are initialized
             mMaxItemId = initializeMaxItemId(db);
             mMaxScreenId = initializeMaxScreenId(db);
-
             return count;
         }
     }
@@ -1080,22 +956,14 @@ public class LauncherProvider extends ContentProvider {
     /**
      * @return the max _id in the provided table.
      */
-    @Thunk static int getMaxId(SQLiteDatabase db, String table) {
-        Cursor c = db.rawQuery("SELECT MAX(_id) FROM " + table, null);
-        // get the result
-        int id = -1;
-        if (c != null && c.moveToNext()) {
-            id = c.getInt(0);
-        }
-        if (c != null) {
-            c.close();
-        }
-
-        if (id == -1) {
-            throw new RuntimeException("Error: could not query max id in " + table);
-        }
-
-        return id;
+    @Thunk static int getMaxId(SQLiteDatabase db, String query, Object... args) {
+        int max = (int) DatabaseUtils.longForQuery(db,
+                String.format(Locale.ENGLISH, query, args),
+                null);
+        if (max < 0) {
+            throw new RuntimeException("Error: could not query max id");
+        }
+        return max;
     }
 
     static class SqlArguments {
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 0b12b15f9..79c820873 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -26,16 +26,17 @@ import android.provider.BaseColumns;
  * Settings related utilities.
  */
 public class LauncherSettings {
-    /** Columns required on table staht will be subject to backup and restore. */
-    static interface ChangeLogColumns extends BaseColumns {
+
+    /**
+     * Favorites.
+     */
+    public static final class Favorites implements BaseColumns {
         /**
          * The time of the last update to this row.
          * <P>Type: INTEGER</P>
          */
         public static final String MODIFIED = "modified";
-    }
 
-    static public interface BaseLauncherColumns extends ChangeLogColumns {
         /**
          * Descriptive name of the gesture that can be displayed to the user.
          * <P>Type: TEXT</P>
@@ -84,34 +85,6 @@ public class LauncherSettings {
          * <P>Type: BLOB</P>
          */
         public static final String ICON = "icon";
-    }
-
-    /**
-     * Workspace Screens.
-     *
-     * Tracks the order of workspace screens.
-     */
-    public static final class WorkspaceScreens implements ChangeLogColumns {
-
-        public static final String TABLE_NAME = "workspaceScreens";
-
-        /**
-         * The content:// style URL for this table
-         */
-        public static final Uri CONTENT_URI = Uri.parse("content://" +
-                LauncherProvider.AUTHORITY + "/" + TABLE_NAME);
-
-        /**
-         * The rank of this screen -- ie. how it is ordered relative to the other screens.
-         * <P>Type: INTEGER</P>
-         */
-        public static final String SCREEN_RANK = "screenRank";
-    }
-
-    /**
-     * Favorites.
-     */
-    public static final class Favorites implements BaseLauncherColumns {
 
         public static final String TABLE_NAME = "favorites";
 
diff --git a/src/com/android/launcher3/SecondaryDropTarget.java b/src/com/android/launcher3/SecondaryDropTarget.java
index 6083415d5..0cf6e44c6 100644
--- a/src/com/android/launcher3/SecondaryDropTarget.java
+++ b/src/com/android/launcher3/SecondaryDropTarget.java
@@ -150,7 +150,7 @@ public class SecondaryDropTarget extends ButtonDropTarget implements OnAlarmList
         Intent intent = null;
         UserHandle user = null;
         if (item != null &&
-                item.itemType == LauncherSettings.BaseLauncherColumns.ITEM_TYPE_APPLICATION) {
+                item.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPLICATION) {
             intent = item.getIntent();
             user = item.user;
         }
diff --git a/src/com/android/launcher3/ShortcutInfo.java b/src/com/android/launcher3/ShortcutInfo.java
index 19c647fd7..c2c328745 100644
--- a/src/com/android/launcher3/ShortcutInfo.java
+++ b/src/com/android/launcher3/ShortcutInfo.java
@@ -89,7 +89,7 @@ public class ShortcutInfo extends ItemInfoWithIcon {
     private int mInstallProgress;
 
     public ShortcutInfo() {
-        itemType = LauncherSettings.BaseLauncherColumns.ITEM_TYPE_SHORTCUT;
+        itemType = LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
     }
 
     public ShortcutInfo(ShortcutInfo info) {
@@ -114,24 +114,23 @@ public class ShortcutInfo extends ItemInfoWithIcon {
     @TargetApi(Build.VERSION_CODES.N)
     public ShortcutInfo(ShortcutInfoCompat shortcutInfo, Context context) {
         user = shortcutInfo.getUserHandle();
-        itemType = LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
+        itemType = Favorites.ITEM_TYPE_DEEP_SHORTCUT;
         updateFromDeepShortcutInfo(shortcutInfo, context);
     }
 
     @Override
     public void onAddToDatabase(ContentWriter writer) {
         super.onAddToDatabase(writer);
-        writer.put(LauncherSettings.BaseLauncherColumns.TITLE, title)
-                .put(LauncherSettings.BaseLauncherColumns.INTENT, getIntent())
-                .put(LauncherSettings.Favorites.RESTORED, status);
+        writer.put(Favorites.TITLE, title)
+                .put(Favorites.INTENT, getIntent())
+                .put(Favorites.RESTORED, status);
 
         if (!usingLowResIcon()) {
             writer.putIcon(iconBitmap, user);
         }
         if (iconResource != null) {
-            writer.put(LauncherSettings.BaseLauncherColumns.ICON_PACKAGE, iconResource.packageName)
-                    .put(LauncherSettings.BaseLauncherColumns.ICON_RESOURCE,
-                            iconResource.resourceName);
+            writer.put(Favorites.ICON_PACKAGE, iconResource.packageName)
+                    .put(Favorites.ICON_RESOURCE, iconResource.resourceName);
         }
     }
 
@@ -189,7 +188,7 @@ public class ShortcutInfo extends ItemInfoWithIcon {
     @Override
     public ComponentName getTargetComponent() {
         ComponentName cn = super.getTargetComponent();
-        if (cn == null && (itemType == LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT
+        if (cn == null && (itemType == Favorites.ITEM_TYPE_SHORTCUT
                 || hasStatusFlag(FLAG_SUPPORTS_WEB_UI))) {
             // Legacy shortcuts and promise icons with web UI may not have a componentName but just
             // a packageName. In that case create a dummy componentName instead of adding additional
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 47ebb1795..56aca02ed 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -619,9 +619,6 @@ public class Workspace extends PagedView<WorkspacePageIndicator>
             // if this is the last screen, convert it to the empty screen
             mWorkspaceScreens.put(EXTRA_EMPTY_SCREEN_ID, finalScreen);
             mScreenOrder.add(EXTRA_EMPTY_SCREEN_ID);
-
-            // Update the model if we have changed any screens
-            LauncherModel.updateWorkspaceScreenOrder(mLauncher, mScreenOrder);
         }
     }
 
@@ -728,9 +725,6 @@ public class Workspace extends PagedView<WorkspacePageIndicator>
         mWorkspaceScreens.put(newId, cl);
         mScreenOrder.add(newId);
 
-        // Update the model for the new screen
-        LauncherModel.updateWorkspaceScreenOrder(mLauncher, mScreenOrder);
-
         return newId;
     }
 
@@ -828,13 +822,6 @@ public class Workspace extends PagedView<WorkspacePageIndicator>
             }
         }
 
-        if (!removeScreens.isEmpty()) {
-            // Update the model if we have changed any screens
-            mLauncher.getModelWriter().enqueueDeleteRunnable(
-                    () -> LauncherModel.updateWorkspaceScreenOrder(mLauncher, mScreenOrder));
-
-        }
-
         if (pageShift >= 0) {
             setCurrentPage(currentPage - pageShift);
         }
diff --git a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
index 756d44e81..5fc555172 100644
--- a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
+++ b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
@@ -15,11 +15,11 @@
  */
 package com.android.launcher3.model;
 
-import android.content.Context;
 import android.content.Intent;
 import android.os.UserHandle;
 import android.util.LongSparseArray;
 import android.util.Pair;
+
 import com.android.launcher3.AllAppsList;
 import com.android.launcher3.AppInfo;
 import com.android.launcher3.FolderInfo;
@@ -27,7 +27,6 @@ import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetInfo;
-import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherModel.CallbackTask;
 import com.android.launcher3.LauncherModel.Callbacks;
 import com.android.launcher3.LauncherSettings;
@@ -58,16 +57,12 @@ public class AddWorkspaceItemsTask extends BaseModelUpdateTask {
         if (mItemList.isEmpty()) {
             return;
         }
-        Context context = app.getContext();
 
         final ArrayList<ItemInfo> addedItemsFinal = new ArrayList<>();
         final IntArray addedWorkspaceScreensFinal = new IntArray();
 
-        // Get the list of workspace screens.  We need to append to this list and
-        // can not use sBgWorkspaceScreens because loadWorkspace() may not have been
-        // called.
-        IntArray workspaceScreens = LauncherModel.loadWorkspaceScreensDb(context);
         synchronized(dataModel) {
+            IntArray workspaceScreens = dataModel.workspaceScreens.clone();
 
             List<ItemInfo> filteredItems = new ArrayList<>();
             for (Pair<ItemInfo, Object> entry : mItemList) {
@@ -116,9 +111,6 @@ public class AddWorkspaceItemsTask extends BaseModelUpdateTask {
             }
         }
 
-        // Update the workspace screens
-        updateScreens(context, workspaceScreens);
-
         if (!addedItemsFinal.isEmpty()) {
             scheduleCallbackTask(new CallbackTask() {
                 @Override
@@ -143,10 +135,6 @@ public class AddWorkspaceItemsTask extends BaseModelUpdateTask {
         }
     }
 
-    protected void updateScreens(Context context, IntArray workspaceScreens) {
-        LauncherModel.updateWorkspaceScreenOrder(context, workspaceScreens);
-    }
-
     /**
      * Returns true if the shortcuts already exists on the workspace. This must be called after
      * the workspace has been loaded. We identify a shortcut by its intent.
diff --git a/src/com/android/launcher3/model/DbDowngradeHelper.java b/src/com/android/launcher3/model/DbDowngradeHelper.java
index cd86b728b..e5c44d1aa 100644
--- a/src/com/android/launcher3/model/DbDowngradeHelper.java
+++ b/src/com/android/launcher3/model/DbDowngradeHelper.java
@@ -21,6 +21,7 @@ import android.database.sqlite.SQLiteException;
 import android.util.Log;
 import android.util.SparseArray;
 
+import com.android.launcher3.R;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.util.IOUtils;
 
@@ -87,8 +88,7 @@ public class DbDowngradeHelper {
         return helper;
     }
 
-    public static void updateSchemaFile(File schemaFile, int expectedVersion,
-            Context context, int schemaResId) {
+    public static void updateSchemaFile(File schemaFile, int expectedVersion, Context context) {
         try {
             if (DbDowngradeHelper.parse(schemaFile).version >= expectedVersion) {
                 return;
@@ -99,7 +99,7 @@ public class DbDowngradeHelper {
 
         // Write the updated schema
         try (FileOutputStream fos = new FileOutputStream(schemaFile);
-            InputStream in = context.getResources().openRawResource(schemaResId)) {
+            InputStream in = context.getResources().openRawResource(R.raw.downgrade_schema)) {
             IOUtils.copy(in, fos);
         } catch (IOException e) {
             Log.e(TAG, "Error writing schema file", e);
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTask.java b/src/com/android/launcher3/model/GridSizeMigrationTask.java
index 289de1a4b..a9ddccb67 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -13,14 +13,12 @@ import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.database.Cursor;
 import android.graphics.Point;
-import android.net.Uri;
 import android.util.Log;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -31,12 +29,14 @@ import com.android.launcher3.compat.PackageInstallerCompat;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.IntSparseArrayMap;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
-import java.util.Locale;
+
+import androidx.annotation.VisibleForTesting;
 
 /**
  * This class takes care of shrinking the workspace (by maximum of one row and one column), as a
@@ -180,11 +180,25 @@ public class GridSizeMigrationTask {
         return applyOperations();
     }
 
+    @VisibleForTesting
+    static IntArray getWorkspaceScreenIds(Context context) {
+        IntSet set = new IntSet();
+        try (Cursor c = context.getContentResolver().query(Favorites.CONTENT_URI,
+                new String[] {Favorites.SCREEN},
+                Favorites.CONTAINER + " = " + Favorites.CONTAINER_DESKTOP,
+                null, Favorites.SCREEN)) {
+            while (c.moveToNext()) {
+                set.add(c.getInt(0));
+            }
+        }
+        return set.getArray();
+    }
+
     /**
      * @return true if any DB change was made
      */
     protected boolean migrateWorkspace() throws Exception {
-        IntArray allScreens = LauncherModel.loadWorkspaceScreensDb(mContext);
+        IntArray allScreens = getWorkspaceScreenIds(mContext);
         if (allScreens.isEmpty()) {
             throw new Exception("Unable to get workspace screens");
         }
@@ -216,9 +230,8 @@ public class GridSizeMigrationTask {
                     int newScreenId = LauncherSettings.Settings.call(
                             mContext.getContentResolver(),
                             LauncherSettings.Settings.METHOD_NEW_SCREEN_ID)
-                            .getInt(LauncherSettings.Settings.EXTRA_VALUE);
+                                .getInt(LauncherSettings.Settings.EXTRA_VALUE);
 
-                    allScreens.add(newScreenId);
                     for (DbEntry item : placement.finalPlacedItems) {
                         if (!mCarryOver.remove(itemMap.get(item.id))) {
                             throw new Exception("Unable to find matching items");
@@ -231,19 +244,6 @@ public class GridSizeMigrationTask {
                 }
 
             } while (!mCarryOver.isEmpty());
-
-            // Update screens
-            final Uri uri = LauncherSettings.WorkspaceScreens.CONTENT_URI;
-            mUpdateOperations.add(ContentProviderOperation.newDelete(uri).build());
-            int count = allScreens.size();
-            for (int i = 0; i < count; i++) {
-                ContentValues v = new ContentValues();
-                int screenId = allScreens.get(i);
-                v.put(LauncherSettings.WorkspaceScreens._ID, screenId);
-                v.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
-                mUpdateOperations.add(ContentProviderOperation.newInsert(uri).withValues(
-                        v).build());
-            }
         }
         return applyOperations();
     }
diff --git a/src/com/android/launcher3/model/LoaderCursor.java b/src/com/android/launcher3/model/LoaderCursor.java
index 8efebf14b..d104a8b22 100644
--- a/src/com/android/launcher3/model/LoaderCursor.java
+++ b/src/com/android/launcher3/model/LoaderCursor.java
@@ -380,7 +380,7 @@ public class LoaderCursor extends CursorWrapper {
      * otherwise marks it for deletion.
      */
     public void checkAndAddItem(ItemInfo info, BgDataModel dataModel) {
-        if (checkItemPlacement(info, dataModel.workspaceScreens)) {
+        if (checkItemPlacement(info)) {
             dataModel.addItem(mContext, info, false);
         } else {
             markDeleted("Item position overlap");
@@ -390,7 +390,7 @@ public class LoaderCursor extends CursorWrapper {
     /**
      * check & update map of what's occupied; used to discard overlapping/invalid items
      */
-    protected boolean checkItemPlacement(ItemInfo item, IntArray workspaceScreens) {
+    protected boolean checkItemPlacement(ItemInfo item) {
         int containerIndex = item.screenId;
         if (item.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT) {
             final GridOccupancy hotseatOccupancy =
@@ -420,12 +420,7 @@ public class LoaderCursor extends CursorWrapper {
                 occupied.put(LauncherSettings.Favorites.CONTAINER_HOTSEAT, occupancy);
                 return true;
             }
-        } else if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
-            if (!workspaceScreens.contains(item.screenId)) {
-                // The item has an invalid screen id.
-                return false;
-            }
-        } else {
+        } else if (item.container != LauncherSettings.Favorites.CONTAINER_DESKTOP) {
             // Skip further checking if it is not the hotseat or workspace container
             return true;
         }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index f711787cf..a4fe57086 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -72,6 +72,7 @@ import com.android.launcher3.shortcuts.ShortcutInfoCompat;
 import com.android.launcher3.shortcuts.ShortcutKey;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LooperIdleLock;
 import com.android.launcher3.util.MultiHashMap;
 import com.android.launcher3.util.PackageManagerHelper;
@@ -293,7 +294,6 @@ public class LoaderTask implements Runnable {
             final HashMap<String, SessionInfo> installingPkgs =
                     mPackageInstaller.updateAndGetActiveSessionCache();
             mFirstScreenBroadcast = new FirstScreenBroadcast(installingPkgs);
-            mBgDataModel.workspaceScreens.addAll(LauncherModel.loadWorkspaceScreensDb(context));
 
             Map<ShortcutKey, ShortcutInfoCompat> shortcutKeyToPinnedShortcuts = new HashMap<>();
             final LoaderCursor c = new LoaderCursor(contentResolver.query(
@@ -780,21 +780,15 @@ public class LoaderTask implements Runnable {
                         new Handler(LauncherModel.getWorkerLooper()));
             }
 
-            // Remove any empty screens
-            IntArray unusedScreens = mBgDataModel.workspaceScreens.clone();
+            // Initialize the screens array. Using an InstSet ensures that the screen ids
+            // are sorted.
+            IntSet screenSet = new IntSet();
             for (ItemInfo item: mBgDataModel.itemsIdMap) {
-                int screenId = item.screenId;
-                if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP &&
-                        unusedScreens.contains(screenId)) {
-                    unusedScreens.removeValue(screenId);
+                if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
+                    screenSet.add(item.screenId);
                 }
             }
-
-            // If there are any empty screens remove them, and update.
-            if (unusedScreens.size() != 0) {
-                mBgDataModel.workspaceScreens.removeAllValues(unusedScreens);
-                LauncherModel.updateWorkspaceScreenOrder(context, mBgDataModel.workspaceScreens);
-            }
+            mBgDataModel.workspaceScreens.addAll(screenSet.getArray());
         }
     }
 
diff --git a/src/com/android/launcher3/model/ModelWriter.java b/src/com/android/launcher3/model/ModelWriter.java
index f7961d5d8..ca5428d94 100644
--- a/src/com/android/launcher3/model/ModelWriter.java
+++ b/src/com/android/launcher3/model/ModelWriter.java
@@ -333,7 +333,7 @@ public class ModelWriter {
      * {@link #commitDelete()} is called (or abandoned if {@link #abortDelete()} is called).
      * Otherwise, we run the Runnable immediately.
      */
-    public void enqueueDeleteRunnable(Runnable r) {
+    private void enqueueDeleteRunnable(Runnable r) {
         if (mPreparingToUndo) {
             mDeleteRunnables.add(r);
         } else {
diff --git a/src/com/android/launcher3/provider/ImportDataTask.java b/src/com/android/launcher3/provider/ImportDataTask.java
index b389b145d..86fcc0680 100644
--- a/src/com/android/launcher3/provider/ImportDataTask.java
+++ b/src/com/android/launcher3/provider/ImportDataTask.java
@@ -33,7 +33,6 @@ import android.os.Process;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.SparseBooleanArray;
-import android.util.SparseIntArray;
 
 import com.android.launcher3.AutoInstallsLayout.LayoutParserCallback;
 import com.android.launcher3.DefaultLayoutParser;
@@ -43,7 +42,6 @@ import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.LauncherSettings.Settings;
-import com.android.launcher3.LauncherSettings.WorkspaceScreens;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.compat.UserManagerCompat;
@@ -72,7 +70,6 @@ public class ImportDataTask {
 
     private final Context mContext;
 
-    private final Uri mOtherScreensUri;
     private final Uri mOtherFavoritesUri;
 
     private int mHotseatSize;
@@ -81,41 +78,14 @@ public class ImportDataTask {
 
     private ImportDataTask(Context context, String sourceAuthority) {
         mContext = context;
-        mOtherScreensUri = Uri.parse("content://" +
-                sourceAuthority + "/" + WorkspaceScreens.TABLE_NAME);
         mOtherFavoritesUri = Uri.parse("content://" + sourceAuthority + "/" + Favorites.TABLE_NAME);
     }
 
     public boolean importWorkspace() throws Exception {
-        IntArray allScreens = LauncherDbUtils.getScreenIdsFromCursor(
-                mContext.getContentResolver().query(mOtherScreensUri, null, null, null,
-                        LauncherSettings.WorkspaceScreens.SCREEN_RANK));
         FileLog.d(TAG, "Importing DB from " + mOtherFavoritesUri);
 
-        // During import we reset the screen IDs to 0-indexed values.
-        if (allScreens.isEmpty()) {
-            // No thing to migrate
-            FileLog.e(TAG, "No data found to import");
-            return false;
-        }
-
         mHotseatSize = mMaxGridSizeX = mMaxGridSizeY = 0;
-
-        // Build screen update
-        ArrayList<ContentProviderOperation> screenOps = new ArrayList<>();
-        int count = allScreens.size();
-        SparseIntArray screenIdMap = new SparseIntArray(count);
-        for (int i = 0; i < count; i++) {
-            ContentValues v = new ContentValues();
-            v.put(LauncherSettings.WorkspaceScreens._ID, i);
-            v.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
-            screenIdMap.put(allScreens.get(i), i);
-            screenOps.add(ContentProviderOperation.newInsert(
-                    LauncherSettings.WorkspaceScreens.CONTENT_URI).withValues(v).build());
-        }
-        mContext.getContentResolver().applyBatch(LauncherProvider.AUTHORITY, screenOps);
-        importWorkspaceItems(allScreens.get(0), screenIdMap);
-
+        importWorkspaceItems();
         GridSizeMigrationTask.markForMigration(mContext, mMaxGridSizeX, mMaxGridSizeY, mHotseatSize);
 
         // Create empty DB flag.
@@ -129,17 +99,17 @@ public class ImportDataTask {
      * 2) For home screen entries, maps the screen id based on {@param screenIdMap}
      * 3) In the end fills any holes in hotseat with items from default hotseat layout.
      */
-    private void importWorkspaceItems(
-            int firstScreenId, SparseIntArray screenIdMap) throws Exception {
+    private void importWorkspaceItems() throws Exception {
         String profileId = Long.toString(UserManagerCompat.getInstance(mContext)
                 .getSerialNumberForUser(Process.myUserHandle()));
 
         boolean createEmptyRowOnFirstScreen;
         if (FeatureFlags.QSB_ON_FIRST_SCREEN.get()) {
             try (Cursor c = mContext.getContentResolver().query(mOtherFavoritesUri, null,
-                    // get items on the first row of the first screen
-                    "profileId = ? AND container = -100 AND screen = ? AND cellY = 0",
-                    new String[]{profileId, Integer.toString(firstScreenId)},
+                    // get items on the first row of the first screen (min screen id)
+                    "profileId = ? AND container = -100 AND cellY = 0 AND screen = " +
+                    "(SELECT MIN(screen) FROM favorites WHERE container = -100)",
+                    new String[]{profileId},
                     null)) {
                 // First row of first screen is not empty
                 createEmptyRowOnFirstScreen = c.moveToNext();
@@ -163,7 +133,7 @@ public class ImportDataTask {
                         Favorites.PROFILE_ID + " = ?", new String[]{profileId},
                         // Get the items sorted by container, so that the folders are loaded
                         // before the corresponding items.
-                        Favorites.CONTAINER)) {
+                        Favorites.CONTAINER + " , " + Favorites.SCREEN)) {
 
             // various columns we expect to exist.
             final int idIndex = c.getColumnIndexOrThrow(Favorites._ID);
@@ -185,6 +155,7 @@ public class ImportDataTask {
             SparseBooleanArray mValidFolders = new SparseBooleanArray();
             ContentValues values = new ContentValues();
 
+            Integer firstScreenId = null;
             while (c.moveToNext()) {
                 values.clear();
                 int id = c.getInt(idIndex);
@@ -201,16 +172,21 @@ public class ImportDataTask {
 
                 switch (container) {
                     case Favorites.CONTAINER_DESKTOP: {
-                        Integer newScreenId = screenIdMap.get(screen);
-                        if (newScreenId == null) {
-                            FileLog.d(TAG, String.format("Skipping item %d, type %d not on a valid screen %d", id, type, screen));
+                        if (screen < Workspace.FIRST_SCREEN_ID) {
+                            FileLog.d(TAG, String.format(
+                                    "Skipping item %d, type %d not on a valid screen %d",
+                                    id, type, screen));
                             continue;
                         }
+                        if (firstScreenId == null) {
+                            firstScreenId = screen;
+                        }
                         // Reset the screen to 0-index value
-                        screen = newScreenId;
-                        if (createEmptyRowOnFirstScreen && screen == Workspace.FIRST_SCREEN_ID) {
+                        if (createEmptyRowOnFirstScreen && firstScreenId.equals(screen)) {
                             // Shift items by 1.
                             cellY++;
+                            // Change the screen id to first screen
+                            screen = Workspace.FIRST_SCREEN_ID;
                         }
 
                         mMaxGridSizeX = Math.max(mMaxGridSizeX, cellX + spanX);
@@ -218,7 +194,7 @@ public class ImportDataTask {
                         break;
                     }
                     case Favorites.CONTAINER_HOTSEAT: {
-                        mHotseatSize = Math.max(mHotseatSize, (int) screen + 1);
+                        mHotseatSize = Math.max(mHotseatSize, screen + 1);
                         break;
                     }
                     default:
diff --git a/src/com/android/launcher3/provider/LauncherDbUtils.java b/src/com/android/launcher3/provider/LauncherDbUtils.java
index ab0703fed..b79478ac6 100644
--- a/src/com/android/launcher3/provider/LauncherDbUtils.java
+++ b/src/com/android/launcher3/provider/LauncherDbUtils.java
@@ -25,11 +25,9 @@ import android.util.Log;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.LauncherSettings.WorkspaceScreens;
 import com.android.launcher3.util.IntArray;
 
-import java.util.ArrayList;
-import java.util.Collection;
+import java.util.Locale;
 
 /**
  * A set of utility methods for Launcher DB used for DB updates and migration.
@@ -47,26 +45,25 @@ public class LauncherDbUtils {
      */
     public static boolean prepareScreenZeroToHostQsb(Context context, SQLiteDatabase db) {
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            // Get the existing screens
-            IntArray screenIds = getScreenIdsFromCursor(db.query(WorkspaceScreens.TABLE_NAME,
-                    null, null, null, null, null, WorkspaceScreens.SCREEN_RANK));
+            // Get the first screen
+            final int firstScreenId;
+            try (Cursor c = db.rawQuery(String.format(Locale.ENGLISH,
+                    "SELECT MIN(%1$s) from %2$s where %3$s = %4$d",
+                    Favorites.SCREEN, Favorites.TABLE_NAME, Favorites.CONTAINER,
+                    Favorites.CONTAINER_DESKTOP), null)) {
+
+                if (!c.moveToNext()) {
+                    // No update needed
+                    t.commit();
+                    return true;
+                }
 
-            if (screenIds.isEmpty()) {
-                // No update needed
-                t.commit();
-                return true;
+                firstScreenId = c.getInt(0);
             }
-            if (screenIds.get(0) != 0) {
-                // First screen is not 0, we need to rename screens
-                if (screenIds.contains(0)) {
-                    // There is already a screen 0. First rename it to a different screen.
-                    int newScreenId = 1;
-                    while (screenIds.contains(newScreenId)) newScreenId++;
-                    renameScreen(db, 0, newScreenId);
-                }
 
+            if (firstScreenId != 0) {
                 // Rename the first screen to 0.
-                renameScreen(db, screenIds.get(0), 0);
+                renameScreen(db, firstScreenId, 0);
             }
 
             // Check if the first row is empty
@@ -89,31 +86,19 @@ public class LauncherDbUtils {
 
     private static void renameScreen(SQLiteDatabase db, int oldScreen, int newScreen) {
         String[] whereParams = new String[] { Integer.toString(oldScreen) };
-
         ContentValues values = new ContentValues();
-        values.put(WorkspaceScreens._ID, newScreen);
-        db.update(WorkspaceScreens.TABLE_NAME, values, "_id = ?", whereParams);
-
-        values.clear();
         values.put(Favorites.SCREEN, newScreen);
         db.update(Favorites.TABLE_NAME, values, "container = -100 and screen = ?", whereParams);
     }
 
-    /**
-     * Parses the cursor containing workspace screens table and returns the list of screen IDs
-     */
-    public static IntArray getScreenIdsFromCursor(Cursor sc) {
-        try {
-            return iterateCursor(sc,
-                    sc.getColumnIndexOrThrow(WorkspaceScreens._ID), new IntArray());
-        } finally {
-            sc.close();
-        }
-    }
-
-    public static IntArray iterateCursor(Cursor c, int columnIndex, IntArray out) {
-        while (c.moveToNext()) {
-            out.add(c.getInt(columnIndex));
+    public static IntArray queryIntArray(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)) {
+            while (c.moveToNext()) {
+                out.add(c.getInt(0));
+            }
         }
         return out;
     }
diff --git a/src/com/android/launcher3/util/IntSet.java b/src/com/android/launcher3/util/IntSet.java
index 63499b06d..2459fad32 100644
--- a/src/com/android/launcher3/util/IntSet.java
+++ b/src/com/android/launcher3/util/IntSet.java
@@ -48,4 +48,15 @@ public class IntSet {
     public int size() {
         return mArray.size();
     }
+
+    public IntArray getArray() {
+        return mArray;
+    }
+
+    public static IntSet wrap(IntArray array) {
+        IntSet set = new IntSet();
+        set.mArray.addAll(array);
+        Arrays.sort(set.mArray.mValues, 0, set.mArray.mSize);
+        return set;
+    }
 }