Merge "Change KeyguardUpdateMonitorLogger to accept nullable params" into tm-qpr-dev
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java b/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
index 2e51b51..b69afeb 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeLog.java
@@ -287,8 +287,8 @@
     /**
      * Appends sensor event dropped event to logs
      */
-    public void traceSensorEventDropped(int sensorEvent, String reason) {
-        mLogger.logSensorEventDropped(sensorEvent, reason);
+    public void traceSensorEventDropped(@Reason int pulseReason, String reason) {
+        mLogger.logSensorEventDropped(pulseReason, reason);
     }
 
     /**
@@ -386,6 +386,47 @@
         mLogger.logSetAodDimmingScrim((long) scrimOpacity);
     }
 
+    /**
+     * Appends sensor attempted to register and whether it was a successful registration.
+     */
+    public void traceSensorRegisterAttempt(String sensorName, boolean successfulRegistration) {
+        mLogger.logSensorRegisterAttempt(sensorName, successfulRegistration);
+    }
+
+    /**
+     * Appends sensor attempted to unregister and whether it was successfully unregistered.
+     */
+    public void traceSensorUnregisterAttempt(String sensorInfo, boolean successfullyUnregistered) {
+        mLogger.logSensorUnregisterAttempt(sensorInfo, successfullyUnregistered);
+    }
+
+    /**
+     * Appends sensor attempted to unregister and whether it was successfully unregistered
+     * with a reason the sensor is being unregistered.
+     */
+    public void traceSensorUnregisterAttempt(String sensorInfo, boolean successfullyUnregistered,
+            String reason) {
+        mLogger.logSensorUnregisterAttempt(sensorInfo, successfullyUnregistered, reason);
+    }
+
+    /**
+     * Appends the event of skipping a sensor registration since it's already registered.
+     */
+    public void traceSkipRegisterSensor(String sensorInfo) {
+        mLogger.logSkipSensorRegistration(sensorInfo);
+    }
+
+    /**
+     * Appends a plugin sensor was registered or unregistered event.
+     */
+    public void tracePluginSensorUpdate(boolean registered) {
+        if (registered) {
+            mLogger.log("register plugin sensor");
+        } else {
+            mLogger.log("unregister plugin sensor");
+        }
+    }
+
     private class SummaryStats {
         private int mCount;
 
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt b/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt
index f8e2566..18c8e01 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeLogger.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.plugins.log.LogLevel.INFO
 import com.android.systemui.log.dagger.DozeLog
 import com.android.systemui.statusbar.policy.DevicePostureController
+import com.google.errorprone.annotations.CompileTimeConstant
 import java.text.SimpleDateFormat
 import java.util.Date
 import java.util.Locale
@@ -324,6 +325,50 @@
             "Doze car mode started"
         })
     }
+
+    fun logSensorRegisterAttempt(sensorInfo: String, successfulRegistration: Boolean) {
+        buffer.log(TAG, INFO, {
+            str1 = sensorInfo
+            bool1 = successfulRegistration
+        }, {
+            "Register sensor. Success=$bool1 sensor=$str1"
+        })
+    }
+
+    fun logSensorUnregisterAttempt(sensorInfo: String, successfulUnregister: Boolean) {
+        buffer.log(TAG, INFO, {
+            str1 = sensorInfo
+            bool1 = successfulUnregister
+        }, {
+            "Unregister sensor. Success=$bool1 sensor=$str1"
+        })
+    }
+
+    fun logSensorUnregisterAttempt(
+            sensorInfo: String,
+            successfulUnregister: Boolean,
+            reason: String
+    ) {
+        buffer.log(TAG, INFO, {
+            str1 = sensorInfo
+            bool1 = successfulUnregister
+            str2 = reason
+        }, {
+            "Unregister sensor. reason=$str2. Success=$bool1 sensor=$str1"
+        })
+    }
+
+    fun logSkipSensorRegistration(sensor: String) {
+        buffer.log(TAG, DEBUG, {
+            str1 = sensor
+        }, {
+            "Skipping sensor registration because its already registered. sensor=$str1"
+        })
+    }
+
+    fun log(@CompileTimeConstant msg: String) {
+        buffer.log(TAG, DEBUG, msg)
+    }
 }
 
 private const val TAG = "DozeLog"
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java b/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
index da6c163..9a091e7 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
@@ -23,7 +23,6 @@
 
 import android.annotation.AnyThread;
 import android.app.ActivityManager;
-import android.content.Context;
 import android.database.ContentObserver;
 import android.hardware.Sensor;
 import android.hardware.SensorManager;
@@ -37,7 +36,6 @@
 import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.IndentingPrintWriter;
-import android.util.Log;
 import android.view.Display;
 
 import androidx.annotation.NonNull;
@@ -88,12 +86,9 @@
  * trigger callbacks on the provided {@link mProxCallback}.
  */
 public class DozeSensors {
-
-    private static final boolean DEBUG = DozeService.DEBUG;
     private static final String TAG = "DozeSensors";
     private static final UiEventLogger UI_EVENT_LOGGER = new UiEventLoggerImpl();
 
-    private final Context mContext;
     private final AsyncSensorManager mSensorManager;
     private final AmbientDisplayConfiguration mConfig;
     private final WakeLock mWakeLock;
@@ -144,7 +139,6 @@
     }
 
     DozeSensors(
-            Context context,
             AsyncSensorManager sensorManager,
             DozeParameters dozeParameters,
             AmbientDisplayConfiguration config,
@@ -157,7 +151,6 @@
             AuthController authController,
             DevicePostureController devicePostureController
     ) {
-        mContext = context;
         mSensorManager = sensorManager;
         mConfig = config;
         mWakeLock = wakeLock;
@@ -605,10 +598,7 @@
             // cancel the previous sensor:
             if (mRegistered) {
                 final boolean rt = mSensorManager.cancelTriggerSensor(this, oldSensor);
-                if (DEBUG) {
-                    Log.d(TAG, "posture changed, cancelTriggerSensor[" + oldSensor + "] "
-                            + rt);
-                }
+                mDozeLog.traceSensorUnregisterAttempt(oldSensor.toString(), rt, "posture changed");
                 mRegistered = false;
             }
 
@@ -654,19 +644,13 @@
             if (mRequested && !mDisabled && (enabledBySetting() || mIgnoresSetting)) {
                 if (!mRegistered) {
                     mRegistered = mSensorManager.requestTriggerSensor(this, sensor);
-                    if (DEBUG) {
-                        Log.d(TAG, "requestTriggerSensor[" + sensor + "] " + mRegistered);
-                    }
+                    mDozeLog.traceSensorRegisterAttempt(sensor.toString(), mRegistered);
                 } else {
-                    if (DEBUG) {
-                        Log.d(TAG, "requestTriggerSensor[" + sensor + "] already registered");
-                    }
+                    mDozeLog.traceSkipRegisterSensor(sensor.toString());
                 }
             } else if (mRegistered) {
                 final boolean rt = mSensorManager.cancelTriggerSensor(this, sensor);
-                if (DEBUG) {
-                    Log.d(TAG, "cancelTriggerSensor[" + sensor + "] " + rt);
-                }
+                mDozeLog.traceSensorUnregisterAttempt(sensor.toString(), rt);
                 mRegistered = false;
             }
         }
@@ -704,7 +688,6 @@
             final Sensor sensor = mSensors[mPosture];
             mDozeLog.traceSensor(mPulseReason);
             mHandler.post(mWakeLock.wrap(() -> {
-                if (DEBUG) Log.d(TAG, "onTrigger: " + triggerEventToString(event));
                 if (sensor != null && sensor.getType() == Sensor.TYPE_PICK_UP_GESTURE) {
                     UI_EVENT_LOGGER.log(DozeSensorsUiEvent.ACTION_AMBIENT_GESTURE_PICKUP);
                 }
@@ -776,11 +759,11 @@
                     && !mRegistered) {
                 asyncSensorManager.registerPluginListener(mPluginSensor, this);
                 mRegistered = true;
-                if (DEBUG) Log.d(TAG, "registerPluginListener");
+                mDozeLog.tracePluginSensorUpdate(true /* registered */);
             } else if (mRegistered) {
                 asyncSensorManager.unregisterPluginListener(mPluginSensor, this);
                 mRegistered = false;
-                if (DEBUG) Log.d(TAG, "unregisterPluginListener");
+                mDozeLog.tracePluginSensorUpdate(false /* registered */);
             }
         }
 
@@ -813,10 +796,9 @@
             mHandler.post(mWakeLock.wrap(() -> {
                 final long now = SystemClock.uptimeMillis();
                 if (now < mDebounceFrom + mDebounce) {
-                    Log.d(TAG, "onSensorEvent dropped: " + triggerEventToString(event));
+                    mDozeLog.traceSensorEventDropped(mPulseReason, "debounce");
                     return;
                 }
-                if (DEBUG) Log.d(TAG, "onSensorEvent: " + triggerEventToString(event));
                 mSensorCallback.onSensorPulse(mPulseReason, -1, -1, event.getValues());
             }));
         }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
index 97a2179..32cb1c0 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
@@ -198,7 +198,7 @@
         mAllowPulseTriggers = true;
         mSessionTracker = sessionTracker;
 
-        mDozeSensors = new DozeSensors(context, mSensorManager, dozeParameters,
+        mDozeSensors = new DozeSensors(mSensorManager, dozeParameters,
                 config, wakeLock, this::onSensor, this::onProximityFar, dozeLog, proximitySensor,
                 secureSettings, authController, devicePostureController);
         mDockManager = dockManager;
diff --git a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
index 7451a76..e3311db 100644
--- a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
@@ -43,7 +43,7 @@
     @SysUISingleton
     @DozeLog
     public static LogBuffer provideDozeLogBuffer(LogBufferFactory factory) {
-        return factory.create("DozeLog", 120);
+        return factory.create("DozeLog", 150);
     }
 
     /** Provides a logging buffer for all logs related to the data layer of notifications. */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
index 832a739..0380fff 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
@@ -20,8 +20,9 @@
 /** Describes usage of a notification. */
 data class NotificationMemoryUsage(
     val packageName: String,
-    val notificationId: String,
+    val notificationKey: String,
     val objectUsage: NotificationObjectUsage,
+    val viewUsage: List<NotificationViewUsage>
 )
 
 /**
@@ -39,3 +40,26 @@
     val extender: Int,
     val hasCustomView: Boolean,
 )
+
+enum class ViewType {
+    PUBLIC_VIEW,
+    PRIVATE_CONTRACTED_VIEW,
+    PRIVATE_EXPANDED_VIEW,
+    PRIVATE_HEADS_UP_VIEW,
+    TOTAL
+}
+
+/**
+ * Describes current memory of a notification view hierarchy.
+ *
+ * The values are in bytes.
+ */
+data class NotificationViewUsage(
+    val viewType: ViewType,
+    val smallIcon: Int,
+    val largeIcon: Int,
+    val systemIcons: Int,
+    val style: Int,
+    val customViews: Int,
+    val softwareBitmapsPenalty: Int,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt
new file mode 100644
index 0000000..7d39e18
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt
@@ -0,0 +1,212 @@
+package com.android.systemui.statusbar.notification.logging
+
+import android.app.Notification
+import android.app.Person
+import android.graphics.Bitmap
+import android.graphics.drawable.Icon
+import android.os.Bundle
+import android.os.Parcel
+import android.os.Parcelable
+import androidx.annotation.WorkerThread
+import com.android.systemui.statusbar.notification.NotificationUtils
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+
+/** Calculates estimated memory usage of [Notification] and [NotificationEntry] objects. */
+internal object NotificationMemoryMeter {
+
+    private const val CAR_EXTENSIONS = "android.car.EXTENSIONS"
+    private const val CAR_EXTENSIONS_LARGE_ICON = "large_icon"
+    private const val TV_EXTENSIONS = "android.tv.EXTENSIONS"
+    private const val WEARABLE_EXTENSIONS = "android.wearable.EXTENSIONS"
+    private const val WEARABLE_EXTENSIONS_BACKGROUND = "background"
+
+    /** Returns a list of memory use entries for currently shown notifications. */
+    @WorkerThread
+    fun notificationMemoryUse(
+        notifications: Collection<NotificationEntry>,
+    ): List<NotificationMemoryUsage> {
+        return notifications
+            .asSequence()
+            .map { entry ->
+                val packageName = entry.sbn.packageName
+                val notificationObjectUsage =
+                    notificationMemoryUse(entry.sbn.notification, hashSetOf())
+                val notificationViewUsage = NotificationMemoryViewWalker.getViewUsage(entry.row)
+                NotificationMemoryUsage(
+                    packageName,
+                    NotificationUtils.logKey(entry.sbn.key),
+                    notificationObjectUsage,
+                    notificationViewUsage
+                )
+            }
+            .toList()
+    }
+
+    @WorkerThread
+    fun notificationMemoryUse(
+        entry: NotificationEntry,
+        seenBitmaps: HashSet<Int> = hashSetOf(),
+    ): NotificationMemoryUsage {
+        return NotificationMemoryUsage(
+            entry.sbn.packageName,
+            NotificationUtils.logKey(entry.sbn.key),
+            notificationMemoryUse(entry.sbn.notification, seenBitmaps),
+            NotificationMemoryViewWalker.getViewUsage(entry.row)
+        )
+    }
+
+    /**
+     * Computes the estimated memory usage of a given [Notification] object. It'll attempt to
+     * inspect Bitmaps in the object and provide summary of memory usage.
+     */
+    @WorkerThread
+    fun notificationMemoryUse(
+        notification: Notification,
+        seenBitmaps: HashSet<Int> = hashSetOf(),
+    ): NotificationObjectUsage {
+        val extras = notification.extras
+        val smallIconUse = computeIconUse(notification.smallIcon, seenBitmaps)
+        val largeIconUse = computeIconUse(notification.getLargeIcon(), seenBitmaps)
+
+        // Collect memory usage of extra styles
+
+        // Big Picture
+        val bigPictureIconUse =
+            computeParcelableUse(extras, Notification.EXTRA_LARGE_ICON_BIG, seenBitmaps)
+        val bigPictureUse =
+            computeParcelableUse(extras, Notification.EXTRA_PICTURE, seenBitmaps) +
+                computeParcelableUse(extras, Notification.EXTRA_PICTURE_ICON, seenBitmaps)
+
+        // People
+        val peopleList = extras.getParcelableArrayList<Person>(Notification.EXTRA_PEOPLE_LIST)
+        val peopleUse =
+            peopleList?.sumOf { person -> computeIconUse(person.icon, seenBitmaps) } ?: 0
+
+        // Calling
+        val callingPersonUse =
+            computeParcelableUse(extras, Notification.EXTRA_CALL_PERSON, seenBitmaps)
+        val verificationIconUse =
+            computeParcelableUse(extras, Notification.EXTRA_VERIFICATION_ICON, seenBitmaps)
+
+        // Messages
+        val messages =
+            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
+                extras.getParcelableArray(Notification.EXTRA_MESSAGES)
+            )
+        val messagesUse =
+            messages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
+        val historicMessages =
+            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
+                extras.getParcelableArray(Notification.EXTRA_HISTORIC_MESSAGES)
+            )
+        val historyicMessagesUse =
+            historicMessages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
+
+        // Extenders
+        val carExtender = extras.getBundle(CAR_EXTENSIONS)
+        val carExtenderSize = carExtender?.let { computeBundleSize(it) } ?: 0
+        val carExtenderIcon =
+            computeParcelableUse(carExtender, CAR_EXTENSIONS_LARGE_ICON, seenBitmaps)
+
+        val tvExtender = extras.getBundle(TV_EXTENSIONS)
+        val tvExtenderSize = tvExtender?.let { computeBundleSize(it) } ?: 0
+
+        val wearExtender = extras.getBundle(WEARABLE_EXTENSIONS)
+        val wearExtenderSize = wearExtender?.let { computeBundleSize(it) } ?: 0
+        val wearExtenderBackground =
+            computeParcelableUse(wearExtender, WEARABLE_EXTENSIONS_BACKGROUND, seenBitmaps)
+
+        val style = notification.notificationStyle
+        val hasCustomView = notification.contentView != null || notification.bigContentView != null
+        val extrasSize = computeBundleSize(extras)
+
+        return NotificationObjectUsage(
+            smallIcon = smallIconUse,
+            largeIcon = largeIconUse,
+            extras = extrasSize,
+            style = style?.simpleName,
+            styleIcon =
+                bigPictureIconUse +
+                    peopleUse +
+                    callingPersonUse +
+                    verificationIconUse +
+                    messagesUse +
+                    historyicMessagesUse,
+            bigPicture = bigPictureUse,
+            extender =
+                carExtenderSize +
+                    carExtenderIcon +
+                    tvExtenderSize +
+                    wearExtenderSize +
+                    wearExtenderBackground,
+            hasCustomView = hasCustomView
+        )
+    }
+
+    /**
+     * Calculates size of the bundle data (excluding FDs and other shared objects like ashmem
+     * bitmaps). Can be slow.
+     */
+    private fun computeBundleSize(extras: Bundle): Int {
+        val parcel = Parcel.obtain()
+        try {
+            extras.writeToParcel(parcel, 0)
+            return parcel.dataSize()
+        } finally {
+            parcel.recycle()
+        }
+    }
+
+    /**
+     * Deserializes [Icon], [Bitmap] or [Person] from extras and computes its memory use. Returns 0
+     * if the key does not exist in extras.
+     */
+    private fun computeParcelableUse(extras: Bundle?, key: String, seenBitmaps: HashSet<Int>): Int {
+        return when (val parcelable = extras?.getParcelable<Parcelable>(key)) {
+            is Bitmap -> computeBitmapUse(parcelable, seenBitmaps)
+            is Icon -> computeIconUse(parcelable, seenBitmaps)
+            is Person -> computeIconUse(parcelable.icon, seenBitmaps)
+            else -> 0
+        }
+    }
+
+    /**
+     * Calculates the byte size of bitmaps or data in the Icon object. Returns 0 if the icon is
+     * defined via Uri or a resource.
+     *
+     * @return memory usage in bytes or 0 if the icon is Uri/Resource based
+     */
+    private fun computeIconUse(icon: Icon?, seenBitmaps: HashSet<Int>) =
+        when (icon?.type) {
+            Icon.TYPE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
+            Icon.TYPE_ADAPTIVE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
+            Icon.TYPE_DATA -> computeDataUse(icon, seenBitmaps)
+            else -> 0
+        }
+
+    /**
+     * Returns the amount of memory a given bitmap is using. If the bitmap reference is part of
+     * seenBitmaps set, this method returns 0 to avoid double counting.
+     *
+     * @return memory usage of the bitmap in bytes
+     */
+    private fun computeBitmapUse(bitmap: Bitmap, seenBitmaps: HashSet<Int>? = null): Int {
+        val refId = System.identityHashCode(bitmap)
+        if (seenBitmaps?.contains(refId) == true) {
+            return 0
+        }
+
+        seenBitmaps?.add(refId)
+        return bitmap.allocationByteCount
+    }
+
+    private fun computeDataUse(icon: Icon, seenBitmaps: HashSet<Int>): Int {
+        val refId = System.identityHashCode(icon.dataBytes)
+        if (seenBitmaps.contains(refId)) {
+            return 0
+        }
+
+        seenBitmaps.add(refId)
+        return icon.dataLength
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
index 958978e..c09cc43 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
@@ -17,22 +17,11 @@
 
 package com.android.systemui.statusbar.notification.logging
 
-import android.app.Notification
-import android.app.Person
-import android.graphics.Bitmap
-import android.graphics.drawable.Icon
-import android.os.Bundle
-import android.os.Parcel
-import android.os.Parcelable
 import android.util.Log
-import androidx.annotation.WorkerThread
-import androidx.core.util.contains
 import com.android.systemui.Dumpable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.statusbar.notification.NotificationUtils
 import com.android.systemui.statusbar.notification.collection.NotifPipeline
-import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import java.io.PrintWriter
 import javax.inject.Inject
 
@@ -46,12 +35,7 @@
 ) : Dumpable {
 
     companion object {
-        private const val TAG = "NotificationMemMonitor"
-        private const val CAR_EXTENSIONS = "android.car.EXTENSIONS"
-        private const val CAR_EXTENSIONS_LARGE_ICON = "large_icon"
-        private const val TV_EXTENSIONS = "android.tv.EXTENSIONS"
-        private const val WEARABLE_EXTENSIONS = "android.wearable.EXTENSIONS"
-        private const val WEARABLE_EXTENSIONS_BACKGROUND = "background"
+        private const val TAG = "NotificationMemory"
     }
 
     fun init() {
@@ -60,184 +44,123 @@
     }
 
     override fun dump(pw: PrintWriter, args: Array<out String>) {
-        currentNotificationMemoryUse().forEach { use -> pw.println(use.toString()) }
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(notificationPipeline.allNotifs)
+                .sortedWith(compareBy({ it.packageName }, { it.notificationKey }))
+        dumpNotificationObjects(pw, memoryUse)
+        dumpNotificationViewUsage(pw, memoryUse)
     }
 
-    @WorkerThread
-    fun currentNotificationMemoryUse(): List<NotificationMemoryUsage> {
-        return notificationMemoryUse(notificationPipeline.allNotifs)
-    }
-
-    /** Returns a list of memory use entries for currently shown notifications. */
-    @WorkerThread
-    fun notificationMemoryUse(
-        notifications: Collection<NotificationEntry>
-    ): List<NotificationMemoryUsage> {
-        return notifications
-            .asSequence()
-            .map { entry ->
-                val packageName = entry.sbn.packageName
-                val notificationObjectUsage =
-                    computeNotificationObjectUse(entry.sbn.notification, hashSetOf())
-                NotificationMemoryUsage(
-                    packageName,
-                    NotificationUtils.logKey(entry.sbn.key),
-                    notificationObjectUsage
-                )
-            }
-            .toList()
-    }
-
-    /**
-     * Computes the estimated memory usage of a given [Notification] object. It'll attempt to
-     * inspect Bitmaps in the object and provide summary of memory usage.
-     */
-    private fun computeNotificationObjectUse(
-        notification: Notification,
-        seenBitmaps: HashSet<Int>
-    ): NotificationObjectUsage {
-        val extras = notification.extras
-        val smallIconUse = computeIconUse(notification.smallIcon, seenBitmaps)
-        val largeIconUse = computeIconUse(notification.getLargeIcon(), seenBitmaps)
-
-        // Collect memory usage of extra styles
-
-        // Big Picture
-        val bigPictureIconUse =
-            computeParcelableUse(extras, Notification.EXTRA_PICTURE_ICON, seenBitmaps) +
-                computeParcelableUse(extras, Notification.EXTRA_LARGE_ICON_BIG, seenBitmaps)
-        val bigPictureUse =
-            computeParcelableUse(extras, Notification.EXTRA_PICTURE, seenBitmaps) +
-                computeParcelableUse(extras, Notification.EXTRA_PICTURE_ICON, seenBitmaps)
-
-        // People
-        val peopleList = extras.getParcelableArrayList<Person>(Notification.EXTRA_PEOPLE_LIST)
-        val peopleUse =
-            peopleList?.sumOf { person -> computeIconUse(person.icon, seenBitmaps) } ?: 0
-
-        // Calling
-        val callingPersonUse =
-            computeParcelableUse(extras, Notification.EXTRA_CALL_PERSON, seenBitmaps)
-        val verificationIconUse =
-            computeParcelableUse(extras, Notification.EXTRA_VERIFICATION_ICON, seenBitmaps)
-
-        // Messages
-        val messages =
-            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
-                extras.getParcelableArray(Notification.EXTRA_MESSAGES)
-            )
-        val messagesUse =
-            messages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
-        val historicMessages =
-            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
-                extras.getParcelableArray(Notification.EXTRA_HISTORIC_MESSAGES)
-            )
-        val historyicMessagesUse =
-            historicMessages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
-
-        // Extenders
-        val carExtender = extras.getBundle(CAR_EXTENSIONS)
-        val carExtenderSize = carExtender?.let { computeBundleSize(it) } ?: 0
-        val carExtenderIcon =
-            computeParcelableUse(carExtender, CAR_EXTENSIONS_LARGE_ICON, seenBitmaps)
-
-        val tvExtender = extras.getBundle(TV_EXTENSIONS)
-        val tvExtenderSize = tvExtender?.let { computeBundleSize(it) } ?: 0
-
-        val wearExtender = extras.getBundle(WEARABLE_EXTENSIONS)
-        val wearExtenderSize = wearExtender?.let { computeBundleSize(it) } ?: 0
-        val wearExtenderBackground =
-            computeParcelableUse(wearExtender, WEARABLE_EXTENSIONS_BACKGROUND, seenBitmaps)
-
-        val style = notification.notificationStyle
-        val hasCustomView = notification.contentView != null || notification.bigContentView != null
-        val extrasSize = computeBundleSize(extras)
-
-        return NotificationObjectUsage(
-            smallIconUse,
-            largeIconUse,
-            extrasSize,
-            style?.simpleName,
-            bigPictureIconUse +
-                peopleUse +
-                callingPersonUse +
-                verificationIconUse +
-                messagesUse +
-                historyicMessagesUse,
-            bigPictureUse,
-            carExtenderSize +
-                carExtenderIcon +
-                tvExtenderSize +
-                wearExtenderSize +
-                wearExtenderBackground,
-            hasCustomView
+    /** Renders a table of notification object usage into passed [PrintWriter]. */
+    private fun dumpNotificationObjects(pw: PrintWriter, memoryUse: List<NotificationMemoryUsage>) {
+        pw.println("Notification Object Usage")
+        pw.println("-----------")
+        pw.println(
+            "Package".padEnd(35) +
+                "\t\tSmall\tLarge\t${"Style".padEnd(15)}\t\tStyle\tBig\tExtend.\tExtras\tCustom"
         )
+        pw.println("".padEnd(35) + "\t\tIcon\tIcon\t${"".padEnd(15)}\t\tIcon\tPicture\t \t \tView")
+        pw.println()
+
+        memoryUse.forEach { use ->
+            pw.println(
+                use.packageName.padEnd(35) +
+                    "\t\t" +
+                    "${use.objectUsage.smallIcon}\t${use.objectUsage.largeIcon}\t" +
+                    (use.objectUsage.style?.take(15) ?: "").padEnd(15) +
+                    "\t\t${use.objectUsage.styleIcon}\t" +
+                    "${use.objectUsage.bigPicture}\t${use.objectUsage.extender}\t" +
+                    "${use.objectUsage.extras}\t${use.objectUsage.hasCustomView}\t" +
+                    use.notificationKey
+            )
+        }
+
+        // Calculate totals for easily glanceable summary.
+        data class Totals(
+            var smallIcon: Int = 0,
+            var largeIcon: Int = 0,
+            var styleIcon: Int = 0,
+            var bigPicture: Int = 0,
+            var extender: Int = 0,
+            var extras: Int = 0,
+        )
+
+        val totals =
+            memoryUse.fold(Totals()) { t, usage ->
+                t.smallIcon += usage.objectUsage.smallIcon
+                t.largeIcon += usage.objectUsage.largeIcon
+                t.styleIcon += usage.objectUsage.styleIcon
+                t.bigPicture += usage.objectUsage.bigPicture
+                t.extender += usage.objectUsage.extender
+                t.extras += usage.objectUsage.extras
+                t
+            }
+
+        pw.println()
+        pw.println("TOTALS")
+        pw.println(
+            "".padEnd(35) +
+                "\t\t" +
+                "${toKb(totals.smallIcon)}\t${toKb(totals.largeIcon)}\t" +
+                "".padEnd(15) +
+                "\t\t${toKb(totals.styleIcon)}\t" +
+                "${toKb(totals.bigPicture)}\t${toKb(totals.extender)}\t" +
+                toKb(totals.extras)
+        )
+        pw.println()
     }
 
-    /**
-     * Calculates size of the bundle data (excluding FDs and other shared objects like ashmem
-     * bitmaps). Can be slow.
-     */
-    private fun computeBundleSize(extras: Bundle): Int {
-        val parcel = Parcel.obtain()
-        try {
-            extras.writeToParcel(parcel, 0)
-            return parcel.dataSize()
-        } finally {
-            parcel.recycle()
-        }
+    /** Renders a table of notification view usage into passed [PrintWriter] */
+    private fun dumpNotificationViewUsage(
+        pw: PrintWriter,
+        memoryUse: List<NotificationMemoryUsage>,
+    ) {
+
+        data class Totals(
+            var smallIcon: Int = 0,
+            var largeIcon: Int = 0,
+            var style: Int = 0,
+            var customViews: Int = 0,
+            var softwareBitmapsPenalty: Int = 0,
+        )
+
+        val totals = Totals()
+        pw.println("Notification View Usage")
+        pw.println("-----------")
+        pw.println("View Type".padEnd(24) + "\tSmall\tLarge\tStyle\tCustom\tSoftware")
+        pw.println("".padEnd(24) + "\tIcon\tIcon\tUse\tView\tBitmaps")
+        pw.println()
+        memoryUse
+            .filter { it.viewUsage.isNotEmpty() }
+            .forEach { use ->
+                pw.println(use.packageName + " " + use.notificationKey)
+                use.viewUsage.forEach { view ->
+                    pw.println(
+                        "  ${view.viewType.toString().padEnd(24)}\t${view.smallIcon}" +
+                            "\t${view.largeIcon}\t${view.style}" +
+                            "\t${view.customViews}\t${view.softwareBitmapsPenalty}"
+                    )
+
+                    if (view.viewType == ViewType.TOTAL) {
+                        totals.smallIcon += view.smallIcon
+                        totals.largeIcon += view.largeIcon
+                        totals.style += view.style
+                        totals.customViews += view.customViews
+                        totals.softwareBitmapsPenalty += view.softwareBitmapsPenalty
+                    }
+                }
+            }
+        pw.println()
+        pw.println("TOTALS")
+        pw.println(
+            "  ${"".padEnd(24)}\t${toKb(totals.smallIcon)}" +
+                "\t${toKb(totals.largeIcon)}\t${toKb(totals.style)}" +
+                "\t${toKb(totals.customViews)}\t${toKb(totals.softwareBitmapsPenalty)}"
+        )
+        pw.println()
     }
 
-    /**
-     * Deserializes [Icon], [Bitmap] or [Person] from extras and computes its memory use. Returns 0
-     * if the key does not exist in extras.
-     */
-    private fun computeParcelableUse(extras: Bundle?, key: String, seenBitmaps: HashSet<Int>): Int {
-        return when (val parcelable = extras?.getParcelable<Parcelable>(key)) {
-            is Bitmap -> computeBitmapUse(parcelable, seenBitmaps)
-            is Icon -> computeIconUse(parcelable, seenBitmaps)
-            is Person -> computeIconUse(parcelable.icon, seenBitmaps)
-            else -> 0
-        }
-    }
-
-    /**
-     * Calculates the byte size of bitmaps or data in the Icon object. Returns 0 if the icon is
-     * defined via Uri or a resource.
-     *
-     * @return memory usage in bytes or 0 if the icon is Uri/Resource based
-     */
-    private fun computeIconUse(icon: Icon?, seenBitmaps: HashSet<Int>) =
-        when (icon?.type) {
-            Icon.TYPE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
-            Icon.TYPE_ADAPTIVE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
-            Icon.TYPE_DATA -> computeDataUse(icon, seenBitmaps)
-            else -> 0
-        }
-
-    /**
-     * Returns the amount of memory a given bitmap is using. If the bitmap reference is part of
-     * seenBitmaps set, this method returns 0 to avoid double counting.
-     *
-     * @return memory usage of the bitmap in bytes
-     */
-    private fun computeBitmapUse(bitmap: Bitmap, seenBitmaps: HashSet<Int>? = null): Int {
-        val refId = System.identityHashCode(bitmap)
-        if (seenBitmaps?.contains(refId) == true) {
-            return 0
-        }
-
-        seenBitmaps?.add(refId)
-        return bitmap.allocationByteCount
-    }
-
-    private fun computeDataUse(icon: Icon, seenBitmaps: HashSet<Int>): Int {
-        val refId = System.identityHashCode(icon.dataBytes)
-        if (seenBitmaps.contains(refId)) {
-            return 0
-        }
-
-        seenBitmaps.add(refId)
-        return icon.dataLength
+    private fun toKb(bytes: Int): String {
+        return (bytes / 1024).toString() + " KB"
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
new file mode 100644
index 0000000..a0bee15
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
@@ -0,0 +1,173 @@
+package com.android.systemui.statusbar.notification.logging
+
+import android.graphics.Bitmap
+import android.graphics.drawable.BitmapDrawable
+import android.graphics.drawable.Drawable
+import android.util.Log
+import android.view.View
+import android.view.ViewGroup
+import android.widget.ImageView
+import com.android.internal.R
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
+import com.android.systemui.util.children
+
+/** Walks view hiearchy of a given notification to estimate its memory use. */
+internal object NotificationMemoryViewWalker {
+
+    private const val TAG = "NotificationMemory"
+
+    /** Builder for [NotificationViewUsage] objects. */
+    private class UsageBuilder {
+        private var smallIcon: Int = 0
+        private var largeIcon: Int = 0
+        private var systemIcons: Int = 0
+        private var style: Int = 0
+        private var customViews: Int = 0
+        private var softwareBitmaps = 0
+
+        fun addSmallIcon(smallIconUse: Int) = apply { smallIcon += smallIconUse }
+        fun addLargeIcon(largeIconUse: Int) = apply { largeIcon += largeIconUse }
+        fun addSystem(systemIconUse: Int) = apply { systemIcons += systemIconUse }
+        fun addStyle(styleUse: Int) = apply { style += styleUse }
+        fun addSoftwareBitmapPenalty(softwareBitmapUse: Int) = apply {
+            softwareBitmaps += softwareBitmapUse
+        }
+
+        fun addCustomViews(customViewsUse: Int) = apply { customViews += customViewsUse }
+
+        fun build(viewType: ViewType): NotificationViewUsage {
+            return NotificationViewUsage(
+                viewType = viewType,
+                smallIcon = smallIcon,
+                largeIcon = largeIcon,
+                systemIcons = systemIcons,
+                style = style,
+                customViews = customViews,
+                softwareBitmapsPenalty = softwareBitmaps,
+            )
+        }
+    }
+
+    /**
+     * Returns memory usage of public and private views contained in passed
+     * [ExpandableNotificationRow]
+     */
+    fun getViewUsage(row: ExpandableNotificationRow?): List<NotificationViewUsage> {
+        if (row == null) {
+            return listOf()
+        }
+
+        // The ordering here is significant since it determines deduplication of seen drawables.
+        return listOf(
+            getViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, row.privateLayout?.expandedChild),
+            getViewUsage(ViewType.PRIVATE_CONTRACTED_VIEW, row.privateLayout?.contractedChild),
+            getViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, row.privateLayout?.headsUpChild),
+            getViewUsage(ViewType.PUBLIC_VIEW, row.publicLayout),
+            getTotalUsage(row)
+        )
+    }
+
+    /**
+     * Calculate total usage of all views - we need to do a separate traversal to make sure we don't
+     * double count fields.
+     */
+    private fun getTotalUsage(row: ExpandableNotificationRow): NotificationViewUsage {
+        val totalUsage = UsageBuilder()
+        val seenObjects = hashSetOf<Int>()
+
+        row.publicLayout?.let { computeViewHierarchyUse(it, totalUsage, seenObjects) }
+        row.privateLayout?.let { child ->
+            for (view in listOf(child.expandedChild, child.contractedChild, child.headsUpChild)) {
+                (view as? ViewGroup)?.let { v ->
+                    computeViewHierarchyUse(v, totalUsage, seenObjects)
+                }
+            }
+        }
+        return totalUsage.build(ViewType.TOTAL)
+    }
+
+    private fun getViewUsage(
+        type: ViewType,
+        rootView: View?,
+        seenObjects: HashSet<Int> = hashSetOf()
+    ): NotificationViewUsage {
+        val usageBuilder = UsageBuilder()
+        (rootView as? ViewGroup)?.let { computeViewHierarchyUse(it, usageBuilder, seenObjects) }
+        return usageBuilder.build(type)
+    }
+
+    private fun computeViewHierarchyUse(
+        rootView: ViewGroup,
+        builder: UsageBuilder,
+        seenObjects: HashSet<Int> = hashSetOf(),
+    ) {
+        for (child in rootView.children) {
+            if (child is ViewGroup) {
+                computeViewHierarchyUse(child, builder, seenObjects)
+            } else {
+                computeViewUse(child, builder, seenObjects)
+            }
+        }
+    }
+
+    private fun computeViewUse(view: View, builder: UsageBuilder, seenObjects: HashSet<Int>) {
+        if (view !is ImageView) return
+        val drawable = view.drawable ?: return
+        val drawableRef = System.identityHashCode(drawable)
+        if (seenObjects.contains(drawableRef)) return
+        val drawableUse = computeDrawableUse(drawable, seenObjects)
+        // TODO(b/235451049): We need to make sure we traverse large icon before small icon -
+        // sometimes the large icons are assigned to small icon views and we want to
+        // attribute them to large view in those cases.
+        when (view.id) {
+            R.id.left_icon,
+            R.id.icon,
+            R.id.conversation_icon -> builder.addSmallIcon(drawableUse)
+            R.id.right_icon -> builder.addLargeIcon(drawableUse)
+            R.id.big_picture -> builder.addStyle(drawableUse)
+            // Elements that are part of platform with resources
+            R.id.phishing_alert,
+            R.id.feedback,
+            R.id.alerted_icon,
+            R.id.expand_button_icon,
+            R.id.remote_input_send -> builder.addSystem(drawableUse)
+            // Custom view ImageViews
+            else -> {
+                if (Log.isLoggable(TAG, Log.DEBUG)) {
+                    Log.d(TAG, "Custom view: ${identifierForView(view)}")
+                }
+                builder.addCustomViews(drawableUse)
+            }
+        }
+
+        if (isDrawableSoftwareBitmap(drawable)) {
+            builder.addSoftwareBitmapPenalty(drawableUse)
+        }
+
+        seenObjects.add(drawableRef)
+    }
+
+    private fun computeDrawableUse(drawable: Drawable, seenObjects: HashSet<Int>): Int =
+        when (drawable) {
+            is BitmapDrawable -> {
+                val ref = System.identityHashCode(drawable.bitmap)
+                if (seenObjects.contains(ref)) {
+                    0
+                } else {
+                    seenObjects.add(ref)
+                    drawable.bitmap.allocationByteCount
+                }
+            }
+            else -> 0
+        }
+
+    private fun isDrawableSoftwareBitmap(drawable: Drawable) =
+        drawable is BitmapDrawable && drawable.bitmap.config != Bitmap.Config.HARDWARE
+
+    private fun identifierForView(view: View) =
+        if (view.id == View.NO_ID) {
+            "no-id"
+        } else {
+            view.resources.getResourceName(view.id)
+        }
+}
diff --git a/packages/SystemUI/tests/res/layout/custom_view_dark.xml b/packages/SystemUI/tests/res/layout/custom_view_dark.xml
index 9e460a5..112d73d2 100644
--- a/packages/SystemUI/tests/res/layout/custom_view_dark.xml
+++ b/packages/SystemUI/tests/res/layout/custom_view_dark.xml
@@ -14,6 +14,7 @@
     limitations under the License.
 -->
 <ImageView xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/custom_view_dark_image"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:background="#ff000000"
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeSensorsTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeSensorsTest.java
index 9ffc5a5..c40c187 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeSensorsTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeSensorsTest.java
@@ -423,7 +423,7 @@
 
     @Test
     public void testGesturesAllInitiallyRespectSettings() {
-        DozeSensors dozeSensors = new DozeSensors(getContext(), mSensorManager, mDozeParameters,
+        DozeSensors dozeSensors = new DozeSensors(mSensorManager, mDozeParameters,
                 mAmbientDisplayConfiguration, mWakeLock, mCallback, mProxCallback, mDozeLog,
                 mProximitySensor, mFakeSettings, mAuthController,
                 mDevicePostureController);
@@ -435,7 +435,7 @@
 
     private class TestableDozeSensors extends DozeSensors {
         TestableDozeSensors() {
-            super(getContext(), mSensorManager, mDozeParameters,
+            super(mSensorManager, mDozeParameters,
                     mAmbientDisplayConfiguration, mWakeLock, mCallback, mProxCallback, mDozeLog,
                     mProximitySensor, mFakeSettings, mAuthController,
                     mDevicePostureController);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeterTest.kt
similarity index 83%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitorTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeterTest.kt
index 16e2441..f69839b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeterTest.kt
@@ -28,30 +28,21 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.statusbar.notification.NotificationUtils
-import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.MockitoAnnotations
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
-class NotificationMemoryMonitorTest : SysuiTestCase() {
-
-    @Before
-    fun setUp() {
-        MockitoAnnotations.initMocks(this)
-    }
+class NotificationMemoryMeterTest : SysuiTestCase() {
 
     @Test
     fun currentNotificationMemoryUse_plainNotification() {
         val notification = createBasicNotification().build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -69,8 +60,8 @@
     fun currentNotificationMemoryUse_plainNotification_dontDoubleCountSameBitmap() {
         val icon = Icon.createWithBitmap(Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888))
         val notification = createBasicNotification().setLargeIcon(icon).setSmallIcon(icon).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -92,8 +83,8 @@
                     RemoteViews(context.packageName, android.R.layout.list_content)
                 )
                 .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -112,8 +103,8 @@
         val dataIcon = Icon.createWithData(ByteArray(444444), 0, 444444)
         val notification =
             createBasicNotification().setLargeIcon(dataIcon).setSmallIcon(dataIcon).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = 444444,
@@ -141,8 +132,8 @@
                         .bigLargeIcon(bigPictureIcon)
                 )
                 .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -167,8 +158,8 @@
             createBasicNotification()
                 .setStyle(Notification.CallStyle.forIncomingCall(person, fakeIntent, fakeIntent))
                 .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -203,8 +194,8 @@
                         .addHistoricMessage(historicMessage)
                 )
                 .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -225,8 +216,8 @@
         val carIcon = Bitmap.createBitmap(432, 322, Bitmap.Config.ARGB_8888)
         val extender = Notification.CarExtender().setLargeIcon(carIcon)
         val notification = createBasicNotification().extend(extender).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -246,8 +237,8 @@
         val wearBackground = Bitmap.createBitmap(443, 433, Bitmap.Config.ARGB_8888)
         val wearExtender = Notification.WearableExtender().setBackground(wearBackground)
         val notification = createBasicNotification().extend(tvExtender).extend(wearExtender).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
         assertNotificationObjectSizes(
             memoryUse = memoryUse,
             smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -283,10 +274,10 @@
         extender: Int,
         style: String?,
         styleIcon: Int,
-        hasCustomView: Boolean
+        hasCustomView: Boolean,
     ) {
         assertThat(memoryUse.packageName).isEqualTo("test_pkg")
-        assertThat(memoryUse.notificationId)
+        assertThat(memoryUse.notificationKey)
             .isEqualTo(NotificationUtils.logKey("0|test_pkg|0|test|0"))
         assertThat(memoryUse.objectUsage.smallIcon).isEqualTo(smallIcon)
         assertThat(memoryUse.objectUsage.largeIcon).isEqualTo(largeIcon)
@@ -301,21 +292,14 @@
     }
 
     private fun getUseObject(
-        singleItemUseList: List<NotificationMemoryUsage>
+        singleItemUseList: List<NotificationMemoryUsage>,
     ): NotificationMemoryUsage {
         assertThat(singleItemUseList).hasSize(1)
         return singleItemUseList[0]
     }
 
-    private fun createNMMWithNotifications(
-        notifications: List<Notification>
-    ): NotificationMemoryMonitor {
-        val notifPipeline: NotifPipeline = mock()
-        val notificationEntries =
-            notifications.map { n ->
-                NotificationEntryBuilder().setTag("test").setNotification(n).build()
-            }
-        whenever(notifPipeline.allNotifs).thenReturn(notificationEntries)
-        return NotificationMemoryMonitor(notifPipeline, mock())
-    }
+    private fun createNotificationEntry(
+        notification: Notification,
+    ): NotificationEntry =
+        NotificationEntryBuilder().setTag("test").setNotification(notification).build()
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt
new file mode 100644
index 0000000..3a16fb3
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt
@@ -0,0 +1,148 @@
+package com.android.systemui.statusbar.notification.logging
+
+import android.app.Notification
+import android.graphics.Bitmap
+import android.graphics.drawable.Icon
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import android.widget.RemoteViews
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.notification.row.NotificationTestHelper
+import com.android.systemui.tests.R
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@TestableLooper.RunWithLooper
+class NotificationMemoryViewWalkerTest : SysuiTestCase() {
+
+    private lateinit var testHelper: NotificationTestHelper
+
+    @Before
+    fun setUp() {
+        allowTestableLooperAsMainThread()
+        testHelper = NotificationTestHelper(mContext, mDependency, TestableLooper.get(this))
+    }
+
+    @Test
+    fun testViewWalker_nullRow_returnsEmptyView() {
+        val result = NotificationMemoryViewWalker.getViewUsage(null)
+        assertThat(result).isNotNull()
+        assertThat(result).isEmpty()
+    }
+
+    @Test
+    fun testViewWalker_plainNotification() {
+        val row = testHelper.createRow()
+        val result = NotificationMemoryViewWalker.getViewUsage(row)
+        assertThat(result).hasSize(5)
+        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_CONTRACTED_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, 0, 0, 0, 0, 0, 0))
+    }
+
+    @Test
+    fun testViewWalker_bigPictureNotification() {
+        val bigPicture = Bitmap.createBitmap(300, 300, Bitmap.Config.ARGB_8888)
+        val icon = Icon.createWithBitmap(Bitmap.createBitmap(20, 20, Bitmap.Config.ARGB_8888))
+        val largeIcon = Icon.createWithBitmap(Bitmap.createBitmap(60, 60, Bitmap.Config.ARGB_8888))
+        val row =
+            testHelper.createRow(
+                Notification.Builder(mContext)
+                    .setContentText("Test")
+                    .setContentTitle("title")
+                    .setSmallIcon(icon)
+                    .setLargeIcon(largeIcon)
+                    .setStyle(Notification.BigPictureStyle().bigPicture(bigPicture))
+                    .build()
+            )
+        val result = NotificationMemoryViewWalker.getViewUsage(row)
+        assertThat(result).hasSize(5)
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_EXPANDED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    largeIcon.bitmap.allocationByteCount,
+                    0,
+                    bigPicture.allocationByteCount,
+                    0,
+                    bigPicture.allocationByteCount +
+                        icon.bitmap.allocationByteCount +
+                        largeIcon.bitmap.allocationByteCount
+                )
+            )
+
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_CONTRACTED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    largeIcon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    icon.bitmap.allocationByteCount + largeIcon.bitmap.allocationByteCount
+                )
+            )
+        // Due to deduplication, this should all be 0.
+        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+    }
+
+    @Test
+    fun testViewWalker_customView() {
+        val icon = Icon.createWithBitmap(Bitmap.createBitmap(20, 20, Bitmap.Config.ARGB_8888))
+        val bitmap = Bitmap.createBitmap(300, 300, Bitmap.Config.ARGB_8888)
+
+        val views = RemoteViews(mContext.packageName, R.layout.custom_view_dark)
+        views.setImageViewBitmap(R.id.custom_view_dark_image, bitmap)
+        val row =
+            testHelper.createRow(
+                Notification.Builder(mContext)
+                    .setContentText("Test")
+                    .setContentTitle("title")
+                    .setSmallIcon(icon)
+                    .setCustomContentView(views)
+                    .setCustomBigContentView(views)
+                    .build()
+            )
+        val result = NotificationMemoryViewWalker.getViewUsage(row)
+        assertThat(result).hasSize(5)
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_CONTRACTED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    bitmap.allocationByteCount,
+                    bitmap.allocationByteCount + icon.bitmap.allocationByteCount
+                )
+            )
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_EXPANDED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    bitmap.allocationByteCount,
+                    bitmap.allocationByteCount + icon.bitmap.allocationByteCount
+                )
+            )
+        // Due to deduplication, this should all be 0.
+        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+    }
+}