Merge "Ensure each dot contains only relevant notification to the shortcut and displays system shortcut when long click on deep shortcut." into ub-launcher3-qt-qpr1-dev
diff --git a/go/src/com/android/launcher3/shortcuts/DeepShortcutManager.java b/go/src/com/android/launcher3/shortcuts/DeepShortcutManager.java
index 73adaa1..7fd4a9d 100644
--- a/go/src/com/android/launcher3/shortcuts/DeepShortcutManager.java
+++ b/go/src/com/android/launcher3/shortcuts/DeepShortcutManager.java
@@ -53,11 +53,14 @@
         return false;
     }
 
-    public static boolean supportsNotificationDots(
-            ItemInfo info, List<NotificationKeyData> notifications) {
+    public static boolean supportsDeepShortcuts(ItemInfo info) {
         return false;
     }
 
+    public static String getShortcutIdIfApplicable(ItemInfo info) {
+        return null;
+    }
+
     public boolean wasLastCallSuccess() {
         return false;
     }
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index dd496b0..6e4883e 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -130,7 +131,8 @@
         for (PackageUserKey packageUserKey : mPackageUserToDotInfos.keySet()) {
             DotInfo prevDot = updatedDots.get(packageUserKey);
             DotInfo newDot = mPackageUserToDotInfos.get(packageUserKey);
-            if (prevDot == null) {
+            if (prevDot == null
+                    || prevDot.getNotificationCount() != newDot.getNotificationCount()) {
                 updatedDots.put(packageUserKey, newDot);
             } else {
                 // No need to update the dot if it already existed (no visual change).
@@ -156,7 +158,7 @@
     }
 
     public int getShortcutCountForItem(ItemInfo info) {
-        if (!DeepShortcutManager.supportsShortcuts(info)) {
+        if (!DeepShortcutManager.supportsDeepShortcuts(info)) {
             return 0;
         }
         ComponentName component = info.getTargetComponent();
@@ -169,10 +171,16 @@
     }
 
     public @Nullable DotInfo getDotInfoForItem(@NonNull ItemInfo info) {
+        if (!DeepShortcutManager.supportsShortcuts(info)) {
+            return null;
+        }
         DotInfo dotInfo = mPackageUserToDotInfos.get(PackageUserKey.fromItemInfo(info));
-        List<NotificationKeyData> notifications =
-                dotInfo == null ? Collections.EMPTY_LIST : dotInfo.getNotificationKeys();
-        if (!DeepShortcutManager.supportsNotificationDots(info, notifications)) {
+        if (dotInfo == null) {
+            return null;
+        }
+        List<NotificationKeyData> notifications = getNotificationsForItem(
+                info, dotInfo.getNotificationKeys());
+        if (notifications.isEmpty()) {
             return null;
         }
         return dotInfo;
@@ -180,7 +188,8 @@
 
     public @NonNull List<NotificationKeyData> getNotificationKeysForItem(ItemInfo info) {
         DotInfo dotInfo = getDotInfoForItem(info);
-        return dotInfo == null ? Collections.EMPTY_LIST : dotInfo.getNotificationKeys();
+        return dotInfo == null ? Collections.EMPTY_LIST
+                : getNotificationsForItem(info, dotInfo.getNotificationKeys());
     }
 
     /** This makes a potentially expensive binder call and should be run on a background thread. */
@@ -229,6 +238,20 @@
         return null;
     }
 
+    /**
+     * Returns a list of notifications that are relevant to given ItemInfo.
+     */
+    public static @NonNull List<NotificationKeyData> getNotificationsForItem(
+            @NonNull ItemInfo info, @NonNull List<NotificationKeyData> notifications) {
+        String shortcutId = DeepShortcutManager.getShortcutIdIfApplicable(info);
+        if (shortcutId == null) {
+            return notifications;
+        }
+        return notifications.stream().filter((NotificationKeyData notification) ->
+                shortcutId.equals(notification.shortcutId)
+        ).collect(Collectors.toList());
+    }
+
     public interface PopupDataChangeListener {
 
         PopupDataChangeListener INSTANCE = new PopupDataChangeListener() { };
diff --git a/src_shortcuts_overrides/com/android/launcher3/shortcuts/DeepShortcutManager.java b/src_shortcuts_overrides/com/android/launcher3/shortcuts/DeepShortcutManager.java
index f42bafe..aec2bad 100644
--- a/src_shortcuts_overrides/com/android/launcher3/shortcuts/DeepShortcutManager.java
+++ b/src_shortcuts_overrides/com/android/launcher3/shortcuts/DeepShortcutManager.java
@@ -30,7 +30,6 @@
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.WorkspaceItemInfo;
-import com.android.launcher3.notification.NotificationKeyData;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -65,16 +64,16 @@
     }
 
     public static boolean supportsShortcuts(ItemInfo info) {
+        return isActive(info) && (isApp(info) || isPinnedShortcut(info));
+    }
+
+    public static boolean supportsDeepShortcuts(ItemInfo info) {
         return isActive(info) && isApp(info);
     }
 
-    public static boolean supportsNotificationDots(
-            ItemInfo info, List<NotificationKeyData> notifications) {
-        if (!isActive(info)) {
-            return false;
-        }
-        return isApp(info) || (isPinnedShortcut(info)
-                && shouldShowNotificationDotForPinnedShortcut(info, notifications));
+    public static String getShortcutIdIfApplicable(ItemInfo info) {
+        return isActive(info) && isPinnedShortcut(info) ?
+                ShortcutKey.fromItemInfo(info).getId() : null;
     }
 
     private static boolean isApp(ItemInfo info) {
@@ -87,20 +86,6 @@
                 && info instanceof WorkspaceItemInfo;
     }
 
-    private static boolean shouldShowNotificationDotForPinnedShortcut(
-            ItemInfo info, List<NotificationKeyData> notifications) {
-        String shortcutId = ((WorkspaceItemInfo) info).getDeepShortcutId();
-        if (shortcutId == null) {
-            return false;
-        }
-        for (NotificationKeyData notification : notifications) {
-            if (shortcutId.equals(notification.shortcutId)) {
-                return true;
-            }
-        }
-        return false;
-    }
-
     public boolean wasLastCallSuccess() {
         return mWasLastCallSuccess;
     }