Initial support for restore workspace from last stable db entry.

(see go/play-launcher-plan-launcher-implementation)

1. When Launcher launches for the first time, creates a backup
   of the workspace before sanitizing db entries.
2. Creates a new path in LauncherProvider that triggers workspace
   restore using last stable db entry of the same grid size.
3. When restore from backup created this way, the table will be
   sanitized afterward.

Test:
1. apply on master, build & refresh on physical device
2. factory reset, go through SuW and perform restore
3. exit SuW without signing into Work Profile
4. run following commands in console
adb root
adb remount
adb pull
/data/data/com.google.android.apps.nexuslauncher/databases/launcher.db
sqlite3 ./launcher.db
.tables
SELECT * FROM favorites_bakup;

Bug: 141472083
Change-Id: I8032866a97eb333946d4f62352595d180364126b
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 900c966..b589560 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -387,6 +387,11 @@
                         tableExists(mOpenHelper.getReadableDatabase(), Favorites.BACKUP_TABLE_NAME);
                 return null;
             }
+            case LauncherSettings.Settings.METHOD_RESTORE_BACKUP_TABLE: {
+                RestoreDbTask.restoreIfPossible(
+                        getContext(), mOpenHelper, new BackupManager(getContext()));
+                return null;
+            }
         }
         return null;
     }
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index ec307db..4c5c61c 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -300,6 +300,8 @@
 
         public static final String METHOD_REFRESH_BACKUP_TABLE = "refresh_backup_table";
 
+        public static final String METHOD_RESTORE_BACKUP_TABLE = "restore_backup_table";
+
         public static final String EXTRA_VALUE = "value";
 
         public static Bundle call(ContentResolver cr, String method) {
diff --git a/src/com/android/launcher3/model/GridBackupTable.java b/src/com/android/launcher3/model/GridBackupTable.java
index 11d4edd..fc9948e 100644
--- a/src/com/android/launcher3/model/GridBackupTable.java
+++ b/src/com/android/launcher3/model/GridBackupTable.java
@@ -27,6 +27,8 @@
 import android.os.Process;
 import android.util.Log;
 
+import androidx.annotation.IntDef;
+
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.LauncherSettings.Settings;
 import com.android.launcher3.pm.UserCache;
@@ -45,6 +47,19 @@
     private static final String KEY_GRID_Y_SIZE = Favorites.SPANY;
     private static final String KEY_DB_VERSION = Favorites.RANK;
 
+    public static final int OPTION_REQUIRES_SANITIZATION = 1;
+
+    /** STATE_NOT_FOUND indicates backup doesn't exist in the db. */
+    private static final int STATE_NOT_FOUND = 0;
+    /**
+     *  STATE_RAW indicates the backup has not yet been sanitized. This implies it might still
+     *  posses app info that doesn't exist in the workspace and needed to be sanitized before
+     *  put into use.
+     */
+    private static final int STATE_RAW = 1;
+    /** STATE_SANITIZED indicates the backup has already been sanitized, thus can be used as-is. */
+    private static final int STATE_SANITIZED = 2;
+
     private final Context mContext;
     private final SQLiteDatabase mDb;
 
@@ -56,6 +71,9 @@
     private int mRestoredGridX;
     private int mRestoredGridY;
 
+    @IntDef({STATE_NOT_FOUND, STATE_RAW, STATE_SANITIZED})
+    private @interface BackupState { }
+
     public GridBackupTable(Context context, SQLiteDatabase db,
             int hotseatSize, int gridX, int gridY) {
         mContext = context;
@@ -66,6 +84,10 @@
         mOldGridY = gridY;
     }
 
+    /**
+     * Create a backup from current workspace layout if one isn't created already (Note backup
+     * created this way is always sanitized). Otherwise restore from the backup instead.
+     */
     public boolean backupOrRestoreAsNeeded() {
         // Check if backup table exists
         if (!tableExists(mDb, BACKUP_TABLE_NAME)) {
@@ -74,16 +96,16 @@
                 // No need to copy if empty DB was created.
                 return false;
             }
-
-            copyTable(Favorites.TABLE_NAME, BACKUP_TABLE_NAME);
-            encodeDBProperties();
+            doBackup(UserCache.INSTANCE.get(mContext).getSerialNumberForUser(
+                    Process.myUserHandle()), 0);
             return false;
         }
-
-        if (!loadDbProperties()) {
+        if (loadDBProperties() != STATE_SANITIZED) {
             return false;
         }
-        copyTable(BACKUP_TABLE_NAME, Favorites.TABLE_NAME);
+        long userSerial = UserCache.INSTANCE.get(mContext).getSerialNumberForUser(
+                Process.myUserHandle());
+        copyTable(BACKUP_TABLE_NAME, Favorites.TABLE_NAME, userSerial);
         Log.d(TAG, "Backup table found");
         return true;
     }
@@ -93,43 +115,84 @@
         return mRestoredHotseatSize;
     }
 
-    private void copyTable(String from, String to) {
-        long userSerial = UserCache.INSTANCE.get(mContext).getSerialNumberForUser(
-                Process.myUserHandle());
+    /**
+     * Copy valid grid entries from one table to another.
+     */
+    private void copyTable(String from, String to, long userSerial) {
         dropTable(mDb, to);
         Favorites.addTableToDb(mDb, userSerial, false, to);
         mDb.execSQL("INSERT INTO " + to + " SELECT * FROM " + from + " where _id > " + ID_PROPERTY);
     }
 
-    private void encodeDBProperties() {
+    private void encodeDBProperties(int options) {
         ContentValues values = new ContentValues();
         values.put(Favorites._ID, ID_PROPERTY);
         values.put(KEY_DB_VERSION, mDb.getVersion());
         values.put(KEY_GRID_X_SIZE, mOldGridX);
         values.put(KEY_GRID_Y_SIZE, mOldGridY);
         values.put(KEY_HOTSEAT_SIZE, mOldHotseatSize);
+        values.put(Favorites.OPTIONS, options);
         mDb.insert(BACKUP_TABLE_NAME, null, values);
     }
 
-    private boolean loadDbProperties() {
+    /**
+     * Load DB properties from grid backup table.
+     */
+    public @BackupState int loadDBProperties() {
         try (Cursor c = mDb.query(BACKUP_TABLE_NAME, new String[] {
-                        KEY_DB_VERSION,     // 0
-                        KEY_GRID_X_SIZE,    // 1
-                        KEY_GRID_Y_SIZE,    // 2
-                        KEY_HOTSEAT_SIZE},  // 3
+                KEY_DB_VERSION,     // 0
+                KEY_GRID_X_SIZE,    // 1
+                KEY_GRID_Y_SIZE,    // 2
+                KEY_HOTSEAT_SIZE,   // 3
+                Favorites.OPTIONS}, // 4
                 "_id=" + ID_PROPERTY, null, null, null, null)) {
             if (!c.moveToNext()) {
                 Log.e(TAG, "Meta data not found in backup table");
-                return false;
+                return STATE_NOT_FOUND;
             }
-            if (mDb.getVersion() != c.getInt(0)) {
-                return false;
+            if (!validateDBVersion(mDb.getVersion(), c.getInt(0))) {
+                return STATE_NOT_FOUND;
             }
 
             mRestoredGridX = c.getInt(1);
             mRestoredGridY = c.getInt(2);
             mRestoredHotseatSize = c.getInt(3);
-            return true;
+            boolean isSanitized = (c.getInt(4) & OPTION_REQUIRES_SANITIZATION) == 0;
+            return isSanitized ? STATE_SANITIZED : STATE_RAW;
         }
     }
+
+    /**
+     * Restore workspace from raw backup if available.
+     */
+    public boolean restoreFromRawBackupIfAvailable(long oldProfileId) {
+        if (!tableExists(mDb, Favorites.BACKUP_TABLE_NAME)
+                || loadDBProperties() != STATE_RAW
+                || mOldHotseatSize != mRestoredHotseatSize
+                || mOldGridX != mRestoredGridX
+                || mOldGridY != mRestoredGridY) {
+            // skip restore if dimensions in backup table differs from current setup.
+            return false;
+        }
+        copyTable(Favorites.BACKUP_TABLE_NAME, Favorites.TABLE_NAME, oldProfileId);
+        Log.d(TAG, "Backup restored");
+        return true;
+    }
+
+    /**
+     * Performs a backup on the workspace layout.
+     */
+    public void doBackup(long profileId, int options) {
+        copyTable(Favorites.TABLE_NAME, Favorites.BACKUP_TABLE_NAME, profileId);
+        encodeDBProperties(options);
+    }
+
+    private static boolean validateDBVersion(int expected, int actual) {
+        if (expected != actual) {
+            Log.e(TAG, String.format("Launcher.db version mismatch, expecting %d but %d was found",
+                    expected, actual));
+            return false;
+        }
+        return true;
+    }
 }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 5037c9d..7ee30cc 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -32,12 +32,15 @@
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.AppWidgetsRestoredReceiver;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.WorkspaceItemInfo;
 import com.android.launcher3.logging.FileLog;
+import com.android.launcher3.model.GridBackupTable;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.LogConfig;
@@ -67,6 +70,7 @@
         SQLiteDatabase db = helper.getWritableDatabase();
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
             RestoreDbTask task = new RestoreDbTask();
+            task.backupWorkspace(context, db);
             task.sanitizeDB(helper, db, backupManager);
             task.restoreAppWidgetIdsIfExists(context);
             t.commit();
@@ -78,6 +82,45 @@
     }
 
     /**
+     * Restore the workspace if backup is available.
+     */
+    public static boolean restoreIfPossible(@NonNull Context context,
+            @NonNull DatabaseHelper helper, @NonNull BackupManager backupManager) {
+        final SQLiteDatabase db = helper.getWritableDatabase();
+        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
+            RestoreDbTask task = new RestoreDbTask();
+            task.restoreWorkspace(context, db, helper, backupManager);
+            task.restoreAppWidgetIdsIfExists(context);
+            t.commit();
+            return true;
+        } catch (Exception e) {
+            FileLog.e(TAG, "Failed to restore db", e);
+            return false;
+        }
+    }
+
+    /**
+     * Backup the workspace so that if things go south in restore, we can recover these entries.
+     */
+    private void backupWorkspace(Context context, SQLiteDatabase db) throws Exception {
+        InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
+        new GridBackupTable(context, db, idp.numHotseatIcons, idp.numColumns, idp.numRows)
+                .doBackup(getDefaultProfileId(db), GridBackupTable.OPTION_REQUIRES_SANITIZATION);
+    }
+
+    private void restoreWorkspace(@NonNull Context context, @NonNull SQLiteDatabase db,
+            @NonNull DatabaseHelper helper, @NonNull BackupManager backupManager)
+            throws Exception {
+        final InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
+        GridBackupTable backupTable = new GridBackupTable(context, db,
+                idp.numHotseatIcons, idp.numColumns, idp.numRows);
+        if (backupTable.restoreFromRawBackupIfAvailable(getDefaultProfileId(db))) {
+            sanitizeDB(helper, db, backupManager);
+            LauncherAppState.getInstance(context).getModel().forceReload();
+        }
+    }
+
+    /**
      * Makes the following changes in the provider DB.
      *   1. Removes all entries belonging to any profiles that were not restored.
      *   2. Marks all entries as restored. The flags are updated during first load or as
@@ -126,15 +169,16 @@
         db.update(Favorites.TABLE_NAME, values, null, null);
 
         // Mark widgets with appropriate restore flag.
-        values.put(Favorites.RESTORED,  LauncherAppWidgetInfo.FLAG_ID_NOT_VALID |
-                LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY |
-                LauncherAppWidgetInfo.FLAG_UI_NOT_READY |
-                (keepAllIcons ? LauncherAppWidgetInfo.FLAG_RESTORE_STARTED : 0));
+        values.put(Favorites.RESTORED,  LauncherAppWidgetInfo.FLAG_ID_NOT_VALID
+                | LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY
+                | LauncherAppWidgetInfo.FLAG_UI_NOT_READY
+                | (keepAllIcons ? LauncherAppWidgetInfo.FLAG_RESTORE_STARTED : 0));
         db.update(Favorites.TABLE_NAME, values, "itemType = ?",
                 new String[]{Integer.toString(Favorites.ITEM_TYPE_APPWIDGET)});
 
-        // Migrate ids. To avoid any overlap, we initially move conflicting ids to a temp location.
-        // Using Long.MIN_VALUE since profile ids can not be negative, so there will be no overlap.
+        // Migrate ids. To avoid any overlap, we initially move conflicting ids to a temp
+        // location. Using Long.MIN_VALUE since profile ids can not be negative, so there will
+        // be no overlap.
         final long tempLocationOffset = Long.MIN_VALUE;
         SparseLongArray tempMigratedIds = new SparseLongArray(profileMapping.size());
         int numTempMigrations = 0;
@@ -192,10 +236,10 @@
     private LongSparseArray<Long> getManagedProfileIds(SQLiteDatabase db, long defaultProfileId) {
         LongSparseArray<Long> ids = new LongSparseArray<>();
         try (Cursor c = db.rawQuery("SELECT profileId from favorites WHERE profileId != ? "
-                + "GROUP BY profileId", new String[] {Long.toString(defaultProfileId)})){
-                while (c.moveToNext()) {
-                    ids.put(c.getLong(c.getColumnIndex(Favorites.PROFILE_ID)), null);
-                }
+                + "GROUP BY profileId", new String[] {Long.toString(defaultProfileId)})) {
+            while (c.moveToNext()) {
+                ids.put(c.getLong(c.getColumnIndex(Favorites.PROFILE_ID)), null);
+            }
         }
         return ids;
     }
@@ -216,7 +260,7 @@
      * Returns the profile id used in the favorites table of the provided db.
      */
     protected long getDefaultProfileId(SQLiteDatabase db) throws Exception {
-        try (Cursor c = db.rawQuery("PRAGMA table_info (favorites)", null)){
+        try (Cursor c = db.rawQuery("PRAGMA table_info (favorites)", null)) {
             int nameIndex = c.getColumnIndex(INFO_COLUMN_NAME);
             while (c.moveToNext()) {
                 if (Favorites.PROFILE_ID.equals(c.getString(nameIndex))) {