Adding support for private/clone profile badge in Launcher3.

Bug: 289224457
Flag: None
Test: Flash on device and Presubmit
Change-Id: Icebb5b5129a2f0d909481f00d51e047431a23a15
diff --git a/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java b/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
index 2d32407..4a26559 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
@@ -21,6 +21,7 @@
 import android.content.Context;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
+import android.content.pm.LauncherUserInfo;
 import android.content.pm.ShortcutInfo;
 import android.graphics.drawable.ColorDrawable;
 import android.os.UserHandle;
@@ -28,6 +29,7 @@
 import android.util.ArrayMap;
 import android.window.RemoteTransition;
 
+import com.android.launcher3.Flags;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.util.UserIconInfo;
 import com.android.quickstep.util.FadeOutRemoteTransition;
@@ -69,17 +71,32 @@
         List<UserHandle> usersActual = um.getUserProfiles();
         if (usersActual != null) {
             for (UserHandle user : usersActual) {
-                long serial = um.getSerialNumberForUser(user);
+                if (android.os.Flags.allowPrivateProfile() && Flags.enablePrivateSpace()) {
+                    LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
+                    LauncherUserInfo launcherUserInfo = launcherApps.getLauncherUserInfo(user);
+                    // UserTypes not supported in Launcher are deemed to be the current
+                    // Foreground User.
+                    int userType = switch (launcherUserInfo.getUserType()) {
+                        case UserManager.USER_TYPE_PROFILE_MANAGED -> UserIconInfo.TYPE_WORK;
+                        case UserManager.USER_TYPE_PROFILE_CLONE -> UserIconInfo.TYPE_CLONED;
+                        case UserManager.USER_TYPE_PROFILE_PRIVATE -> UserIconInfo.TYPE_PRIVATE;
+                        default -> UserIconInfo.TYPE_MAIN;
+                    };
+                    long serial = launcherUserInfo.getUserSerialNumber();
+                    users.put(user, new UserIconInfo(user, userType, serial));
+                } else {
+                    long serial = um.getSerialNumberForUser(user);
 
-                // Simple check to check if the provided user is work profile
-                // TODO: Migrate to a better platform API
-                NoopDrawable d = new NoopDrawable();
-                boolean isWork = (d != context.getPackageManager().getUserBadgedIcon(d, user));
-                UserIconInfo info = new UserIconInfo(
-                        user,
-                        isWork ? UserIconInfo.TYPE_WORK : UserIconInfo.TYPE_MAIN,
-                        serial);
-                users.put(user, info);
+                    // Simple check to check if the provided user is work profile
+                    // TODO: Migrate to a better platform API
+                    NoopDrawable d = new NoopDrawable();
+                    boolean isWork = (d != context.getPackageManager().getUserBadgedIcon(d, user));
+                    UserIconInfo info = new UserIconInfo(
+                            user,
+                            isWork ? UserIconInfo.TYPE_WORK : UserIconInfo.TYPE_MAIN,
+                            serial);
+                    users.put(user, info);
+                }
             }
         }
         return users;
diff --git a/quickstep/src/com/android/quickstep/TaskIconCache.java b/quickstep/src/com/android/quickstep/TaskIconCache.java
index 164a366..f0f8aec 100644
--- a/quickstep/src/com/android/quickstep/TaskIconCache.java
+++ b/quickstep/src/com/android/quickstep/TaskIconCache.java
@@ -40,6 +40,7 @@
 import com.android.launcher3.icons.BaseIconFactory.IconOptions;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.IconProvider;
+import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.DisplayController.DisplayInfoChangeListener;
 import com.android.launcher3.util.DisplayController.Info;
@@ -233,7 +234,9 @@
             } else {
                 try (BaseIconFactory li = getIconFactory()) {
                     BitmapInfo info = mDefaultIconBase.withFlags(
-                            li.getBitmapFlagOp(new IconOptions().setUser(UserHandle.of(userId))));
+                            li.getBitmapFlagOp(new IconOptions()
+                                    .setUser(UserCache.INSTANCE.get(mContext)
+                                            .getUserInfo(UserHandle.of(userId)))));
                     mDefaultIcons.put(userId, info);
                     return info.newIcon(mContext);
                 }
@@ -249,7 +252,9 @@
 
             // User version code O, so that the icon is always wrapped in an adaptive icon container
             return bif.createBadgedIconBitmap(drawable,
-                    new IconOptions().setUser(UserHandle.of(userId))
+                    new IconOptions()
+                            .setUser(UserCache.INSTANCE.get(mContext)
+                                    .getUserInfo(UserHandle.of(userId)))
                             .setInstantApp(isInstantApp)
                             .setExtractedColor(0));
         }
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index e797147..8ba6d2c 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -56,7 +56,6 @@
 import android.os.DeadObjectException;
 import android.os.Handler;
 import android.os.Message;
-import android.os.Process;
 import android.os.TransactionTooLargeException;
 import android.text.Spannable;
 import android.text.SpannableString;
@@ -80,12 +79,15 @@
 
 import com.android.launcher3.dragndrop.FolderAdaptiveIcon;
 import com.android.launcher3.graphics.TintedDrawableSpan;
+import com.android.launcher3.icons.BaseIconFactory;
+import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.icons.ShortcutCachingLogic;
 import com.android.launcher3.icons.ThemedIconDrawable;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
+import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.shortcuts.ShortcutKey;
 import com.android.launcher3.shortcuts.ShortcutRequest;
 import com.android.launcher3.testing.shared.ResourceUtils;
@@ -674,11 +676,15 @@
         }
 
         if (badge == null) {
-            badge = Process.myUserHandle().equals(info.user)
-                    ? new ColorDrawable(Color.TRANSPARENT)
-                    : context.getDrawable(useTheme
-                            ? R.drawable.ic_work_app_badge_themed
-                            : R.drawable.ic_work_app_badge);
+            try (LauncherIcons li = LauncherIcons.obtain(context)) {
+                badge = BitmapInfo.LOW_RES_INFO.withFlags(
+                                li.getBitmapFlagOp(new BaseIconFactory.IconOptions().setUser(
+                                        UserCache.INSTANCE.get(context).getUserInfo(info.user))))
+                        .getBadgeDrawable(context, useTheme);
+            }
+            if (badge == null) {
+                badge = new ColorDrawable(Color.TRANSPARENT);
+            }
         }
         return Pair.create(result, badge);
     }