Attach work badge to Weather and Battery widgets

Moving the part where widget icon is generated from WidgetsListHeader (UI thread) to IconCache (backend)

Test: Open widget picker -> switch to work widget picker -> verify that Battery and Weather are badged
Fix: 226132413, 209995894
Change-Id: I3d649f2b26d7d8e7b756129b5bae4433ea344d43

src/com/android/launcher3/icons/IconCache.java

@@ -42,8 +42,10 @@ import android.os.Trace;
 import android.os.UserHandle;
 import android.text.TextUtils;
 import android.util.Log;
+import android.util.SparseArray;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.core.util.Pair;
 
 import com.android.launcher3.InvariantDeviceProfile;
@@ -94,6 +96,8 @@ public class IconCache extends BaseIconCache {
     private final InstantAppResolver mInstantAppResolver;
     private final IconProvider mIconProvider;
 
+    private final SparseArray<BitmapInfo> mWidgetCategoryBitmapInfos;
+
     private int mPendingIconRequestCount = 0;
 
     public IconCache(Context context, InvariantDeviceProfile idp) {
@@ -111,6 +115,7 @@ public class IconCache extends BaseIconCache {
         mUserManager = UserCache.INSTANCE.get(mContext);
         mInstantAppResolver = InstantAppResolver.newInstance(mContext);
         mIconProvider = iconProvider;
+        mWidgetCategoryBitmapInfos = new SparseArray<>();
     }
 
     @Override
@@ -477,13 +482,39 @@ public class IconCache extends BaseIconCache {
         CacheEntry entry = getEntryForPackageLocked(
                 infoInOut.packageName, infoInOut.user, useLowResIcon);
         applyCacheEntry(entry, infoInOut);
-        if (infoInOut.widgetCategory != NO_CATEGORY) {
-            WidgetSection widgetSection = WidgetSections.getWidgetSections(mContext)
-                    .get(infoInOut.widgetCategory);
-            infoInOut.title = mContext.getString(widgetSection.mSectionTitle);
-            infoInOut.contentDescription = mPackageManager.getUserBadgedLabel(
-                    infoInOut.title, infoInOut.user);
+        if (infoInOut.widgetCategory == NO_CATEGORY) {
+            return;
         }
+
+        WidgetSection widgetSection = WidgetSections.getWidgetSections(mContext)
+                .get(infoInOut.widgetCategory);
+        infoInOut.title = mContext.getString(widgetSection.mSectionTitle);
+        infoInOut.contentDescription = mPackageManager.getUserBadgedLabel(
+                infoInOut.title, infoInOut.user);
+        final BitmapInfo cachedBitmap = mWidgetCategoryBitmapInfos.get(infoInOut.widgetCategory);
+        if (cachedBitmap != null) {
+            infoInOut.bitmap = getBadgedIcon(cachedBitmap, infoInOut.user);
+            return;
+        }
+
+        try (LauncherIcons li = LauncherIcons.obtain(mContext)) {
+            final BitmapInfo tempBitmap = li.createBadgedIconBitmap(
+                    mContext.getDrawable(widgetSection.mSectionDrawable),
+                    new BaseIconFactory.IconOptions().setShrinkNonAdaptiveIcons(false));
+            mWidgetCategoryBitmapInfos.put(infoInOut.widgetCategory, tempBitmap);
+            infoInOut.bitmap = getBadgedIcon(tempBitmap, infoInOut.user);
+        } catch (Exception e) {
+            Log.e(TAG, "Error initializing bitmap for icons with widget category", e);
+        }
+
+    }
+
+    private synchronized BitmapInfo getBadgedIcon(@Nullable final BitmapInfo bitmap,
+            @NonNull final UserHandle user) {
+        if (bitmap == null) {
+            return getDefaultIcon(user);
+        }
+        return bitmap.withFlags(getUserFlagOpLocked(user));
     }
 
     protected void applyCacheEntry(CacheEntry entry, ItemInfoWithIcon info) {

src/com/android/launcher3/widget/PendingAppWidgetHostView.java

@@ -18,7 +18,6 @@ package com.android.launcher3.widget;
 
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
 import static com.android.launcher3.icons.FastBitmapDrawable.getDisabledColorFilter;
-import static com.android.launcher3.widget.WidgetSections.getWidgetSections;
 
 import android.content.Context;
 import android.graphics.Canvas;
@@ -341,8 +340,6 @@ public class PendingAppWidgetHostView extends LauncherAppWidgetHostView
         if (mInfo.pendingItemInfo.widgetCategory == WidgetSections.NO_CATEGORY) {
             return null;
         }
-        Context context = getContext();
-        return context.getDrawable(getWidgetSections(context).get(
-                mInfo.pendingItemInfo.widgetCategory).mSectionDrawable);
+        return mInfo.pendingItemInfo.newIcon(getContext());
     }
 }

src/com/android/launcher3/widget/picker/WidgetsListHeader.java

@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.widget.picker;
 
-import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
 
 import android.content.Context;
 import android.content.res.Resources;
@@ -43,8 +42,6 @@ import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.PluralMessageFormat;
 import com.android.launcher3.views.ActivityContext;
-import com.android.launcher3.widget.WidgetSections;
-import com.android.launcher3.widget.WidgetSections.WidgetSection;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
@@ -177,13 +174,7 @@ public final class WidgetsListHeader extends LinearLayout implements ItemInfoUpd
 
     private void setIcon(PackageItemInfo info) {
         Drawable icon;
-        if (info.widgetCategory == NO_CATEGORY) {
-            icon = info.newIcon(getContext());
-        } else {
-            WidgetSection widgetSection = WidgetSections.getWidgetSections(getContext())
-                    .get(info.widgetCategory);
-            icon = getContext().getDrawable(widgetSection.mSectionDrawable);
-        }
+        icon = info.newIcon(getContext());
         applyDrawables(icon);
         mIconDrawable = icon;
         if (mIconDrawable != null) {