Adding support for storing container based item list in the model

These items get updated automatically during various model tasks.
Also simplifying the pinned shortcut state, by calculating the
list of ppined shortcut on demand, instead of storing a refCount.

Bug: 160748731
Change-Id: I3169d293552b05b4f4d6c529397fbc761887a282
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 8c0a2d7..f3cc164 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2495,7 +2495,7 @@
      * @param updated list of shortcuts which have changed.
      */
     @Override
-    public void bindWorkspaceItemsChanged(ArrayList<WorkspaceItemInfo> updated) {
+    public void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) {
         if (!updated.isEmpty()) {
             mWorkspace.updateShortcuts(updated);
         }
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index a6283ff..15e0daa 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -115,6 +115,7 @@
 
 import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.List;
 import java.util.function.Predicate;
 
 /**
@@ -3087,7 +3088,7 @@
         return false;
     }
 
-    void updateShortcuts(ArrayList<WorkspaceItemInfo> shortcuts) {
+    void updateShortcuts(List<WorkspaceItemInfo> shortcuts) {
         final HashSet<WorkspaceItemInfo> updates = new HashSet<>(shortcuts);
         ItemOperator op = (info, v) -> {
             if (v instanceof BubbleTextView && updates.contains(info)) {
diff --git a/src/com/android/launcher3/model/BaseModelUpdateTask.java b/src/com/android/launcher3/model/BaseModelUpdateTask.java
index 9013cba..d1e5017 100644
--- a/src/com/android/launcher3/model/BaseModelUpdateTask.java
+++ b/src/com/android/launcher3/model/BaseModelUpdateTask.java
@@ -22,7 +22,9 @@
 import com.android.launcher3.LauncherModel.CallbackTask;
 import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.model.BgDataModel.Callbacks;
+import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.ItemInfoMatcher;
@@ -30,7 +32,10 @@
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.Executor;
+import java.util.stream.Collectors;
 
 /**
  * Extension of {@link ModelUpdateTask} with some utility methods
@@ -88,11 +93,27 @@
         return mModel.getWriter(false /* hasVerticalHotseat */, false /* verifyChanges */);
     }
 
-
-    public void bindUpdatedWorkspaceItems(final ArrayList<WorkspaceItemInfo> updatedShortcuts) {
-        if (!updatedShortcuts.isEmpty()) {
-            scheduleCallbackTask(c -> c.bindWorkspaceItemsChanged(updatedShortcuts));
+    public void bindUpdatedWorkspaceItems(List<WorkspaceItemInfo> allUpdates) {
+        // Bind workspace items
+        List<WorkspaceItemInfo> workspaceUpdates = allUpdates.stream()
+                .filter(info -> info.id != ItemInfo.NO_ID)
+                .collect(Collectors.toList());
+        if (!workspaceUpdates.isEmpty()) {
+            scheduleCallbackTask(c -> c.bindWorkspaceItemsChanged(workspaceUpdates));
         }
+
+        // Bind extra items if any
+        allUpdates.stream()
+                .mapToInt(info -> info.container)
+                .distinct()
+                .mapToObj(mDataModel.extraItems::get)
+                .filter(Objects::nonNull)
+                .forEach(this::bindExtraContainerItems);
+    }
+
+    public void bindExtraContainerItems(FixedContainerItems item) {
+        FixedContainerItems copy = item.clone();
+        scheduleCallbackTask(c -> c.bindExtraContainerItems(copy));
     }
 
     public void bindDeepShortcuts(BgDataModel dataModel) {
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index 7524920..dfdc138 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -15,19 +15,25 @@
  */
 package com.android.launcher3.model;
 
+import static android.content.pm.LauncherApps.ShortcutQuery.FLAG_GET_KEY_FIELDS_ONLY;
+
 import static com.android.launcher3.model.WidgetsModel.GO_DISABLE_WIDGETS;
 import static com.android.launcher3.shortcuts.ShortcutRequest.PINNED;
 
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.mapping;
+
 import android.content.Context;
 import android.content.pm.LauncherApps;
 import android.content.pm.ShortcutInfo;
 import android.os.UserHandle;
 import android.text.TextUtils;
+import android.util.ArraySet;
 import android.util.Log;
-import android.util.MutableInt;
 
 import com.android.launcher3.InstallShortcutReceiver;
 import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.data.AppInfo;
@@ -36,8 +42,10 @@
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.PromiseAppInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.shortcuts.ShortcutKey;
 import com.android.launcher3.shortcuts.ShortcutRequest;
+import com.android.launcher3.shortcuts.ShortcutRequest.QueryResult;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
@@ -50,14 +58,16 @@
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.function.BiConsumer;
+import java.util.Set;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * All the data stored in-memory and managed by the LauncherModel
@@ -89,9 +99,9 @@
     public final IntSparseArrayMap<FolderInfo> folders = new IntSparseArrayMap<>();
 
     /**
-     * Map of ShortcutKey to the number of times it is pinned.
+     * Extra container based items
      */
-    public final Map<ShortcutKey, MutableInt> pinnedShortcutCounts = new HashMap<>();
+    public final IntSparseArrayMap<FixedContainerItems> extraItems = new IntSparseArrayMap<>();
 
     /**
      * List of all cached predicted items visible on home screen
@@ -128,8 +138,8 @@
         appWidgets.clear();
         folders.clear();
         itemsIdMap.clear();
-        pinnedShortcutCounts.clear();
         deepShortcutMap.clear();
+        extraItems.clear();
     }
 
     /**
@@ -182,6 +192,7 @@
     }
 
     public synchronized void removeItem(Context context, Iterable<? extends ItemInfo> items) {
+        ArraySet<UserHandle> updatedDeepShortcuts = new ArraySet<>();
         for (ItemInfo item : items) {
             switch (item.itemType) {
                 case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
@@ -200,14 +211,7 @@
                     workspaceItems.remove(item);
                     break;
                 case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT: {
-                    // Decrement pinned shortcut count
-                    ShortcutKey pinnedShortcut = ShortcutKey.fromItemInfo(item);
-                    MutableInt count = pinnedShortcutCounts.get(pinnedShortcut);
-                    if ((count == null || --count.value == 0)
-                            && !InstallShortcutReceiver.getPendingShortcuts(context)
-                                .contains(pinnedShortcut)) {
-                        unpinShortcut(context, pinnedShortcut);
-                    }
+                    updatedDeepShortcuts.add(item.user);
                     // Fall through.
                 }
                 case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
@@ -221,6 +225,7 @@
             }
             itemsIdMap.remove(item.id);
         }
+        updatedDeepShortcuts.forEach(user -> updateShortcutPinnedState(context, user));
     }
 
     public synchronized void addItem(Context context, ItemInfo item, boolean newItem) {
@@ -230,23 +235,7 @@
                 folders.put(item.id, (FolderInfo) item);
                 workspaceItems.add(item);
                 break;
-            case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT: {
-                // Increment the count for the given shortcut
-                ShortcutKey pinnedShortcut = ShortcutKey.fromItemInfo(item);
-                MutableInt count = pinnedShortcutCounts.get(pinnedShortcut);
-                if (count == null) {
-                    count = new MutableInt(1);
-                    pinnedShortcutCounts.put(pinnedShortcut, count);
-                } else {
-                    count.value++;
-                }
-
-                // Since this is a new item, pin the shortcut in the system server.
-                if (newItem && count.value == 1) {
-                    updatePinnedShortcuts(context, pinnedShortcut, List::add);
-                }
-                // Fall through
-            }
+            case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT:
             case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
             case LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT:
                 if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP ||
@@ -271,36 +260,87 @@
                 appWidgets.add((LauncherAppWidgetInfo) item);
                 break;
         }
+        if (newItem && item.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
+            updateShortcutPinnedState(context, item.user);
+        }
     }
 
     /**
-     * Removes the given shortcut from the current list of pinned shortcuts.
-     * (Runs on background thread)
+     * Updates the deep shortucts state in system to match out internal model, pinning any missing
+     * shortcuts and unpinning any extra shortcuts.
      */
-    public void unpinShortcut(Context context, ShortcutKey key) {
-        updatePinnedShortcuts(context, key, List::remove);
+    public void updateShortcutPinnedState(Context context) {
+        for (UserHandle user : UserCache.INSTANCE.get(context).getUserProfiles()) {
+            updateShortcutPinnedState(context, user);
+        }
     }
 
-    private void updatePinnedShortcuts(Context context, ShortcutKey key,
-            BiConsumer<List<String>, String> idOp) {
+    /**
+     * Updates the deep shortucts state in system to match out internal model, pinning any missing
+     * shortcuts and unpinning any extra shortcuts.
+     */
+    public synchronized void updateShortcutPinnedState(Context context, UserHandle user) {
         if (GO_DISABLE_WIDGETS) {
             return;
         }
-        String packageName = key.componentName.getPackageName();
-        String id = key.getId();
-        UserHandle user = key.user;
-        List<String> pinnedIds = new ShortcutRequest(context, user)
-                .forPackage(packageName)
-                .query(PINNED)
-                .stream()
-                .map(ShortcutInfo::getId)
-                .collect(Collectors.toCollection(ArrayList::new));
-        idOp.accept(pinnedIds, id);
-        try {
-            context.getSystemService(LauncherApps.class).pinShortcuts(packageName, pinnedIds, user);
-        } catch (SecurityException | IllegalStateException e) {
-            Log.w(TAG, "Failed to pin shortcut", e);
+
+        // Collect all system shortcuts
+        QueryResult result = new ShortcutRequest(context, user)
+                .query(PINNED | FLAG_GET_KEY_FIELDS_ONLY);
+        if (!result.wasSuccess()) {
+            return;
         }
+        // Map of packageName to shortcutIds that are currently in the system
+        Map<String, Set<String>> systemMap = result.stream()
+                .collect(groupingBy(ShortcutInfo::getPackage,
+                        mapping(ShortcutInfo::getId, Collectors.toSet())));
+
+        // Collect all model shortcuts
+        Stream.Builder<WorkspaceItemInfo> itemStream = Stream.builder();
+        forAllWorkspaceItemInfos(user, itemStream::accept);
+        // Map of packageName to shortcutIds that are currently in our model
+        Map<String, Set<String>> modelMap = Stream.concat(
+                    // Model shortcuts
+                    itemStream.build()
+                        .filter(wi -> wi.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT)
+                        .map(ShortcutKey::fromItemInfo),
+                    // Pending shortcuts
+                    InstallShortcutReceiver.getPendingShortcuts(context)
+                        .stream().filter(si -> si.user.equals(user)))
+                .collect(groupingBy(ShortcutKey::getPackageName,
+                        mapping(ShortcutKey::getId, Collectors.toSet())));
+
+        // Check for diff
+        for (Map.Entry<String, Set<String>> entry : modelMap.entrySet()) {
+            Set<String> modelShortcuts = entry.getValue();
+            Set<String> systemShortcuts = systemMap.remove(entry.getKey());
+            if (systemShortcuts == null) {
+                systemShortcuts = Collections.emptySet();
+            }
+
+            // Do not use .equals as it can vary based on the type of set
+            if (systemShortcuts.size() != modelShortcuts.size()
+                    || !systemShortcuts.containsAll(modelShortcuts)) {
+                // Update system state for this package
+                try {
+                    context.getSystemService(LauncherApps.class).pinShortcuts(
+                            entry.getKey(), new ArrayList<>(modelShortcuts), user);
+                } catch (SecurityException | IllegalStateException e) {
+                    Log.w(TAG, "Failed to pin shortcut", e);
+                }
+            }
+        }
+
+        // If there are any extra pinned shortcuts, remove them
+        systemMap.keySet().forEach(packageName -> {
+            // Update system state
+            try {
+                context.getSystemService(LauncherApps.class).pinShortcuts(
+                        packageName, Collections.emptyList(), user);
+            } catch (SecurityException | IllegalStateException e) {
+                Log.w(TAG, "Failed to unpin shortcut", e);
+            }
+        });
     }
 
     /**
@@ -360,8 +400,40 @@
                 op.accept((WorkspaceItemInfo) info);
             }
         }
+
+        for (int i = extraItems.size() - 1; i >= 0; i--) {
+            for (ItemInfo info : extraItems.valueAt(i).items) {
+                if (info instanceof WorkspaceItemInfo && userHandle.equals(info.user)) {
+                    op.accept((WorkspaceItemInfo) info);
+                }
+            }
+        }
     }
 
+    /**
+     * An object containing items corresponding to a fixed container
+     */
+    public static class FixedContainerItems {
+
+        public final int containerId;
+        public final List<ItemInfo> items;
+
+        public FixedContainerItems(int containerId) {
+            this(containerId, new ArrayList<>());
+        }
+
+        public FixedContainerItems(int containerId, List<ItemInfo> items) {
+            this.containerId = containerId;
+            this.items = items;
+        }
+
+        @Override
+        public FixedContainerItems clone() {
+            return new FixedContainerItems(containerId, Collections.unmodifiableList(items));
+        }
+    }
+
+
     public interface Callbacks {
         // If the launcher has permission to access deep shortcuts.
         int FLAG_HAS_SHORTCUT_PERMISSION = 1 << 0;
@@ -384,7 +456,7 @@
         void bindAppsAdded(IntArray newScreens,
                 ArrayList<ItemInfo> addNotAnimated, ArrayList<ItemInfo> addAnimated);
         void bindPromiseAppProgressUpdated(PromiseAppInfo app);
-        void bindWorkspaceItemsChanged(ArrayList<WorkspaceItemInfo> updated);
+        void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated);
         void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets);
         void bindRestoreItemsChange(HashSet<ItemInfo> updates);
         void bindWorkspaceComponentsRemoved(ItemInfoMatcher matcher);
@@ -393,6 +465,11 @@
         void executeOnNextDraw(ViewOnDrawExecutor executor);
         void bindDeepShortcutMap(HashMap<ComponentKey, Integer> deepShortcutMap);
 
+        /**
+         * Binds extra item provided any external source
+         */
+        default void bindExtraContainerItems(FixedContainerItems item) { }
+
         void bindAllApplications(AppInfo[] apps, int flags);
 
         /**
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 4a64522..bea0086 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -46,12 +46,10 @@
 import android.text.TextUtils;
 import android.util.Log;
 import android.util.LongSparseArray;
-import android.util.MutableInt;
 import android.util.TimingLogger;
 
 import androidx.annotation.WorkerThread;
 
-import com.android.launcher3.InstallShortcutReceiver;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherSettings;
@@ -94,7 +92,6 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CancellationException;
@@ -794,17 +791,8 @@
                         LauncherSettings.Settings.METHOD_REMOVE_GHOST_WIDGETS);
             }
 
-            // Unpin shortcuts that don't exist on the workspace.
-            HashSet<ShortcutKey> pendingShortcuts =
-                    InstallShortcutReceiver.getPendingShortcuts(context);
-            for (ShortcutKey key : shortcutKeyToPinnedShortcuts.keySet()) {
-                MutableInt numTimesPinned = mBgDataModel.pinnedShortcutCounts.get(key);
-                if ((numTimesPinned == null || numTimesPinned.value == 0)
-                        && !pendingShortcuts.contains(key)) {
-                    // Shortcut is pinned but doesn't exist on the workspace; unpin it.
-                    mBgDataModel.unpinShortcut(context, key);
-                }
-            }
+            // Update pinned state of model shortcuts
+            mBgDataModel.updateShortcutPinnedState(context);
 
             // Sort the folder items, update ranks, and make sure all preview items are high res.
             FolderGridOrganizer verifier =
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
index a013312..c996748 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
@@ -222,7 +222,7 @@
     }
 
     @Override
-    public void bindWorkspaceItemsChanged(ArrayList<WorkspaceItemInfo> updated) { }
+    public void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) { }
 
     @Override
     public void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets) { }
diff --git a/src/com/android/launcher3/shortcuts/ShortcutKey.java b/src/com/android/launcher3/shortcuts/ShortcutKey.java
index 3ca9490..0c6d675 100644
--- a/src/com/android/launcher3/shortcuts/ShortcutKey.java
+++ b/src/com/android/launcher3/shortcuts/ShortcutKey.java
@@ -30,6 +30,10 @@
         return componentName.getClassName();
     }
 
+    public String getPackageName() {
+        return componentName.getPackageName();
+    }
+
     /**
      * Creates a {@link ShortcutRequest} for this key
      */