2 files changed, 41 insertions, 10 deletions

diff --git a/src/com/android/launcher3/model/GridSizeMigrationTask.java b/src/com/android/launcher3/model/GridSizeMigrationTask.java
index 600768e5d..fd647c79a 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -974,6 +974,26 @@ public class GridSizeMigrationTask {
     }
 
     /**
+     * Removes any broken item from the hotseat.
+     * @return a map with occupied hotseat position set to non-null value.
+     */
+    public static LongArrayMap<Object> removeBrokenHotseatItems(Context context) throws Exception {
+        GridSizeMigrationTask task = new GridSizeMigrationTask(context,
+                LauncherAppState.getInstance().getInvariantDeviceProfile(),
+                getValidPackages(context), Integer.MAX_VALUE, Integer.MAX_VALUE);
+
+        // Load all the valid entries
+        ArrayList<DbEntry> items = task.loadHotseatEntries();
+        // Delete any entry marked for deletion by above load.
+        task.applyOperations();
+        LongArrayMap<Object> positions = new LongArrayMap<>();
+        for (DbEntry item : items) {
+            positions.put(item.screenId, item);
+        }
+        return positions;
+    }
+
+    /**
      * Task to run grid migration in multiple steps when the size difference is more than 1.
      */
     protected static class MultiStepMigrationTask {
diff --git a/src/com/android/launcher3/provider/ImportDataTask.java b/src/com/android/launcher3/provider/ImportDataTask.java
index 233c3edf1..5cb34e856 100644
--- a/src/com/android/launcher3/provider/ImportDataTask.java
+++ b/src/com/android/launcher3/provider/ImportDataTask.java
@@ -148,7 +148,6 @@ public class ImportDataTask {
 
         // Set of package names present in hotseat
         final HashSet<String> hotseatTargetApps = new HashSet<>();
-        final LongArrayMap<Intent> hotseatItems = new LongArrayMap<>();
         int maxId = 0;
 
         // Number of imported items on workspace and hotseat
@@ -270,7 +269,6 @@ public class ImportDataTask {
                     if (intent.getComponent() != null) {
                         intent.setPackage(intent.getComponent().getPackageName());
                     }
-                    hotseatItems.put(screen, intent);
                     hotseatTargetApps.add(getPackage(intent));
                 }
 
@@ -299,7 +297,13 @@ public class ImportDataTask {
         if (totalItemsOnWorkspace < MIN_ITEM_COUNT_FOR_SUCCESSFUL_MIGRATION) {
             throw new Exception("Insufficient data");
         }
+        if (!insertOperations.isEmpty()) {
+            mContext.getContentResolver().applyBatch(ProviderConfig.AUTHORITY,
+                    insertOperations);
+            insertOperations.clear();
+        }
 
+        LongArrayMap<Object> hotseatItems = GridSizeMigrationTask.removeBrokenHotseatItems(mContext);
         int myHotseatCount = LauncherAppState.getInstance().getInvariantDeviceProfile().numHotseatIcons;
         if (!FeatureFlags.NO_ALL_APPS_ICON) {
             myHotseatCount--;
@@ -307,14 +311,15 @@ public class ImportDataTask {
         if (hotseatItems.size() < myHotseatCount) {
             // Insufficient hotseat items. Add a few more.
             HotseatParserCallback parserCallback = new HotseatParserCallback(
-                    hotseatTargetApps, hotseatItems, insertOperations, maxId + 1);
+                    hotseatTargetApps, hotseatItems, insertOperations, maxId + 1, myHotseatCount);
             new HotseatLayoutParser(mContext,
                     parserCallback).loadLayout(null, new ArrayList<Long>());
             mHotseatSize = (int) hotseatItems.keyAt(hotseatItems.size() - 1) + 1;
-        }
-        if (!insertOperations.isEmpty()) {
-            mContext.getContentResolver().applyBatch(ProviderConfig.AUTHORITY,
-                    insertOperations);
+
+            if (!insertOperations.isEmpty()) {
+                mContext.getContentResolver().applyBatch(ProviderConfig.AUTHORITY,
+                        insertOperations);
+            }
         }
     }
 
@@ -404,16 +409,18 @@ public class ImportDataTask {
      */
     private static class HotseatParserCallback implements LayoutParserCallback {
         private final HashSet<String> mExisitingApps;
-        private final LongArrayMap<Intent> mExistingItems;
+        private final LongArrayMap<Object> mExistingItems;
         private final ArrayList<ContentProviderOperation> mOutOps;
+        private final int mRequiredSize;
         private int mStartItemId;
 
         HotseatParserCallback(
-                HashSet<String> existingApps, LongArrayMap<Intent> existingItems,
-                ArrayList<ContentProviderOperation> outOps, int startItemId) {
+                HashSet<String> existingApps, LongArrayMap<Object> existingItems,
+                ArrayList<ContentProviderOperation> outOps, int startItemId, int requiredSize) {
             mExisitingApps = existingApps;
             mExistingItems = existingItems;
             mOutOps = outOps;
+            mRequiredSize = requiredSize;
             mStartItemId = startItemId;
         }
 
@@ -424,6 +431,10 @@ public class ImportDataTask {
 
         @Override
         public long insertAndCheck(SQLiteDatabase db, ContentValues values) {
+            if (mExistingItems.size() >= mRequiredSize) {
+                // No need to add more items.
+                return 0;
+            }
             Intent intent;
             try {
                 intent = Intent.parseUri(values.getAsString(Favorites.INTENT), 0);