1 files changed, 13 insertions, 3 deletions
diff --git a/src/com/android/launcher3/LauncherBackupAgentHelper.java b/src/com/android/launcher3/LauncherBackupAgentHelper.java
index a92a889f9..8eb4e6369 100644
--- a/src/com/android/launcher3/LauncherBackupAgentHelper.java
+++ b/src/com/android/launcher3/LauncherBackupAgentHelper.java
@@ -24,6 +24,8 @@ import android.database.Cursor;
 import android.os.ParcelFileDescriptor;
 import android.util.Log;
 
+import com.android.launcher3.model.MigrateFromRestoreTask;
+
 import java.io.IOException;
 
 public class LauncherBackupAgentHelper extends BackupAgentHelper {
@@ -63,7 +65,7 @@ public class LauncherBackupAgentHelper extends BackupAgentHelper {
     @Override
     public void onRestore(BackupDataInput data, int appVersionCode, ParcelFileDescriptor newState)
             throws IOException {
-        if (!Utilities.isLmpOrAbove()) {
+        if (!Utilities.ATLEAST_LOLLIPOP) {
             // No restore for old devices.
             Log.i(TAG, "You shall not pass!!!");
             Log.d(TAG, "Restore is only supported on devices running Lollipop and above.");
@@ -91,11 +93,19 @@ public class LauncherBackupAgentHelper extends BackupAgentHelper {
             LauncherAppState.getLauncherProvider().clearFlagEmptyDbCreated();
             LauncherClings.synchonouslyMarkFirstRunClingDismissed(this);
 
-            // TODO: Update the backup set to include rank.
+            // Rank was added in v4.
             if (mHelper.restoredBackupVersion <= 3) {
                 LauncherAppState.getLauncherProvider().updateFolderItemsRank();
-                LauncherAppState.getLauncherProvider().convertShortcutsToLauncherActivities();
             }
+
+            if (MigrateFromRestoreTask.ENABLED && mHelper.shouldAttemptWorkspaceMigration()) {
+                MigrateFromRestoreTask.markForMigration(getApplicationContext(),
+                        (int) mHelper.migrationCompatibleProfileData.desktopCols,
+                        (int) mHelper.migrationCompatibleProfileData.desktopRows,
+                        mHelper.widgetSizes);
+            }
+
+            LauncherAppState.getLauncherProvider().convertShortcutsToLauncherActivities();
         } else {
             if (VERBOSE) Log.v(TAG, "Nothing was restored, clearing DB");
             LauncherAppState.getLauncherProvider().createEmptyDB();