Use notification counts (instead of assuming all 1)

We still ignore group summary headers, which means, for example,
we won't get Gmail unread count. But single notifications that
have numbers associated, such as messages from a single contact
will be included in the badge count. So if you have 2 hangounts
threads, one with 10 messages and one with 8, the badge would
say 18.

Bug: 34939841
Change-Id: I20b9a857d91715e10c0da400a1cee209d7b837b8
diff --git a/src/com/android/launcher3/badge/BadgeInfo.java b/src/com/android/launcher3/badge/BadgeInfo.java
index d7aa4e8..08d8ad4 100644
--- a/src/com/android/launcher3/badge/BadgeInfo.java
+++ b/src/com/android/launcher3/badge/BadgeInfo.java
@@ -36,6 +36,8 @@
  */
 public class BadgeInfo {
 
+    public static final int MAX_COUNT = 999;
+
     /** Used to link this BadgeInfo to icons on the workspace and all apps */
     private PackageUserKey mPackageUserKey;
 
@@ -45,6 +47,12 @@
      */
     private List<NotificationKeyData> mNotificationKeys;
 
+    /**
+     * The current sum of the counts in {@link #mNotificationKeys},
+     * updated whenever a key is added or removed.
+     */
+    private int mTotalCount;
+
     /** This will only be initialized if the badge should display the notification icon. */
     private NotificationInfo mNotificationInfo;
 
@@ -60,20 +68,38 @@
     }
 
     /**
-     * Returns whether the notification was added (false if it already existed).
+     * Returns whether the notification was added or its count changed.
      */
-    public boolean addNotificationKeyIfNotExists(NotificationKeyData notificationKey) {
-        if (mNotificationKeys.contains(notificationKey)) {
-            return false;
+    public boolean addOrUpdateNotificationKey(NotificationKeyData notificationKey) {
+        int indexOfPrevKey = mNotificationKeys.indexOf(notificationKey);
+        NotificationKeyData prevKey = indexOfPrevKey == -1 ? null
+                : mNotificationKeys.get(indexOfPrevKey);
+        if (prevKey != null) {
+            if (prevKey.count == notificationKey.count) {
+                return false;
+            }
+            // Notification was updated with a new count.
+            mTotalCount -= prevKey.count;
+            mTotalCount += notificationKey.count;
+            prevKey.count = notificationKey.count;
+            return true;
         }
-        return mNotificationKeys.add(notificationKey);
+        boolean added = mNotificationKeys.add(notificationKey);
+        if (added) {
+            mTotalCount += notificationKey.count;
+        }
+        return added;
     }
 
     /**
      * Returns whether the notification was removed (false if it didn't exist).
      */
     public boolean removeNotificationKey(NotificationKeyData notificationKey) {
-        return mNotificationKeys.remove(notificationKey);
+        boolean removed = mNotificationKeys.remove(notificationKey);
+        if (removed) {
+            mTotalCount -= notificationKey.count;
+        }
+        return removed;
     }
 
     public List<NotificationKeyData> getNotificationKeys() {
@@ -81,7 +107,7 @@
     }
 
     public int getNotificationCount() {
-        return mNotificationKeys.size();
+        return Math.min(mTotalCount, MAX_COUNT);
     }
 
     public void setNotificationToShow(@Nullable NotificationInfo notificationInfo) {
diff --git a/src/com/android/launcher3/badge/FolderBadgeInfo.java b/src/com/android/launcher3/badge/FolderBadgeInfo.java
index 4d1e5c2..f7c64aa 100644
--- a/src/com/android/launcher3/badge/FolderBadgeInfo.java
+++ b/src/com/android/launcher3/badge/FolderBadgeInfo.java
@@ -18,8 +18,6 @@
 
 import com.android.launcher3.Utilities;
 
-import static com.android.launcher3.Utilities.boundToRange;
-
 /**
  * Subclass of BadgeInfo that only contains the badge count,
  * which is the sum of all the Folder's items' counts.
@@ -27,7 +25,6 @@
 public class FolderBadgeInfo extends BadgeInfo {
 
     private static final int MIN_COUNT = 0;
-    private static final int MAX_COUNT = 999;
 
     private int mTotalNotificationCount;
 
@@ -41,7 +38,7 @@
         }
         mTotalNotificationCount += badgeToAdd.getNotificationCount();
         mTotalNotificationCount = Utilities.boundToRange(
-                mTotalNotificationCount, MIN_COUNT, MAX_COUNT);
+                mTotalNotificationCount, MIN_COUNT, BadgeInfo.MAX_COUNT);
     }
 
     public void subtractBadgeInfo(BadgeInfo badgeToSubtract) {
@@ -50,7 +47,7 @@
         }
         mTotalNotificationCount -= badgeToSubtract.getNotificationCount();
         mTotalNotificationCount = Utilities.boundToRange(
-                mTotalNotificationCount, MIN_COUNT, MAX_COUNT);
+                mTotalNotificationCount, MIN_COUNT, BadgeInfo.MAX_COUNT);
     }
 
     @Override
diff --git a/src/com/android/launcher3/notification/NotificationKeyData.java b/src/com/android/launcher3/notification/NotificationKeyData.java
index b3ff8da..bf7ae1a 100644
--- a/src/com/android/launcher3/notification/NotificationKeyData.java
+++ b/src/com/android/launcher3/notification/NotificationKeyData.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.notification;
 
+import android.app.Notification;
 import android.service.notification.StatusBarNotification;
 import android.support.annotation.NonNull;
 
@@ -31,14 +32,17 @@
 public class NotificationKeyData {
     public final String notificationKey;
     public final String shortcutId;
+    public int count;
 
-    private NotificationKeyData(String notificationKey, String shortcutId) {
+    private NotificationKeyData(String notificationKey, String shortcutId, int count) {
         this.notificationKey = notificationKey;
         this.shortcutId = shortcutId;
+        this.count = Math.max(1, count);
     }
 
     public static NotificationKeyData fromNotification(StatusBarNotification sbn) {
-        return new NotificationKeyData(sbn.getKey(), sbn.getNotification().getShortcutId());
+        Notification notif = sbn.getNotification();
+        return new NotificationKeyData(sbn.getKey(), notif.getShortcutId(), notif.number);
     }
 
     public static List<String> extractKeysOnly(@NonNull List<NotificationKeyData> notificationKeys) {
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index 43b2b52..9c4faed 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -63,26 +63,26 @@
     public void onNotificationPosted(PackageUserKey postedPackageUserKey,
             NotificationKeyData notificationKey, boolean shouldBeFilteredOut) {
         BadgeInfo badgeInfo = mPackageUserToBadgeInfos.get(postedPackageUserKey);
-        boolean notificationWasAddedOrRemoved; // As opposed to updated.
+        boolean badgeShouldBeRefreshed;
         if (badgeInfo == null) {
             if (!shouldBeFilteredOut) {
                 BadgeInfo newBadgeInfo = new BadgeInfo(postedPackageUserKey);
-                newBadgeInfo.addNotificationKeyIfNotExists(notificationKey);
+                newBadgeInfo.addOrUpdateNotificationKey(notificationKey);
                 mPackageUserToBadgeInfos.put(postedPackageUserKey, newBadgeInfo);
-                notificationWasAddedOrRemoved = true;
+                badgeShouldBeRefreshed = true;
             } else {
-                notificationWasAddedOrRemoved = false;
+                badgeShouldBeRefreshed = false;
             }
         } else {
-            notificationWasAddedOrRemoved = shouldBeFilteredOut
+            badgeShouldBeRefreshed = shouldBeFilteredOut
                     ? badgeInfo.removeNotificationKey(notificationKey)
-                    : badgeInfo.addNotificationKeyIfNotExists(notificationKey);
+                    : badgeInfo.addOrUpdateNotificationKey(notificationKey);
             if (badgeInfo.getNotificationCount() == 0) {
                 mPackageUserToBadgeInfos.remove(postedPackageUserKey);
             }
         }
         updateLauncherIconBadges(Utilities.singletonHashSet(postedPackageUserKey),
-                notificationWasAddedOrRemoved);
+                badgeShouldBeRefreshed);
     }
 
     @Override
@@ -115,7 +115,7 @@
                 badgeInfo = new BadgeInfo(packageUserKey);
                 mPackageUserToBadgeInfos.put(packageUserKey, badgeInfo);
             }
-            badgeInfo.addNotificationKeyIfNotExists(NotificationKeyData
+            badgeInfo.addOrUpdateNotificationKey(NotificationKeyData
                     .fromNotification(notification));
         }
 
@@ -150,17 +150,17 @@
      * Updates the icons on launcher (workspace, folders, all apps) to refresh their badges.
      * @param updatedBadges The packages whose badges should be refreshed (either a notification was
      *                      added or removed, or the badge should show the notification icon).
-     * @param addedOrRemoved An optional parameter that will allow us to only refresh badges that
-     *                       updated (not added/removed) that have icons. If a badge updated
-     *                       but it doesn't have an icon, then the badge number doesn't change.
+     * @param shouldRefresh An optional parameter that will allow us to only refresh badges that
+     *                      have actually changed. If a notification updated its content but not
+     *                      its count or icon, then the badge doesn't change.
      */
     private void updateLauncherIconBadges(Set<PackageUserKey> updatedBadges,
-            boolean addedOrRemoved) {
+            boolean shouldRefresh) {
         Iterator<PackageUserKey> iterator = updatedBadges.iterator();
         while (iterator.hasNext()) {
             BadgeInfo badgeInfo = mPackageUserToBadgeInfos.get(iterator.next());
-            if (badgeInfo != null && !updateBadgeIcon(badgeInfo) && !addedOrRemoved) {
-                // The notification icon isn't used, and the badge wasn't added or removed
+            if (badgeInfo != null && !updateBadgeIcon(badgeInfo) && !shouldRefresh) {
+                // The notification icon isn't used, and the badge hasn't changed
                 // so there is no update to be made.
                 iterator.remove();
             }