Merge "BatteryStats receives custom measured energies" into sc-dev
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 6e41b3f..0da2998 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -7173,6 +7173,20 @@
                 .getAccumulatedStandardBucketEnergy(MeasuredEnergyStats.ENERGY_BUCKET_SCREEN_DOZE);
     }
 
+    /**
+     * Returns the energy in microjoules that the given custom energy bucket consumed.
+     * Will return {@link #ENERGY_DATA_UNAVAILABLE} if data is unavailable
+     *
+     * @param customEnergyBucket custom energy bucket of interest
+     * @return energy (in microjoules) used by this uid for this energy bucket
+     */
+    public long getCustomMeasuredEnergyMicroJoules(int customEnergyBucket) {
+        if (mGlobalMeasuredEnergyStats == null) {
+            return ENERGY_DATA_UNAVAILABLE;
+        }
+        return mGlobalMeasuredEnergyStats.getAccumulatedCustomBucketEnergy(customEnergyBucket);
+    }
+
     @Override public long getStartClockTime() {
         final long currentTimeMs = System.currentTimeMillis();
         if ((currentTimeMs > MILLISECONDS_IN_YEAR
@@ -7941,6 +7955,13 @@
                     .updateStandardBucket(energyBucket, energyDeltaUJ, accumulate);
         }
 
+        /** Adds the given energy to the given custom energy bucket for this uid. */
+        private void addEnergyToCustomBucketLocked(long energyDeltaUJ, int energyBucket,
+                boolean accumulate) {
+            getOrCreateMeasuredEnergyStatsLocked()
+                    .updateCustomBucket(energyBucket, energyDeltaUJ, accumulate);
+        }
+
         /**
          * Returns the energy used by this uid for a standard energy bucket of interest.
          * @param bucket standard energy bucket of interest
@@ -7958,6 +7979,22 @@
         }
 
         /**
+         * Returns the energy used by this uid for a custom energy bucket of interest.
+         * @param customEnergyBucket custom energy bucket of interest
+         * @return energy (in microjoules) used by this uid for this energy bucket
+         */
+        public long getCustomMeasuredEnergyMicroJoules(int customEnergyBucket) {
+            if (mBsi.mGlobalMeasuredEnergyStats == null
+                    || !mBsi.mGlobalMeasuredEnergyStats.isValidCustomBucket(customEnergyBucket)) {
+                return ENERGY_DATA_UNAVAILABLE;
+            }
+            if (mUidMeasuredEnergyStats == null) {
+                return 0L; // It is supported, but was never filled, so it must be 0
+            }
+            return mUidMeasuredEnergyStats.getAccumulatedCustomBucketEnergy(customEnergyBucket);
+        }
+
+        /**
          * Gets the minimum of the uid's foreground activity time and its PROCESS_STATE_TOP time
          * since last marked. Also sets the mark time for both these timers.
          *
@@ -12465,6 +12502,42 @@
     }
 
     /**
+     * Accumulate Custom energy bucket energy, globally and for each app.
+     *
+     * @param totalEnergyUJ energy (microjoules) used for this bucket since this was last called.
+     * @param uidEnergies map of uid->energy (microjoules) for this bucket since last called.
+     *                    Data inside uidEnergies will not be modified (treated immutable).
+     */
+    public void updateCustomMeasuredEnergyDataLocked(int customEnergyBucket,
+            long totalEnergyUJ, @Nullable SparseLongArray uidEnergies) {
+        if (DEBUG_ENERGY) {
+            Slog.d(TAG, "Updating attributed measured energy stats for custom bucket "
+                    + customEnergyBucket
+                    + " with total energy " + totalEnergyUJ
+                    + " and uid energies " + String.valueOf(uidEnergies));
+        }
+        if (mGlobalMeasuredEnergyStats == null) return;
+        if (!mOnBatteryInternal || mIgnoreNextExternalStats || totalEnergyUJ <= 0) return;
+
+        mGlobalMeasuredEnergyStats.updateCustomBucket(customEnergyBucket, totalEnergyUJ, true);
+
+        if (uidEnergies == null) return;
+        final int numUids = uidEnergies.size();
+        for (int i = 0; i < numUids; i++) {
+            final int uidInt = mapUid(uidEnergies.keyAt(i));
+            final long uidEnergyUJ = uidEnergies.valueAt(i);
+            if (uidEnergyUJ == 0) continue;
+            // TODO: Worry about uids not in BSI currently, including uninstalled uids 'coming back'
+            //  Specifically: What if the uid had been removed? We'll re-create it now.
+            //  And if we instead use getAvailableUidStatsLocked() and chec for null, then we might
+            //  not create a Uid even when we should be (say, the app's first event, somehow, was to
+            //  use GPU). I guess that CPU/kernel data might already have this problem?
+            final Uid uidObj = getUidStatsLocked(uidInt);
+            uidObj.addEnergyToCustomBucketLocked(uidEnergyUJ, customEnergyBucket, true);
+        }
+    }
+
+    /**
      * Read and record Rail Energy data.
      */
     public void updateRailStatsLocked() {
diff --git a/core/java/com/android/internal/power/MeasuredEnergyArray.java b/core/java/com/android/internal/power/MeasuredEnergyArray.java
deleted file mode 100644
index 1f6dc26..0000000
--- a/core/java/com/android/internal/power/MeasuredEnergyArray.java
+++ /dev/null
@@ -1,66 +0,0 @@
-/*
- * Copyright (C) 2020 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.internal.power;
-
-
-import android.annotation.IntDef;
-
-import com.android.internal.os.RailStats;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-
-/**
- * Interface to provide subsystem energy data.
- * TODO: replace this and {@link RailStats} once b/173077356 is done
- */
-public interface MeasuredEnergyArray {
-    int SUBSYSTEM_UNKNOWN = -1;
-    int SUBSYSTEM_DISPLAY = 0;
-    int NUMBER_SUBSYSTEMS = 1;
-    String[] SUBSYSTEM_NAMES = {"display"};
-
-
-    @IntDef(prefix = { "SUBSYSTEM_" }, value = {
-            SUBSYSTEM_UNKNOWN,
-            SUBSYSTEM_DISPLAY,
-    })
-    @Retention(RetentionPolicy.SOURCE)
-    @interface MeasuredEnergySubsystem {}
-
-    /**
-     * Get the subsystem at an index in array.
-     *
-     * @param index into the array.
-     * @return subsystem.
-     */
-    @MeasuredEnergySubsystem
-    int getSubsystem(int index);
-
-    /**
-     * Get the energy (in microjoules) consumed since boot of the subsystem at an index.
-     *
-     * @param index into the array.
-     * @return energy (in microjoules) consumed since boot.
-     */
-    long getEnergy(int index);
-
-    /**
-     * Return number of subsystems in the array.
-     */
-    int size();
-}
diff --git a/core/java/com/android/internal/power/MeasuredEnergyStats.java b/core/java/com/android/internal/power/MeasuredEnergyStats.java
index e310f8d..38ef55c 100644
--- a/core/java/com/android/internal/power/MeasuredEnergyStats.java
+++ b/core/java/com/android/internal/power/MeasuredEnergyStats.java
@@ -28,7 +28,6 @@
 import android.view.Display;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.power.MeasuredEnergyArray.MeasuredEnergySubsystem;
 
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
@@ -247,7 +246,7 @@
     }
 
     /**
-     * Map {@link MeasuredEnergySubsystem} and device state to Display {@link StandardEnergyBucket}.
+     * Map {@link android.view.Display} STATE_ to corresponding {@link StandardEnergyBucket}.
      */
     public static @StandardEnergyBucket int getDisplayEnergyBucket(int screenState) {
         if (Display.isOnState(screenState)) {
@@ -450,7 +449,8 @@
         return bucket >= 0 && bucket < NUMBER_STANDARD_ENERGY_BUCKETS;
     }
 
-    private boolean isValidCustomBucket(int customBucket) {
+    /** Returns whether the given custom bucket is valid (exists) on this device. */
+    public boolean isValidCustomBucket(int customBucket) {
         return customBucket >= 0
                 && customBucketToIndex(customBucket) < mAccumulatedEnergiesMicroJoules.length;
     }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
index 97c07ea..4c52848 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
@@ -24,6 +24,7 @@
 import android.os.BatteryStats.HistoryItem;
 import android.os.BatteryStats.Uid.Sensor;
 import android.os.WorkSource;
+import android.util.SparseLongArray;
 import android.view.Display;
 
 import androidx.test.filters.SmallTest;
@@ -583,6 +584,95 @@
         checkMeasuredEnergy("H", uid1, blame1, uid2, blame2, globalDoze, bi);
     }
 
+    @SmallTest
+    public void testUpdateCustomMeasuredEnergyDataLocked_neverCalled() {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        final MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        bi.setOnBatteryInternal(true);
+
+        final int uid1 = 11500;
+        final int uid2 = 11501;
+
+        // Initially, all custom buckets report energy of 0.
+        checkCustomMeasuredEnergy("0", 0, 0, uid1, 0, 0, uid2, 0, 0, bi);
+    }
+
+    @SmallTest
+    public void testUpdateCustomMeasuredEnergyDataLocked() {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        final MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+        final int bucketA = 0; // Custom bucket 0
+        final int bucketB = 1; // Custom bucket 1
+
+        long totalBlameA = 0; // Total energy consumption for bucketA (may exceed sum of uids)
+        long totalBlameB = 0; // Total energy consumption for bucketB (may exceed sum of uids)
+
+        final int uid1 = 10500;
+        long blame1A = 0; // Blame for uid1 in bucketA
+        long blame1B = 0; // Blame for uid1 in bucketB
+
+        final int uid2 = 10501;
+        long blame2A = 0; // Blame for uid2 in bucketA
+        long blame2B = 0; // Blame for uid2 in bucketB
+
+        final SparseLongArray newEnergiesA = new SparseLongArray(2);
+        final SparseLongArray newEnergiesB = new SparseLongArray(2);
+
+
+        // ----- Case A: battery off (so blame does not increase)
+        bi.setOnBatteryInternal(false);
+
+        newEnergiesA.put(uid1, 20_000);
+        // Implicit newEnergiesA.put(uid2, 0);
+        bi.updateCustomMeasuredEnergyDataLocked(bucketA, 500_000, newEnergiesA);
+
+        newEnergiesB.put(uid1, 60_000);
+        // Implicit newEnergiesB.put(uid2, 0);
+        bi.updateCustomMeasuredEnergyDataLocked(bucketB, 700_000, newEnergiesB);
+
+        checkCustomMeasuredEnergy(
+                "A", totalBlameA, totalBlameB, uid1, blame1A, blame1B, uid2, blame2A, blame2B, bi);
+
+
+        // ----- Case B: battery on
+        bi.setOnBatteryInternal(true);
+
+        newEnergiesA.put(uid1, 7_000); blame1A += 7_000;
+        // Implicit newEnergiesA.put(uid2, 0); blame2A += 0;
+        bi.updateCustomMeasuredEnergyDataLocked(bucketA, 310_000, newEnergiesA);
+        totalBlameA += 310_000;
+
+        newEnergiesB.put(uid1, 63_000); blame1B += 63_000;
+        newEnergiesB.put(uid2, 15_000); blame2B += 15_000;
+        bi.updateCustomMeasuredEnergyDataLocked(bucketB, 790_000, newEnergiesB);
+        totalBlameB += 790_000;
+
+        checkCustomMeasuredEnergy(
+                "B", totalBlameA, totalBlameB, uid1, blame1A, blame1B, uid2, blame2A, blame2B, bi);
+
+
+        // ----- Case C: battery still on
+        newEnergiesA.delete(uid1); blame1A += 0;
+        newEnergiesA.put(uid2, 16_000); blame2A += 16_000;
+        bi.updateCustomMeasuredEnergyDataLocked(bucketA, 560_000, newEnergiesA);
+        totalBlameA += 560_000;
+
+        bi.updateCustomMeasuredEnergyDataLocked(bucketB, 10_000, null);
+        totalBlameB += 10_000;
+
+        checkCustomMeasuredEnergy(
+                "C", totalBlameA, totalBlameB, uid1, blame1A, blame1B, uid2, blame2A, blame2B, bi);
+
+
+        // ----- Case D: battery still on
+        bi.updateCustomMeasuredEnergyDataLocked(bucketA, 0, newEnergiesA);
+        bi.updateCustomMeasuredEnergyDataLocked(bucketB, 15_000, new SparseLongArray(1));
+        totalBlameB += 15_000;
+        checkCustomMeasuredEnergy(
+                "D", totalBlameA, totalBlameB, uid1, blame1A, blame1B, uid2, blame2A, blame2B, bi);
+    }
+
     private void setFgState(int uid, boolean fgOn, MockBatteryStatsImpl bi) {
         // Note that noteUidProcessStateLocked uses ActivityManager process states.
         if (fgOn) {
@@ -610,4 +700,29 @@
         assertEquals("Wrong doze for Case " + caseName, globalDoze,
                 bi.getScreenDozeEnergy());
     }
+
+    private void checkCustomMeasuredEnergy(String caseName,
+            long totalBlameA, long totalBlameB,
+            int uid1, long blame1A, long blame1B,
+            int uid2, long blame2A, long blame2B,
+            MockBatteryStatsImpl bi) {
+
+        assertEquals("Wrong total blame in bucket 0 for Case " + caseName, totalBlameA,
+                bi.getCustomMeasuredEnergyMicroJoules(0));
+
+        assertEquals("Wrong total blame in bucket 1 for Case " + caseName, totalBlameB,
+                bi.getCustomMeasuredEnergyMicroJoules(1));
+
+        assertEquals("Wrong uid1 blame in bucket 0 for Case " + caseName, blame1A,
+                bi.getUidStatsLocked(uid1).getCustomMeasuredEnergyMicroJoules(0));
+
+        assertEquals("Wrong uid1 blame in bucket 1 for Case " + caseName, blame1B,
+                bi.getUidStatsLocked(uid1).getCustomMeasuredEnergyMicroJoules(1));
+
+        assertEquals("Wrong uid2 blame in bucket 0 for Case " + caseName, blame2A,
+                bi.getUidStatsLocked(uid2).getCustomMeasuredEnergyMicroJoules(0));
+
+        assertEquals("Wrong uid2 blame in bucket 1 for Case " + caseName, blame2B,
+                bi.getUidStatsLocked(uid2).getCustomMeasuredEnergyMicroJoules(1));
+    }
 }
diff --git a/core/tests/coretests/src/com/android/internal/power/MeasuredEnergyStatsTest.java b/core/tests/coretests/src/com/android/internal/power/MeasuredEnergyStatsTest.java
index b9908f4..1679774 100644
--- a/core/tests/coretests/src/com/android/internal/power/MeasuredEnergyStatsTest.java
+++ b/core/tests/coretests/src/com/android/internal/power/MeasuredEnergyStatsTest.java
@@ -387,6 +387,24 @@
     }
 
     @Test
+    public void testIsValidCustomBucket() {
+        final MeasuredEnergyStats stats
+                = new MeasuredEnergyStats(new boolean[NUMBER_STANDARD_ENERGY_BUCKETS], 3);
+        assertFalse(stats.isValidCustomBucket(-1));
+        assertTrue(stats.isValidCustomBucket(0));
+        assertTrue(stats.isValidCustomBucket(1));
+        assertTrue(stats.isValidCustomBucket(2));
+        assertFalse(stats.isValidCustomBucket(3));
+        assertFalse(stats.isValidCustomBucket(4));
+
+        final MeasuredEnergyStats boringStats
+                = new MeasuredEnergyStats(new boolean[NUMBER_STANDARD_ENERGY_BUCKETS], 0);
+        assertFalse(boringStats.isValidCustomBucket(-1));
+        assertFalse(boringStats.isValidCustomBucket(0));
+        assertFalse(boringStats.isValidCustomBucket(1));
+    }
+
+    @Test
     public void testReset() {
         final boolean[] supportedStandardBuckets = new boolean[NUMBER_STANDARD_ENERGY_BUCKETS];
         final int numCustomBuckets = 2;
diff --git a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
index c6947c2d..b994389 100644
--- a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
@@ -15,8 +15,6 @@
  */
 package com.android.server.am;
 
-import static com.android.internal.power.MeasuredEnergyArray.SUBSYSTEM_DISPLAY;
-
 import android.annotation.Nullable;
 import android.bluetooth.BluetoothActivityEnergyInfo;
 import android.bluetooth.BluetoothAdapter;
@@ -39,13 +37,12 @@
 import android.telephony.TelephonyManager;
 import android.util.IntArray;
 import android.util.Slog;
-import android.util.SparseIntArray;
+import android.util.SparseArray;
 import android.util.SparseLongArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.BatteryStatsImpl;
-import com.android.internal.power.MeasuredEnergyArray;
 import com.android.internal.power.MeasuredEnergyStats;
 import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.function.pooled.PooledLambda;
@@ -148,13 +145,13 @@
     private WifiActivityEnergyInfo mLastWifiInfo =
             new WifiActivityEnergyInfo(0, 0, 0, 0, 0, 0);
 
-    /** Maps the EnergyConsumer id to it's corresponding {@link MeasuredEnergySubsystem} */
+    /**
+     * Maps an {@link EnergyConsumerType} to it's corresponding {@link EnergyConsumer#id}s,
+     * unless it is of {@link EnergyConsumer#type}=={@link EnergyConsumerType#OTHER}
+     */
+    // TODO: Hook this up (it isn't used yet)
     @GuardedBy("mWorkerLock")
-    private @Nullable SparseIntArray mEnergyConsumerToSubsystemMap = null;
-
-    /** Maps a {@link MeasuredEnergySubsystem} to it's corresponding EnergyConsumer id */
-    @GuardedBy("mWorkerLock")
-    private @Nullable SparseIntArray mSubsystemToEnergyConsumerMap = null;
+    private @Nullable SparseArray<int[]> mEnergyConsumerTypeToIdMap = null;
 
     /** Snapshot of measured energies, or null if no measured energies are supported. */
     @GuardedBy("mWorkerLock")
@@ -204,18 +201,26 @@
             mWifiManager = wm;
             mTelephony = tm;
             mPowerStatsInternal = psi;
+
+            boolean[] supportedStdBuckets = null;
+            int numCustomBuckets = 0;
             if (mPowerStatsInternal != null) {
-                populateEnergyConsumerSubsystemMapsLocked();
-                final MeasuredEnergyArray initialMeasuredEnergies = getEnergyConsumptionData();
-                mMeasuredEnergySnapshot = initialMeasuredEnergies == null
-                        ? null : new MeasuredEnergySnapshot(initialMeasuredEnergies);
-                final boolean[] supportedStdBuckets
-                        = getSupportedEnergyBuckets(initialMeasuredEnergies);
-                final int numCustomBuckets = 0; // TODO: Get this from initialMeasuredEnergies
-                synchronized (mStats) {
-                    mStats.initMeasuredEnergyStatsLocked(supportedStdBuckets, numCustomBuckets);
+                final SparseArray<EnergyConsumer> idToConsumer
+                        = populateEnergyConsumerSubsystemMapsLocked();
+                if (idToConsumer != null) {
+                    mMeasuredEnergySnapshot = new MeasuredEnergySnapshot(idToConsumer);
+                    final EnergyConsumerResult[] initialEcrs = getEnergyConsumptionData();
+                    // According to spec, initialEcrs will include 0s for consumers that haven't
+                    // used any energy yet, as long as they are supported; however, attributed uid
+                    // energies will be absent if their energy is 0.
+                    mMeasuredEnergySnapshot.updateAndGetDelta(initialEcrs);
+                    numCustomBuckets = mMeasuredEnergySnapshot.getNumOtherOrdinals();
+                    supportedStdBuckets = getSupportedEnergyBuckets(idToConsumer);
                 }
             }
+            synchronized (mStats) {
+                mStats.initMeasuredEnergyStatsLocked(supportedStdBuckets, numCustomBuckets);
+            }
         }
     }
 
@@ -568,7 +573,9 @@
         } catch (ExecutionException e) {
             Slog.w(TAG, "exception reading modem stats: " + e.getCause());
         }
-        final SparseLongArray energyDeltas = mMeasuredEnergySnapshot == null ? null :
+
+        final MeasuredEnergySnapshot.MeasuredEnergyDeltaData measuredEnergyDeltas =
+                mMeasuredEnergySnapshot == null ? null :
                 mMeasuredEnergySnapshot.updateAndGetDelta(getMeasuredEnergyLocked(updateFlags));
 
         final long elapsedRealtime = SystemClock.elapsedRealtime();
@@ -576,6 +583,7 @@
         final long elapsedRealtimeUs = elapsedRealtime * 1000;
         final long uptimeUs = uptime * 1000;
 
+        // Now that we have finally received all the data, we can tell mStats about it.
         synchronized (mStats) {
             mStats.addHistoryEventLocked(
                     elapsedRealtime,
@@ -601,10 +609,21 @@
             }
 
             // Inform mStats about each applicable measured energy.
-            if (energyDeltas != null) {
-                final long displayEnergy = energyDeltas.get(SUBSYSTEM_DISPLAY, 0L);
-                // Always pass in what BatteryExternalStatsWorker thinks screenState is.
-                mStats.updateDisplayEnergyLocked(displayEnergy, screenState, elapsedRealtime);
+            if (measuredEnergyDeltas != null) {
+                final long displayEnergy = measuredEnergyDeltas.displayEnergyUJ;
+                if (displayEnergy != MeasuredEnergySnapshot.UNAVAILABLE) {
+                    // If updating, pass in what BatteryExternalStatsWorker thinks screenState is.
+                    mStats.updateDisplayEnergyLocked(displayEnergy, screenState, elapsedRealtime);
+                }
+            }
+            // Inform mStats about each applicable custom energy bucket.
+            if (measuredEnergyDeltas != null && measuredEnergyDeltas.otherTotalEnergyUJ != null) {
+                // Iterate over the custom (EnergyConsumerType.OTHER) ordinals.
+                for (int ord = 0; ord < measuredEnergyDeltas.otherTotalEnergyUJ.length; ord++) {
+                    long totalEnergy = measuredEnergyDeltas.otherTotalEnergyUJ[ord];
+                    SparseLongArray uidEnergies = measuredEnergyDeltas.otherUidEnergiesUJ[ord];
+                    mStats.updateCustomMeasuredEnergyDataLocked(ord, totalEnergy, uidEnergies);
+                }
             }
 
             if (bluetoothInfo != null) {
@@ -621,7 +640,8 @@
 
         if (wifiInfo != null) {
             if (wifiInfo.isValid()) {
-                // TODO: wifiEnergyDelta = energyDeltas.get(MeasuredEnergyArray.SUBSYSTEM_WIFI, 0L);
+                // TODO: wifiEnergyDelta = measuredEnergyDeltas.consumerTypeEnergyUJ
+                //               .get(EnergyConsumerType.WIFI, MeasuredEnergySnapshot.UNAVAILABLE)
                 mStats.updateWifiState(extractDeltaLocked(wifiInfo)
                         /*, TODO: wifiEnergyDelta */, elapsedRealtime, uptime);
             } else {
@@ -740,21 +760,23 @@
     }
 
     /**
-     * Map the {@link MeasuredEnergyArray.MeasuredEnergySubsystem}s in the given energyArray to
+     * Map the {@link EnergyConsumerType}s in the given energyArray to
      * their corresponding {@link MeasuredEnergyStats.StandardEnergyBucket}s.
      * Does not include custom energy buckets (which are always, by definition, supported).
      *
      * @return array with true for index i if standard energy bucket i is supported.
      */
-    private static @Nullable boolean[] getSupportedEnergyBuckets(MeasuredEnergyArray energyArray) {
-        if (energyArray == null) {
+    private static @Nullable boolean[] getSupportedEnergyBuckets(
+            SparseArray<EnergyConsumer> idToConsumer) {
+        if (idToConsumer == null) {
             return null;
         }
         final boolean[] buckets = new boolean[MeasuredEnergyStats.NUMBER_STANDARD_ENERGY_BUCKETS];
-        final int size = energyArray.size();
-        for (int energyIdx = 0; energyIdx < size; energyIdx++) {
-            switch (energyArray.getSubsystem(energyIdx)) {
-                case MeasuredEnergyArray.SUBSYSTEM_DISPLAY:
+        final int size = idToConsumer.size();
+        for (int idx = 0; idx < size; idx++) {
+            final EnergyConsumer consumer = idToConsumer.valueAt(idx);
+            switch (consumer.type) {
+                case EnergyConsumerType.DISPLAY:
                     buckets[MeasuredEnergyStats.ENERGY_BUCKET_SCREEN_ON] = true;
                     buckets[MeasuredEnergyStats.ENERGY_BUCKET_SCREEN_DOZE] = true;
                     buckets[MeasuredEnergyStats.ENERGY_BUCKET_SCREEN_OTHER] = true;
@@ -764,71 +786,22 @@
         return buckets;
     }
 
-    /**
-     * Get a {@link MeasuredEnergyArray} with the latest
-     * {@link MeasuredEnergyArray.MeasuredEnergySubsystem} energy usage since boot.
-     *
-     * TODO(b/176988041): Replace {@link MeasuredEnergyArray} usage with {@link
-     * EnergyConsumerResult}[]
-     */
+    /** Get {@link EnergyConsumerResult}s with the latest energy usage since boot. */
     @GuardedBy("mWorkerLock")
-    @VisibleForTesting
-    public @Nullable MeasuredEnergyArray getEnergyConsumptionData() {
-        final EnergyConsumerResult[] results;
+    private @Nullable EnergyConsumerResult[] getEnergyConsumptionData() {
         try {
-            results = mPowerStatsInternal.getEnergyConsumedAsync(new int[0])
+            return mPowerStatsInternal.getEnergyConsumedAsync(new int[0])
                     .get(EXTERNAL_STATS_SYNC_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
         } catch (Exception e) {
             Slog.e(TAG, "Failed to getEnergyConsumedAsync", e);
             return null;
         }
-        if (results == null) return null;
-        final int size = results.length;
-        final int[] subsystems = new int[size];
-        final long[] energyUJ = new long[size];
-
-        int count = 0;
-        for (int i = 0; i < size; i++) {
-            final EnergyConsumerResult consumer = results[i];
-            final int subsystem = mEnergyConsumerToSubsystemMap.get(consumer.id,
-                    MeasuredEnergyArray.SUBSYSTEM_UNKNOWN);
-            if (subsystem == MeasuredEnergyArray.SUBSYSTEM_UNKNOWN) continue;
-            subsystems[count] = subsystem;
-            energyUJ[count] = consumer.energyUWs;
-            count++;
-        }
-        final int arraySize = count;
-        return new MeasuredEnergyArray() {
-            @Override
-            public int getSubsystem(int index) {
-                if (index >= size()) {
-                    throw new IllegalArgumentException(
-                            "Out of bounds subsystem index! index : " + index + ", size : "
-                                    + size());
-                }
-                return subsystems[index];
-            }
-
-            @Override
-            public long getEnergy(int index) {
-                if (index >= size()) {
-                    throw new IllegalArgumentException(
-                            "Out of bounds subsystem index! index : " + index + ", size : "
-                                    + size());
-                }
-                return energyUJ[index];
-            }
-
-            @Override
-            public int size() {
-                return arraySize;
-            }
-        };
     }
 
-    /** Fetch MeasuredEnergyArray for supported subsystems based on the given updateFlags. */
+    /** Fetch EnergyConsumerResult[] for supported subsystems based on the given updateFlags. */
     @GuardedBy("mWorkerLock")
-    private @Nullable MeasuredEnergyArray getMeasuredEnergyLocked(@ExternalUpdateFlag int flags) {
+    private @Nullable EnergyConsumerResult[] getMeasuredEnergyLocked(@ExternalUpdateFlag int flags)
+    {
         if (mMeasuredEnergySnapshot == null || mPowerStatsInternal == null) return null;
 
         if (flags == UPDATE_ALL) {
@@ -836,12 +809,13 @@
             return getEnergyConsumptionData();
         }
 
-        final List<Integer> energySubsystems = new ArrayList<>();
+        final List<Integer> energyConsumerIds = new ArrayList<>();
         if ((flags & UPDATE_DISPLAY) != 0) {
-            addEnergyConsumerIdLocked(energySubsystems, SUBSYSTEM_DISPLAY);
+            addEnergyConsumerIdLocked(energyConsumerIds, EnergyConsumerType.DISPLAY);
         }
         // TODO: Wifi, Bluetooth, etc., go here
-        if (energySubsystems.isEmpty()) {
+
+        if (energyConsumerIds.isEmpty()) {
             return null;
         }
         // TODO: Query *specific* subsystems from HAL based on energyConsumerIds.toArray()
@@ -849,59 +823,48 @@
     }
 
     @GuardedBy("mWorkerLock")
-    private void addEnergyConsumerIdLocked(List<Integer> energyConsumerIds,
-            @MeasuredEnergyArray.MeasuredEnergySubsystem int consumerId) {
-        if (mMeasuredEnergySnapshot.hasSubsystem(consumerId)) {
-            energyConsumerIds.add(consumerId);
-        }
+    private void addEnergyConsumerIdLocked(
+            List<Integer> energyConsumerIds, @EnergyConsumerType int type) {
+        final int consumerId = 0; // TODO: Use mEnergyConsumerTypeToIdMap to get this
+        energyConsumerIds.add(consumerId);
     }
 
+    /** Populates the cached type->ids map, and returns the (inverse) id->EnergyConsumer map. */
     @GuardedBy("mWorkerLock")
-    private void populateEnergyConsumerSubsystemMapsLocked() {
+    private @Nullable SparseArray<EnergyConsumer> populateEnergyConsumerSubsystemMapsLocked() {
         if (mPowerStatsInternal == null) {
-            // PowerStatsInternal unavailable, don't bother populating maps.
-            mEnergyConsumerToSubsystemMap = null;
-            mSubsystemToEnergyConsumerMap = null;
-            return;
+            return null;
         }
         final EnergyConsumer[] energyConsumers = mPowerStatsInternal.getEnergyConsumerInfo();
-        if (energyConsumers == null) {
-            // EnergyConsumer data unavailable, don't bother populating maps.
-            mEnergyConsumerToSubsystemMap = null;
-            mSubsystemToEnergyConsumerMap = null;
-            return;
+        if (energyConsumers == null || energyConsumers.length == 0) {
+            return null;
         }
 
-        final int length = energyConsumers.length;
-        if (length == 0) {
-            // EnergyConsumer array empty, don't bother populating maps.
-            mEnergyConsumerToSubsystemMap = null;
-            mSubsystemToEnergyConsumerMap = null;
-            return;
-        } else {
-            mEnergyConsumerToSubsystemMap = new SparseIntArray(length);
-            mSubsystemToEnergyConsumerMap = new SparseIntArray(length);
-        }
+        // TODO: Initialize typeToIds
+        // Maps type -> {ids} (1:n map, since multiple ids might have the same type)
+        // final SparseArray<SparseIntArray> typeToIds = new SparseArray<>();
+
+        // Maps id -> EnergyConsumer (1:1 map)
+        final SparseArray<EnergyConsumer> idToConsumer = new SparseArray<>(energyConsumers.length);
 
         // Add all expected EnergyConsumers to the maps
-        for (int i = 0; i < length; i++) {
-            final EnergyConsumer consumer = energyConsumers[i];
-            switch (consumer.type) {
-                case EnergyConsumerType.DISPLAY:
-                    if (consumer.ordinal == 0) {
-                        mEnergyConsumerToSubsystemMap.put(consumer.id,
-                                MeasuredEnergyArray.SUBSYSTEM_DISPLAY);
-                        mSubsystemToEnergyConsumerMap.put(MeasuredEnergyArray.SUBSYSTEM_DISPLAY,
-                                consumer.id);
-                    } else {
-                        Slog.w(TAG, "Unexpected ordinal (" + consumer.ordinal
-                                + ") for EnergyConsumerType.DISPLAY");
-                    }
-                    break;
-                default:
-                    Slog.w(TAG, "Unexpected EnergyConsumerType (" + consumer.type + ")");
+        for (final EnergyConsumer consumer : energyConsumers) {
+            // Check for inappropriate ordinals
+            if (consumer.ordinal != 0) {
+                switch (consumer.type) {
+                    case EnergyConsumerType.OTHER:
+                    case EnergyConsumerType.CPU_CLUSTER:
+                        break;
+                    default:
+                        Slog.w(TAG, "EnergyConsumer '" + consumer.name + "' has unexpected ordinal "
+                                + consumer.ordinal + " for type " + consumer.type);
+                        continue; // Ignore this consumer
+                }
             }
-
+            idToConsumer.put(consumer.id, consumer);
+            // TODO: Also populate typeToIds map
         }
+        // TODO: Store typeToIds in mEnergyConsumerTypeToIdMap.
+        return idToConsumer;
     }
 }
diff --git a/services/core/java/com/android/server/am/MeasuredEnergySnapshot.java b/services/core/java/com/android/server/am/MeasuredEnergySnapshot.java
index b915c0c..9e0aa32 100644
--- a/services/core/java/com/android/server/am/MeasuredEnergySnapshot.java
+++ b/services/core/java/com/android/server/am/MeasuredEnergySnapshot.java
@@ -17,132 +17,260 @@
 package com.android.server.am;
 
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.hardware.power.stats.EnergyConsumer;
+import android.hardware.power.stats.EnergyConsumerAttribution;
+import android.hardware.power.stats.EnergyConsumerResult;
+import android.hardware.power.stats.EnergyConsumerType;
 import android.util.Slog;
+import android.util.SparseArray;
 import android.util.SparseLongArray;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.power.MeasuredEnergyArray;
-import com.android.internal.power.MeasuredEnergyArray.MeasuredEnergySubsystem;
 
 import java.io.PrintWriter;
-import java.util.Arrays;
 
 /**
- * Keeps snapshots of data from previously pulled MeasuredEnergyArrays.
+ * Keeps snapshots of data from previously pulled EnergyConsumerResults.
  */
 @VisibleForTesting
 public class MeasuredEnergySnapshot {
     private static final String TAG = "MeasuredEnergySnapshot";
 
-    private static final long UNAVAILABLE = -1;
+    public static final long UNAVAILABLE = -1L;
+
+    /** Map of {@link EnergyConsumer#id} to its corresponding {@link EnergyConsumer}. */
+    private final SparseArray<EnergyConsumer> mEnergyConsumers;
+
+    /** Number of ordinals for {@link EnergyConsumerType#OTHER}. */
+    private final int mNumOtherOrdinals;
 
     /**
-     * Energy snapshots from the last time each {@link MeasuredEnergySubsystem} was updated.
+     * Energy snapshots, mapping {@link EnergyConsumer#id} to energy (UJ) from the last time
+     * each {@link EnergyConsumer} was updated.
      *
-     * Note that the snapshots for different subsystems may have been taken at different times.
+     * Note that the snapshots for different ids may have been taken at different times.
+     * Note that energies for all existing ids are stored here, including each ordinal of type
+     * {@link EnergyConsumerType#OTHER} (tracking their total energy usage).
      *
-     * A snapshot is {@link #UNAVAILABLE} if the subsystem has never been updated (ie. unsupported).
+     * If an id is not present yet, it is treated as uninitialized (energy {@link #UNAVAILABLE}).
      */
-    private final long[] mMeasuredEnergySnapshots;
+    private final SparseLongArray mMeasuredEnergySnapshots;
 
     /**
-     * Constructor that initializes to the given energyArray;
-     * all subsystems not mentioned in initialEnergyArray are set to UNAVAILABLE.
+     * Energy snapshots <b>per uid</b> from the last time each {@link EnergyConsumer} of type
+     * {@link EnergyConsumerType#OTHER} was updated.
+     * It maps each OTHER {@link EnergyConsumer#id} to a SparseLongArray, which itself maps each
+     * uid to an energy (UJ). That is,
+     * mAttributionSnapshots.get(consumerId).get(uid) = energy used by uid for this consumer.
+     *
+     * If an id is not present yet, it is treated as uninitialized (i.e. each uid is unavailable).
+     * If an id is present but a uid is not present, that uid's energy is 0.
      */
-    public MeasuredEnergySnapshot(MeasuredEnergyArray initialEnergyArray) {
-        this(MeasuredEnergyArray.NUMBER_SUBSYSTEMS, initialEnergyArray);
+    private final SparseArray<SparseLongArray> mAttributionSnapshots;
+
+    /**
+     * Constructor that initializes to the given id->EnergyConsumer map, indicating which consumers
+     * exist and what their details are.
+     */
+    MeasuredEnergySnapshot(@NonNull SparseArray<EnergyConsumer> idToConsumerMap) {
+        mEnergyConsumers = idToConsumerMap;
+        mMeasuredEnergySnapshots = new SparseLongArray(mEnergyConsumers.size());
+
+        mNumOtherOrdinals = calculateNumOtherOrdinals(idToConsumerMap);
+        mAttributionSnapshots = new SparseArray<>(mNumOtherOrdinals);
     }
 
     /**
-     * Constructor (for testing) that initializes to the given energyArray and numSubsystems;
-     * all subsystems not mentioned in initialEnergyArray are set to UNAVAILABLE.
+     * Returns the number of ordinals for {@link EnergyConsumerType#OTHER}, i.e. the number of
+     * custom energy buckets supported by the device.
      */
-    @VisibleForTesting
-    MeasuredEnergySnapshot(int numSubsystems, MeasuredEnergyArray initialEnergyArray) {
-        if (initialEnergyArray.size() > numSubsystems) {
-            throw new IllegalArgumentException("Energy array contains " + initialEnergyArray.size()
-                    + " subsystems, which exceeds the maximum allowed of " + numSubsystems);
-        }
-        mMeasuredEnergySnapshots = new long[numSubsystems];
-        Arrays.fill(mMeasuredEnergySnapshots, UNAVAILABLE);
-        fillGivenSubsystems(initialEnergyArray);
+    public int getNumOtherOrdinals() {
+        return mNumOtherOrdinals;
     }
 
-    /**
-     * For the subsystems present in energyArray, overwrites mMeasuredEnergySnapshots with their
-     * energy values from energyArray.
-     */
-    private void fillGivenSubsystems(MeasuredEnergyArray energyArray) {
-        final int size = energyArray.size();
-        for (int i = 0; i < size; i++) {
-            final int subsystem = energyArray.getSubsystem(i);
-            mMeasuredEnergySnapshots[subsystem] = energyArray.getEnergy(i);
-        }
+    /** Class for returning measured energy delta data. */
+    static class MeasuredEnergyDeltaData {
+        /** The energyUJ for {@link EnergyConsumerType#DISPLAY}. */
+        public long displayEnergyUJ = UNAVAILABLE;
+
+        /** Map of {@link EnergyConsumerType#OTHER} ordinals to their total energyUJ. */
+        public @Nullable long[] otherTotalEnergyUJ = null;
+
+        /** Map of {@link EnergyConsumerType#OTHER} ordinals to their {uid->energyUJ} maps. */
+        public @Nullable SparseLongArray[] otherUidEnergiesUJ = null;
     }
 
     /**
      * Update with the some freshly measured energies and return the difference (delta)
      * between the previously stored values and the passed-in values.
      *
-     * @param energyArray measured energy array for some (possibly not all) subsystems.
+     * @param ecrs EnergyConsumerResults for some (possibly not all) {@link EnergyConsumer}s.
+     *             Consumers that are not present are ignored (they are *not* treated as 0).
      *
-     * @return a map from the updated subsystems to their corresponding energy deltas.
-     *         Subsystems not present in energyArray will not appear.
-     *         Subsystems with no difference in energy will not appear.
-     *         Returns null, if energyArray is null.
+     * @return a MeasuredEnergyDeltaData, containing maps from the updated consumers to
+     *         their corresponding energy deltas.
+     *         Fields with no interesting data (consumers not present in ecrs or with no energy
+     *         difference) will generally be left as their default values.
+     *         otherTotalEnergyUJ and otherUidEnergiesUJ are always either both null or both of
+     *         length {@link #getNumOtherOrdinals()}.
+     *         Returns null, if ecrs is null or empty.
      */
-    public @Nullable SparseLongArray updateAndGetDelta(MeasuredEnergyArray energyArray) {
-        if (energyArray == null) {
+    public @Nullable MeasuredEnergyDeltaData updateAndGetDelta(EnergyConsumerResult[] ecrs) {
+        if (ecrs == null || ecrs.length == 0) {
             return null;
         }
-        final SparseLongArray delta = new SparseLongArray();
-        final int size = energyArray.size();
-        for (int i = 0; i < size; i++) {
-            final int updatedSubsystem = energyArray.getSubsystem(i);
-            final long newEnergyUJ = energyArray.getEnergy(i);
-            final long oldEnergyUJ = mMeasuredEnergySnapshots[updatedSubsystem];
+        final MeasuredEnergyDeltaData output = new MeasuredEnergyDeltaData();
 
-            // If this is the first valid energy, there is no delta to take.
-            if (oldEnergyUJ < 0) continue;
-            final long deltaUJ = newEnergyUJ - oldEnergyUJ;
-            if (deltaUJ == 0) continue;
-            if (deltaUJ < 0) {
-                Slog.e(TAG, "For subsystem " + updatedSubsystem + ", new energy (" + newEnergyUJ
-                        + ") is less than old energy (" + oldEnergyUJ + "). Skipping. ");
+        for (final EnergyConsumerResult ecr : ecrs) {
+            // Extract the new energy data for the current consumer.
+            final int consumerId = ecr.id;
+            final long newEnergyUJ = ecr.energyUWs;
+            final EnergyConsumerAttribution[] newAttributions = ecr.attribution;
+
+            // Look up the static information about this consumer.
+            final EnergyConsumer consumer = mEnergyConsumers.get(consumerId, null);
+            if (consumer == null) {
+                Slog.e(TAG, "updateAndGetDelta given invalid consumerId " + consumerId);
                 continue;
             }
-            delta.put(updatedSubsystem, deltaUJ);
+            final int type = consumer.type;
+            final int ordinal = consumer.ordinal;
+
+            // Look up, and update, the old energy information about this consumer.
+            final long oldEnergyUJ = mMeasuredEnergySnapshots.get(consumerId, UNAVAILABLE);
+            mMeasuredEnergySnapshots.put(consumerId, newEnergyUJ);
+            final SparseLongArray otherUidEnergies
+                    = updateAndGetDeltaForTypeOther(consumer, newAttributions);
+
+            // Everything is fully done being updated. We now calculate the delta for returning.
+
+            // NB: Since sum(attribution.energyUWs)<=energyUWs we assume that if deltaEnergy==0
+            // there's no attribution either. Technically that isn't enforced at the HAL, but we
+            // can't really trust data like that anyway.
+
+            if (oldEnergyUJ < 0) continue; // Generally happens only on initialization.
+            if (newEnergyUJ == oldEnergyUJ) continue;
+            final long deltaUJ = newEnergyUJ - oldEnergyUJ;
+            if (deltaUJ < 0) {
+                Slog.e(TAG, "EnergyConsumer " + consumer.name + ": new energy (" + newEnergyUJ
+                        + ") < old energy (" + oldEnergyUJ + "). Skipping. ");
+                continue;
+            }
+
+            switch (type) {
+                case EnergyConsumerType.DISPLAY:
+                    output.displayEnergyUJ = deltaUJ;
+                    break;
+                case EnergyConsumerType.OTHER:
+                    if (output.otherTotalEnergyUJ == null) {
+                        output.otherTotalEnergyUJ = new long[getNumOtherOrdinals()];
+                        output.otherUidEnergiesUJ = new SparseLongArray[getNumOtherOrdinals()];
+                    }
+                    output.otherTotalEnergyUJ[ordinal] = deltaUJ;
+                    output.otherUidEnergiesUJ[ordinal] = otherUidEnergies;
+                    break;
+                default:
+                    Slog.w(TAG, "Ignoring consumer " + consumer.name + " of unknown type " + type);
+
+            }
         }
-
-        fillGivenSubsystems(energyArray);
-
-        return delta;
+        return output;
     }
 
     /**
-     * Check if a subsystem's measured energy is available.
-     * @param subsystem which subsystem.
-     * @return true if subsystem is available.
+     * For a consumer of type {@link EnergyConsumerType#OTHER}, updates
+     * {@link #mAttributionSnapshots} with freshly measured energies (per uid) and returns the
+     * difference (delta) between the previously stored values and the passed-in values.
+     *
+     * @param consumerInfo a consumer of type {@link EnergyConsumerType#OTHER}.
+     * @param newAttributions Record of uids and their new energyUJ values.
+     *                        Any uid not present is treated as having energy 0.
+     *                        If null or empty, all uids are treated as having energy 0.
+     * @return A map (in the sense of {@link MeasuredEnergyDeltaData#otherUidEnergiesUJ} for this
+     *         consumer) of uid -> energyDelta, with all uids that have a non-zero energyDelta.
+     *         Returns null if no delta available to calculate.
      */
-    public boolean hasSubsystem(@MeasuredEnergySubsystem int subsystem) {
-        return mMeasuredEnergySnapshots[subsystem] != UNAVAILABLE;
+    private @Nullable SparseLongArray updateAndGetDeltaForTypeOther(
+            @NonNull EnergyConsumer consumerInfo,
+            @Nullable EnergyConsumerAttribution[] newAttributions) {
+
+        if (consumerInfo.type != EnergyConsumerType.OTHER) {
+            return null;
+        }
+        if (newAttributions == null) {
+            // Treat null as empty (i.e. all uids have 0 energy).
+            newAttributions = new EnergyConsumerAttribution[0];
+        }
+
+        // SparseLongArray mapping uid -> energyUJ (for this particular consumerId)
+        SparseLongArray uidOldEnergyMap = mAttributionSnapshots.get(consumerInfo.id, null);
+
+        // If uidOldEnergyMap wasn't present, each uid was UNAVAILABLE, so update data and return.
+        if (uidOldEnergyMap == null) {
+            uidOldEnergyMap = new SparseLongArray(newAttributions.length);
+            mAttributionSnapshots.put(consumerInfo.id, uidOldEnergyMap);
+            for (EnergyConsumerAttribution newAttribution : newAttributions) {
+                uidOldEnergyMap.put(newAttribution.uid, newAttribution.energyUWs);
+            }
+            return null;
+        }
+
+        // Map uid -> energyDelta. No initial capacity since many deltas might be 0.
+        final SparseLongArray uidEnergyDeltas = new SparseLongArray();
+
+        for (EnergyConsumerAttribution newAttribution : newAttributions) {
+            final int uid = newAttribution.uid;
+            final long newEnergyUJ = newAttribution.energyUWs;
+            // uidOldEnergyMap was present. So any particular uid that wasn't present, had 0 energy.
+            final long oldEnergyUJ = uidOldEnergyMap.get(uid, 0L);
+            uidOldEnergyMap.put(uid, newEnergyUJ);
+
+            // Everything is fully done being updated. We now calculate the delta for returning.
+            if (oldEnergyUJ < 0) continue;
+            if (newEnergyUJ == oldEnergyUJ) continue;
+            final long deltaUJ = newEnergyUJ - oldEnergyUJ;
+            if (deltaUJ < 0) {
+                Slog.e(TAG, "EnergyConsumer " + consumerInfo.name + ": new energy (" + newEnergyUJ
+                        + ") but old energy (" + oldEnergyUJ + "). Skipping. ");
+                continue;
+            }
+            uidEnergyDeltas.put(uid, deltaUJ);
+        }
+        return uidEnergyDeltas;
     }
 
     /** Dump debug data. */
     public void dump(PrintWriter pw) {
-        pw.println("Measured energy snapshot (microjoules):");
-        pw.print("   ");
-        for (int i = 0; i < MeasuredEnergyArray.NUMBER_SUBSYSTEMS; i++) {
-            final long energyUJ = mMeasuredEnergySnapshots[i];
-            if (energyUJ == UNAVAILABLE) continue;
-            pw.print(MeasuredEnergyArray.SUBSYSTEM_NAMES[i]);
-            pw.print(" : ");
-            pw.print(energyUJ);
-            if (i != MeasuredEnergyArray.NUMBER_SUBSYSTEMS - 1) {
-                pw.print(", ");
-            }
+        pw.println("Measured energy snapshot");
+        pw.println("List of EnergyConsumers:");
+        for (int i = 0; i < mEnergyConsumers.size(); i++) {
+            final int id = mEnergyConsumers.keyAt(i);
+            final EnergyConsumer consumer = mEnergyConsumers.valueAt(i);
+            pw.println(String.format("    Consumer %d is {id=%d, ordinal=%d, type=%d, name=%s}", id,
+                    consumer.id, consumer.ordinal, consumer.type, consumer.name));
         }
+        pw.println("Map of consumerIds to energy (in microjoules):");
+        for (int i = 0; i < mMeasuredEnergySnapshots.size(); i++) {
+            final int id = mMeasuredEnergySnapshots.keyAt(i);
+            final long energyUJ = mMeasuredEnergySnapshots.valueAt(i);
+            pw.println(String.format("    Consumer %d has energy %d uJ}", id, energyUJ));
+        }
+        pw.println("List of the " + mNumOtherOrdinals + " OTHER EnergyConsumers:");
+        pw.println("    " + mAttributionSnapshots);
         pw.println();
     }
+
+    /** Determines the number of ordinals for {@link EnergyConsumerType#OTHER}. */
+    private static int calculateNumOtherOrdinals(SparseArray<EnergyConsumer> idToConsumer) {
+        if (idToConsumer == null) return 0;
+        int numOtherOrdinals = 0;
+        final int size = idToConsumer.size();
+        for (int idx = 0; idx < size; idx++) {
+            final EnergyConsumer consumer = idToConsumer.valueAt(idx);
+            if (consumer.type == EnergyConsumerType.OTHER) numOtherOrdinals++;
+        }
+        return numOtherOrdinals;
+    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/am/BatteryExternalStatsWorkerTest.java b/services/tests/servicestests/src/com/android/server/am/BatteryExternalStatsWorkerTest.java
index fdf5095..a946534 100644
--- a/services/tests/servicestests/src/com/android/server/am/BatteryExternalStatsWorkerTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/BatteryExternalStatsWorkerTest.java
@@ -16,8 +16,6 @@
 
 package com.android.server.am;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 
 import android.content.Context;
@@ -30,15 +28,12 @@
 import android.hardware.power.stats.StateResidencyResult;
 import android.power.PowerStatsInternal;
 import android.util.SparseArray;
-import android.util.SparseLongArray;
 
 import androidx.test.InstrumentationRegistry;
 
 import com.android.internal.os.BatteryStatsImpl;
-import com.android.internal.power.MeasuredEnergyArray;
 
 import org.junit.Before;
-import org.junit.Test;
 
 import java.util.concurrent.CompletableFuture;
 
@@ -63,44 +58,6 @@
                 mBatteryStatsImpl);
     }
 
-    @Test
-    public void getEnergyConsumptionData() {
-        SparseLongArray expectSubsystems = new SparseLongArray();
-        // Add some energy consumers used by BatteryExternalStatsWorker.
-        final int displayId = mPowerStatsInternal.addEnergyConsumer(EnergyConsumerType.DISPLAY, 0,
-                "display");
-        mPowerStatsInternal.incrementEnergyConsumption(displayId, 12345);
-        expectSubsystems.put(MeasuredEnergyArray.SUBSYSTEM_DISPLAY, 12345);
-
-        // Add an arbitrary energy consumer unused by BatteryExternalStatsWorker.
-        // Must be changed if '154' ever becomes an EnergyConsumerType used by BESW.
-        final int someId = mPowerStatsInternal.addEnergyConsumer((byte) 154, 0, "some_consumer");
-        mPowerStatsInternal.incrementEnergyConsumption(someId, 34567);
-
-        // Inform BESW that PowerStatsInternal is ready to query
-        mBatteryExternalStatsWorker.systemServicesReady();
-
-        MeasuredEnergyArray energies = mBatteryExternalStatsWorker.getEnergyConsumptionData();
-
-        assertEquals(expectSubsystems.size(), energies.size());
-        final int size = expectSubsystems.size();
-
-        for (int i = 0; i < size; i++) {
-            int subsystem = expectSubsystems.keyAt(i);
-            // find the subsystem in the returned MeasuredEnergyArray
-            int subsystemIndex = -1;
-            for (int j = 0; j < size; j++) {
-                if (subsystem == energies.getSubsystem(i)) {
-                    subsystemIndex = i;
-                    break;
-                }
-            }
-            assertNotEquals("Subsystem " + subsystem + " not found in MeasuredEnergyArray", -1,
-                    subsystemIndex);
-            assertEquals(expectSubsystems.valueAt(i), energies.getEnergy(subsystemIndex));
-        }
-    }
-
     public class TestInjector extends BatteryExternalStatsWorker.Injector {
         public TestInjector(Context context) {
             super(context);
diff --git a/services/tests/servicestests/src/com/android/server/am/MeasuredEnergySnapshotTest.java b/services/tests/servicestests/src/com/android/server/am/MeasuredEnergySnapshotTest.java
index 67d379a..1efce39 100644
--- a/services/tests/servicestests/src/com/android/server/am/MeasuredEnergySnapshotTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/MeasuredEnergySnapshotTest.java
@@ -16,18 +16,23 @@
 
 package com.android.server.am;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
+import static com.android.server.am.MeasuredEnergySnapshot.UNAVAILABLE;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+
+import android.hardware.power.stats.EnergyConsumer;
+import android.hardware.power.stats.EnergyConsumerAttribution;
+import android.hardware.power.stats.EnergyConsumerResult;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.util.SparseArray;
 import android.util.SparseLongArray;
 
 import androidx.test.filters.SmallTest;
 
-import com.android.internal.power.MeasuredEnergyArray;
+import com.android.server.am.MeasuredEnergySnapshot.MeasuredEnergyDeltaData;
 
-import org.junit.Before;
 import org.junit.Test;
 
 /**
@@ -38,134 +43,198 @@
  */
 @SmallTest
 public final class MeasuredEnergySnapshotTest {
-    private static final int NUMBER_SUBSYSTEMS = 3;
-    private static final int SUBSYSTEM_DISPLAY = 0;
-    private static final int SUBSYSTEM_NEVER_USED = 1;
-    private static final int SUBSYSTEM_CATAPULT = 2;
+    private static final EnergyConsumer CONSUMER_DISPLAY = createEnergyConsumer(
+            0, 0, EnergyConsumerType.DISPLAY, "Display");
+    private static final  EnergyConsumer CONSUMER_OTHER_0 = createEnergyConsumer(
+            47, 0, EnergyConsumerType.OTHER, "GPU");
+    private static final  EnergyConsumer CONSUMER_OTHER_1 = createEnergyConsumer(
+            1, 1, EnergyConsumerType.OTHER, "HPU");
+    private static final  EnergyConsumer CONSUMER_OTHER_2 = createEnergyConsumer(
+            436, 2, EnergyConsumerType.OTHER, "IPU");
 
-    private MeasuredEnergySnapshot mSnapshot;
+    private static final SparseArray<EnergyConsumer> ALL_ID_CONSUMER_MAP = createIdToConsumerMap(
+            CONSUMER_DISPLAY, CONSUMER_OTHER_0, CONSUMER_OTHER_1, CONSUMER_OTHER_2);
+    private static final SparseArray<EnergyConsumer> SOME_ID_CONSUMER_MAP = createIdToConsumerMap(
+            CONSUMER_DISPLAY);
 
-    // Basic MeasuredEnergyArray that supports all the subsystems. Out of order on purpose.
-    private final int[] mAllSubsystems =
-            {SUBSYSTEM_DISPLAY, SUBSYSTEM_CATAPULT, SUBSYSTEM_NEVER_USED};
-    // E.g. mAllSubsystems[mSubsystemIndices[SUBSYSTEM_CATAPULT]]=SUBSYSTEM_CATAPULT
-    private final int[] mSubsystemIndices = {0, 2, 1};
-    private final long[] mCurrentSubsystemEnergyUJ = {111, 0, 0};
-    private final MeasuredEnergyArray mOmniEnergyArray = new MeasuredEnergyArray() {
-        @Override
-        public int getSubsystem(int index) {
-            return mAllSubsystems[index];
-        }
-
-        @Override
-        public long getEnergy(int index) {
-            return mCurrentSubsystemEnergyUJ[index];
-        }
-
-        @Override
-        public int size() {
-            return mAllSubsystems.length;
-        }
+    // Elements in each results are purposefully out of order.
+    private static final  EnergyConsumerResult[] RESULTS_0 = new EnergyConsumerResult[] {
+        createEnergyConsumerResult(CONSUMER_OTHER_0.id, 90, new int[] {47, 3}, new long[] {14, 13}),
+        createEnergyConsumerResult(CONSUMER_DISPLAY.id, 14, null, null),
+        createEnergyConsumerResult(CONSUMER_OTHER_1.id, 0, null, null),
+        // No CONSUMER_OTHER_2
     };
-    private final MeasuredEnergyArray mJustDisplayEnergyArray = new MeasuredEnergyArray() {
-        @Override
-        public int getSubsystem(int index) {
-            return mAllSubsystems[0];
-        }
-
-        @Override
-        public long getEnergy(int index) {
-            return mCurrentSubsystemEnergyUJ[0];
-        }
-
-        @Override
-        public int size() {
-            return 1;
-        }
+    private static final  EnergyConsumerResult[] RESULTS_1 = new EnergyConsumerResult[] {
+        createEnergyConsumerResult(CONSUMER_DISPLAY.id, 24, null, null),
+        createEnergyConsumerResult(CONSUMER_OTHER_0.id, 90, new int[] {47, 3}, new long[] {14, 13}),
+        createEnergyConsumerResult(CONSUMER_OTHER_2.id, 12, new int[] {6}, new long[] {10}),
+        createEnergyConsumerResult(CONSUMER_OTHER_1.id, 12_000, null, null),
+    };
+    private static final  EnergyConsumerResult[] RESULTS_2 = new EnergyConsumerResult[] {
+        createEnergyConsumerResult(CONSUMER_DISPLAY.id, 36, null, null),
+        // No CONSUMER_OTHER_0
+        // No CONSUMER_OTHER_1
+        // No CONSUMER_OTHER_2
+    };
+    private static final  EnergyConsumerResult[] RESULTS_3 = new EnergyConsumerResult[] {
+        // No CONSUMER_DISPLAY
+        createEnergyConsumerResult(CONSUMER_OTHER_2.id, 13, new int[] {6}, new long[] {10}),
+        createEnergyConsumerResult(
+                CONSUMER_OTHER_0.id, 190, new int[] {2, 3, 47, 7}, new long[] {9, 18, 14, 6}),
+        createEnergyConsumerResult(CONSUMER_OTHER_1.id, 12_000, null, null),
+    };
+    private static final  EnergyConsumerResult[] RESULTS_4 = new EnergyConsumerResult[] {
+        createEnergyConsumerResult(CONSUMER_DISPLAY.id, 43, null, null),
+        createEnergyConsumerResult(
+                CONSUMER_OTHER_0.id, 290, new int[] {7, 47, 3, 2}, new long[] {6, 14, 18, 11}),
+        // No CONSUMER_OTHER_1
+        createEnergyConsumerResult(CONSUMER_OTHER_2.id, 165, new int[] {6, 47}, new long[] {10, 8}),
     };
 
-    @Before
-    public void setUp() {
-        mSnapshot = new MeasuredEnergySnapshot(NUMBER_SUBSYSTEMS, mOmniEnergyArray);
+    @Test
+    public void testUpdateAndGetDelta_empty() {
+        final MeasuredEnergySnapshot snapshot = new MeasuredEnergySnapshot(ALL_ID_CONSUMER_MAP);
+        assertNull(snapshot.updateAndGetDelta(null));
+        assertNull(snapshot.updateAndGetDelta(new EnergyConsumerResult[0]));
     }
 
     @Test
     public void testUpdateAndGetDelta() {
-        SparseLongArray result;
+        final MeasuredEnergySnapshot snapshot = new MeasuredEnergySnapshot(ALL_ID_CONSUMER_MAP);
 
-        // Increment DISPLAY by 15
-        incrementEnergyOfSubsystem(SUBSYSTEM_DISPLAY, 15);
-        result = mSnapshot.updateAndGetDelta(mOmniEnergyArray);
-        assertEquals(1, result.size());
-        assertEquals(15, result.get(SUBSYSTEM_DISPLAY));
+        // results0
+        MeasuredEnergyDeltaData delta = snapshot.updateAndGetDelta(RESULTS_0);
+        if (delta != null) { // null is fine here. If non-null, it better be uninteresting though.
+            assertEquals(UNAVAILABLE, delta.displayEnergyUJ);
+            assertNull(delta.otherTotalEnergyUJ);
+            assertNull(delta.otherUidEnergiesUJ);
+        }
 
-        // Increment DISPLAY by 7
-        // Increment CATAPULT by 5. But do NOT include (pull) it in the passed in energy array.
-        incrementEnergyOfSubsystem(SUBSYSTEM_DISPLAY, 7);
-        incrementEnergyOfSubsystem(SUBSYSTEM_CATAPULT, 5);
-        result = mSnapshot.updateAndGetDelta(mJustDisplayEnergyArray); // Just pull display.
-        assertEquals(1, result.size());
-        assertEquals(7, result.get(SUBSYSTEM_DISPLAY));
+        // results1
+        delta = snapshot.updateAndGetDelta(RESULTS_1);
+        assertNotNull(delta);
+        assertEquals(24 - 14, delta.displayEnergyUJ);
 
-        // Increment CATAPULT by 64 (in addition to the previous increase of 5)
-        incrementEnergyOfSubsystem(SUBSYSTEM_CATAPULT, 64);
-        result = mSnapshot.updateAndGetDelta(mOmniEnergyArray);
-        assertEquals(1, result.size());
-        assertEquals(5 + 64, result.get(SUBSYSTEM_CATAPULT));
+        assertNotNull(delta.otherTotalEnergyUJ);
+        assertEquals(90 - 90, delta.otherTotalEnergyUJ[0]);
+        assertEquals(12_000 - 0, delta.otherTotalEnergyUJ[1]);
+        assertEquals(0, delta.otherTotalEnergyUJ[2]); // First good pull. Treat delta as 0.
 
-        // Do nothing
-        result = mSnapshot.updateAndGetDelta(mOmniEnergyArray);
-        assertEquals("0 results should not appear at all", 0, result.size());
+        assertNotNull(delta.otherUidEnergiesUJ);
+        assertNullOrEmpty(delta.otherUidEnergiesUJ[0]); // No change in uid energies
+        assertNullOrEmpty(delta.otherUidEnergiesUJ[1]);
+        assertNullOrEmpty(delta.otherUidEnergiesUJ[2]);
 
-        // Increment DISPLAY by 42
-        incrementEnergyOfSubsystem(SUBSYSTEM_DISPLAY, 42);
-        result = mSnapshot.updateAndGetDelta(mOmniEnergyArray);
-        assertEquals(1, result.size());
-        assertEquals(42, result.get(SUBSYSTEM_DISPLAY));
+        // results2
+        delta = snapshot.updateAndGetDelta(RESULTS_2);
+        assertNotNull(delta);
+        assertEquals(36 - 24, delta.displayEnergyUJ);
+        assertNull(delta.otherUidEnergiesUJ);
+        assertNull(delta.otherTotalEnergyUJ);
 
-        // Increment DISPLAY by 106 and CATAPULT by 13
-        incrementEnergyOfSubsystem(SUBSYSTEM_DISPLAY, 106);
-        incrementEnergyOfSubsystem(SUBSYSTEM_CATAPULT, 13);
-        result = mSnapshot.updateAndGetDelta(mOmniEnergyArray);
-        assertEquals(2, result.size());
-        assertEquals(106, result.get(SUBSYSTEM_DISPLAY));
-        assertEquals(13, result.get(SUBSYSTEM_CATAPULT));
+        // results3
+        delta = snapshot.updateAndGetDelta(RESULTS_3);
+        assertNotNull(delta);
+        assertEquals(UNAVAILABLE, delta.displayEnergyUJ);
+
+        assertNotNull(delta.otherTotalEnergyUJ);
+        assertEquals(190 - 90, delta.otherTotalEnergyUJ[0]);
+        assertEquals(12_000 - 12_000, delta.otherTotalEnergyUJ[1]);
+        assertEquals(13 - 12, delta.otherTotalEnergyUJ[2]);
+
+        assertNotNull(delta.otherUidEnergiesUJ);
+        assertEquals(3, delta.otherUidEnergiesUJ[0].size());
+        assertEquals(9 - 0, delta.otherUidEnergiesUJ[0].get(2));
+        assertEquals(18 - 13, delta.otherUidEnergiesUJ[0].get(3));
+        assertEquals(6 - 0, delta.otherUidEnergiesUJ[0].get(7));
+        assertNullOrEmpty(delta.otherUidEnergiesUJ[1]);
+        assertNullOrEmpty(delta.otherUidEnergiesUJ[2]);
+
+        // results4
+        delta = snapshot.updateAndGetDelta(RESULTS_4);
+        assertNotNull(delta);
+        assertEquals(43 - 36, delta.displayEnergyUJ);
+
+        assertNotNull(delta.otherTotalEnergyUJ);
+        assertEquals(290 - 190, delta.otherTotalEnergyUJ[0]);
+        assertEquals(0, delta.otherTotalEnergyUJ[1]); // Not present (e.g. missing data)
+        assertEquals(165 - 13, delta.otherTotalEnergyUJ[2]);
+
+        assertNotNull(delta.otherUidEnergiesUJ);
+        assertEquals(1, delta.otherUidEnergiesUJ[0].size());
+        assertEquals(11 - 9, delta.otherUidEnergiesUJ[0].get(2));
+        assertNullOrEmpty(delta.otherUidEnergiesUJ[1]); // Not present
+        assertEquals(1, delta.otherUidEnergiesUJ[2].size());
+        assertEquals(8, delta.otherUidEnergiesUJ[2].get(47));
     }
 
-    private void incrementEnergyOfSubsystem(int subsystem, long energy) {
-        mCurrentSubsystemEnergyUJ[mSubsystemIndices[subsystem]] += energy;
+    /** Test updateAndGetDelta() when the results have consumers absent from idToConsumerMap. */
+    @Test
+    public void testUpdateAndGetDelta_some() {
+        final MeasuredEnergySnapshot snapshot = new MeasuredEnergySnapshot(SOME_ID_CONSUMER_MAP);
+
+        // results0
+        MeasuredEnergyDeltaData delta = snapshot.updateAndGetDelta(RESULTS_0);
+        if (delta != null) { // null is fine here. If non-null, it better be uninteresting though.
+            assertEquals(UNAVAILABLE, delta.displayEnergyUJ);
+            assertNull(delta.otherTotalEnergyUJ);
+            assertNull(delta.otherUidEnergiesUJ);
+        }
+
+        // results1
+        delta = snapshot.updateAndGetDelta(RESULTS_1);
+        assertNotNull(delta);
+        assertEquals(24 - 14, delta.displayEnergyUJ);
+        assertNull(delta.otherTotalEnergyUJ); // Although in the results, they're not in the idMap
+        assertNull(delta.otherUidEnergiesUJ);
     }
 
     @Test
-    public void testUpdateAndGetDelta_null() {
-        assertNull(mSnapshot.updateAndGetDelta(null));
+    public void testGetNumOtherOrdinals() {
+        final MeasuredEnergySnapshot snapshot = new MeasuredEnergySnapshot(ALL_ID_CONSUMER_MAP);
+        assertEquals(3, snapshot.getNumOtherOrdinals());
     }
 
     @Test
-    public void testHasSubsystem() {
-        // Setup MeasuredEnergySnapshot which reported some of the subsystems.
-        final int[] subsystems = {SUBSYSTEM_DISPLAY, SUBSYSTEM_CATAPULT};
-        MeasuredEnergyArray measuredEnergyArray = new MeasuredEnergyArray() {
-            @Override
-            public int getSubsystem(int index) {
-                return subsystems[index];
-            }
+    public void testGetNumOtherOrdinals_none() {
+        final MeasuredEnergySnapshot snapshot = new MeasuredEnergySnapshot(SOME_ID_CONSUMER_MAP);
+        assertEquals(0, snapshot.getNumOtherOrdinals());
+    }
 
-            @Override
-            public long getEnergy(int index) {
-                return 0; // Irrelevant for this test.
-            }
+    private static EnergyConsumer createEnergyConsumer(int id, int ord, byte type, String name) {
+        final EnergyConsumer ec = new EnergyConsumer();
+        ec.id = id;
+        ec.ordinal = ord;
+        ec.type = type;
+        ec.name = name;
+        return ec;
+    }
 
-            @Override
-            public int size() {
-                return subsystems.length;
-            }
-        };
-        final MeasuredEnergySnapshot snapshot =
-                new MeasuredEnergySnapshot(NUMBER_SUBSYSTEMS, measuredEnergyArray);
+    private static SparseArray<EnergyConsumer> createIdToConsumerMap(EnergyConsumer ... ecs) {
+        final SparseArray<EnergyConsumer> map = new SparseArray<>();
+        for (EnergyConsumer ec : ecs) {
+            map.put(ec.id, ec);
+        }
+        return map;
+    }
 
-        assertTrue(snapshot.hasSubsystem(SUBSYSTEM_DISPLAY));
-        assertTrue(snapshot.hasSubsystem(SUBSYSTEM_CATAPULT));
-        assertFalse(snapshot.hasSubsystem(SUBSYSTEM_NEVER_USED));
+    private static EnergyConsumerResult createEnergyConsumerResult(
+            int id, long energyUWs, int[] uids, long[] uidEnergies) {
+        final EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.id = id;
+        ecr.energyUWs = energyUWs;
+        if (uids != null) {
+            ecr.attribution = new EnergyConsumerAttribution[uids.length];
+            for (int i = 0; i < uids.length; i++) {
+                ecr.attribution[i] = new EnergyConsumerAttribution();
+                ecr.attribution[i].uid = uids[i];
+                ecr.attribution[i].energyUWs = uidEnergies[i];
+            }
+        }
+        return ecr;
+    }
+
+    private void assertNullOrEmpty(SparseLongArray a) {
+        if (a != null) assertEquals("Array should be null or empty", 0, a.size());
     }
 }