Merge changes Ie7eef8c7,I6ef2411e,I50e9806a,I632daa93 into main

* changes:
  Implement hit-ratio dumping for NotifCollectionCache
  Introduce NotifCollectionCache
  Notif redesign: reduce conversation icon margin
  Notif redesign: Don't apply effects to app icons
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCacheTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCacheTest.kt
new file mode 100644
index 0000000..d2a3a19
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCacheTest.kt
@@ -0,0 +1,232 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.notification.collection
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.time.FakeSystemClock
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.assertNotNull
+import kotlin.test.assertNull
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class NotifCollectionCacheTest : SysuiTestCase() {
+    companion object {
+        const val A = "a"
+        const val B = "b"
+        const val C = "c"
+    }
+
+    val systemClock = FakeSystemClock()
+    val underTest =
+        NotifCollectionCache<String>(purgeTimeoutMillis = 200L, systemClock = systemClock)
+
+    @After
+    fun cleanUp() {
+        underTest.clear()
+    }
+
+    @Test
+    fun fetch_isOnlyCalledOncePerEntry() {
+        val fetchList = mutableListOf<String>()
+        val fetch = { key: String ->
+            fetchList.add(key)
+            key
+        }
+
+        // Construct the cache and make sure fetch is called
+        assertThat(underTest.getOrFetch(A, fetch)).isEqualTo(A)
+        assertThat(underTest.getOrFetch(B, fetch)).isEqualTo(B)
+        assertThat(underTest.getOrFetch(C, fetch)).isEqualTo(C)
+        assertThat(fetchList).containsExactly(A, B, C).inOrder()
+
+        // Verify that further calls don't trigger fetch again
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+        assertThat(fetchList).containsExactly(A, B, C).inOrder()
+
+        // Verify that fetch gets called again if the entries are cleared
+        underTest.clear()
+        underTest.getOrFetch(A, fetch)
+        assertThat(fetchList).containsExactly(A, B, C, A).inOrder()
+    }
+
+    @Test
+    fun purge_beforeTimeout_doesNothing() {
+        // Populate cache
+        val fetch = { key: String -> key }
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+
+        // B starts off with ♥ ︎♥︎
+        assertThat(underTest.getLives(B)).isEqualTo(2)
+        // First purge run removes a ︎♥︎
+        underTest.purge(listOf(A, C))
+        assertNotNull(underTest.cache[B])
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+        // Second purge run done too early does nothing to B
+        systemClock.advanceTime(100L)
+        underTest.purge(listOf(A, C))
+        assertNotNull(underTest.cache[B])
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+        // Purge done after timeout (200ms) clears B
+        systemClock.advanceTime(100L)
+        underTest.purge(listOf(A, C))
+        assertNull(underTest.cache[B])
+    }
+
+    @Test
+    fun get_resetsLives() {
+        // Populate cache
+        val fetch = { key: String -> key }
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+
+        // Bring B down to one ︎♥︎
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+
+        // Get should restore B to ♥ ︎♥︎
+        underTest.getOrFetch(B, fetch)
+        assertThat(underTest.getLives(B)).isEqualTo(2)
+
+        // Subsequent purge should remove a life regardless of timing
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+    }
+
+    @Test
+    fun purge_resetsLives() {
+        // Populate cache
+        val fetch = { key: String -> key }
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+
+        // Bring B down to one ︎♥︎
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+
+        // When B is back to wantedKeys, it is restored to to ♥ ︎♥ ︎︎
+        underTest.purge(listOf(B))
+        assertThat(underTest.getLives(B)).isEqualTo(2)
+        assertThat(underTest.getLives(A)).isEqualTo(1)
+        assertThat(underTest.getLives(C)).isEqualTo(1)
+
+        // Subsequent purge should remove a life regardless of timing
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+    }
+
+    @Test
+    fun purge_worksWithMoreLives() {
+        val multiLivesCache =
+            NotifCollectionCache<String>(
+                retainCount = 3,
+                purgeTimeoutMillis = 100L,
+                systemClock = systemClock,
+            )
+
+        // Populate cache
+        val fetch = { key: String -> key }
+        multiLivesCache.getOrFetch(A, fetch)
+        multiLivesCache.getOrFetch(B, fetch)
+        multiLivesCache.getOrFetch(C, fetch)
+
+        // B starts off with ♥ ︎♥︎ ♥ ︎♥︎
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(4)
+        // First purge run removes a ︎♥︎
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(3)
+        // Second purge run done too early does nothing to B
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(3)
+        // Staggered purge runs remove further ︎♥︎
+        systemClock.advanceTime(100L)
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(2)
+        systemClock.advanceTime(100L)
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(1)
+        systemClock.advanceTime(100L)
+        multiLivesCache.purge(listOf(A, C))
+        assertNull(multiLivesCache.cache[B])
+    }
+
+    @Test
+    fun purge_worksWithNoLives() {
+        val noLivesCache =
+            NotifCollectionCache<String>(
+                retainCount = 0,
+                purgeTimeoutMillis = 0L,
+                systemClock = systemClock,
+            )
+
+        val fetch = { key: String -> key }
+        noLivesCache.getOrFetch(A, fetch)
+        noLivesCache.getOrFetch(B, fetch)
+        noLivesCache.getOrFetch(C, fetch)
+
+        // Purge immediately removes entry
+        noLivesCache.purge(listOf(A, C))
+
+        assertNotNull(noLivesCache.cache[A])
+        assertNull(noLivesCache.cache[B])
+        assertNotNull(noLivesCache.cache[C])
+    }
+
+    @Test
+    fun hitsAndMisses_areAccurate() {
+        val fetch = { key: String -> key }
+
+        // Construct the cache
+        assertThat(underTest.getOrFetch(A, fetch)).isEqualTo(A)
+        assertThat(underTest.getOrFetch(B, fetch)).isEqualTo(B)
+        assertThat(underTest.getOrFetch(C, fetch)).isEqualTo(C)
+        assertThat(underTest.hits.get()).isEqualTo(0)
+        assertThat(underTest.misses.get()).isEqualTo(3)
+
+        // Verify that further calls count as hits
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+        assertThat(underTest.hits.get()).isEqualTo(4)
+        assertThat(underTest.misses.get()).isEqualTo(3)
+
+        // Verify that a miss is counted again if the entries are cleared
+        underTest.clear()
+        underTest.getOrFetch(A, fetch)
+        assertThat(underTest.hits.get()).isEqualTo(4)
+        assertThat(underTest.misses.get()).isEqualTo(4)
+    }
+
+    private fun <V> NotifCollectionCache<V>.getLives(key: String) = this.cache[key]?.lives
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt
index 9d990b1..9a6a699 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt
@@ -26,6 +26,7 @@
 import com.android.internal.widget.MessagingGroup
 import com.android.internal.widget.MessagingImageMessage
 import com.android.internal.widget.MessagingLinearLayout
+import com.android.internal.widget.NotificationRowIconView
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper
@@ -90,7 +91,7 @@
 
     private fun fakeConversationLayout(
         mockDrawableGroupMessage: AnimatedImageDrawable,
-        mockDrawableImageMessage: AnimatedImageDrawable
+        mockDrawableImageMessage: AnimatedImageDrawable,
     ): View {
         val mockMessagingImageMessage: MessagingImageMessage =
             mock<MessagingImageMessage>().apply {
@@ -126,6 +127,7 @@
                 whenever(requireViewById<CachingIconView>(R.id.conversation_icon))
                     .thenReturn(mock())
                 whenever(findViewById<CachingIconView>(R.id.icon)).thenReturn(mock())
+                whenever(requireViewById<NotificationRowIconView>(R.id.icon)).thenReturn(mock())
                 whenever(requireViewById<View>(R.id.conversation_icon_badge_bg)).thenReturn(mock())
                 whenever(requireViewById<View>(R.id.expand_button)).thenReturn(mock())
                 whenever(requireViewById<View>(R.id.expand_button_container)).thenReturn(mock())
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
index 8d3f728..30f564f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar;
 
+import android.app.Flags;
 import android.app.Notification;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.Icon;
@@ -60,20 +61,6 @@
             return row.getEntry().getSbn().getNotification();
         }
     };
-    private static final IconComparator ICON_VISIBILITY_COMPARATOR = new IconComparator() {
-        public boolean compare(View parent, View child, Object parentData,
-                Object childData) {
-            return hasSameIcon(parentData, childData)
-                    && hasSameColor(parentData, childData);
-        }
-    };
-    private static final IconComparator GREY_COMPARATOR = new IconComparator() {
-        public boolean compare(View parent, View child, Object parentData,
-                Object childData) {
-            return !hasSameIcon(parentData, childData)
-                    || hasSameColor(parentData, childData);
-        }
-    };
     private static final ResultApplicator GREY_APPLICATOR = new ResultApplicator() {
         @Override
         public void apply(View parent, View view, boolean apply, boolean reset) {
@@ -90,34 +77,58 @@
 
     public NotificationGroupingUtil(ExpandableNotificationRow row) {
         mRow = row;
+
+        final IconComparator iconVisibilityComparator = new IconComparator(mRow) {
+            public boolean compare(View parent, View child, Object parentData,
+                    Object childData) {
+                return hasSameIcon(parentData, childData)
+                        && hasSameColor(parentData, childData);
+            }
+        };
+        final IconComparator greyComparator = new IconComparator(mRow) {
+            public boolean compare(View parent, View child, Object parentData,
+                    Object childData) {
+                if (Flags.notificationsRedesignAppIcons() && mRow.isShowingAppIcon()) {
+                    return false;
+                }
+                return !hasSameIcon(parentData, childData)
+                        || hasSameColor(parentData, childData);
+            }
+        };
+
         // To hide the icons if they are the same and the color is the same
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.icon,
                 ICON_EXTRACTOR,
-                ICON_VISIBILITY_COMPARATOR,
+                iconVisibilityComparator,
                 VISIBILITY_APPLICATOR));
-        // To grey them out the icons and expand button when the icons are not the same
+        // To grey out the icons when they are not the same, or they have the same color
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.status_bar_latest_event_content,
                 ICON_EXTRACTOR,
-                GREY_COMPARATOR,
+                greyComparator,
                 GREY_APPLICATOR));
+        // To show the large icon on the left side instead if all the small icons are the same
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.status_bar_latest_event_content,
                 ICON_EXTRACTOR,
-                ICON_VISIBILITY_COMPARATOR,
+                iconVisibilityComparator,
                 LEFT_ICON_APPLICATOR));
+        // To only show the work profile icon in the group header
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.profile_badge,
                 null /* Extractor */,
                 BADGE_COMPARATOR,
                 VISIBILITY_APPLICATOR));
+        // To hide the app name in group children
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.app_name_text,
                 null,
                 APP_NAME_COMPARATOR,
                 APP_NAME_APPLICATOR));
+        // To hide the header text if it's the same
         mProcessors.add(Processor.forTextView(mRow, com.android.internal.R.id.header_text));
+
         mDividers.add(com.android.internal.R.id.header_text_divider);
         mDividers.add(com.android.internal.R.id.header_text_secondary_divider);
         mDividers.add(com.android.internal.R.id.time_divider);
@@ -261,6 +272,7 @@
             mParentData = mExtractor == null ? null : mExtractor.extractData(mParentRow);
             mApply = !mComparator.isEmpty(mParentView);
         }
+
         public void compareToGroupParent(ExpandableNotificationRow row) {
             if (!mApply) {
                 return;
@@ -356,12 +368,21 @@
     }
 
     private abstract static class IconComparator implements ViewComparator {
+        private final ExpandableNotificationRow mRow;
+
+        IconComparator(ExpandableNotificationRow row) {
+            mRow = row;
+        }
+
         @Override
         public boolean compare(View parent, View child, Object parentData, Object childData) {
             return false;
         }
 
         protected boolean hasSameIcon(Object parentData, Object childData) {
+            if (Flags.notificationsRedesignAppIcons() && mRow.isShowingAppIcon()) {
+                return true;
+            }
             Icon parentIcon = getIcon((Notification) parentData);
             Icon childIcon = getIcon((Notification) childData);
             return parentIcon.sameAs(childIcon);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCache.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCache.kt
new file mode 100644
index 0000000..2ee1dffd
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCache.kt
@@ -0,0 +1,188 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.notification.collection
+
+import android.annotation.SuppressLint
+import com.android.internal.annotations.VisibleForTesting
+import com.android.systemui.Dumpable
+import com.android.systemui.util.asIndenting
+import com.android.systemui.util.printCollection
+import com.android.systemui.util.time.SystemClock
+import com.android.systemui.util.time.SystemClockImpl
+import com.android.systemui.util.withIncreasedIndent
+import java.io.PrintWriter
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicInteger
+
+/**
+ * A cache in which entries can "survive" getting purged [retainCount] times, given consecutive
+ * [purge] calls made at least [purgeTimeoutMillis] apart. See also [purge].
+ *
+ * This cache is safe for multithreaded usage, and is recommended for objects that take a while to
+ * resolve (such as drawables, or things that require binder calls). As such, [getOrFetch] is
+ * recommended to be run on a background thread, while [purge] can be done from any thread.
+ */
+@SuppressLint("DumpableNotRegistered") // this will be dumped by container classes
+class NotifCollectionCache<V>(
+    private val retainCount: Int = 1,
+    private val purgeTimeoutMillis: Long = 1000L,
+    private val systemClock: SystemClock = SystemClockImpl(),
+) : Dumpable {
+    @get:VisibleForTesting val cache = ConcurrentHashMap<String, CacheEntry>()
+
+    // Counters for cache hits and misses to be used to calculate and dump the hit ratio
+    @get:VisibleForTesting val misses = AtomicInteger(0)
+    @get:VisibleForTesting val hits = AtomicInteger(0)
+
+    init {
+        if (retainCount < 0) {
+            throw IllegalArgumentException("retainCount cannot be negative")
+        }
+    }
+
+    inner class CacheEntry(val key: String, val value: V) {
+        /**
+         * The "lives" represent how many times the entry will remain in the cache when purging it
+         * is attempted.
+         */
+        @get:VisibleForTesting var lives: Int = retainCount + 1
+        /**
+         * The last time this entry lost a "life". Starts at a negative value chosen so that the
+         * first purge is always considered "valid".
+         */
+        private var lastValidPurge: Long = -purgeTimeoutMillis
+
+        fun resetLives() {
+            // Lives/timeouts don't matter if retainCount is 0
+            if (retainCount == 0) {
+                return
+            }
+
+            synchronized(key) {
+                lives = retainCount + 1
+                lastValidPurge = -purgeTimeoutMillis
+            }
+            // Add it to the cache again just in case it was deleted before we could reset the lives
+            cache[key] = this
+        }
+
+        fun tryPurge(): Boolean {
+            // Lives/timeouts don't matter if retainCount is 0
+            if (retainCount == 0) {
+                return true
+            }
+
+            // Using uptimeMillis since it's guaranteed to be monotonic, as we don't want a
+            // timezone/clock change to break us
+            val now = systemClock.uptimeMillis()
+
+            // Cannot purge the same entry from two threads simultaneously
+            synchronized(key) {
+                if (now - lastValidPurge < purgeTimeoutMillis) {
+                    return false
+                }
+                lastValidPurge = now
+                return --lives <= 0
+            }
+        }
+    }
+
+    /**
+     * Get value from cache, or fetch it and add it to cache if not found. This can be called from
+     * any thread, but is usually expected to be called from the background.
+     *
+     * @param key key for the object to be obtained
+     * @param fetch method to fetch the object and add it to the cache if not present; note that
+     *   there is no guarantee that two [fetch] cannot run in parallel for the same [key] (if
+     *   [getOrFetch] is called simultaneously from different threads), so be mindful of potential
+     *   side effects
+     */
+    fun getOrFetch(key: String, fetch: (String) -> V): V {
+        val entry = cache[key]
+        if (entry != null) {
+            hits.incrementAndGet()
+            // Refresh lives on access
+            entry.resetLives()
+            return entry.value
+        }
+
+        misses.incrementAndGet()
+        val value = fetch(key)
+        cache[key] = CacheEntry(key, value)
+        return value
+    }
+
+    /**
+     * Clear entries that are NOT in [wantedKeys] if appropriate. This can be called from any
+     * thread.
+     *
+     * If retainCount > 0, a given entry will need to not be present in [wantedKeys] for
+     * ([retainCount] + 1) consecutive [purge] calls made within at least [purgeTimeoutMillis] of
+     * each other in order to be cleared. This count will be reset for any given entry 1) if
+     * [getOrFetch] is called for the entry or 2) if the entry is present in [wantedKeys] in a
+     * subsequent [purge] call. We prioritize keeping the entry if possible, so if [purge] is called
+     * simultaneously with [getOrFetch] on different threads for example, we will try to keep it in
+     * the cache, although it is not guaranteed. If avoiding cache misses is a concern, consider
+     * increasing the [retainCount] or [purgeTimeoutMillis].
+     *
+     * For example, say [retainCount] = 1 and [purgeTimeoutMillis] = 1000 and we start with entries
+     * (a, b, c) in the cache:
+     * ```kotlin
+     * purge((a, c)); // marks b for deletion
+     * Thread.sleep(500)
+     * purge((a, c)); // does nothing as it was called earlier than the min 1s
+     * Thread.sleep(500)
+     * purge((b, c)); // b is no longer marked for deletion, but now a is
+     * Thread.sleep(1000);
+     * purge((c));    // deletes a from the cache and marks b for deletion, etc.
+     * ```
+     */
+    fun purge(wantedKeys: List<String>) {
+        for ((key, entry) in cache) {
+            if (key in wantedKeys) {
+                entry.resetLives()
+            } else if (entry.tryPurge()) {
+                cache.remove(key)
+            }
+        }
+    }
+
+    /** Clear all entries from the cache. */
+    fun clear() {
+        cache.clear()
+    }
+
+    override fun dump(pwOrig: PrintWriter, args: Array<out String>) {
+        val pw = pwOrig.asIndenting()
+
+        pw.println("$TAG(retainCount = $retainCount, purgeTimeoutMillis = $purgeTimeoutMillis)")
+        pw.withIncreasedIndent {
+            pw.printCollection("keys present in cache", cache.keys.stream().sorted().toList())
+
+            val misses = misses.get()
+            val hits = hits.get()
+            pw.println(
+                "cache hit ratio = ${(hits.toFloat() / (hits + misses)) * 100}% " +
+                    "($hits hits, $misses misses)"
+            )
+        }
+    }
+
+    companion object {
+        const val TAG = "NotifCollectionCache"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index 38e6609..933f793 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -262,6 +262,11 @@
      */
     private boolean mIsHeadsUp;
 
+    /**
+     * Whether or not the notification is showing the app icon instead of the small icon.
+     */
+    private boolean mIsShowingAppIcon;
+
     private boolean mLastChronometerRunning = true;
     private ViewStub mChildrenContainerStub;
     private GroupMembershipManager mGroupMembershipManager;
@@ -816,6 +821,20 @@
         }
     }
 
+    /**
+     * Indicate that the notification is showing the app icon instead of the small icon.
+     */
+    public void setIsShowingAppIcon(boolean isShowingAppIcon) {
+        mIsShowingAppIcon = isShowingAppIcon;
+    }
+
+    /**
+     * Whether or not the notification is showing the app icon instead of the small icon.
+     */
+    public boolean isShowingAppIcon() {
+        return mIsShowingAppIcon;
+    }
+
     @Override
     public boolean showingPulsing() {
         return isHeadsUpState() && (isDozing() || (mOnKeyguard && isBypassEnabled()));
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 79defd2..7b85bfd 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
@@ -21,6 +21,7 @@
 import android.util.AttributeSet
 import android.view.View
 import com.android.internal.widget.NotificationRowIconView
+import com.android.internal.widget.NotificationRowIconView.NotificationIconProvider
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.NotifRemoteViewsFactory
 import com.android.systemui.statusbar.notification.row.NotificationRowContentBinder
@@ -47,20 +48,27 @@
         return when (name) {
             NotificationRowIconView::class.java.name ->
                 NotificationRowIconView(context, attrs).also { view ->
-                    val sbn = row.entry.sbn
-                    view.setIconProvider(
-                        object : NotificationRowIconView.NotificationIconProvider {
-                            override fun shouldShowAppIcon(): Boolean {
-                                return iconStyleProvider.shouldShowAppIcon(row.entry.sbn, context)
-                            }
-
-                            override fun getAppIcon(): Drawable {
-                                return appIconProvider.getOrFetchAppIcon(sbn.packageName, context)
-                            }
-                        }
-                    )
+                    view.setIconProvider(createIconProvider(row, context))
                 }
             else -> null
         }
     }
+
+    private fun createIconProvider(
+        row: ExpandableNotificationRow,
+        context: Context,
+    ): NotificationIconProvider {
+        val sbn = row.entry.sbn
+        return object : NotificationIconProvider {
+            override fun shouldShowAppIcon(): Boolean {
+                val shouldShowAppIcon = iconStyleProvider.shouldShowAppIcon(row.entry.sbn, context)
+                row.setIsShowingAppIcon(shouldShowAppIcon)
+                return shouldShowAppIcon
+            }
+
+            override fun getAppIcon(): Drawable {
+                return appIconProvider.getOrFetchAppIcon(sbn.packageName, context)
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt
index b4411f1..f8aff69 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt
@@ -16,15 +16,18 @@
 
 package com.android.systemui.statusbar.notification.row.wrapper
 
+import android.app.Flags
 import android.content.Context
 import android.graphics.drawable.AnimatedImageDrawable
 import android.view.View
 import android.view.ViewGroup
+import android.view.ViewGroup.MarginLayoutParams
 import com.android.internal.widget.CachingIconView
 import com.android.internal.widget.ConversationLayout
 import com.android.internal.widget.MessagingGroup
 import com.android.internal.widget.MessagingImageMessage
 import com.android.internal.widget.MessagingLinearLayout
+import com.android.internal.widget.NotificationRowIconView
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.notification.NotificationFadeAware
 import com.android.systemui.statusbar.notification.NotificationUtils
@@ -32,23 +35,23 @@
 import com.android.systemui.statusbar.notification.row.wrapper.NotificationMessagingTemplateViewWrapper.setCustomImageMessageTransform
 import com.android.systemui.util.children
 
-/**
- * Wraps a notification containing a conversation template
- */
-class NotificationConversationTemplateViewWrapper constructor(
+/** Wraps a notification containing a conversation template */
+class NotificationConversationTemplateViewWrapper(
     ctx: Context,
     view: View,
-    row: ExpandableNotificationRow
+    row: ExpandableNotificationRow,
 ) : NotificationTemplateViewWrapper(ctx, view, row) {
 
-    private val minHeightWithActions: Int = NotificationUtils.getFontScaledHeight(
+    private val minHeightWithActions: Int =
+        NotificationUtils.getFontScaledHeight(
             ctx,
-            R.dimen.notification_messaging_actions_min_height
-    )
+            R.dimen.notification_messaging_actions_min_height,
+        )
     private val conversationLayout: ConversationLayout = view as ConversationLayout
 
     private lateinit var conversationIconContainer: View
     private lateinit var conversationIconView: CachingIconView
+    private lateinit var badgeIconView: NotificationRowIconView
     private lateinit var conversationBadgeBg: View
     private lateinit var expandBtn: View
     private lateinit var expandBtnContainer: View
@@ -68,10 +71,13 @@
         messageContainers = conversationLayout.messagingGroups
         with(conversationLayout) {
             conversationIconContainer =
-                    requireViewById(com.android.internal.R.id.conversation_icon_container)
+                requireViewById(com.android.internal.R.id.conversation_icon_container)
             conversationIconView = requireViewById(com.android.internal.R.id.conversation_icon)
+            if (Flags.notificationsRedesignAppIcons()) {
+                badgeIconView = requireViewById(com.android.internal.R.id.icon)
+            }
             conversationBadgeBg =
-                    requireViewById(com.android.internal.R.id.conversation_icon_badge_bg)
+                requireViewById(com.android.internal.R.id.conversation_icon_badge_bg)
             expandBtn = requireViewById(com.android.internal.R.id.expand_button)
             expandBtnContainer = requireViewById(com.android.internal.R.id.expand_button_container)
             importanceRing = requireViewById(com.android.internal.R.id.conversation_icon_badge_ring)
@@ -80,7 +86,7 @@
             facePileTop = findViewById(com.android.internal.R.id.conversation_face_pile_top)
             facePileBottom = findViewById(com.android.internal.R.id.conversation_face_pile_bottom)
             facePileBottomBg =
-                    findViewById(com.android.internal.R.id.conversation_face_pile_bottom_background)
+                findViewById(com.android.internal.R.id.conversation_face_pile_bottom_background)
         }
     }
 
@@ -88,6 +94,13 @@
         // Reinspect the notification. Before the super call, because the super call also updates
         // the transformation types and we need to have our values set by then.
         resolveViews()
+        if (Flags.notificationsRedesignAppIcons() && row.isShowingAppIcon) {
+            // Override the margins to be 2dp instead of 4dp according to the new design if we're
+            // showing the app icon.
+            val lp = badgeIconView.layoutParams as MarginLayoutParams
+            lp.setMargins(2, 2, 2, 2)
+            badgeIconView.layoutParams = lp
+        }
         super.onContentUpdated(row)
     }
 
@@ -96,56 +109,50 @@
         super.updateTransformedTypes()
 
         mTransformationHelper.addTransformedView(TRANSFORMING_VIEW_TITLE, conversationTitleView)
-        addTransformedViews(
-                messagingLinearLayout,
-                appName
-        )
+        addTransformedViews(messagingLinearLayout, appName)
 
         setCustomImageMessageTransform(mTransformationHelper, imageMessageContainer)
 
         addViewsTransformingToSimilar(
-                conversationIconView,
-                conversationBadgeBg,
-                expandBtn,
-                importanceRing,
-                facePileTop,
-                facePileBottom,
-                facePileBottomBg
+            conversationIconView,
+            conversationBadgeBg,
+            expandBtn,
+            importanceRing,
+            facePileTop,
+            facePileBottom,
+            facePileBottomBg,
         )
     }
 
     override fun getShelfTransformationTarget(): View? =
-            if (conversationLayout.isImportantConversation)
-                if (conversationIconView.visibility != View.GONE)
-                    conversationIconView
-                else
-                    // A notification with a fallback icon was set to important. Currently
-                    // the transformation doesn't work for these and needs to be fixed.
-                    // In the meantime those are using the icon.
-                    super.getShelfTransformationTarget()
+        if (conversationLayout.isImportantConversation)
+            if (conversationIconView.visibility != View.GONE) conversationIconView
             else
-                super.getShelfTransformationTarget()
+            // A notification with a fallback icon was set to important. Currently
+            // the transformation doesn't work for these and needs to be fixed.
+            // In the meantime those are using the icon.
+            super.getShelfTransformationTarget()
+        else super.getShelfTransformationTarget()
 
     override fun setRemoteInputVisible(visible: Boolean) =
-            conversationLayout.showHistoricMessages(visible)
+        conversationLayout.showHistoricMessages(visible)
 
     override fun updateExpandability(
         expandable: Boolean,
         onClickListener: View.OnClickListener,
-        requestLayout: Boolean
+        requestLayout: Boolean,
     ) = conversationLayout.updateExpandability(expandable, onClickListener)
 
     override fun disallowSingleClick(x: Float, y: Float): Boolean {
-        val isOnExpandButton = expandBtnContainer.visibility == View.VISIBLE &&
-                isOnView(expandBtnContainer, x, y)
+        val isOnExpandButton =
+            expandBtnContainer.visibility == View.VISIBLE && isOnView(expandBtnContainer, x, y)
         return isOnExpandButton || super.disallowSingleClick(x, y)
     }
 
     override fun getMinLayoutHeight(): Int =
-            if (mActionsContainer != null && mActionsContainer.visibility != View.GONE)
-                minHeightWithActions
-            else
-                super.getMinLayoutHeight()
+        if (mActionsContainer != null && mActionsContainer.visibility != View.GONE)
+            minHeightWithActions
+        else super.getMinLayoutHeight()
 
     override fun setNotificationFaded(faded: Boolean) {
         // Do not call super
@@ -157,16 +164,17 @@
     override fun setAnimationsRunning(running: Boolean) {
         // We apply to both the child message containers in a conversation group,
         // and the top level image message container.
-        val containers = messageContainers.asSequence().map { it.messageContainer } +
+        val containers =
+            messageContainers.asSequence().map { it.messageContainer } +
                 sequenceOf(imageMessageContainer)
         val drawables =
-                containers
-                        .flatMap { it.children }
-                        .mapNotNull { child ->
-                            (child as? MessagingImageMessage)?.let { imageMessage ->
-                                imageMessage.drawable as? AnimatedImageDrawable
-                            }
-                        }
+            containers
+                .flatMap { it.children }
+                .mapNotNull { child ->
+                    (child as? MessagingImageMessage)?.let { imageMessage ->
+                        imageMessage.drawable as? AnimatedImageDrawable
+                    }
+                }
         drawables.toSet().forEach {
             when {
                 running -> it.start()