Merge "Notif redesign: Improve the way we display the work badge" into main
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/AppIconProvider.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/AppIconProvider.kt
index 0ddf9f72..172b76c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/AppIconProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/AppIconProvider.kt
@@ -22,12 +22,14 @@
 import android.content.Context
 import android.content.pm.PackageManager.NameNotFoundException
 import android.graphics.Color
-import android.graphics.drawable.BitmapDrawable
 import android.graphics.drawable.ColorDrawable
 import android.graphics.drawable.Drawable
+import android.os.UserHandle
 import android.util.Log
 import com.android.internal.R
 import com.android.launcher3.icons.BaseIconFactory
+import com.android.launcher3.icons.BaseIconFactory.IconOptions
+import com.android.launcher3.util.UserIconInfo
 import com.android.systemui.Dumpable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dump.DumpManager
@@ -48,7 +50,11 @@
      */
     @Throws(NameNotFoundException::class)
     @WorkerThread
-    fun getOrFetchAppIcon(packageName: String, context: Context): Drawable
+    fun getOrFetchAppIcon(
+        packageName: String,
+        context: Context,
+        withWorkProfileBadge: Boolean = false,
+    ): Drawable
 
     /**
      * Mark all the entries in the cache that are NOT in [wantedPackages] to be cleared. If they're
@@ -81,21 +87,52 @@
 
     private val cache = NotifCollectionCache<Drawable>()
 
-    override fun getOrFetchAppIcon(packageName: String, context: Context): Drawable {
-        return cache.getOrFetch(packageName) { fetchAppIcon(packageName, context) }
+    override fun getOrFetchAppIcon(
+        packageName: String,
+        context: Context,
+        withWorkProfileBadge: Boolean,
+    ): Drawable {
+        // Add a suffix to distinguish the app installed on the work profile, since the icon will
+        // be different.
+        val key = packageName + if (withWorkProfileBadge) WORK_SUFFIX else ""
+
+        return cache.getOrFetch(key) { fetchAppIcon(packageName, context, withWorkProfileBadge) }
     }
 
     @WorkerThread
-    private fun fetchAppIcon(packageName: String, context: Context): BitmapDrawable {
-        val icon = context.packageManager.getApplicationIcon(packageName)
-        return BitmapDrawable(
-            context.resources,
-            iconFactory.createScaledBitmap(icon, BaseIconFactory.MODE_HARDWARE),
+    private fun fetchAppIcon(
+        packageName: String,
+        context: Context,
+        withWorkProfileBadge: Boolean,
+    ): Drawable {
+        val pm = context.packageManager
+        val icon = pm.getApplicationInfo(packageName, 0).loadUnbadgedIcon(pm)
+
+        val options =
+            IconOptions().apply {
+                setUser(userIconInfo(context, withWorkProfileBadge))
+                setBitmapGenerationMode(BaseIconFactory.MODE_HARDWARE)
+                // This color is not used since we're not showing the themed icons. We're just
+                // setting it so that the icon factory doesn't try to extract colors from our bitmap
+                // (since it won't work, given it's a hardware bitmap).
+                setExtractedColor(Color.BLUE)
+            }
+        val badgedIcon = iconFactory.createBadgedIconBitmap(icon, options)
+        return badgedIcon.newIcon(sysuiContext)
+    }
+
+    private fun userIconInfo(context: Context, withWorkProfileBadge: Boolean): UserIconInfo {
+        val userId = context.userId
+        return UserIconInfo(
+            UserHandle.of(userId),
+            if (withWorkProfileBadge) UserIconInfo.TYPE_WORK else UserIconInfo.TYPE_MAIN,
         )
     }
 
     override fun purgeCache(wantedPackages: Collection<String>) {
-        cache.purge(wantedPackages)
+        // We don't know from the packages if it's the work profile app or not, so let's just keep
+        // both if they're present in the cache.
+        cache.purge(wantedPackages.flatMap { listOf(it, "$it$WORK_SUFFIX") })
     }
 
     override fun dump(pwOrig: PrintWriter, args: Array<out String>) {
@@ -114,6 +151,7 @@
 
     companion object {
         const val TAG = "AppIconProviderImpl"
+        const val WORK_SUFFIX = "|WORK"
     }
 }
 
@@ -122,7 +160,11 @@
         const val TAG = "NoOpIconProvider"
     }
 
-    override fun getOrFetchAppIcon(packageName: String, context: Context): Drawable {
+    override fun getOrFetchAppIcon(
+        packageName: String,
+        context: Context,
+        withWorkProfileBadge: Boolean,
+    ): Drawable {
         Log.wtf(TAG, "NoOpIconProvider should not be used anywhere.")
         return ColorDrawable(Color.WHITE)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProvider.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProvider.kt
index 35e38c2..7177a7b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProvider.kt
@@ -20,6 +20,7 @@
 import android.app.Flags
 import android.content.Context
 import android.content.pm.ApplicationInfo
+import android.os.UserManager
 import android.service.notification.StatusBarNotification
 import android.util.Log
 import com.android.systemui.Dumpable
@@ -47,6 +48,12 @@
     fun shouldShowAppIcon(notification: StatusBarNotification, context: Context): Boolean
 
     /**
+     * Whether the [notification] is coming from a work profile app, and therefore should display
+     * the briefcase badge.
+     */
+    fun shouldShowWorkProfileBadge(notification: StatusBarNotification, context: Context): Boolean
+
+    /**
      * Mark all the entries in the cache that are NOT in [wantedPackages] to be cleared. If they're
      * still not needed on the next call of this method (made after a timeout of 1s, in case they
      * happen more frequently than that), they will be purged. This can be done from any thread.
@@ -55,7 +62,9 @@
 }
 
 @SysUISingleton
-class NotificationIconStyleProviderImpl @Inject constructor(dumpManager: DumpManager) :
+class NotificationIconStyleProviderImpl
+@Inject
+constructor(private val userManager: UserManager, dumpManager: DumpManager) :
     NotificationIconStyleProvider, Dumpable {
     init {
         dumpManager.registerNormalDumpable(TAG, this)
@@ -89,6 +98,15 @@
         }
     }
 
+    override fun shouldShowWorkProfileBadge(
+        notification: StatusBarNotification,
+        context: Context,
+    ): Boolean {
+        val packageContext = notification.getPackageContext(context)
+        // UserManager already caches this, so we don't need to.
+        return userManager.isManagedProfile(packageContext.userId)
+    }
+
     override fun purgeCache(wantedPackages: Collection<String>) {
         cache.purge(wantedPackages)
     }
@@ -114,6 +132,14 @@
         return true
     }
 
+    override fun shouldShowWorkProfileBadge(
+        notification: StatusBarNotification,
+        context: Context,
+    ): Boolean {
+        Log.wtf(TAG, "NoOpIconStyleProvider should not be used anywhere.")
+        return false
+    }
+
     override fun purgeCache(wantedPackages: Collection<String>) {
         Log.wtf(TAG, "NoOpIconStyleProvider should not be used anywhere.")
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt
index 7b85bfd..64fdf6f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt
@@ -50,6 +50,7 @@
                 NotificationRowIconView(context, attrs).also { view ->
                     view.setIconProvider(createIconProvider(row, context))
                 }
+
             else -> null
         }
     }
@@ -61,13 +62,19 @@
         val sbn = row.entry.sbn
         return object : NotificationIconProvider {
             override fun shouldShowAppIcon(): Boolean {
-                val shouldShowAppIcon = iconStyleProvider.shouldShowAppIcon(row.entry.sbn, context)
+                val shouldShowAppIcon = iconStyleProvider.shouldShowAppIcon(sbn, context)
                 row.setIsShowingAppIcon(shouldShowAppIcon)
                 return shouldShowAppIcon
             }
 
             override fun getAppIcon(): Drawable {
-                return appIconProvider.getOrFetchAppIcon(sbn.packageName, context)
+                val withWorkProfileBadge =
+                    iconStyleProvider.shouldShowWorkProfileBadge(sbn, context)
+                return appIconProvider.getOrFetchAppIcon(
+                    sbn.packageName,
+                    context,
+                    withWorkProfileBadge,
+                )
             }
         }
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
index c3996e40..31d3908 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
@@ -23,6 +23,7 @@
 import android.content.Context
 import android.content.pm.LauncherApps
 import android.os.UserHandle
+import android.os.UserManager
 import android.provider.DeviceConfig
 import androidx.core.os.bundleOf
 import com.android.internal.config.sysui.SystemUiDeviceConfigFlags
@@ -109,6 +110,7 @@
     private val mKeyguardBypassController: KeyguardBypassController
     private val mGroupMembershipManager: GroupMembershipManager
     private val mGroupExpansionManager: GroupExpansionManager
+    private val mUserManager: UserManager
     private val mHeadsUpManager: HeadsUpManager
     private val mIconManager: IconManager
     private val mContentBinder: NotificationRowContentBinder
@@ -143,6 +145,7 @@
         mSmartReplyController = Mockito.mock(SmartReplyController::class.java, STUB_ONLY)
 
         mGroupExpansionManager = GroupExpansionManagerImpl(mDumpManager, mGroupMembershipManager)
+        mUserManager = Mockito.mock(UserManager::class.java, STUB_ONLY)
         mHeadsUpManager = Mockito.mock(HeadsUpManager::class.java, STUB_ONLY)
         mIconManager =
             IconManager(
@@ -289,7 +292,7 @@
             { Mockito.mock(NotificationViewFlipperFactory::class.java) },
             NotificationRowIconViewInflaterFactory(
                 AppIconProviderImpl(context, mDumpManager),
-                NotificationIconStyleProviderImpl(mDumpManager),
+                NotificationIconStyleProviderImpl(mUserManager, mDumpManager),
             ),
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProviderKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProviderKosmos.kt
index 0fe84fb..b4fb7dd 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProviderKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/icon/NotificationIconStyleProviderKosmos.kt
@@ -16,8 +16,9 @@
 
 package com.android.systemui.statusbar.notification.row.icon
 
+import android.os.userManager
 import com.android.systemui.dump.dumpManager
 import com.android.systemui.kosmos.Kosmos
 
 val Kosmos.notificationIconStyleProvider by
-    Kosmos.Fixture { NotificationIconStyleProviderImpl(dumpManager) }
+    Kosmos.Fixture { NotificationIconStyleProviderImpl(userManager, dumpManager) }