1 files changed, 31 insertions, 19 deletions

diff --git a/iconloaderlib/src/com/android/launcher3/icons/BitmapInfo.java b/iconloaderlib/src/com/android/launcher3/icons/BitmapInfo.java
index d1ef6f7..86db5b8 100644
--- a/iconloaderlib/src/com/android/launcher3/icons/BitmapInfo.java
+++ b/iconloaderlib/src/com/android/launcher3/icons/BitmapInfo.java
@@ -16,10 +16,10 @@
 package com.android.launcher3.icons;
 
 import android.content.Context;
-import android.content.res.Configuration;
 import android.graphics.Bitmap;
 import android.graphics.Bitmap.Config;
 import android.graphics.Canvas;
+import android.graphics.drawable.Drawable;
 
 import androidx.annotation.IntDef;
 import androidx.annotation.NonNull;
@@ -29,13 +29,15 @@ import com.android.launcher3.util.FlagOp;
 
 public class BitmapInfo {
 
-    static final int FLAG_WORK = 1 << 0;
-    static final int FLAG_INSTANT = 1 << 1;
-    static final int FLAG_CLONE = 1 << 2;
+    public static final int FLAG_WORK = 1 << 0;
+    public static final int FLAG_INSTANT = 1 << 1;
+    public static final int FLAG_CLONE = 1 << 2;
+    public static final int FLAG_PRIVATE = 1 << 3;
     @IntDef(flag = true, value = {
             FLAG_WORK,
             FLAG_INSTANT,
-            FLAG_CLONE
+            FLAG_CLONE,
+            FLAG_PRIVATE
     })
     @interface BitmapInfoFlags {}
 
@@ -151,25 +153,35 @@ public class BitmapInfo {
     protected void applyFlags(Context context, FastBitmapDrawable drawable,
             @DrawableCreationFlags int creationFlags) {
         drawable.mDisabledAlpha = GraphicsUtils.getFloat(context, R.attr.disabledIconAlpha, 1f);
+        drawable.mCreationFlags = creationFlags;
         if ((creationFlags & FLAG_NO_BADGE) == 0) {
-            if (badgeInfo != null) {
-                drawable.setBadge(badgeInfo.newIcon(context, creationFlags));
-            } else if ((flags & FLAG_INSTANT) != 0) {
-                drawable.setBadge(context.getDrawable(drawable.isThemed()
-                        ? R.drawable.ic_instant_app_badge_themed
-                        : R.drawable.ic_instant_app_badge));
-            } else if ((flags & FLAG_WORK) != 0) {
-                drawable.setBadge(context.getDrawable(drawable.isThemed()
-                        ? R.drawable.ic_work_app_badge_themed
-                        : R.drawable.ic_work_app_badge));
-            } else if ((flags & FLAG_CLONE) != 0) {
-                drawable.setBadge(context.getDrawable(drawable.isThemed()
-                        ? R.drawable.ic_clone_app_badge_themed
-                        : R.drawable.ic_clone_app_badge));
+            Drawable badge = getBadgeDrawable(context, (creationFlags & FLAG_THEMED) != 0);
+            if (badge != null) {
+                drawable.setBadge(badge);
             }
         }
     }
 
+    /**
+     * Returns a drawable representing the badge for this info
+     */
+    @Nullable
+    public Drawable getBadgeDrawable(Context context, boolean isThemed) {
+        if (badgeInfo != null) {
+            return badgeInfo.newIcon(context, isThemed ? FLAG_THEMED : 0);
+        } else if ((flags & FLAG_INSTANT) != 0) {
+            return new UserBadgeDrawable(context, R.drawable.ic_instant_app_badge, isThemed);
+        } else if ((flags & FLAG_WORK) != 0) {
+            return new UserBadgeDrawable(context, R.drawable.ic_work_app_badge, isThemed);
+        } else if ((flags & FLAG_CLONE) != 0) {
+            return new UserBadgeDrawable(context, R.drawable.ic_clone_app_badge, isThemed);
+        } else if ((flags & FLAG_PRIVATE) != 0) {
+            return new UserBadgeDrawable(
+                    context, R.drawable.ic_private_profile_app_badge, isThemed);
+        }
+        return null;
+    }
+
     public static BitmapInfo fromBitmap(@NonNull Bitmap bitmap) {
         return of(bitmap, 0);
     }