Ensure app widget ids are restored after database is sanitized.

Previously, it was possible for AppWidgetsRestoredReceiver to
start the restore process before work profile has finished restoring which
resulted in the work profile items being removed from the workspace.

Bug: 131315856
Change-Id: I2f295a1ca91f1996522bcc8052aa139979526e3b
diff --git a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
index 0250c36..d949141 100644
--- a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
+++ b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
@@ -8,7 +8,6 @@
 import android.content.Context;
 import android.content.Intent;
 import android.database.Cursor;
-import android.os.Handler;
 import android.util.Log;
 
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -34,16 +33,8 @@
 
             final int[] oldIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS);
             final int[] newIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS);
-            if (oldIds.length == newIds.length) {
-                final PendingResult asyncResult = goAsync();
-                new Handler(LauncherModel.getWorkerLooper())
-                        .postAtFrontOfQueue(new Runnable() {
-                            @Override
-                            public void run() {
-                                restoreAppWidgetIds(context, oldIds, newIds);
-                                asyncResult.finish();
-                            }
-                        });
+            if (oldIds != null && newIds != null && oldIds.length == newIds.length) {
+                RestoreDbTask.setRestoredAppWidgetIds(context, oldIds, newIds);
             } else {
                 Log.e(TAG, "Invalid host restored received");
             }
@@ -54,7 +45,7 @@
      * Updates the app widgets whose id has changed during the restore process.
      */
     @WorkerThread
-    static void restoreAppWidgetIds(Context context, int[] oldWidgetIds, int[] newWidgetIds) {
+    public static void restoreAppWidgetIds(Context context, int[] oldWidgetIds, int[] newWidgetIds) {
         AppWidgetHost appWidgetHost = new LauncherAppWidgetHost(context);
         if (FeatureFlags.GO_DISABLE_WIDGETS) {
             Log.e(TAG, "Skipping widget ID remap as widgets not supported");
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index f830301..6ad5c36 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -158,7 +158,8 @@
             mOpenHelper = new DatabaseHelper(getContext(), mListenerHandler);
 
             if (RestoreDbTask.isPending(getContext())) {
-                if (!RestoreDbTask.performRestore(mOpenHelper, new BackupManager(getContext()))) {
+                if (!RestoreDbTask.performRestore(getContext(), mOpenHelper,
+                        new BackupManager(getContext()))) {
                     mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
                 }
                 // Set is pending to false irrespective of the result, so that it doesn't get
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 5cfd95c..02fc84b 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -74,6 +74,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Locale;
+import java.util.StringTokenizer;
 import java.util.concurrent.Executor;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -652,4 +653,23 @@
             return null;
         }
     }
+
+    public static int[] getIntArrayFromString(String tokenized) {
+        StringTokenizer tokenizer = new StringTokenizer(tokenized, ",");
+        int[] array = new int[tokenizer.countTokens()];
+        int count = 0;
+        while (tokenizer.hasMoreTokens()) {
+            array[count] = Integer.parseInt(tokenizer.nextToken());
+            count++;
+        }
+        return array;
+    }
+
+    public static String getStringFromIntArray(int[] array) {
+        StringBuilder str = new StringBuilder();
+        for (int value : array) {
+            str.append(value).append(",");
+        }
+        return str.toString();
+    }
 }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 053c493..3c0c5fd 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -16,18 +16,23 @@
 
 package com.android.launcher3.provider;
 
+import static com.android.launcher3.Utilities.getIntArrayFromString;
+import static com.android.launcher3.Utilities.getStringFromIntArray;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 
 import android.app.backup.BackupManager;
 import android.content.ContentValues;
 import android.content.Context;
+import android.content.SharedPreferences;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
-import android.os.Build;
 import android.os.UserHandle;
 import android.util.LongSparseArray;
 import android.util.SparseLongArray;
 
+import androidx.annotation.NonNull;
+
+import com.android.launcher3.AppWidgetsRestoredReceiver;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -53,10 +58,16 @@
     private static final String INFO_COLUMN_NAME = "name";
     private static final String INFO_COLUMN_DEFAULT_VALUE = "dflt_value";
 
-    public static boolean performRestore(DatabaseHelper helper, BackupManager backupManager) {
+    private static final String APPWIDGET_OLD_IDS = "appwidget_old_ids";
+    private static final String APPWIDGET_IDS = "appwidget_ids";
+
+    public static boolean performRestore(Context context, DatabaseHelper helper,
+            BackupManager backupManager) {
         SQLiteDatabase db = helper.getWritableDatabase();
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            new RestoreDbTask().sanitizeDB(helper, db, backupManager);
+            RestoreDbTask task = new RestoreDbTask();
+            task.sanitizeDB(helper, db, backupManager);
+            task.restoreAppWidgetIdsIfExists(context);
             t.commit();
             return true;
         } catch (Exception e) {
@@ -230,4 +241,27 @@
         FileLog.d(TAG, "Restore data received through full backup " + isPending);
         Utilities.getPrefs(context).edit().putBoolean(RESTORE_TASK_PENDING, isPending).commit();
     }
+
+    private void restoreAppWidgetIdsIfExists(Context context) {
+        SharedPreferences prefs = Utilities.getPrefs(context);
+        if (prefs.contains(APPWIDGET_OLD_IDS) && prefs.contains(APPWIDGET_IDS)) {
+            AppWidgetsRestoredReceiver.restoreAppWidgetIds(context,
+                    getIntArrayFromString(prefs.getString(APPWIDGET_OLD_IDS, "")),
+                    getIntArrayFromString(prefs.getString(APPWIDGET_IDS, "")));
+        } else {
+            FileLog.d(TAG, "No app widget ids to restore.");
+        }
+
+        prefs.edit().remove(APPWIDGET_OLD_IDS)
+                .remove(APPWIDGET_IDS).apply();
+    }
+
+    public static void setRestoredAppWidgetIds(Context context, @NonNull int[] oldIds,
+            @NonNull int[] newIds) {
+        Utilities.getPrefs(context).edit()
+                .putString(APPWIDGET_OLD_IDS, getStringFromIntArray(oldIds))
+                .putString(APPWIDGET_IDS, getStringFromIntArray(newIds))
+                .commit();
+    }
+
 }