1 files changed, 60 insertions, 2 deletions

diff --git a/src/com/android/launcher3/badge/BadgeInfo.java b/src/com/android/launcher3/badge/BadgeInfo.java
index 77355c75e..532396cdd 100644
--- a/src/com/android/launcher3/badge/BadgeInfo.java
+++ b/src/com/android/launcher3/badge/BadgeInfo.java
@@ -16,6 +16,14 @@
 
 package com.android.launcher3.badge;
 
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.graphics.BitmapShader;
+import android.graphics.Canvas;
+import android.graphics.Shader;
+import android.graphics.drawable.Drawable;
+import android.support.annotation.Nullable;
+
 import com.android.launcher3.notification.NotificationInfo;
 import com.android.launcher3.util.PackageUserKey;
 
@@ -29,12 +37,22 @@ public class BadgeInfo {
 
     /** Used to link this BadgeInfo to icons on the workspace and all apps */
     private PackageUserKey mPackageUserKey;
+
     /**
      * The keys of the notifications that this badge represents. These keys can later be
      * used to retrieve {@link NotificationInfo}'s.
      */
     private List<String> mNotificationKeys;
 
+    /** This will only be initialized if the badge should display the notification icon. */
+    private NotificationInfo mNotificationInfo;
+
+    /**
+     * When retrieving the notification icon, we draw it into this shader, which can be clipped
+     * as necessary when drawn in a badge.
+     */
+    private Shader mNotificationIcon;
+
     public BadgeInfo(PackageUserKey packageUserKey) {
         mPackageUserKey = packageUserKey;
         mNotificationKeys = new ArrayList<>();
@@ -65,15 +83,55 @@ public class BadgeInfo {
         return mNotificationKeys.size();
     }
 
+    public void setNotificationToShow(@Nullable NotificationInfo notificationInfo) {
+        mNotificationInfo = notificationInfo;
+        mNotificationIcon = null;
+    }
+
+    public boolean hasNotificationToShow() {
+        return mNotificationInfo != null;
+    }
+
+    /**
+     * Returns a shader to set on a Paint that will draw the notification icon in a badge.
+     *
+     * The shader is cached until {@link #setNotificationToShow(NotificationInfo)} is called.
+     */
+    public @Nullable Shader getNotificationIconForBadge(Context context, int badgeColor,
+            int badgeSize, int badgePadding) {
+        if (mNotificationInfo == null) {
+            return null;
+        }
+        if (mNotificationIcon == null) {
+            Drawable icon = mNotificationInfo.getIconForBackground(context, badgeColor)
+                    .getConstantState().newDrawable();
+            int iconSize = badgeSize - badgePadding * 2;
+            icon.setBounds(0, 0, iconSize, iconSize);
+            Bitmap iconBitmap = Bitmap.createBitmap(badgeSize, badgeSize, Bitmap.Config.ARGB_8888);
+            Canvas canvas = new Canvas(iconBitmap);
+            canvas.translate(badgePadding, badgePadding);
+            icon.draw(canvas);
+            mNotificationIcon = new BitmapShader(iconBitmap, Shader.TileMode.CLAMP,
+                    Shader.TileMode.CLAMP);
+        }
+        return mNotificationIcon;
+    }
+
+    public boolean isIconLarge() {
+        return mNotificationInfo != null && mNotificationInfo.isIconLarge();
+    }
+
     /**
      * Whether newBadge represents the same PackageUserKey as this badge, and icons with
      * this badge should be invalidated. So, for instance, if a badge has 3 notifications
      * and one of those notifications is updated, this method should return false because
      * the badge still says "3" and the contents of those notifications are only retrieved
-     * upon long-click. This method always returns true when adding or removing notifications.
+     * upon long-click. This method always returns true when adding or removing notifications,
+     * or if the badge has a notification icon to show.
      */
     public boolean shouldBeInvalidated(BadgeInfo newBadge) {
         return mPackageUserKey.equals(newBadge.mPackageUserKey)
-                && getNotificationCount() != newBadge.getNotificationCount();
+                && (getNotificationCount() != newBadge.getNotificationCount()
+                    || hasNotificationToShow());
     }
 }