Fixing drawable cloning does not preserve badge

Bug: 209503720
Test: Manual
Change-Id: Iacfc5f2fdbdeb4d76a639b3e187b1ef7854f0d5b
diff --git a/iconloaderlib/src/com/android/launcher3/icons/ClockDrawableWrapper.java b/iconloaderlib/src/com/android/launcher3/icons/ClockDrawableWrapper.java
index 9667cae..cbb02cc 100644
--- a/iconloaderlib/src/com/android/launcher3/icons/ClockDrawableWrapper.java
+++ b/iconloaderlib/src/com/android/launcher3/icons/ClockDrawableWrapper.java
@@ -419,24 +419,22 @@
         }
 
         @Override
-        public ConstantState getConstantState() {
-            return new ClockConstantState(mInfo, isDisabled());
+        public FastBitmapConstantState newConstantState() {
+            return new ClockConstantState(mInfo);
         }
 
         private static class ClockConstantState extends FastBitmapConstantState {
 
             private final ClockBitmapInfo mInfo;
 
-            ClockConstantState(ClockBitmapInfo info, boolean isDisabled) {
-                super(info.icon, info.color, isDisabled);
+            ClockConstantState(ClockBitmapInfo info) {
+                super(info.icon, info.color);
                 mInfo = info;
             }
 
             @Override
-            public FastBitmapDrawable newDrawable() {
-                ClockIconDrawable drawable = new ClockIconDrawable(mInfo);
-                drawable.setIsDisabled(mIsDisabled);
-                return drawable;
+            public FastBitmapDrawable createDrawable() {
+                return new ClockIconDrawable(mInfo);
             }
         }
     }
diff --git a/iconloaderlib/src/com/android/launcher3/icons/FastBitmapDrawable.java b/iconloaderlib/src/com/android/launcher3/icons/FastBitmapDrawable.java
index 2632519..8a1ad44 100644
--- a/iconloaderlib/src/com/android/launcher3/icons/FastBitmapDrawable.java
+++ b/iconloaderlib/src/com/android/launcher3/icons/FastBitmapDrawable.java
@@ -87,14 +87,9 @@
     }
 
     protected FastBitmapDrawable(Bitmap b, int iconColor) {
-        this(b, iconColor, false);
-    }
-
-    protected FastBitmapDrawable(Bitmap b, int iconColor, boolean isDisabled) {
         mBitmap = b;
         mIconColor = iconColor;
         setFilterBitmap(true);
-        setIsDisabled(isDisabled);
     }
 
     @Override
@@ -290,9 +285,18 @@
         invalidateSelf();
     }
 
+    protected FastBitmapConstantState newConstantState() {
+        return new FastBitmapConstantState(mBitmap, mIconColor);
+    }
+
     @Override
-    public ConstantState getConstantState() {
-        return new FastBitmapConstantState(mBitmap, mIconColor, mIsDisabled);
+    public final ConstantState getConstantState() {
+        FastBitmapConstantState cs = newConstantState();
+        cs.mIsDisabled = mIsDisabled;
+        if (mBadge != null) {
+            cs.mBadgeConstantState = mBadge.getConstantState();
+        }
+        return cs;
     }
 
     public static ColorFilter getDisabledColorFilter() {
@@ -349,17 +353,29 @@
     protected static class FastBitmapConstantState extends ConstantState {
         protected final Bitmap mBitmap;
         protected final int mIconColor;
-        protected final boolean mIsDisabled;
 
-        public FastBitmapConstantState(Bitmap bitmap, int color, boolean isDisabled) {
+        // These are initialized later so that subclasses don't need to
+        // pass everything in constructor
+        protected boolean mIsDisabled;
+        private ConstantState mBadgeConstantState;
+
+        public FastBitmapConstantState(Bitmap bitmap, int color) {
             mBitmap = bitmap;
             mIconColor = color;
-            mIsDisabled = isDisabled;
+        }
+
+        protected FastBitmapDrawable createDrawable() {
+            return new FastBitmapDrawable(mBitmap, mIconColor);
         }
 
         @Override
-        public FastBitmapDrawable newDrawable() {
-            return new FastBitmapDrawable(mBitmap, mIconColor, mIsDisabled);
+        public final FastBitmapDrawable newDrawable() {
+            FastBitmapDrawable drawable = createDrawable();
+            drawable.setIsDisabled(mIsDisabled);
+            if (mBadgeConstantState != null) {
+                drawable.setBadge(mBadgeConstantState.newDrawable());
+            }
+            return drawable;
         }
 
         @Override
diff --git a/iconloaderlib/src/com/android/launcher3/icons/ThemedIconDrawable.java b/iconloaderlib/src/com/android/launcher3/icons/ThemedIconDrawable.java
index 7398530..1a2d0d6 100644
--- a/iconloaderlib/src/com/android/launcher3/icons/ThemedIconDrawable.java
+++ b/iconloaderlib/src/com/android/launcher3/icons/ThemedIconDrawable.java
@@ -62,7 +62,7 @@
     private final AdaptiveIconDrawable mBgWrapper;
 
     protected ThemedIconDrawable(ThemedConstantState constantState) {
-        super(constantState.mBitmap, constantState.colorFg, constantState.mIsDisabled);
+        super(constantState.mBitmap, constantState.colorFg);
         bitmapInfo = constantState.bitmapInfo;
         colorBg = constantState.colorBg;
         colorFg = constantState.colorFg;
@@ -95,8 +95,8 @@
     }
 
     @Override
-    public ConstantState getConstantState() {
-        return new ThemedConstantState(bitmapInfo, colorBg, colorFg, mIsDisabled);
+    public FastBitmapConstantState newConstantState() {
+        return new ThemedConstantState(bitmapInfo, colorBg, colorFg);
     }
 
     static class ThemedConstantState extends FastBitmapConstantState {
@@ -105,15 +105,15 @@
         final int colorFg, colorBg;
 
         public ThemedConstantState(ThemedBitmapInfo bitmapInfo,
-                int colorBg, int colorFg, boolean isDisabled) {
-            super(bitmapInfo.icon, bitmapInfo.color, isDisabled);
+                int colorBg, int colorFg) {
+            super(bitmapInfo.icon, bitmapInfo.color);
             this.bitmapInfo = bitmapInfo;
             this.colorBg = colorBg;
             this.colorFg = colorFg;
         }
 
         @Override
-        public FastBitmapDrawable newDrawable() {
+        public FastBitmapDrawable createDrawable() {
             return new ThemedIconDrawable(this);
         }
     }
@@ -136,7 +136,7 @@
             if ((creationFlags & FLAG_THEMED) != 0) {
                 int[] colors = getColors(context);
                 FastBitmapDrawable drawable =
-                        new ThemedConstantState(this, colors[0], colors[1], false).newDrawable();
+                        new ThemedConstantState(this, colors[0], colors[1]).newDrawable();
                 applyFlags(context, drawable, creationFlags);
                 return drawable;
             }