Add Power State and Screen State to BatteryUsageStats

Bug: 352835319
Test: atest PowerStatsTestsRavenwood; atest PowerStatsTests
Flag: com.android.server.power.optimization.battery_usage_stats_by_power_and_screen_state
Change-Id: I60f4b688880304da0879e682b153a278f8a19682

17 files changed, 1065 insertions, 415 deletions

diff --git a/core/java/android/os/AggregateBatteryConsumer.java b/core/java/android/os/AggregateBatteryConsumer.java
index 67e21957839a..c7f8878f104e 100644
--- a/core/java/android/os/AggregateBatteryConsumer.java
+++ b/core/java/android/os/AggregateBatteryConsumer.java
@@ -55,7 +55,7 @@ public final class AggregateBatteryConsumer extends BatteryConsumer {
 
     @Override
     public void dump(PrintWriter pw, boolean skipEmptyComponents) {
-        mPowerComponents.dump(pw, skipEmptyComponents);
+        mPowerComponents.dump(pw, SCREEN_STATE_ANY, POWER_STATE_ANY, skipEmptyComponents);
     }
 
     @Override
diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index 744f6a8d61ac..2447ff93fdbc 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -19,14 +19,19 @@ package android.os;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.database.Cursor;
 import android.database.CursorWindow;
+import android.util.IntArray;
 import android.util.Slog;
+import android.util.SparseArray;
 import android.util.proto.ProtoOutputStream;
 
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 
 /**
  * Interface for objects containing battery attribution data.
@@ -192,31 +197,106 @@ public abstract class BatteryConsumer {
         sProcessStateNames[PROCESS_STATE_CACHED] = "cached";
     }
 
-    private static final int[] SUPPORTED_POWER_COMPONENTS_PER_PROCESS_STATE = {
-            POWER_COMPONENT_CPU,
-            POWER_COMPONENT_MOBILE_RADIO,
-            POWER_COMPONENT_WIFI,
-            POWER_COMPONENT_BLUETOOTH,
-            POWER_COMPONENT_AUDIO,
-            POWER_COMPONENT_VIDEO,
-            POWER_COMPONENT_FLASHLIGHT,
-            POWER_COMPONENT_CAMERA,
-            POWER_COMPONENT_GNSS,
+    private static final IntArray SUPPORTED_POWER_COMPONENTS_PER_PROCESS_STATE;
+    static {
+        int[] supportedPowerComponents = {
+                POWER_COMPONENT_CPU,
+                POWER_COMPONENT_MOBILE_RADIO,
+                POWER_COMPONENT_WIFI,
+                POWER_COMPONENT_BLUETOOTH,
+                POWER_COMPONENT_AUDIO,
+                POWER_COMPONENT_VIDEO,
+                POWER_COMPONENT_FLASHLIGHT,
+                POWER_COMPONENT_CAMERA,
+                POWER_COMPONENT_GNSS};
+        Arrays.sort(supportedPowerComponents);
+        SUPPORTED_POWER_COMPONENTS_PER_PROCESS_STATE = IntArray.wrap(supportedPowerComponents);
     };
 
     static final int COLUMN_INDEX_BATTERY_CONSUMER_TYPE = 0;
     static final int COLUMN_COUNT = 1;
 
     /**
+     * Identifiers of consumed power aggregations per SCREEN state.
+     *
+     * @hide
+     */
+    @IntDef(prefix = {"SCREEN_STATE_"}, value = {
+            SCREEN_STATE_UNSPECIFIED,
+            SCREEN_STATE_ANY,
+            SCREEN_STATE_ON,
+            SCREEN_STATE_OTHER,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface ScreenState {
+    }
+
+    public static final int SCREEN_STATE_UNSPECIFIED = 0;
+    public static final int SCREEN_STATE_ANY = SCREEN_STATE_UNSPECIFIED;
+    public static final int SCREEN_STATE_ON = 1;
+    public static final int SCREEN_STATE_OTHER = 2;  // Off, doze etc
+
+    public static final int SCREEN_STATE_COUNT = 3;
+
+    private static final String[] sScreenStateNames = new String[SCREEN_STATE_COUNT];
+
+    static {
+        // Assign individually to avoid future mismatch
+        sScreenStateNames[SCREEN_STATE_UNSPECIFIED] = "unspecified";
+        sScreenStateNames[SCREEN_STATE_ON] = "on";
+        sScreenStateNames[SCREEN_STATE_OTHER] = "off/doze";
+    }
+
+    /**
+     * Identifiers of consumed power aggregations per POWER state.
+     *
+     * @hide
+     */
+    @IntDef(prefix = {"POWER_STATE_"}, value = {
+            POWER_STATE_UNSPECIFIED,
+            POWER_STATE_ANY,
+            POWER_STATE_BATTERY,
+            POWER_STATE_OTHER,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface PowerState {
+    }
+
+    public static final int POWER_STATE_UNSPECIFIED = 0;
+    public static final int POWER_STATE_ANY = POWER_STATE_UNSPECIFIED;
+    public static final int POWER_STATE_BATTERY = 1;
+    public static final int POWER_STATE_OTHER = 2;   // Plugged in, or on wireless charger, etc.
+
+    public static final int POWER_STATE_COUNT = 3;
+
+    private static final String[] sPowerStateNames = new String[POWER_STATE_COUNT];
+
+    static {
+        // Assign individually to avoid future mismatch
+        sPowerStateNames[POWER_STATE_UNSPECIFIED] = "unspecified";
+        sPowerStateNames[POWER_STATE_BATTERY] = "on battery";
+        sPowerStateNames[POWER_STATE_OTHER] = "not on battery";
+    }
+
+    /**
      * Identifies power attribution dimensions that a caller is interested in.
      */
     public static final class Dimensions {
         public final @PowerComponent int powerComponent;
         public final @ProcessState int processState;
+        public final @ScreenState int screenState;
+        public final @PowerState int powerState;
 
-        public Dimensions(int powerComponent, int processState) {
+        public Dimensions(@PowerComponent int powerComponent, @ProcessState int processState) {
+            this(powerComponent, processState, SCREEN_STATE_UNSPECIFIED, POWER_STATE_UNSPECIFIED);
+        }
+
+        public Dimensions(@PowerComponent int powerComponent, int processState,
+                @ScreenState int screenState, @PowerState int powerState) {
             this.powerComponent = powerComponent;
             this.processState = processState;
+            this.screenState = screenState;
+            this.powerState = powerState;
         }
 
         @Override
@@ -234,6 +314,20 @@ public abstract class BatteryConsumer {
                 sb.append("processState=").append(sProcessStateNames[processState]);
                 dimensionSpecified = true;
             }
+            if (screenState != SCREEN_STATE_ANY) {
+                if (dimensionSpecified) {
+                    sb.append(", ");
+                }
+                sb.append("screenState=").append(screenStateToString(screenState));
+                dimensionSpecified = true;
+            }
+            if (powerState != POWER_STATE_ANY) {
+                if (dimensionSpecified) {
+                    sb.append(", ");
+                }
+                sb.append("powerState=").append(powerStateToString(powerState));
+                dimensionSpecified = true;
+            }
             if (!dimensionSpecified) {
                 sb.append("any components and process states");
             }
@@ -242,7 +336,8 @@ public abstract class BatteryConsumer {
     }
 
     public static final Dimensions UNSPECIFIED_DIMENSIONS =
-            new Dimensions(POWER_COMPONENT_ANY, PROCESS_STATE_ANY);
+            new Dimensions(POWER_COMPONENT_ANY, PROCESS_STATE_ANY, SCREEN_STATE_ANY,
+                    POWER_STATE_ANY);
 
     /**
      * Identifies power attribution dimensions that are captured by a data element of
@@ -258,52 +353,93 @@ public abstract class BatteryConsumer {
     public static final class Key {
         public final @PowerComponent int powerComponent;
         public final @ProcessState int processState;
+        public final @ScreenState int screenState;
+        public final @PowerState int powerState;
 
         final int mPowerModelColumnIndex;
         final int mPowerColumnIndex;
         final int mDurationColumnIndex;
-        private String mShortString;
 
-        private Key(int powerComponent, int processState, int powerModelColumnIndex,
+        private Key(@PowerComponent int powerComponent, @ProcessState int processState,
+                @ScreenState int screenState, @PowerState int powerState, int powerModelColumnIndex,
                 int powerColumnIndex, int durationColumnIndex) {
             this.powerComponent = powerComponent;
             this.processState = processState;
+            this.screenState = screenState;
+            this.powerState = powerState;
 
             mPowerModelColumnIndex = powerModelColumnIndex;
             mPowerColumnIndex = powerColumnIndex;
             mDurationColumnIndex = durationColumnIndex;
         }
 
+        /**
+         * Returns true if this key should be included in an enumeration parameterized with
+         * the supplied dimensions.
+         */
+        boolean matches(@PowerComponent int powerComponent, @ProcessState int processState,
+                @ScreenState int screenState, @PowerState int powerState) {
+            if (powerComponent != POWER_COMPONENT_ANY && this.powerComponent != powerComponent) {
+                return false;
+            }
+            if (processState != PROCESS_STATE_ANY && this.processState != processState) {
+                return false;
+            }
+            if (screenState != SCREEN_STATE_ANY && this.screenState != screenState) {
+                return false;
+            }
+            if (powerState != POWER_STATE_ANY && this.powerState != powerState) {
+                return false;
+            }
+            return true;
+        }
+
         @SuppressWarnings("EqualsUnsafeCast")
         @Override
         public boolean equals(Object o) {
             // Skipping null and class check for performance
             final Key key = (Key) o;
             return powerComponent == key.powerComponent
-                && processState == key.processState;
+                    && processState == key.processState
+                    && screenState == key.screenState
+                    && powerState == key.powerState;
         }
 
         @Override
         public int hashCode() {
             int result = powerComponent;
             result = 31 * result + processState;
+            result = 31 * result + screenState;
+            result = 31 * result + powerState;
             return result;
         }
 
         /**
          * Returns a string suitable for use in dumpsys.
          */
-        public String toShortString() {
-            if (mShortString == null) {
-                StringBuilder sb = new StringBuilder();
-                sb.append(powerComponentIdToString(powerComponent));
-                if (processState != PROCESS_STATE_UNSPECIFIED) {
-                    sb.append(':');
-                    sb.append(processStateToString(processState));
-                }
-                mShortString = sb.toString();
+        public static String toString(@PowerComponent int powerComponent,
+                @ProcessState int processState, @ScreenState int screenState,
+                @PowerState int powerState) {
+            StringBuilder sb = new StringBuilder();
+            sb.append(powerComponentIdToString(powerComponent));
+            if (processState != PROCESS_STATE_UNSPECIFIED) {
+                sb.append(':');
+                sb.append(processStateToString(processState));
+            }
+            if (screenState != SCREEN_STATE_UNSPECIFIED) {
+                sb.append(":scr-");
+                sb.append(sScreenStateNames[screenState]);
+            }
+            if (powerState != POWER_STATE_UNSPECIFIED) {
+                sb.append(":pwr-");
+                sb.append(sPowerStateNames[powerState]);
             }
-            return mShortString;
+            return sb.toString();
+        }
+
+        @Override
+        public String toString() {
+            return toString(powerComponent, processState, screenState, powerState);
         }
     }
 
@@ -335,11 +471,18 @@ public abstract class BatteryConsumer {
     }
 
     /**
+     * Returns the amount of usage time  aggregated over the specified dimensions, in millis.
+     */
+    public long getUsageDurationMillis(@NonNull Dimensions dimensions) {
+        return mPowerComponents.getUsageDurationMillis(dimensions);
+    }
+
+    /**
      * Returns keys for various power values attributed to the specified component
      * held by this BatteryUsageStats object.
      */
     public Key[] getKeys(@PowerComponent int componentId) {
-        return mData.getKeys(componentId);
+        return mData.layout.getKeys(componentId);
     }
 
     /**
@@ -347,14 +490,16 @@ public abstract class BatteryConsumer {
      * for all values of other dimensions such as process state.
      */
     public Key getKey(@PowerComponent int componentId) {
-        return mData.getKey(componentId, PROCESS_STATE_UNSPECIFIED);
+        return mData.layout.getKey(componentId, PROCESS_STATE_UNSPECIFIED, SCREEN_STATE_UNSPECIFIED,
+                POWER_STATE_UNSPECIFIED);
     }
 
     /**
      * Returns the key for the power attributed to the specified component and process state.
      */
     public Key getKey(@PowerComponent int componentId, @ProcessState int processState) {
-        return mData.getKey(componentId, processState);
+        return mData.layout.getKey(componentId, processState, SCREEN_STATE_UNSPECIFIED,
+                POWER_STATE_UNSPECIFIED);
     }
 
     /**
@@ -365,8 +510,8 @@ public abstract class BatteryConsumer {
      * @return Amount of consumed power in mAh.
      */
     public double getConsumedPower(@PowerComponent int componentId) {
-        return mPowerComponents.getConsumedPower(
-                mData.getKeyOrThrow(componentId, PROCESS_STATE_UNSPECIFIED));
+        return mPowerComponents.getConsumedPower(componentId, PROCESS_STATE_UNSPECIFIED,
+                        SCREEN_STATE_UNSPECIFIED, POWER_STATE_UNSPECIFIED);
     }
 
     /**
@@ -388,7 +533,8 @@ public abstract class BatteryConsumer {
      */
     public @PowerModel int getPowerModel(@BatteryConsumer.PowerComponent int componentId) {
         return mPowerComponents.getPowerModel(
-                mData.getKeyOrThrow(componentId, PROCESS_STATE_UNSPECIFIED));
+                mData.layout.getKeyOrThrow(componentId, PROCESS_STATE_UNSPECIFIED,
+                        SCREEN_STATE_UNSPECIFIED, POWER_STATE_UNSPECIFIED));
     }
 
     /**
@@ -507,6 +653,20 @@ public abstract class BatteryConsumer {
     }
 
     /**
+     * Returns the human-readable name of the specified power state (on battery or not)
+     */
+    public static String powerStateToString(@PowerState int powerState) {
+        return sPowerStateNames[powerState];
+    }
+
+    /**
+     * Returns the human-readable name of the specified screen state (on or off/doze)
+     */
+    public static String screenStateToString(@ScreenState int screenState) {
+        return sScreenStateNames[screenState];
+    }
+
+    /**
      * Prints the stats in a human-readable format.
      */
     public void dump(PrintWriter pw) {
@@ -591,42 +751,11 @@ public abstract class BatteryConsumer {
             return new BatteryConsumerData(cursorWindow, cursorRow, layout);
         }
 
-        public Key[] getKeys(int componentId) {
-            return layout.keys[componentId];
-        }
-
-        Key getKeyOrThrow(int componentId, int processState) {
-            Key key = getKey(componentId, processState);
-            if (key == null) {
-                if (processState == PROCESS_STATE_ANY) {
-                    throw new IllegalArgumentException(
-                            "Unsupported power component ID: " + componentId);
-                } else {
-                    throw new IllegalArgumentException(
-                            "Unsupported power component ID: " + componentId
-                                    + " process state: " + processState);
-                }
-            }
-            return key;
-        }
-
-        Key getKey(int componentId, int processState) {
-            if (componentId >= POWER_COMPONENT_COUNT) {
-                return null;
-            }
-
-            if (processState == PROCESS_STATE_ANY) {
-                // The 0-th key for each component corresponds to the roll-up,
-                // across all dimensions. We might as well skip the iteration over the array.
-                return layout.keys[componentId][0];
-            } else {
-                for (Key key : layout.keys[componentId]) {
-                    if (key.processState == processState) {
-                        return key;
-                    }
-                }
+        boolean hasValue(int columnIndex) {
+            if (mCursorRow == -1) {
+                return false;
             }
-            return null;
+            return mCursorWindow.getType(mCursorRow, columnIndex) != Cursor.FIELD_TYPE_NULL;
         }
 
         void putInt(int columnIndex, int value) {
@@ -693,113 +822,158 @@ public abstract class BatteryConsumer {
         public final int customPowerComponentCount;
         public final boolean powerModelsIncluded;
         public final boolean processStateDataIncluded;
-        public final Key[][] keys;
+        public final boolean screenStateDataIncluded;
+        public final boolean powerStateDataIncluded;
+        public final Key[] keys;
+        public final SparseArray<Key> indexedKeys;
         public final int totalConsumedPowerColumnIndex;
         public final int firstCustomConsumedPowerColumn;
         public final int firstCustomUsageDurationColumn;
         public final int columnCount;
-        public final Key[][] processStateKeys;
+        private Key[][] mPerComponentKeys;
 
         private BatteryConsumerDataLayout(int firstColumn, String[] customPowerComponentNames,
-                boolean powerModelsIncluded, boolean includeProcessStateData) {
+                boolean powerModelsIncluded, boolean includeProcessStateData,
+                boolean includeScreenState, boolean includePowerState) {
             this.customPowerComponentNames = customPowerComponentNames;
             this.customPowerComponentCount = customPowerComponentNames.length;
             this.powerModelsIncluded = powerModelsIncluded;
             this.processStateDataIncluded = includeProcessStateData;
+            this.screenStateDataIncluded = includeScreenState;
+            this.powerStateDataIncluded = includePowerState;
 
             int columnIndex = firstColumn;
 
             totalConsumedPowerColumnIndex = columnIndex++;
 
-            keys = new Key[POWER_COMPONENT_COUNT][];
+            ArrayList<Key> keyList = new ArrayList<>();
+            for (int screenState = 0; screenState < SCREEN_STATE_COUNT; screenState++) {
+                if (!includeScreenState && screenState != SCREEN_STATE_UNSPECIFIED) {
+                    continue;
+                }
+                for (int powerState = 0; powerState < POWER_STATE_COUNT; powerState++) {
+                    if (!includePowerState && powerState != POWER_STATE_UNSPECIFIED) {
+                        continue;
+                    }
+                    for (int componentId = 0; componentId < POWER_COMPONENT_COUNT; componentId++) {
+                        columnIndex = addKeys(keyList, powerModelsIncluded, includeProcessStateData,
+                                componentId, screenState, powerState, columnIndex);
+                    }
+                }
+            }
+
+            firstCustomConsumedPowerColumn = columnIndex;
+            columnIndex += customPowerComponentCount;
+
+            firstCustomUsageDurationColumn = columnIndex;
+            columnIndex += customPowerComponentCount;
+
+            columnCount = columnIndex;
 
-            ArrayList<Key> perComponentKeys = new ArrayList<>();
-            for (int componentId = 0; componentId < POWER_COMPONENT_COUNT; componentId++) {
-                perComponentKeys.clear();
+            keys = keyList.toArray(KEY_ARRAY);
+            indexedKeys = new SparseArray<>(keys.length);
+            for (int i = 0; i < keys.length; i++) {
+                Key key = keys[i];
+                int index = keyIndex(key.powerComponent, key.processState, key.screenState,
+                        key.powerState);
+                indexedKeys.put(index, key);
+            }
+        }
 
-                // Declare the Key for the power component, ignoring other dimensions.
-                perComponentKeys.add(
-                        new Key(componentId, PROCESS_STATE_ANY,
+        private int addKeys(List<Key> keys, boolean powerModelsIncluded,
+                boolean includeProcessStateData, int componentId,
+                int screenState, int powerState, int columnIndex) {
+            keys.add(new Key(componentId, PROCESS_STATE_ANY, screenState, powerState,
+                    powerModelsIncluded
+                            ? columnIndex++
+                            : POWER_MODEL_NOT_INCLUDED,  // power model
+                    columnIndex++,      // power
+                    columnIndex++       // usage duration
+            ));
+
+            // Declare Keys for all process states, if needed
+            if (includeProcessStateData) {
+                boolean isSupported = SUPPORTED_POWER_COMPONENTS_PER_PROCESS_STATE
+                        .binarySearch(componentId) >= 0;
+                if (isSupported) {
+                    for (int processState = 0; processState < PROCESS_STATE_COUNT;
+                            processState++) {
+                        if (processState == PROCESS_STATE_UNSPECIFIED) {
+                            continue;
+                        }
+
+                        keys.add(new Key(componentId, processState, screenState, powerState,
                                 powerModelsIncluded
                                         ? columnIndex++
-                                        : POWER_MODEL_NOT_INCLUDED,  // power model
+                                        : POWER_MODEL_NOT_INCLUDED, // power model
                                 columnIndex++,      // power
                                 columnIndex++       // usage duration
                         ));
-
-                // Declare Keys for all process states, if needed
-                if (includeProcessStateData) {
-                    boolean isSupported = false;
-                    for (int id : SUPPORTED_POWER_COMPONENTS_PER_PROCESS_STATE) {
-                        if (id == componentId) {
-                            isSupported = true;
-                            break;
-                        }
-                    }
-                    if (isSupported) {
-                        for (int processState = 0; processState < PROCESS_STATE_COUNT;
-                                processState++) {
-                            if (processState == PROCESS_STATE_UNSPECIFIED) {
-                                continue;
-                            }
-
-                            perComponentKeys.add(
-                                    new Key(componentId, processState,
-                                            powerModelsIncluded
-                                                    ? columnIndex++
-                                                    : POWER_MODEL_NOT_INCLUDED, // power model
-                                            columnIndex++,      // power
-                                            columnIndex++       // usage duration
-                                    ));
-                        }
                     }
                 }
-
-                keys[componentId] = perComponentKeys.toArray(KEY_ARRAY);
             }
+            return columnIndex;
+        }
 
-            if (includeProcessStateData) {
-                processStateKeys = new Key[BatteryConsumer.PROCESS_STATE_COUNT][];
-                ArrayList<Key> perProcStateKeys = new ArrayList<>();
-                for (int processState = 0; processState < PROCESS_STATE_COUNT; processState++) {
-                    if (processState == PROCESS_STATE_UNSPECIFIED) {
-                        continue;
-                    }
+        Key getKey(@PowerComponent int componentId, @ProcessState int processState,
+                @ScreenState int screenState, @PowerState int powerState) {
+            return indexedKeys.get(keyIndex(componentId, processState, screenState, powerState));
+        }
 
-                    perProcStateKeys.clear();
-                    for (int i = 0; i < keys.length; i++) {
-                        for (int j = 0; j < keys[i].length; j++) {
-                            if (keys[i][j].processState == processState) {
-                                perProcStateKeys.add(keys[i][j]);
-                            }
+        Key getKeyOrThrow(@PowerComponent int componentId, @ProcessState int processState,
+                @ScreenState int screenState, @PowerState int powerState) {
+            Key key = getKey(componentId, processState, screenState, powerState);
+            if (key == null) {
+                throw new IllegalArgumentException(
+                        "Unsupported power component ID: " + Key.toString(componentId, processState,
+                                screenState, powerState));
+            }
+            return key;
+        }
+
+        public Key[] getKeys(@PowerComponent int componentId) {
+            synchronized (this) {
+                if (mPerComponentKeys == null) {
+                    mPerComponentKeys = new Key[BatteryConsumer.POWER_COMPONENT_COUNT][];
+                }
+                Key[] componentKeys = mPerComponentKeys[componentId];
+                if (componentKeys == null) {
+                    ArrayList<Key> out = new ArrayList<>();
+                    for (Key key : keys) {
+                        if (key.powerComponent == componentId) {
+                            out.add(key);
                         }
                     }
-                    processStateKeys[processState] = perProcStateKeys.toArray(KEY_ARRAY);
+                    componentKeys = out.toArray(new Key[out.size()]);
+                    mPerComponentKeys[componentId] = componentKeys;
                 }
-            } else {
-                processStateKeys = null;
+                return componentKeys;
             }
+        }
 
-            firstCustomConsumedPowerColumn = columnIndex;
-            columnIndex += customPowerComponentCount;
-
-            firstCustomUsageDurationColumn = columnIndex;
-            columnIndex += customPowerComponentCount;
-
-            columnCount = columnIndex;
+        private int keyIndex(@PowerComponent int componentId, @ProcessState int processState,
+                @ScreenState int screenState, @PowerState int powerState) {
+            // [CCCCCCPPPSSBB]
+            // C - component ID
+            // P - process state
+            // S - screen state
+            // B - power state
+            return componentId << 7 | processState << 4 | screenState << 2 | powerState;
         }
     }
 
     static BatteryConsumerDataLayout createBatteryConsumerDataLayout(
             String[] customPowerComponentNames, boolean includePowerModels,
-            boolean includeProcessStateData) {
+            boolean includeProcessStateData, boolean includeScreenStateData,
+            boolean includePowerStateData) {
         int columnCount = BatteryConsumer.COLUMN_COUNT;
         columnCount = Math.max(columnCount, AggregateBatteryConsumer.COLUMN_COUNT);
         columnCount = Math.max(columnCount, UidBatteryConsumer.COLUMN_COUNT);
         columnCount = Math.max(columnCount, UserBatteryConsumer.COLUMN_COUNT);
 
         return new BatteryConsumerDataLayout(columnCount, customPowerComponentNames,
-                includePowerModels, includeProcessStateData);
+                includePowerModels, includeProcessStateData, includeScreenStateData,
+                includePowerStateData);
     }
 
     protected abstract static class BaseBuilder<T extends BaseBuilder<?>> {
@@ -816,12 +990,19 @@ public abstract class BatteryConsumer {
 
         @Nullable
         public Key[] getKeys(@PowerComponent int componentId) {
-            return mData.getKeys(componentId);
+            return mData.layout.getKeys(componentId);
         }
 
         @Nullable
         public Key getKey(@PowerComponent int componentId, @ProcessState int processState) {
-            return mData.getKey(componentId, processState);
+            return mData.layout.getKey(componentId, processState, SCREEN_STATE_UNSPECIFIED,
+                    POWER_STATE_UNSPECIFIED);
+        }
+
+        @Nullable
+        public Key getKey(@PowerComponent int componentId, @ProcessState int processState,
+                @ScreenState int screenState, @PowerState int powerState) {
+            return mData.layout.getKey(componentId, processState, screenState, powerState);
         }
 
         /**
diff --git a/core/java/android/os/BatteryUsageStats.java b/core/java/android/os/BatteryUsageStats.java
index 61cc23d994f3..dd484f6bace5 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -102,9 +102,13 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
     static final String XML_ATTR_SCOPE = "scope";
     static final String XML_ATTR_PREFIX_CUSTOM_COMPONENT = "custom_component_";
     static final String XML_ATTR_PREFIX_INCLUDES_PROC_STATE_DATA = "includes_proc_state_data";
+    static final String XML_ATTR_PREFIX_INCLUDES_SCREEN_STATE_DATA = "includes_screen_state_data";
+    static final String XML_ATTR_PREFIX_INCLUDES_POWER_STATE_DATA = "includes_power_state_data";
     static final String XML_ATTR_START_TIMESTAMP = "start_timestamp";
     static final String XML_ATTR_END_TIMESTAMP = "end_timestamp";
     static final String XML_ATTR_PROCESS_STATE = "process_state";
+    static final String XML_ATTR_SCREEN_STATE = "screen_state";
+    static final String XML_ATTR_POWER_STATE = "power_state";
     static final String XML_ATTR_POWER = "power";
     static final String XML_ATTR_DURATION = "duration";
     static final String XML_ATTR_MODEL = "model";
@@ -144,10 +148,13 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
     private final String[] mCustomPowerComponentNames;
     private final boolean mIncludesPowerModels;
     private final boolean mIncludesProcessStateData;
+    private final boolean mIncludesScreenStateData;
+    private final boolean mIncludesPowerStateData;
     private final List<UidBatteryConsumer> mUidBatteryConsumers;
     private final List<UserBatteryConsumer> mUserBatteryConsumers;
     private final AggregateBatteryConsumer[] mAggregateBatteryConsumers;
     private final BatteryStatsHistory mBatteryStatsHistory;
+    private BatteryConsumer.BatteryConsumerDataLayout mBatteryConsumerDataLayout;
     private CursorWindow mBatteryConsumersCursorWindow;
 
     private BatteryUsageStats(@NonNull Builder builder) {
@@ -165,6 +172,9 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
         mCustomPowerComponentNames = builder.mCustomPowerComponentNames;
         mIncludesPowerModels = builder.mIncludePowerModels;
         mIncludesProcessStateData = builder.mIncludesProcessStateData;
+        mIncludesScreenStateData = builder.mIncludesScreenStateData;
+        mIncludesPowerStateData = builder.mIncludesPowerStateData;
+        mBatteryConsumerDataLayout = builder.mBatteryConsumerDataLayout;
         mBatteryConsumersCursorWindow = builder.mBatteryConsumersCursorWindow;
 
         double totalPowerMah = 0;
@@ -347,11 +357,13 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
         mCustomPowerComponentNames = source.readStringArray();
         mIncludesPowerModels = source.readBoolean();
         mIncludesProcessStateData = source.readBoolean();
+        mIncludesScreenStateData = source.readBoolean();
+        mIncludesPowerStateData = source.readBoolean();
 
         mBatteryConsumersCursorWindow = CursorWindow.newFromParcel(source);
-        BatteryConsumer.BatteryConsumerDataLayout dataLayout =
-                BatteryConsumer.createBatteryConsumerDataLayout(mCustomPowerComponentNames,
-                        mIncludesPowerModels, mIncludesProcessStateData);
+        mBatteryConsumerDataLayout = BatteryConsumer.createBatteryConsumerDataLayout(
+                mCustomPowerComponentNames, mIncludesPowerModels, mIncludesProcessStateData,
+                mIncludesScreenStateData, mIncludesPowerStateData);
 
         final int numRows = mBatteryConsumersCursorWindow.getNumRows();
 
@@ -363,7 +375,7 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
         for (int i = 0; i < numRows; i++) {
             final BatteryConsumer.BatteryConsumerData data =
                     new BatteryConsumer.BatteryConsumerData(mBatteryConsumersCursorWindow, i,
-                            dataLayout);
+                            mBatteryConsumerDataLayout);
 
             int consumerType = mBatteryConsumersCursorWindow.getInt(i,
                             BatteryConsumer.COLUMN_INDEX_BATTERY_CONSUMER_TYPE);
@@ -405,6 +417,8 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
         dest.writeStringArray(mCustomPowerComponentNames);
         dest.writeBoolean(mIncludesPowerModels);
         dest.writeBoolean(mIncludesProcessStateData);
+        dest.writeBoolean(mIncludesScreenStateData);
+        dest.writeBoolean(mIncludesPowerStateData);
 
         mBatteryConsumersCursorWindow.writeToParcel(dest, flags);
 
@@ -598,23 +612,16 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
 
         for (int componentId = 0; componentId < BatteryConsumer.POWER_COMPONENT_COUNT;
                 componentId++) {
-            for (BatteryConsumer.Key key : deviceConsumer.getKeys(componentId)) {
-                final double devicePowerMah = deviceConsumer.getConsumedPower(key);
-                final double appsPowerMah = appsConsumer.getConsumedPower(key);
-                if (devicePowerMah == 0 && appsPowerMah == 0) {
-                    continue;
-                }
-
-                String label = BatteryConsumer.powerComponentIdToString(componentId);
-                if (key.processState != BatteryConsumer.PROCESS_STATE_UNSPECIFIED) {
-                    label = label
-                            + "(" + BatteryConsumer.processStateToString(key.processState) + ")";
-                }
-                printPowerComponent(pw, prefix, label, devicePowerMah, appsPowerMah,
-                        mIncludesPowerModels ? deviceConsumer.getPowerModel(key)
-                                : BatteryConsumer.POWER_MODEL_UNDEFINED,
-                        deviceConsumer.getUsageDurationMillis(key));
+            final double devicePowerMah = deviceConsumer.getConsumedPower(componentId);
+            final double appsPowerMah = appsConsumer.getConsumedPower(componentId);
+            if (devicePowerMah == 0 && appsPowerMah == 0) {
+                continue;
             }
+
+            printPowerComponent(pw, prefix, BatteryConsumer.powerComponentIdToString(componentId),
+                    devicePowerMah, appsPowerMah,
+                    BatteryConsumer.POWER_MODEL_UNDEFINED,
+                    deviceConsumer.getUsageDurationMillis(componentId));
         }
 
         for (int componentId = BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
@@ -635,6 +642,59 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
                     deviceConsumer.getUsageDurationForCustomComponentMillis(componentId));
         }
 
+        if (mIncludesScreenStateData || mIncludesPowerStateData) {
+            String prefixPlus = prefix + "  ";
+            StringBuilder stateLabel = new StringBuilder();
+            int screenState = BatteryConsumer.SCREEN_STATE_UNSPECIFIED;
+            int powerState = BatteryConsumer.POWER_STATE_UNSPECIFIED;
+            for (BatteryConsumer.Key key : mBatteryConsumerDataLayout.keys) {
+                if (key.processState != BatteryConsumer.PROCESS_STATE_UNSPECIFIED) {
+                    continue;
+                }
+
+                if (key.screenState == BatteryConsumer.SCREEN_STATE_UNSPECIFIED
+                        && key.powerState == BatteryConsumer.POWER_STATE_UNSPECIFIED) {
+                    // Totals already printed earlier in this method
+                    continue;
+                }
+
+                final double devicePowerMah = deviceConsumer.getConsumedPower(key);
+                final double appsPowerMah = appsConsumer.getConsumedPower(key);
+                if (devicePowerMah == 0 && appsPowerMah == 0) {
+                    continue;
+                }
+
+                if (key.screenState != screenState || key.powerState != powerState) {
+                    screenState = key.screenState;
+                    powerState = key.powerState;
+
+                    boolean empty = true;
+                    stateLabel.setLength(0);
+                    stateLabel.append("      (");
+                    if (powerState != BatteryConsumer.POWER_STATE_UNSPECIFIED) {
+                        stateLabel.append(BatteryConsumer.powerStateToString(powerState));
+                        empty = false;
+                    }
+                    if (screenState != BatteryConsumer.SCREEN_STATE_UNSPECIFIED) {
+                        if (!empty) {
+                            stateLabel.append(", ");
+                        }
+                        stateLabel.append("screen ").append(
+                                BatteryConsumer.screenStateToString(screenState));
+                        empty = false;
+                    }
+                    if (!empty) {
+                        stateLabel.append(")");
+                        pw.println(stateLabel);
+                    }
+                }
+                String label = BatteryConsumer.powerComponentIdToString(key.powerComponent);
+                printPowerComponent(pw, prefixPlus, label, devicePowerMah, appsPowerMah,
+                        mIncludesPowerModels ? deviceConsumer.getPowerModel(key)
+                                : BatteryConsumer.POWER_MODEL_UNDEFINED,
+                        deviceConsumer.getUsageDurationMillis(key));
+            }
+        }
         dumpSortedBatteryConsumers(pw, prefix, getUidBatteryConsumers());
         dumpSortedBatteryConsumers(pw, prefix, getUserBatteryConsumers());
         pw.println();
@@ -643,7 +703,7 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
     private void printPowerComponent(PrintWriter pw, String prefix, String label,
             double devicePowerMah, double appsPowerMah, int powerModel, long durationMs) {
         StringBuilder sb = new StringBuilder();
-        sb.append(prefix).append("      ").append(label).append(": ")
+        sb.append(prefix).append("    ").append(label).append(": ")
                 .append(BatteryStats.formatCharge(devicePowerMah));
         if (powerModel != BatteryConsumer.POWER_MODEL_UNDEFINED
                 && powerModel != BatteryConsumer.POWER_MODEL_POWER_PROFILE) {
@@ -657,7 +717,7 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
             BatteryStats.formatTimeMs(sb, durationMs);
         }
 
-        pw.println(sb.toString());
+        pw.println(sb);
     }
 
     private void dumpSortedBatteryConsumers(PrintWriter pw, String prefix,
@@ -670,9 +730,8 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
                 continue;
             }
             pw.print(prefix);
-            pw.print("    ");
+            pw.print("  ");
             consumer.dump(pw);
-            pw.println();
         }
     }
 
@@ -686,6 +745,10 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
         }
         serializer.attributeBoolean(null, XML_ATTR_PREFIX_INCLUDES_PROC_STATE_DATA,
                 mIncludesProcessStateData);
+        serializer.attributeBoolean(null, XML_ATTR_PREFIX_INCLUDES_SCREEN_STATE_DATA,
+                mIncludesScreenStateData);
+        serializer.attributeBoolean(null, XML_ATTR_PREFIX_INCLUDES_POWER_STATE_DATA,
+                mIncludesPowerStateData);
         serializer.attributeLong(null, XML_ATTR_START_TIMESTAMP, mStatsStartTimestampMs);
         serializer.attributeLong(null, XML_ATTR_END_TIMESTAMP, mStatsEndTimestampMs);
         serializer.attributeLong(null, XML_ATTR_DURATION, mStatsDurationMs);
@@ -732,9 +795,13 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
 
                 final boolean includesProcStateData = parser.getAttributeBoolean(null,
                         XML_ATTR_PREFIX_INCLUDES_PROC_STATE_DATA, false);
+                final boolean includesScreenStateData = parser.getAttributeBoolean(null,
+                        XML_ATTR_PREFIX_INCLUDES_SCREEN_STATE_DATA, false);
+                final boolean includesPowerStateData = parser.getAttributeBoolean(null,
+                        XML_ATTR_PREFIX_INCLUDES_POWER_STATE_DATA, false);
 
                 builder = new Builder(customComponentNames.toArray(new String[0]), true,
-                        includesProcStateData, 0);
+                        includesProcStateData, includesScreenStateData, includesPowerStateData, 0);
 
                 builder.setStatsStartTimestamp(
                         parser.getAttributeLong(null, XML_ATTR_START_TIMESTAMP));
@@ -818,6 +885,8 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
         private final String[] mCustomPowerComponentNames;
         private final boolean mIncludePowerModels;
         private final boolean mIncludesProcessStateData;
+        private final boolean mIncludesScreenStateData;
+        private final boolean mIncludesPowerStateData;
         private final double mMinConsumedPowerThreshold;
         private final BatteryConsumer.BatteryConsumerDataLayout mBatteryConsumerDataLayout;
         private long mStatsStartTimestampMs;
@@ -839,21 +908,24 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
         private BatteryStatsHistory mBatteryStatsHistory;
 
         public Builder(@NonNull String[] customPowerComponentNames) {
-            this(customPowerComponentNames, false, false, 0);
+            this(customPowerComponentNames, false, false, false, false, 0);
         }
 
         public Builder(@NonNull String[] customPowerComponentNames, boolean includePowerModels,
-                boolean includeProcessStateData, double minConsumedPowerThreshold) {
+                boolean includeProcessStateData, boolean includeScreenStateData,
+                boolean includesPowerStateData, double minConsumedPowerThreshold) {
             mBatteryConsumersCursorWindow =
                     new CursorWindow(null, BATTERY_CONSUMER_CURSOR_WINDOW_SIZE);
-            mBatteryConsumerDataLayout =
-                    BatteryConsumer.createBatteryConsumerDataLayout(customPowerComponentNames,
-                            includePowerModels, includeProcessStateData);
+            mBatteryConsumerDataLayout = BatteryConsumer.createBatteryConsumerDataLayout(
+                    customPowerComponentNames, includePowerModels, includeProcessStateData,
+                    includeScreenStateData, includesPowerStateData);
             mBatteryConsumersCursorWindow.setNumColumns(mBatteryConsumerDataLayout.columnCount);
 
             mCustomPowerComponentNames = customPowerComponentNames;
             mIncludePowerModels = includePowerModels;
             mIncludesProcessStateData = includeProcessStateData;
+            mIncludesScreenStateData = includeScreenStateData;
+            mIncludesPowerStateData = includesPowerStateData;
             mMinConsumedPowerThreshold = minConsumedPowerThreshold;
             for (int scope = 0; scope < AGGREGATE_BATTERY_CONSUMER_SCOPE_COUNT; scope++) {
                 final BatteryConsumer.BatteryConsumerData data =
@@ -869,6 +941,14 @@ public final class BatteryUsageStats implements Parcelable, Closeable {
             return mIncludesProcessStateData;
         }
 
+        public boolean isScreenStateDataNeeded() {
+            return mIncludesScreenStateData;
+        }
+
+        public boolean isPowerStateDataNeeded() {
+            return mIncludesPowerStateData;
+        }
+
         /**
          * Returns true if this Builder is configured to hold data for the specified
          * custom power component ID.
diff --git a/core/java/android/os/BatteryUsageStatsQuery.java b/core/java/android/os/BatteryUsageStatsQuery.java
index 203ef47d857e..d0ed297d2bda 100644
--- a/core/java/android/os/BatteryUsageStatsQuery.java
+++ b/core/java/android/os/BatteryUsageStatsQuery.java
@@ -73,6 +73,10 @@ public final class BatteryUsageStatsQuery implements Parcelable {
 
     public static final int FLAG_BATTERY_USAGE_STATS_INCLUDE_VIRTUAL_UIDS = 0x0010;
 
+    public static final int FLAG_BATTERY_USAGE_STATS_INCLUDE_SCREEN_STATE = 0x0020;
+
+    public static final int FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_STATE = 0x0040;
+
     private static final long DEFAULT_MAX_STATS_AGE_MS = 5 * 60 * 1000;
 
     private final int mFlags;
@@ -123,6 +127,14 @@ public final class BatteryUsageStatsQuery implements Parcelable {
         return (mFlags & FLAG_BATTERY_USAGE_STATS_INCLUDE_PROCESS_STATE_DATA) != 0;
     }
 
+    public boolean isScreenStateDataNeeded() {
+        return (mFlags & FLAG_BATTERY_USAGE_STATS_INCLUDE_SCREEN_STATE) != 0;
+    }
+
+    public boolean isPowerStateDataNeeded() {
+        return (mFlags & FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_STATE) != 0;
+    }
+
     /**
      * Returns the power components that should be estimated or null if all power components
      * are being requested.
@@ -297,6 +309,24 @@ public final class BatteryUsageStatsQuery implements Parcelable {
         }
 
         /**
+         * Requests that screen state data (screen-on, screen-other) be included in the
+         * BatteryUsageStats, if available.
+         */
+        public Builder includeScreenStateData() {
+            mFlags |= BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_SCREEN_STATE;
+            return this;
+        }
+
+        /**
+         * Requests that power state data (on-battery, power-other) be included in the
+         * BatteryUsageStats, if available.
+         */
+        public Builder includePowerStateData() {
+            mFlags |= BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_STATE;
+            return this;
+        }
+
+        /**
          * Requests to aggregate stored snapshots between the two supplied timestamps
          * @param fromTimestamp Exclusive starting timestamp, as per System.currentTimeMillis()
          * @param toTimestamp Inclusive ending timestamp, as per System.currentTimeMillis()
diff --git a/core/java/android/os/PowerComponents.java b/core/java/android/os/PowerComponents.java
index b035f123d5db..f22e1eac9643 100644
--- a/core/java/android/os/PowerComponents.java
+++ b/core/java/android/os/PowerComponents.java
@@ -17,8 +17,12 @@ package android.os;
 
 import static android.os.BatteryConsumer.BatteryConsumerDataLayout.POWER_MODEL_NOT_INCLUDED;
 import static android.os.BatteryConsumer.POWER_COMPONENT_ANY;
+import static android.os.BatteryConsumer.POWER_STATE_ANY;
+import static android.os.BatteryConsumer.POWER_STATE_UNSPECIFIED;
 import static android.os.BatteryConsumer.PROCESS_STATE_ANY;
 import static android.os.BatteryConsumer.PROCESS_STATE_UNSPECIFIED;
+import static android.os.BatteryConsumer.SCREEN_STATE_ANY;
+import static android.os.BatteryConsumer.SCREEN_STATE_UNSPECIFIED;
 import static android.os.BatteryConsumer.convertMahToDeciCoulombs;
 
 import android.annotation.NonNull;
@@ -56,23 +60,100 @@ class PowerComponents {
      * Total power consumed by this consumer, aggregated over the specified dimensions, in mAh.
      */
     public double getConsumedPower(@NonNull BatteryConsumer.Dimensions dimensions) {
-        if (dimensions.powerComponent != POWER_COMPONENT_ANY) {
-            return mData.getDouble(mData.getKeyOrThrow(dimensions.powerComponent,
-                    dimensions.processState).mPowerColumnIndex);
-        } else if (dimensions.processState != PROCESS_STATE_ANY) {
-            if (!mData.layout.processStateDataIncluded) {
-                throw new IllegalArgumentException(
-                        "No data included in BatteryUsageStats for " + dimensions);
+        return getConsumedPower(dimensions.powerComponent, dimensions.processState,
+                dimensions.screenState, dimensions.powerState);
+    }
+
+    /**
+     * Total power consumed by this consumer, aggregated over the specified dimensions, in mAh.
+     */
+    public double getConsumedPower(@BatteryConsumer.PowerComponent int powerComponent,
+            @BatteryConsumer.ProcessState int processState,
+            @BatteryConsumer.ScreenState int screenState,
+            @BatteryConsumer.PowerState int powerState) {
+        if (powerComponent == POWER_COMPONENT_ANY && processState == PROCESS_STATE_ANY
+                && screenState == SCREEN_STATE_ANY && powerState == POWER_STATE_ANY) {
+            return mData.getDouble(mData.layout.totalConsumedPowerColumnIndex);
+        }
+
+        if (powerComponent != POWER_COMPONENT_ANY
+                && ((mData.layout.screenStateDataIncluded && screenState != SCREEN_STATE_ANY)
+                || (mData.layout.powerStateDataIncluded && powerState != POWER_STATE_ANY))) {
+            BatteryConsumer.Key key = mData.layout.getKey(powerComponent,
+                    processState, screenState, powerState);
+            if (key != null) {
+                return mData.getDouble(key.mPowerColumnIndex);
             }
-            final BatteryConsumer.Key[] keys =
-                    mData.layout.processStateKeys[dimensions.processState];
-            double totalPowerMah = 0;
-            for (int i = keys.length - 1; i >= 0; i--) {
-                totalPowerMah += mData.getDouble(keys[i].mPowerColumnIndex);
+            return 0;
+        }
+
+        if (mData.layout.processStateDataIncluded || mData.layout.screenStateDataIncluded
+                || mData.layout.powerStateDataIncluded) {
+            double total = 0;
+            for (BatteryConsumer.Key key : mData.layout.keys) {
+                if (key.processState != PROCESS_STATE_UNSPECIFIED
+                        && key.matches(powerComponent, processState, screenState, powerState)) {
+                    total += mData.getDouble(key.mPowerColumnIndex);
+                }
             }
-            return totalPowerMah;
+            if (total != 0) {
+                return total;
+            }
+        }
+
+        BatteryConsumer.Key key = mData.layout.getKey(powerComponent, processState,
+                SCREEN_STATE_UNSPECIFIED, POWER_STATE_UNSPECIFIED);
+        if (key != null) {
+            return mData.getDouble(key.mPowerColumnIndex);
         } else {
-            return mData.getDouble(mData.layout.totalConsumedPowerColumnIndex);
+            return 0;
+        }
+    }
+
+    /**
+     * Total usage duration by this consumer, aggregated over the specified dimensions, in ms.
+     */
+    public long getUsageDurationMillis(@NonNull BatteryConsumer.Dimensions dimensions) {
+        return getUsageDurationMillis(dimensions.powerComponent, dimensions.processState,
+                dimensions.screenState, dimensions.powerState);
+    }
+
+    /**
+     * Total usage duration by this consumer, aggregated over the specified dimensions, in ms.
+     */
+    public long getUsageDurationMillis(@BatteryConsumer.PowerComponent int powerComponent,
+            @BatteryConsumer.ProcessState int processState,
+            @BatteryConsumer.ScreenState int screenState,
+            @BatteryConsumer.PowerState int powerState) {
+        if ((mData.layout.screenStateDataIncluded && screenState != SCREEN_STATE_ANY)
+                || (mData.layout.powerStateDataIncluded && powerState != POWER_STATE_ANY)) {
+            BatteryConsumer.Key key = mData.layout.getKey(powerComponent,
+                    processState, screenState, powerState);
+            if (key != null) {
+                return mData.getLong(key.mDurationColumnIndex);
+            }
+            return 0;
+        }
+
+        if (mData.layout.screenStateDataIncluded || mData.layout.powerStateDataIncluded) {
+            long total = 0;
+            for (BatteryConsumer.Key key : mData.layout.keys) {
+                if (key.processState != PROCESS_STATE_UNSPECIFIED
+                        && key.matches(powerComponent, processState, screenState, powerState)) {
+                    total += mData.getLong(key.mDurationColumnIndex);
+                }
+            }
+            if (total != 0) {
+                return total;
+            }
+        }
+
+        BatteryConsumer.Key key = mData.layout.getKey(powerComponent, processState,
+                SCREEN_STATE_UNSPECIFIED, POWER_STATE_UNSPECIFIED);
+        if (key != null) {
+            return mData.getLong(key.mDurationColumnIndex);
+        } else {
+            return 0;
         }
     }
 
@@ -84,7 +165,11 @@ class PowerComponents {
      * @return Amount of consumed power in mAh.
      */
     public double getConsumedPower(@NonNull BatteryConsumer.Key key) {
-        return mData.getDouble(key.mPowerColumnIndex);
+        if (mData.hasValue(key.mPowerColumnIndex)) {
+            return mData.getDouble(key.mPowerColumnIndex);
+        }
+        return getConsumedPower(key.powerComponent, key.processState, key.screenState,
+                key.powerState);
     }
 
     /**
@@ -135,7 +220,12 @@ class PowerComponents {
      * @return Amount of time in milliseconds.
      */
     public long getUsageDurationMillis(BatteryConsumer.Key key) {
-        return mData.getLong(key.mDurationColumnIndex);
+        if (mData.hasValue(key.mDurationColumnIndex)) {
+            return mData.getLong(key.mDurationColumnIndex);
+        }
+
+        return getUsageDurationMillis(key.powerComponent, key.processState, key.screenState,
+                key.powerState);
     }
 
     /**
@@ -154,51 +244,77 @@ class PowerComponents {
         }
     }
 
-    public void dump(PrintWriter pw, boolean skipEmptyComponents) {
-        String separator = "";
+    void dump(PrintWriter pw, @BatteryConsumer.ScreenState int screenState,
+            @BatteryConsumer.PowerState int powerState, boolean skipEmptyComponents) {
         StringBuilder sb = new StringBuilder();
-
         for (int componentId = 0; componentId < BatteryConsumer.POWER_COMPONENT_COUNT;
                 componentId++) {
-            for (BatteryConsumer.Key key: mData.getKeys(componentId)) {
-                final double componentPower = getConsumedPower(key);
-                final long durationMs = getUsageDurationMillis(key);
-                if (skipEmptyComponents && componentPower == 0 && durationMs == 0) {
-                    continue;
+            dump(sb, componentId, PROCESS_STATE_ANY, screenState, powerState, skipEmptyComponents);
+            if (mData.layout.processStateDataIncluded) {
+                for (int processState = 0; processState < BatteryConsumer.PROCESS_STATE_COUNT;
+                        processState++) {
+                    if (processState == PROCESS_STATE_UNSPECIFIED) {
+                        continue;
+                    }
+                    dump(sb, componentId, processState, screenState, powerState,
+                            skipEmptyComponents);
                 }
+            }
+        }
 
-                sb.append(separator);
-                separator = " ";
-                sb.append(key.toShortString());
-                sb.append("=");
-                sb.append(BatteryStats.formatCharge(componentPower));
-
-                if (durationMs != 0) {
-                    sb.append(" (");
-                    BatteryStats.formatTimeMsNoSpace(sb, durationMs);
-                    sb.append(")");
+        // TODO(b/352835319): take into account screen and power states
+        if (screenState == SCREEN_STATE_ANY && powerState == POWER_STATE_ANY) {
+            final int customComponentCount = mData.layout.customPowerComponentCount;
+            for (int customComponentId = BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
+                    customComponentId < BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID
+                            + customComponentCount;
+                    customComponentId++) {
+                final double customComponentPower =
+                        getConsumedPowerForCustomComponent(customComponentId);
+                if (skipEmptyComponents && customComponentPower == 0) {
+                    continue;
                 }
+                sb.append(getCustomPowerComponentName(customComponentId));
+                sb.append("=");
+                sb.append(BatteryStats.formatCharge(customComponentPower));
+                sb.append(" ");
             }
         }
 
-        final int customComponentCount = mData.layout.customPowerComponentCount;
-        for (int customComponentId = BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
-                customComponentId < BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID
-                        + customComponentCount;
-                customComponentId++) {
-            final double customComponentPower =
-                    getConsumedPowerForCustomComponent(customComponentId);
-            if (skipEmptyComponents && customComponentPower == 0) {
-                continue;
-            }
-            sb.append(separator);
-            separator = " ";
-            sb.append(getCustomPowerComponentName(customComponentId));
-            sb.append("=");
-            sb.append(BatteryStats.formatCharge(customComponentPower));
+        // Remove trailing spaces
+        while (!sb.isEmpty() && Character.isWhitespace(sb.charAt(sb.length() - 1))) {
+            sb.setLength(sb.length() - 1);
         }
 
-        pw.print(sb);
+        pw.println(sb);
+    }
+
+    private void dump(StringBuilder sb, @BatteryConsumer.PowerComponent int powerComponent,
+            @BatteryConsumer.ProcessState int processState,
+            @BatteryConsumer.ScreenState int screenState,
+            @BatteryConsumer.PowerState int powerState, boolean skipEmptyComponents) {
+        final double componentPower = getConsumedPower(powerComponent, processState, screenState,
+                powerState);
+        final long durationMs = getUsageDurationMillis(powerComponent, processState, screenState,
+                powerState);
+        if (skipEmptyComponents && componentPower == 0 && durationMs == 0) {
+            return;
+        }
+
+        sb.append(BatteryConsumer.powerComponentIdToString(powerComponent));
+        if (processState != PROCESS_STATE_UNSPECIFIED) {
+            sb.append(':');
+            sb.append(BatteryConsumer.processStateToString(processState));
+        }
+        sb.append("=");
+        sb.append(BatteryStats.formatCharge(componentPower));
+
+        if (durationMs != 0) {
+            sb.append(" (");
+            BatteryStats.formatTimeMsNoSpace(sb, durationMs);
+            sb.append(")");
+        }
+        sb.append(' ');
     }
 
     /** Returns whether there are any atoms.proto POWER_COMPONENTS data to write to a proto. */
@@ -220,11 +336,13 @@ class PowerComponents {
 
         for (int componentId = 0; componentId < BatteryConsumer.POWER_COMPONENT_COUNT;
                 componentId++) {
-
-            final BatteryConsumer.Key[] keys = mData.getKeys(componentId);
+            final BatteryConsumer.Key[] keys = mData.layout.getKeys(componentId);
             for (BatteryConsumer.Key key : keys) {
-                final long powerDeciCoulombs = convertMahToDeciCoulombs(getConsumedPower(key));
-                final long durationMs = getUsageDurationMillis(key);
+                final long powerDeciCoulombs = convertMahToDeciCoulombs(
+                        getConsumedPower(key.powerComponent, key.processState, key.screenState,
+                                key.powerState));
+                final long durationMs = getUsageDurationMillis(key.powerComponent, key.processState,
+                        key.screenState, key.powerState);
 
                 if (powerDeciCoulombs == 0 && durationMs == 0) {
                     // No interesting data. Make sure not to even write the COMPONENT int.
@@ -329,34 +447,43 @@ class PowerComponents {
 
     void writeToXml(TypedXmlSerializer serializer) throws IOException {
         serializer.startTag(null, BatteryUsageStats.XML_TAG_POWER_COMPONENTS);
-        for (int componentId = 0; componentId < BatteryConsumer.POWER_COMPONENT_COUNT;
-                componentId++) {
-            final BatteryConsumer.Key[] keys = mData.getKeys(componentId);
-            for (BatteryConsumer.Key key : keys) {
-                final double powerMah = getConsumedPower(key);
-                final long durationMs = getUsageDurationMillis(key);
-                if (powerMah == 0 && durationMs == 0) {
-                    continue;
-                }
+        for (BatteryConsumer.Key key : mData.layout.keys) {
+            if (!mData.hasValue(key.mPowerColumnIndex)
+                    && !mData.hasValue(key.mDurationColumnIndex)) {
+                continue;
+            }
 
-                serializer.startTag(null, BatteryUsageStats.XML_TAG_COMPONENT);
-                serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_ID, componentId);
-                if (key.processState != PROCESS_STATE_UNSPECIFIED) {
-                    serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_PROCESS_STATE,
-                            key.processState);
-                }
-                if (powerMah != 0) {
-                    serializer.attributeDouble(null, BatteryUsageStats.XML_ATTR_POWER, powerMah);
-                }
-                if (durationMs != 0) {
-                    serializer.attributeLong(null, BatteryUsageStats.XML_ATTR_DURATION, durationMs);
-                }
-                if (mData.layout.powerModelsIncluded) {
-                    serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_MODEL,
-                            getPowerModel(key));
-                }
-                serializer.endTag(null, BatteryUsageStats.XML_TAG_COMPONENT);
+            final double powerMah = getConsumedPower(key);
+            final long durationMs = getUsageDurationMillis(key);
+            if (powerMah == 0 && durationMs == 0) {
+                continue;
+            }
+
+            serializer.startTag(null, BatteryUsageStats.XML_TAG_COMPONENT);
+            serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_ID, key.powerComponent);
+            if (key.processState != PROCESS_STATE_UNSPECIFIED) {
+                serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_PROCESS_STATE,
+                        key.processState);
+            }
+            if (key.screenState != SCREEN_STATE_UNSPECIFIED) {
+                serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_SCREEN_STATE,
+                        key.screenState);
+            }
+            if (key.powerState != POWER_STATE_UNSPECIFIED) {
+                serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_POWER_STATE,
+                        key.powerState);
+            }
+            if (powerMah != 0) {
+                serializer.attributeDouble(null, BatteryUsageStats.XML_ATTR_POWER, powerMah);
+            }
+            if (durationMs != 0) {
+                serializer.attributeLong(null, BatteryUsageStats.XML_ATTR_DURATION, durationMs);
+            }
+            if (mData.layout.powerModelsIncluded) {
+                serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_MODEL,
+                        getPowerModel(key));
             }
+            serializer.endTag(null, BatteryUsageStats.XML_TAG_COMPONENT);
         }
 
         final int customComponentEnd = BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID
@@ -401,6 +528,8 @@ class PowerComponents {
                     case BatteryUsageStats.XML_TAG_COMPONENT: {
                         int componentId = -1;
                         int processState = PROCESS_STATE_UNSPECIFIED;
+                        int screenState = SCREEN_STATE_UNSPECIFIED;
+                        int powerState = POWER_STATE_UNSPECIFIED;
                         double powerMah = 0;
                         long durationMs = 0;
                         int model = BatteryConsumer.POWER_MODEL_UNDEFINED;
@@ -412,6 +541,12 @@ class PowerComponents {
                                 case BatteryUsageStats.XML_ATTR_PROCESS_STATE:
                                     processState = parser.getAttributeInt(i);
                                     break;
+                                case BatteryUsageStats.XML_ATTR_SCREEN_STATE:
+                                    screenState = parser.getAttributeInt(i);
+                                    break;
+                                case BatteryUsageStats.XML_ATTR_POWER_STATE:
+                                    powerState = parser.getAttributeInt(i);
+                                    break;
                                 case BatteryUsageStats.XML_ATTR_POWER:
                                     powerMah = parser.getAttributeDouble(i);
                                     break;
@@ -423,8 +558,8 @@ class PowerComponents {
                                     break;
                             }
                         }
-                        final BatteryConsumer.Key key =
-                                builder.mData.getKey(componentId, processState);
+                        final BatteryConsumer.Key key = builder.mData.layout.getKey(componentId,
+                                processState, screenState, powerState);
                         builder.setConsumedPower(key, powerMah, model);
                         builder.setUsageDurationMillis(key, durationMs);
                         break;
@@ -468,11 +603,9 @@ class PowerComponents {
         Builder(BatteryConsumer.BatteryConsumerData data, double minConsumedPowerThreshold) {
             mData = data;
             mMinConsumedPowerThreshold = minConsumedPowerThreshold;
-            for (BatteryConsumer.Key[] keys : mData.layout.keys) {
-                for (BatteryConsumer.Key key : keys) {
-                    if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
-                        mData.putInt(key.mPowerModelColumnIndex, POWER_MODEL_UNINITIALIZED);
-                    }
+            for (BatteryConsumer.Key key : mData.layout.keys) {
+                if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
+                    mData.putInt(key.mPowerModelColumnIndex, POWER_MODEL_UNINITIALIZED);
                 }
             }
         }
@@ -572,51 +705,41 @@ class PowerComponents {
                                 + ", expected: " + mData.layout.customPowerComponentCount);
             }
 
-            for (int componentId = BatteryConsumer.POWER_COMPONENT_COUNT - 1; componentId >= 0;
-                    componentId--) {
-                final BatteryConsumer.Key[] keys = mData.layout.keys[componentId];
-                for (BatteryConsumer.Key key: keys) {
-                    BatteryConsumer.Key otherKey = null;
-                    for (BatteryConsumer.Key aKey: otherData.layout.keys[componentId]) {
-                        if (aKey.equals(key)) {
-                            otherKey = aKey;
-                            break;
-                        }
-                    }
-
-                    if (otherKey == null) {
-                        continue;
-                    }
+            for (BatteryConsumer.Key key : mData.layout.keys) {
+                BatteryConsumer.Key otherKey = otherData.layout.getKey(key.powerComponent,
+                        key.processState, key.screenState, key.powerState);
+                if (otherKey == null) {
+                    continue;
+                }
 
-                    mData.putDouble(key.mPowerColumnIndex,
-                            mData.getDouble(key.mPowerColumnIndex)
-                                    + otherData.getDouble(otherKey.mPowerColumnIndex));
-                    mData.putLong(key.mDurationColumnIndex,
-                            mData.getLong(key.mDurationColumnIndex)
-                                    + otherData.getLong(otherKey.mDurationColumnIndex));
+                mData.putDouble(key.mPowerColumnIndex,
+                        mData.getDouble(key.mPowerColumnIndex)
+                                + otherData.getDouble(otherKey.mPowerColumnIndex));
+                mData.putLong(key.mDurationColumnIndex,
+                        mData.getLong(key.mDurationColumnIndex)
+                                + otherData.getLong(otherKey.mDurationColumnIndex));
 
-                    if (key.mPowerModelColumnIndex == POWER_MODEL_NOT_INCLUDED) {
-                        continue;
-                    }
+                if (key.mPowerModelColumnIndex == POWER_MODEL_NOT_INCLUDED) {
+                    continue;
+                }
 
-                    boolean undefined = false;
-                    if (otherKey.mPowerModelColumnIndex == POWER_MODEL_NOT_INCLUDED) {
+                boolean undefined = false;
+                if (otherKey.mPowerModelColumnIndex == POWER_MODEL_NOT_INCLUDED) {
+                    undefined = true;
+                } else {
+                    final int powerModel = mData.getInt(key.mPowerModelColumnIndex);
+                    int otherPowerModel = otherData.getInt(otherKey.mPowerModelColumnIndex);
+                    if (powerModel == POWER_MODEL_UNINITIALIZED) {
+                        mData.putInt(key.mPowerModelColumnIndex, otherPowerModel);
+                    } else if (powerModel != otherPowerModel
+                            && otherPowerModel != POWER_MODEL_UNINITIALIZED) {
                         undefined = true;
-                    } else {
-                        final int powerModel = mData.getInt(key.mPowerModelColumnIndex);
-                        int otherPowerModel = otherData.getInt(otherKey.mPowerModelColumnIndex);
-                        if (powerModel == POWER_MODEL_UNINITIALIZED) {
-                            mData.putInt(key.mPowerModelColumnIndex, otherPowerModel);
-                        } else if (powerModel != otherPowerModel
-                                && otherPowerModel != POWER_MODEL_UNINITIALIZED) {
-                            undefined = true;
-                        }
                     }
+                }
 
-                    if (undefined) {
-                        mData.putInt(key.mPowerModelColumnIndex,
-                                BatteryConsumer.POWER_MODEL_UNDEFINED);
-                    }
+                if (undefined) {
+                    mData.putInt(key.mPowerModelColumnIndex,
+                            BatteryConsumer.POWER_MODEL_UNDEFINED);
                 }
             }
 
@@ -631,10 +754,8 @@ class PowerComponents {
                 final int usageColumnIndex = mData.layout.firstCustomUsageDurationColumn + i;
                 final int otherDurationColumnIndex =
                         otherData.layout.firstCustomUsageDurationColumn + i;
-                mData.putLong(usageColumnIndex,
-                        mData.getLong(usageColumnIndex) + otherData.getLong(
-                                otherDurationColumnIndex)
-                );
+                mData.putLong(usageColumnIndex, mData.getLong(usageColumnIndex)
+                        + otherData.getLong(otherDurationColumnIndex));
             }
         }
 
@@ -647,7 +768,8 @@ class PowerComponents {
             for (int componentId = 0; componentId < BatteryConsumer.POWER_COMPONENT_COUNT;
                     componentId++) {
                 totalPowerMah += mData.getDouble(
-                        mData.getKeyOrThrow(componentId, PROCESS_STATE_ANY).mPowerColumnIndex);
+                        mData.layout.getKeyOrThrow(componentId, PROCESS_STATE_ANY, SCREEN_STATE_ANY,
+                                POWER_STATE_ANY).mPowerColumnIndex);
             }
             for (int i = 0; i < mData.layout.customPowerComponentCount; i++) {
                 totalPowerMah += mData.getDouble(
@@ -661,19 +783,17 @@ class PowerComponents {
          */
         @NonNull
         public PowerComponents build() {
-            for (BatteryConsumer.Key[] keys : mData.layout.keys) {
-                for (BatteryConsumer.Key key : keys) {
-                    if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
-                        if (mData.getInt(key.mPowerModelColumnIndex) == POWER_MODEL_UNINITIALIZED) {
-                            mData.putInt(key.mPowerModelColumnIndex,
-                                    BatteryConsumer.POWER_MODEL_UNDEFINED);
-                        }
+            for (BatteryConsumer.Key key: mData.layout.keys) {
+                if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
+                    if (mData.getInt(key.mPowerModelColumnIndex) == POWER_MODEL_UNINITIALIZED) {
+                        mData.putInt(key.mPowerModelColumnIndex,
+                                BatteryConsumer.POWER_MODEL_UNDEFINED);
                     }
+                }
 
-                    if (mMinConsumedPowerThreshold != 0) {
-                        if (mData.getDouble(key.mPowerColumnIndex) < mMinConsumedPowerThreshold) {
-                            mData.putDouble(key.mPowerColumnIndex, 0);
-                        }
+                if (mMinConsumedPowerThreshold != 0) {
+                    if (mData.getDouble(key.mPowerColumnIndex) < mMinConsumedPowerThreshold) {
+                        mData.putDouble(key.mPowerColumnIndex, 0);
                     }
                 }
             }
diff --git a/core/java/android/os/UidBatteryConsumer.java b/core/java/android/os/UidBatteryConsumer.java
index 53af838cb535..9b5a378d6591 100644
--- a/core/java/android/os/UidBatteryConsumer.java
+++ b/core/java/android/os/UidBatteryConsumer.java
@@ -140,12 +140,50 @@ public final class UidBatteryConsumer extends BatteryConsumer {
                     skipEmptyComponents);
             appendProcessStateData(sb, BatteryConsumer.PROCESS_STATE_CACHED,
                     skipEmptyComponents);
-            pw.print(sb);
+            pw.println(sb);
+        } else {
+            pw.println();
         }
 
-        pw.print(" ( ");
-        mPowerComponents.dump(pw, skipEmptyComponents  /* skipTotalPowerComponent */);
-        pw.print(" ) ");
+        pw.print("      ");
+        mPowerComponents.dump(pw, SCREEN_STATE_ANY, POWER_STATE_ANY, skipEmptyComponents);
+
+        if (mData.layout.powerStateDataIncluded || mData.layout.screenStateDataIncluded) {
+            for (int powerState = 0; powerState < POWER_STATE_COUNT; powerState++) {
+                if (mData.layout.powerStateDataIncluded && powerState == POWER_STATE_UNSPECIFIED) {
+                    continue;
+                }
+
+                for (int screenState = 0; screenState < SCREEN_STATE_COUNT; screenState++) {
+                    if (mData.layout.screenStateDataIncluded
+                            && screenState == POWER_STATE_UNSPECIFIED) {
+                        continue;
+                    }
+
+                    final double consumedPower = mPowerComponents.getConsumedPower(
+                            POWER_COMPONENT_ANY,
+                            PROCESS_STATE_ANY, screenState, powerState);
+                    if (consumedPower == 0) {
+                        continue;
+                    }
+
+                    pw.print("      (");
+                    if (powerState != POWER_STATE_UNSPECIFIED) {
+                        pw.print(BatteryConsumer.powerStateToString(powerState));
+                    }
+                    if (screenState != SCREEN_STATE_UNSPECIFIED) {
+                        if (powerState != POWER_STATE_UNSPECIFIED) {
+                            pw.print(", ");
+                        }
+                        pw.print("screen ");
+                        pw.print(BatteryConsumer.screenStateToString(screenState));
+                    }
+                    pw.print(") ");
+                    mPowerComponents.dump(pw, screenState, powerState,
+                            skipEmptyComponents  /* skipTotalPowerComponent */);
+                }
+            }
+        }
     }
 
     private void appendProcessStateData(StringBuilder sb, @ProcessState int processState,
@@ -160,10 +198,6 @@ public final class UidBatteryConsumer extends BatteryConsumer {
                 .append(BatteryStats.formatCharge(power));
     }
 
-    static UidBatteryConsumer create(BatteryConsumerData data) {
-        return new UidBatteryConsumer(data);
-    }
-
     /** Serializes this object to XML */
     void writeToXml(TypedXmlSerializer serializer) throws IOException {
         if (getConsumedPower() == 0) {
diff --git a/core/java/android/os/UserBatteryConsumer.java b/core/java/android/os/UserBatteryConsumer.java
index 23ba0c635eca..ea2be7b80e94 100644
--- a/core/java/android/os/UserBatteryConsumer.java
+++ b/core/java/android/os/UserBatteryConsumer.java
@@ -60,10 +60,10 @@ public class UserBatteryConsumer extends BatteryConsumer {
         pw.print("User ");
         pw.print(getUserId());
         pw.print(": ");
-        pw.print(BatteryStats.formatCharge(consumedPower));
-        pw.print(" ( ");
-        mPowerComponents.dump(pw, skipEmptyComponents  /* skipTotalPowerComponent */);
-        pw.print(" ) ");
+        pw.println(BatteryStats.formatCharge(consumedPower));
+        pw.print("      ");
+        mPowerComponents.dump(pw, SCREEN_STATE_ANY, POWER_STATE_ANY,
+                skipEmptyComponents  /* skipTotalPowerComponent */);
     }
 
     /** Serializes this object to XML */
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 092ee16f3342..33b87f36f0da 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -3252,6 +3252,9 @@ public final class BatteryStatsService extends IBatteryStats.Stub
                 .setMaxStatsAgeMs(0)
                 .includeProcessStateData()
                 .includePowerModels();
+        if (Flags.batteryUsageStatsByPowerAndScreenState()) {
+            builder.includeScreenStateData().includePowerStateData();
+        }
         if (model == BatteryConsumer.POWER_MODEL_POWER_PROFILE) {
             builder.powerProfileModeledOnly();
         }
@@ -3270,7 +3273,7 @@ public final class BatteryStatsService extends IBatteryStats.Stub
         if (proto) {
             batteryUsageStats.dumpToProto(fd);
         } else {
-            batteryUsageStats.dump(pw, "");
+            batteryUsageStats.dump(pw, "  ");
         }
     }
 
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsDumpHelperImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsDumpHelperImpl.java
index ad146afe16e7..fb54c5de260a 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsDumpHelperImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsDumpHelperImpl.java
@@ -20,6 +20,8 @@ import android.os.BatteryStats;
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
 
+import com.android.server.power.optimization.Flags;
+
 public class BatteryStatsDumpHelperImpl implements BatteryStats.BatteryStatsDumpHelper {
     private final BatteryUsageStatsProvider mBatteryUsageStatsProvider;
 
@@ -33,6 +35,9 @@ public class BatteryStatsDumpHelperImpl implements BatteryStats.BatteryStatsDump
                 .setMaxStatsAgeMs(0);
         if (detailed) {
             builder.includePowerModels().includeProcessStateData().includeVirtualUids();
+            if (Flags.batteryUsageStatsByPowerAndScreenState()) {
+                builder.includePowerStateData().includeScreenStateData();
+            }
         }
         return mBatteryUsageStatsProvider.getBatteryUsageStats((BatteryStatsImpl) batteryStats,
                 builder.build());
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index 8127b8217bb0..ac6896696de6 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -201,7 +201,8 @@ public class BatteryUsageStatsProvider {
 
             batteryUsageStatsBuilder = new BatteryUsageStats.Builder(
                     stats.getCustomEnergyConsumerNames(), includePowerModels,
-                    includeProcessStateData, minConsumedPowerThreshold);
+                    includeProcessStateData, query.isScreenStateDataNeeded(),
+                    query.isPowerStateDataNeeded(), minConsumedPowerThreshold);
 
             // TODO(b/188068523): use a monotonic clock to ensure resilience of order and duration
             // of batteryUsageStats sessions to wall-clock adjustments
@@ -348,6 +349,7 @@ public class BatteryUsageStatsProvider {
         final String[] customEnergyConsumerNames = stats.getCustomEnergyConsumerNames();
         final BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
                 customEnergyConsumerNames, includePowerModels, includeProcessStateData,
+                query.isScreenStateDataNeeded(), query.isPowerStateDataNeeded(),
                 minConsumedPowerThreshold);
         if (mPowerStatsStore == null) {
             Log.e(TAG, "PowerStatsStore is unavailable");
@@ -408,7 +410,6 @@ public class BatteryUsageStatsProvider {
                             + " does not include process state data");
                     continue;
                 }
-
                 builder.add(snapshot);
             }
         }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsExporter.java b/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
index 549a97ea49cd..0f1349287a0f 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
@@ -129,17 +129,55 @@ public class PowerStatsExporter {
         if (descriptor == null) {
             return;
         }
+        boolean isCustomComponent =
+                descriptor.powerComponentId >= BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
 
         PowerStatsLayout layout = new PowerStatsLayout();
         layout.fromExtras(descriptor.extras);
 
         long[] deviceStats = new long[descriptor.statsArrayLength];
+        for (int screenState = 0; screenState < BatteryConsumer.SCREEN_STATE_COUNT; screenState++) {
+            if (batteryUsageStatsBuilder.isScreenStateDataNeeded()) {
+                if (screenState == BatteryConsumer.SCREEN_STATE_UNSPECIFIED) {
+                    continue;
+                }
+            } else if (screenState != BatteryConsumer.SCREEN_STATE_UNSPECIFIED) {
+                continue;
+            }
+
+            for (int powerState = 0; powerState < BatteryConsumer.POWER_STATE_COUNT; powerState++) {
+                if (batteryUsageStatsBuilder.isPowerStateDataNeeded() && !isCustomComponent) {
+                    if (powerState == BatteryConsumer.POWER_STATE_UNSPECIFIED) {
+                        continue;
+                    }
+                } else if (powerState != BatteryConsumer.POWER_STATE_BATTERY) {
+                    continue;
+                }
+
+                populateAggregatedBatteryConsumer(batteryUsageStatsBuilder, powerComponentStats,
+                        layout, deviceStats, screenState, powerState);
+            }
+        }
+        if (layout.isUidPowerAttributionSupported()) {
+            populateBatteryConsumers(batteryUsageStatsBuilder,
+                    powerComponentStats, layout);
+        }
+    }
+
+    private static void populateAggregatedBatteryConsumer(
+            BatteryUsageStats.Builder batteryUsageStatsBuilder,
+            PowerComponentAggregatedPowerStats powerComponentStats, PowerStatsLayout layout,
+            long[] deviceStats, @BatteryConsumer.ScreenState int screenState,
+            @BatteryConsumer.PowerState int powerState) {
+        int powerComponentId = powerComponentStats.powerComponentId;
+        boolean isCustomComponent =
+                powerComponentId >= BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
+
         double[] totalPower = new double[1];
         MultiStateStats.States.forEachTrackedStateCombination(
                 powerComponentStats.getConfig().getDeviceStateConfig(),
                 states -> {
-                    if (states[AggregatedPowerStatsConfig.STATE_POWER]
-                            != AggregatedPowerStatsConfig.POWER_STATE_BATTERY) {
+                    if (!areMatchingStates(states, screenState, powerState)) {
                         return;
                     }
 
@@ -153,24 +191,23 @@ public class PowerStatsExporter {
         AggregateBatteryConsumer.Builder deviceScope =
                 batteryUsageStatsBuilder.getAggregateBatteryConsumerBuilder(
                         BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
-        if (descriptor.powerComponentId >= BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID) {
-            if (batteryUsageStatsBuilder.isSupportedCustomPowerComponent(
-                    descriptor.powerComponentId)) {
-                deviceScope.addConsumedPowerForCustomComponent(descriptor.powerComponentId,
-                        totalPower[0]);
+        if (isCustomComponent) {
+            if (batteryUsageStatsBuilder.isSupportedCustomPowerComponent(powerComponentId)) {
+                deviceScope.addConsumedPowerForCustomComponent(powerComponentId, totalPower[0]);
             }
         } else {
-            deviceScope.addConsumedPower(descriptor.powerComponentId,
-                    totalPower[0], BatteryConsumer.POWER_MODEL_UNDEFINED);
-        }
-
-        if (layout.isUidPowerAttributionSupported()) {
-            populateUidBatteryConsumers(batteryUsageStatsBuilder,
-                    powerComponentStats, layout);
+            BatteryConsumer.Key key = deviceScope.getKey(powerComponentId,
+                    BatteryConsumer.PROCESS_STATE_ANY, screenState, powerState);
+            if (key != null) {
+                deviceScope.addConsumedPower(key, totalPower[0],
+                        BatteryConsumer.POWER_MODEL_UNDEFINED);
+            }
+            deviceScope.addConsumedPower(powerComponentId, totalPower[0],
+                    BatteryConsumer.POWER_MODEL_UNDEFINED);
         }
     }
 
-    private static void populateUidBatteryConsumers(
+    private static void populateBatteryConsumers(
             BatteryUsageStats.Builder batteryUsageStatsBuilder,
             PowerComponentAggregatedPowerStats powerComponentStats,
             PowerStatsLayout layout) {
@@ -185,11 +222,44 @@ public class PowerStatsExporter {
                 .getUidStateConfig()[AggregatedPowerStatsConfig.STATE_PROCESS_STATE].isTracked()
                 && powerComponentId < BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
 
+        ArrayList<Integer> uids = new ArrayList<>();
+        powerComponentStats.collectUids(uids);
+        for (int screenState = 0; screenState < BatteryConsumer.SCREEN_STATE_COUNT; screenState++) {
+            if (batteryUsageStatsBuilder.isScreenStateDataNeeded()) {
+                if (screenState == BatteryConsumer.SCREEN_STATE_UNSPECIFIED) {
+                    continue;
+                }
+            } else if (screenState != BatteryConsumer.SCREEN_STATE_UNSPECIFIED) {
+                continue;
+            }
+
+            for (int powerState = 0; powerState < BatteryConsumer.POWER_STATE_COUNT; powerState++) {
+                if (batteryUsageStatsBuilder.isPowerStateDataNeeded()) {
+                    if (powerState == BatteryConsumer.POWER_STATE_UNSPECIFIED) {
+                        continue;
+                    }
+                } else if (powerState != BatteryConsumer.POWER_STATE_BATTERY) {
+                    continue;
+                }
+
+                populateUidBatteryConsumers(batteryUsageStatsBuilder, powerComponentStats, layout,
+                        uids, powerComponent, uidStats, breakDownByProcState, screenState,
+                        powerState);
+            }
+        }
+    }
+
+    private static void populateUidBatteryConsumers(
+            BatteryUsageStats.Builder batteryUsageStatsBuilder,
+            PowerComponentAggregatedPowerStats powerComponentStats, PowerStatsLayout layout,
+            List<Integer> uids, AggregatedPowerStatsConfig.PowerComponent powerComponent,
+            long[] uidStats, boolean breakDownByProcState,
+            @BatteryConsumer.ScreenState int screenState,
+            @BatteryConsumer.PowerState int powerState) {
+        int powerComponentId = powerComponentStats.powerComponentId;
         double[] powerByProcState =
                 new double[breakDownByProcState ? BatteryConsumer.PROCESS_STATE_COUNT : 1];
         double powerAllApps = 0;
-        ArrayList<Integer> uids = new ArrayList<>();
-        powerComponentStats.collectUids(uids);
         for (int uid : uids) {
             UidBatteryConsumer.Builder builder =
                     batteryUsageStatsBuilder.getOrCreateUidBatteryConsumerBuilder(uid);
@@ -199,8 +269,7 @@ public class PowerStatsExporter {
             MultiStateStats.States.forEachTrackedStateCombination(
                     powerComponent.getUidStateConfig(),
                     states -> {
-                        if (states[AggregatedPowerStatsConfig.STATE_POWER]
-                                != AggregatedPowerStatsConfig.POWER_STATE_BATTERY) {
+                        if (!areMatchingStates(states, screenState, powerState)) {
                             return;
                         }
 
@@ -224,8 +293,17 @@ public class PowerStatsExporter {
                 powerAllProcStates += power;
                 if (breakDownByProcState
                         && procState != BatteryConsumer.PROCESS_STATE_UNSPECIFIED) {
-                    builder.addConsumedPower(builder.getKey(powerComponentId, procState), power,
-                            BatteryConsumer.POWER_MODEL_UNDEFINED);
+                    if (batteryUsageStatsBuilder.isPowerStateDataNeeded()) {
+                        builder.addConsumedPower(
+                                builder.getKey(powerComponentId, procState, screenState,
+                                        powerState),
+                                power, BatteryConsumer.POWER_MODEL_UNDEFINED);
+                    } else {
+                        builder.addConsumedPower(
+                                builder.getKey(powerComponentId, procState, screenState,
+                                        BatteryConsumer.POWER_STATE_UNSPECIFIED),
+                                power, BatteryConsumer.POWER_MODEL_UNDEFINED);
+                    }
                 }
             }
             if (powerComponentId >= BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID) {
@@ -243,8 +321,49 @@ public class PowerStatsExporter {
         if (powerComponentId >= BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID) {
             allAppsScope.addConsumedPowerForCustomComponent(powerComponentId, powerAllApps);
         } else {
+            BatteryConsumer.Key key = allAppsScope.getKey(powerComponentId,
+                    BatteryConsumer.PROCESS_STATE_ANY, screenState, powerState);
+            if (key != null) {
+                allAppsScope.addConsumedPower(key, powerAllApps,
+                        BatteryConsumer.POWER_MODEL_UNDEFINED);
+            }
             allAppsScope.addConsumedPower(powerComponentId, powerAllApps,
                     BatteryConsumer.POWER_MODEL_UNDEFINED);
         }
     }
+
+    private static boolean areMatchingStates(int[] states,
+            @BatteryConsumer.ScreenState int screenState,
+            @BatteryConsumer.PowerState int powerState) {
+        switch (screenState) {
+            case BatteryConsumer.SCREEN_STATE_ON:
+                if (states[AggregatedPowerStatsConfig.STATE_SCREEN]
+                        != AggregatedPowerStatsConfig.SCREEN_STATE_ON) {
+                    return false;
+                }
+                break;
+            case BatteryConsumer.SCREEN_STATE_OTHER:
+                if (states[AggregatedPowerStatsConfig.STATE_SCREEN]
+                        != AggregatedPowerStatsConfig.SCREEN_STATE_OTHER) {
+                    return false;
+                }
+                break;
+        }
+
+        switch (powerState) {
+            case BatteryConsumer.POWER_STATE_BATTERY:
+                if (states[AggregatedPowerStatsConfig.STATE_POWER]
+                        != AggregatedPowerStatsConfig.POWER_STATE_BATTERY) {
+                    return false;
+                }
+                break;
+            case BatteryConsumer.POWER_STATE_OTHER:
+                if (states[AggregatedPowerStatsConfig.STATE_POWER]
+                        != AggregatedPowerStatsConfig.POWER_STATE_OTHER) {
+                    return false;
+                }
+                break;
+        }
+        return true;
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/flags.aconfig b/services/core/java/com/android/server/power/stats/flags.aconfig
index d34498a3764b..05496cff2ee9 100644
--- a/services/core/java/com/android/server/power/stats/flags.aconfig
+++ b/services/core/java/com/android/server/power/stats/flags.aconfig
@@ -54,3 +54,10 @@ flag {
     description: "Adds battery_usage_stats_slice atom"
     bug: "324602949"
 }
+
+flag {
+    name: "battery_usage_stats_by_power_and_screen_state"
+    namespace: "backstage_power"
+    description: "Batterystats dumpsys is enhanced by including power break-down by power s"
+    bug: "352835319"
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
index 62efbc3cfa35..37d8f2f74850 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
@@ -488,6 +488,8 @@ public class BatteryUsageStatsAtomTest {
                 new BatteryUsageStats.Builder(new String[]{"CustomConsumer1", "CustomConsumer2"},
                         /* includePowerModels */ true,
                         /* includeProcessStats */ true,
+                        /* includeScreenStateData */ false,
+                        /* includePowerStateData */ false,
                         /* minConsumedPowerThreshold */ 0)
                         .setDischargePercentage(20)
                         .setDischargedPowerRange(1000, 2000)
@@ -574,7 +576,7 @@ public class BatteryUsageStatsAtomTest {
     @Test
     public void testLargeAtomTruncated() {
         final BatteryUsageStats.Builder builder =
-                new BatteryUsageStats.Builder(new String[0], true, false, 0);
+                new BatteryUsageStats.Builder(new String[0], true, false, false, false, 0);
         // If not truncated, this BatteryUsageStats object would generate a proto buffer
         // significantly larger than 50 Kb
         for (int i = 0; i < 3000; i++) {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
index 6edfedee9e5b..624b18948c49 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
@@ -397,10 +397,14 @@ public class BatteryUsageStatsRule implements TestRule {
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_MODELS) != 0;
         final boolean includeProcessStateData = (query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_PROCESS_STATE_DATA) != 0;
+        final boolean includeScreenStateData = (query.getFlags()
+                & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_SCREEN_STATE) != 0;
+        final boolean includePowerStateData = (query.getFlags()
+                & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_STATE) != 0;
         final double minConsumedPowerThreshold = query.getMinConsumedPowerThreshold();
         BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
                 customPowerComponentNames, includePowerModels, includeProcessStateData,
-                minConsumedPowerThreshold);
+                includeScreenStateData, includePowerStateData, minConsumedPowerThreshold);
         SparseArray<? extends BatteryStats.Uid> uidStats = mBatteryStats.getUidStats();
         for (int i = 0; i < uidStats.size(); i++) {
             builder.getOrCreateUidBatteryConsumerBuilder(uidStats.valueAt(i));
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java
index a3f0770ec8ba..52bb5e839ca2 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsTest.java
@@ -31,6 +31,8 @@ import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import static java.util.regex.Pattern.quote;
+
 import android.os.AggregateBatteryConsumer;
 import android.os.BatteryConsumer;
 import android.os.BatteryUsageStats;
@@ -91,7 +93,7 @@ public class BatteryUsageStatsTest {
         final Parcel parcel = Parcel.obtain();
         parcel.writeParcelable(outBatteryUsageStats, 0);
 
-        assertThat(parcel.dataSize()).isLessThan(12000);
+        assertThat(parcel.dataSize()).isLessThan(100000);
 
         parcel.setDataPosition(0);
 
@@ -161,15 +163,47 @@ public class BatteryUsageStatsTest {
         assertThat(dump).contains("Computed drain: 30000");
         assertThat(dump).contains("actual drain: 1000-2000");
         assertThat(dump).contains("cpu: 20100 apps: 10100 duration: 20s 300ms");
-        assertThat(dump).contains("cpu(fg): 2333 apps: 1333 duration: 3s 332ms");
-        assertThat(dump).contains("cpu(bg): 2444 apps: 1444 duration: 4s 442ms");
-        assertThat(dump).contains("cpu(fgs): 2555 apps: 1555 duration: 5s 552ms");
-        assertThat(dump).contains("cpu(cached): 123 apps: 123 duration: 456ms");
         assertThat(dump).contains("FOO: 20200 apps: 10200 duration: 20s 400ms");
-        assertThat(dump).contains("UID 271: 1200 fg: 1777 bg: 1888 fgs: 1999 cached: 123 "
-                + "( screen=300 cpu=400 (600ms) cpu:fg=1777 (7s 771ms) cpu:bg=1888 (8s 881ms) "
-                + "cpu:fgs=1999 (9s 991ms) cpu:cached=123 (456ms) FOO=500 )");
-        assertThat(dump).contains("User 42: 30.0 ( cpu=10.0 (30ms) FOO=20.0 )");
+        assertThat(dump).containsMatch(quote("(on battery, screen on)") + "\\s*"
+                + "cpu: 2333 apps: 1333 duration: 3s 332ms");
+        assertThat(dump).containsMatch(quote("(not on battery, screen on)") + "\\s*"
+                + "cpu: 2555 apps: 1555 duration: 5s 552ms");
+        assertThat(dump).containsMatch(quote("(on battery, screen off/doze)") + "\\s*"
+                + "cpu: 2444 apps: 1444 duration: 4s 442ms");
+        assertThat(dump).containsMatch(quote("(not on battery, screen off/doze)") + "\\s*"
+                + "cpu: 123 apps: 123 duration: 456ms");
+        assertThat(dump).containsMatch(
+                "UID 271: 1200 fg: 1777 bg: 1888 fgs: 1999 cached: 123\\s*"
+                        + quote("screen=300 cpu=5787 (27s 99ms) cpu:fg=1777 (7s 771ms) "
+                        + "cpu:bg=1888 (8s 881ms) cpu:fgs=1999 (9s 991ms) "
+                        + "cpu:cached=123 (456ms) FOO=500") + "\\s*"
+                        + quote("(on battery, screen on)") + "\\s*"
+                        + quote("cpu:fg=1777 (7s 771ms)"));
+        assertThat(dump).containsMatch("User 42: 30.0\\s*"
+                + quote("cpu=10.0 (30ms) FOO=20.0"));
+    }
+
+    @Test
+    public void testDumpNoScreenOrPowerState() {
+        final BatteryUsageStats stats = buildBatteryUsageStats1(true, false, false).build();
+        final StringWriter out = new StringWriter();
+        try (PrintWriter pw = new PrintWriter(out)) {
+            stats.dump(pw, "  ");
+        }
+        final String dump = out.toString();
+
+        assertThat(dump).contains("Capacity: 4000");
+        assertThat(dump).contains("Computed drain: 30000");
+        assertThat(dump).contains("actual drain: 1000-2000");
+        assertThat(dump).contains("cpu: 20100 apps: 10100 duration: 20s 300ms");
+        assertThat(dump).contains("FOO: 20200 apps: 10200 duration: 20s 400ms");
+        assertThat(dump).containsMatch(
+                "UID 271: 1200 fg: 1777 bg: 1888 fgs: 1999 cached: 123\\s*"
+                        + quote("screen=300 cpu=5787 (600ms) cpu:fg=1777 (7s 771ms) "
+                        + "cpu:bg=1888 (8s 881ms) cpu:fgs=1999 (9s 991ms) "
+                        + "cpu:cached=123 (456ms) FOO=500"));
+        assertThat(dump).containsMatch("User 42: 30.0\\s*"
+                + quote("cpu=10.0 (30ms) FOO=20.0"));
     }
 
     @Test
@@ -186,9 +220,8 @@ public class BatteryUsageStatsTest {
     public void testAdd() {
         final BatteryUsageStats stats1 = buildBatteryUsageStats1(false).build();
         final BatteryUsageStats stats2 = buildBatteryUsageStats2(new String[]{"FOO"}, true).build();
-
         final BatteryUsageStats sum =
-                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, 0)
+                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, true, true, 0)
                         .add(stats1)
                         .add(stats2)
                         .build();
@@ -200,14 +233,14 @@ public class BatteryUsageStatsTest {
         for (UidBatteryConsumer uidBatteryConsumer : uidBatteryConsumers) {
             if (uidBatteryConsumer.getUid() == APP_UID1) {
                 assertUidBatteryConsumer(uidBatteryConsumer, 2124, null,
-                        5321, 6900, 532, 423, BatteryConsumer.POWER_MODEL_POWER_PROFILE, 745,
+                        5321, 6900, 532, 423, BatteryConsumer.POWER_MODEL_POWER_PROFILE, 11772,
                         POWER_MODEL_UNDEFINED,
                         956, 1167, 1478,
                         true, 3554, 3776, 3998, 444, 3554, 15542, 3776, 17762, 3998, 19982,
                         444, 1110);
             } else if (uidBatteryConsumer.getUid() == APP_UID2) {
                 assertUidBatteryConsumer(uidBatteryConsumer, 1332, "bar",
-                        1111, 2220, 2, 333, BatteryConsumer.POWER_MODEL_POWER_PROFILE, 444,
+                        1111, 2220, 2, 333, BatteryConsumer.POWER_MODEL_POWER_PROFILE, 5985,
                         BatteryConsumer.POWER_MODEL_POWER_PROFILE,
                         555, 666, 777,
                         true, 1777, 1888, 1999, 321, 1777, 7771, 1888, 8881, 1999, 9991,
@@ -229,7 +262,7 @@ public class BatteryUsageStatsTest {
     @Test
     public void testAdd_customComponentMismatch() {
         final BatteryUsageStats.Builder builder =
-                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, 0);
+                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, true, true, 0);
         final BatteryUsageStats stats = buildBatteryUsageStats2(new String[]{"BAR"}, false).build();
 
         assertThrows(IllegalArgumentException.class, () -> builder.add(stats));
@@ -238,7 +271,7 @@ public class BatteryUsageStatsTest {
     @Test
     public void testAdd_processStateDataMismatch() {
         final BatteryUsageStats.Builder builder =
-                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, 0);
+                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, true, true, 0);
         final BatteryUsageStats stats = buildBatteryUsageStats2(new String[]{"FOO"}, false).build();
 
         assertThrows(IllegalArgumentException.class, () -> builder.add(stats));
@@ -259,15 +292,23 @@ public class BatteryUsageStatsTest {
         parser.setInput(in, StandardCharsets.UTF_8.name());
         final BatteryUsageStats fromXml = BatteryUsageStats.createFromXml(parser);
 
+        System.out.println("stats = " + stats);
+        System.out.println("fromXml = " + fromXml);
         assertBatteryUsageStats1(fromXml, true);
     }
 
     private BatteryUsageStats.Builder buildBatteryUsageStats1(boolean includeUserBatteryConsumer) {
+        return buildBatteryUsageStats1(includeUserBatteryConsumer, true, true);
+    }
+
+    private BatteryUsageStats.Builder buildBatteryUsageStats1(boolean includeUserBatteryConsumer,
+            boolean includeScreenState, boolean includePowerState) {
         final MockClock clocks = new MockClock();
         final MockBatteryStatsImpl batteryStats = new MockBatteryStatsImpl(clocks);
 
         final BatteryUsageStats.Builder builder =
-                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true, 0)
+                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, true,
+                        includeScreenState, includePowerState, 0)
                         .setBatteryCapacity(4000)
                         .setDischargePercentage(20)
                         .setDischargedPowerRange(1000, 2000)
@@ -312,7 +353,7 @@ public class BatteryUsageStatsTest {
 
         final BatteryUsageStats.Builder builder =
                 new BatteryUsageStats.Builder(customPowerComponentNames, true,
-                        includeProcessStateData, 0);
+                        includeProcessStateData, true, true, 0);
         builder.setDischargePercentage(30)
                 .setDischargedPowerRange(1234, 2345)
                 .setStatsStartTimestamp(2000)
@@ -371,9 +412,15 @@ public class BatteryUsageStatsTest {
                 .setUsageDurationForCustomComponentMillis(
                         BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentDuration);
         if (builder.isProcessStateDataNeeded()) {
-            final BatteryConsumer.Key cpuFgKey = uidBuilder.getKey(
-                    BatteryConsumer.POWER_COMPONENT_CPU,
-                    BatteryConsumer.PROCESS_STATE_FOREGROUND);
+            final BatteryConsumer.Key cpuFgKey = builder.isScreenStateDataNeeded()
+                    ? uidBuilder.getKey(
+                            BatteryConsumer.POWER_COMPONENT_CPU,
+                            BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                            BatteryConsumer.SCREEN_STATE_ON,
+                            BatteryConsumer.POWER_STATE_BATTERY)
+                    : uidBuilder.getKey(
+                            BatteryConsumer.POWER_COMPONENT_CPU,
+                            BatteryConsumer.PROCESS_STATE_FOREGROUND);
             final BatteryConsumer.Key cpuBgKey = uidBuilder.getKey(
                     BatteryConsumer.POWER_COMPONENT_CPU,
                     BatteryConsumer.PROCESS_STATE_BACKGROUND);
@@ -401,9 +448,9 @@ public class BatteryUsageStatsTest {
 
     private void addAggregateBatteryConsumer(BatteryUsageStats.Builder builder, int scope,
             double consumedPower, int cpuPower, int customComponentPower, int cpuDuration,
-            int customComponentDuration, double cpuPowerForeground, long cpuDurationForeground,
-            double cpuPowerBackground, long cpuDurationBackground, double cpuPowerFgs,
-            long cpuDurationFgs, double cpuPowerCached, long cpuDurationCached) {
+            int customComponentDuration, double cpuPowerBatScrOn, long cpuDurationBatScrOn,
+            double cpuPowerBatScrOff, long cpuDurationBatScrOff, double cpuPowerChgScrOn,
+            long cpuDurationChgScrOn, double cpuPowerChgScrOff, long cpuDurationChgScrOff) {
         final AggregateBatteryConsumer.Builder aggBuilder =
                 builder.getAggregateBatteryConsumerBuilder(scope)
                         .setConsumedPower(consumedPower)
@@ -417,32 +464,40 @@ public class BatteryUsageStatsTest {
                         .setUsageDurationForCustomComponentMillis(
                                 BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID,
                                 customComponentDuration);
-        if (builder.isProcessStateDataNeeded()) {
-            final BatteryConsumer.Key cpuFgKey = aggBuilder.getKey(
+        if (builder.isPowerStateDataNeeded() || builder.isScreenStateDataNeeded()) {
+            final BatteryConsumer.Key cpuBatScrOn = aggBuilder.getKey(
                     BatteryConsumer.POWER_COMPONENT_CPU,
-                    BatteryConsumer.PROCESS_STATE_FOREGROUND);
-            final BatteryConsumer.Key cpuBgKey = aggBuilder.getKey(
+                    BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                    BatteryConsumer.SCREEN_STATE_ON,
+                    BatteryConsumer.POWER_STATE_BATTERY);
+            final BatteryConsumer.Key cpuBatScrOff = aggBuilder.getKey(
                     BatteryConsumer.POWER_COMPONENT_CPU,
-                    BatteryConsumer.PROCESS_STATE_BACKGROUND);
-            final BatteryConsumer.Key cpuFgsKey = aggBuilder.getKey(
+                    BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                    BatteryConsumer.SCREEN_STATE_OTHER,
+                    BatteryConsumer.POWER_STATE_BATTERY);
+            final BatteryConsumer.Key cpuChgScrOn = aggBuilder.getKey(
                     BatteryConsumer.POWER_COMPONENT_CPU,
-                    BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE);
-            final BatteryConsumer.Key cpuCachedKey = aggBuilder.getKey(
+                    BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                    BatteryConsumer.SCREEN_STATE_ON,
+                    BatteryConsumer.POWER_STATE_OTHER);
+            final BatteryConsumer.Key cpuChgScrOff = aggBuilder.getKey(
                     BatteryConsumer.POWER_COMPONENT_CPU,
-                    BatteryConsumer.PROCESS_STATE_CACHED);
+                    BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                    BatteryConsumer.SCREEN_STATE_OTHER,
+                    BatteryConsumer.POWER_STATE_OTHER);
             aggBuilder
-                    .setConsumedPower(cpuFgKey, cpuPowerForeground,
+                    .setConsumedPower(cpuBatScrOn, cpuPowerBatScrOn,
                             BatteryConsumer.POWER_MODEL_POWER_PROFILE)
-                    .setUsageDurationMillis(cpuFgKey, cpuDurationForeground)
-                    .setConsumedPower(cpuBgKey, cpuPowerBackground,
+                    .setUsageDurationMillis(cpuBatScrOn, cpuDurationBatScrOn)
+                    .setConsumedPower(cpuBatScrOff, cpuPowerBatScrOff,
                             BatteryConsumer.POWER_MODEL_POWER_PROFILE)
-                    .setUsageDurationMillis(cpuBgKey, cpuDurationBackground)
-                    .setConsumedPower(cpuFgsKey, cpuPowerFgs,
+                    .setUsageDurationMillis(cpuBatScrOff, cpuDurationBatScrOff)
+                    .setConsumedPower(cpuChgScrOn, cpuPowerChgScrOn,
                             BatteryConsumer.POWER_MODEL_POWER_PROFILE)
-                    .setUsageDurationMillis(cpuFgsKey, cpuDurationFgs)
-                    .setConsumedPower(cpuCachedKey, cpuPowerCached,
+                    .setUsageDurationMillis(cpuChgScrOn, cpuDurationChgScrOn)
+                    .setConsumedPower(cpuChgScrOff, cpuPowerChgScrOff,
                             BatteryConsumer.POWER_MODEL_POWER_PROFILE)
-                    .setUsageDurationMillis(cpuCachedKey, cpuDurationCached);
+                    .setUsageDurationMillis(cpuChgScrOff, cpuDurationChgScrOff);
         }
     }
 
@@ -456,7 +511,7 @@ public class BatteryUsageStatsTest {
         for (UidBatteryConsumer uidBatteryConsumer : uidBatteryConsumers) {
             if (uidBatteryConsumer.getUid() == APP_UID1) {
                 assertUidBatteryConsumer(uidBatteryConsumer, 1200, "foo",
-                        1000, 1500, 500, 300, BatteryConsumer.POWER_MODEL_POWER_PROFILE, 400,
+                        1000, 1500, 500, 300, BatteryConsumer.POWER_MODEL_POWER_PROFILE, 5787,
                         BatteryConsumer.POWER_MODEL_POWER_PROFILE,
                         500, 600, 800,
                         true, 1777, 1888, 1999, 123, 1777, 7771, 1888, 8881, 1999, 9991, 123, 456);
@@ -568,54 +623,53 @@ public class BatteryUsageStatsTest {
                     .isEqualTo(totalPowerCached);
         }
 
-        final BatteryConsumer.Key cpuFgKey = uidBatteryConsumer.getKey(
+        final BatteryConsumer.Dimensions cpuFg = new BatteryConsumer.Dimensions(
                 BatteryConsumer.POWER_COMPONENT_CPU, BatteryConsumer.PROCESS_STATE_FOREGROUND);
         if (processStateDataIncluded) {
-            assertThat(cpuFgKey).isNotNull();
-            assertThat(uidBatteryConsumer.getConsumedPower(cpuFgKey))
+            assertThat(uidBatteryConsumer.getConsumedPower(cpuFg))
                     .isEqualTo(cpuPowerForeground);
-            assertThat(uidBatteryConsumer.getUsageDurationMillis(cpuFgKey))
+            assertThat(uidBatteryConsumer.getUsageDurationMillis(cpuFg))
                     .isEqualTo(cpuDurationForeground);
         } else {
-            assertThat(cpuFgKey).isNull();
+            assertThat(uidBatteryConsumer.getConsumedPower(cpuFg)).isEqualTo(0);
         }
 
-        final BatteryConsumer.Key cpuBgKey = uidBatteryConsumer.getKey(
+        final BatteryConsumer.Dimensions cpuBg = new BatteryConsumer.Dimensions(
                 BatteryConsumer.POWER_COMPONENT_CPU, BatteryConsumer.PROCESS_STATE_BACKGROUND);
         if (processStateDataIncluded) {
-            assertThat(cpuBgKey).isNotNull();
-            assertThat(uidBatteryConsumer.getConsumedPower(cpuBgKey))
+            assertThat(uidBatteryConsumer.getConsumedPower(cpuBg))
                     .isEqualTo(cpuPowerBackground);
-            assertThat(uidBatteryConsumer.getUsageDurationMillis(cpuBgKey))
+            assertThat(uidBatteryConsumer.getUsageDurationMillis(cpuBg))
                     .isEqualTo(cpuDurationBackground);
         } else {
-            assertThat(cpuBgKey).isNull();
+            assertThat(uidBatteryConsumer.getConsumedPower(cpuBg))
+                    .isEqualTo(0);
         }
 
-        final BatteryConsumer.Key cpuFgsKey = uidBatteryConsumer.getKey(
+        final BatteryConsumer.Dimensions cpuFgs = new BatteryConsumer.Dimensions(
                 BatteryConsumer.POWER_COMPONENT_CPU,
                 BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE);
         if (processStateDataIncluded) {
-            assertThat(cpuFgsKey).isNotNull();
-            assertThat(uidBatteryConsumer.getConsumedPower(cpuFgsKey))
+            assertThat(uidBatteryConsumer.getConsumedPower(cpuFgs))
                     .isEqualTo(cpuPowerFgs);
-            assertThat(uidBatteryConsumer.getUsageDurationMillis(cpuFgsKey))
+            assertThat(uidBatteryConsumer.getUsageDurationMillis(cpuFgs))
                     .isEqualTo(cpuDurationFgs);
         } else {
-            assertThat(cpuFgsKey).isNotNull();
+            assertThat(uidBatteryConsumer.getConsumedPower(cpuFgs))
+                    .isEqualTo(0);
         }
 
-        final BatteryConsumer.Key cachedKey = uidBatteryConsumer.getKey(
+        final BatteryConsumer.Dimensions cached = new BatteryConsumer.Dimensions(
                 BatteryConsumer.POWER_COMPONENT_CPU,
                 BatteryConsumer.PROCESS_STATE_CACHED);
         if (processStateDataIncluded) {
-            assertThat(cachedKey).isNotNull();
-            assertThat(uidBatteryConsumer.getConsumedPower(cachedKey))
+            assertThat(uidBatteryConsumer.getConsumedPower(cached))
                     .isEqualTo(cpuPowerCached);
-            assertThat(uidBatteryConsumer.getUsageDurationMillis(cachedKey))
+            assertThat(uidBatteryConsumer.getUsageDurationMillis(cached))
                     .isEqualTo(cpuDurationCached);
         } else {
-            assertThat(cpuFgsKey).isNotNull();
+            assertThat(uidBatteryConsumer.getConsumedPower(cached))
+                    .isEqualTo(0);
         }
     }
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
index 32bfb2cd3507..7f7967ba4d7b 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
@@ -19,7 +19,6 @@ package com.android.server.power.stats;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
-import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.mock;
 
 import android.os.AggregateBatteryConsumer;
@@ -131,9 +130,20 @@ public class PowerStatsExporterTest {
 
     @Test
     public void breakdownByProcState_fullRange() throws Exception {
+        breakdownByProcState_fullRange(false, false);
+    }
+
+    @Test
+    public void breakdownByProcStateScreenAndPower_fullRange() throws Exception {
+        breakdownByProcState_fullRange(true, true);
+    }
+
+    private void breakdownByProcState_fullRange(boolean includeScreenStateData,
+            boolean includePowerStateData) throws Exception {
         BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
                 new String[]{"cu570m"}, /* includePowerModels */ false,
-                /* includeProcessStateData */ true, /* powerThreshold */ 0);
+                /* includeProcessStateData */ true, includeScreenStateData,
+                includePowerStateData, /* powerThreshold */ 0);
         exportAggregatedPowerStats(builder, 1000, 10000);
 
         BatteryUsageStats actual = builder.build();
@@ -177,7 +187,7 @@ public class PowerStatsExporterTest {
     public void breakdownByProcState_subRange() throws Exception {
         BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
                 new String[]{"cu570m"}, /* includePowerModels */ false,
-                /* includeProcessStateData */ true, /* powerThreshold */ 0);
+                /* includeProcessStateData */ true, true, true, /* powerThreshold */ 0);
         exportAggregatedPowerStats(builder, 3700, 6700);
 
         BatteryUsageStats actual = builder.build();
@@ -209,7 +219,7 @@ public class PowerStatsExporterTest {
     public void combinedProcessStates() throws Exception {
         BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
                 new String[]{"cu570m"}, /* includePowerModels */ false,
-                /* includeProcessStateData */ false, /* powerThreshold */ 0);
+                /* includeProcessStateData */ false, true, true, /* powerThreshold */ 0);
         exportAggregatedPowerStats(builder, 1000, 10000);
 
         BatteryUsageStats actual = builder.build();
@@ -229,13 +239,13 @@ public class PowerStatsExporterTest {
         UidBatteryConsumer uidScope = actual.getUidBatteryConsumers().stream()
                 .filter(us -> us.getUid() == APP_UID1).findFirst().orElse(null);
         // There shouldn't be any per-procstate data
-        assertThrows(
-                IllegalArgumentException.class,
-                () -> uidScope.getConsumedPower(new BatteryConsumer.Dimensions(
+        for (int procState = 0; procState < BatteryConsumer.PROCESS_STATE_COUNT; procState++) {
+            if (procState != BatteryConsumer.PROCESS_STATE_UNSPECIFIED) {
+                assertThat(uidScope.getConsumedPower(new BatteryConsumer.Dimensions(
                         BatteryConsumer.POWER_COMPONENT_CPU,
-                        BatteryConsumer.PROCESS_STATE_FOREGROUND)));
-
-
+                        BatteryConsumer.PROCESS_STATE_FOREGROUND))).isEqualTo(0);
+            }
+        }
         actual.close();
     }
 
diff --git a/tests/BatteryStatsPerfTest/src/com/android/internal/os/BatteryUsageStatsPerfTest.java b/tests/BatteryStatsPerfTest/src/com/android/internal/os/BatteryUsageStatsPerfTest.java
index 08430f2f2744..4143f595f9a0 100644
--- a/tests/BatteryStatsPerfTest/src/com/android/internal/os/BatteryUsageStatsPerfTest.java
+++ b/tests/BatteryStatsPerfTest/src/com/android/internal/os/BatteryUsageStatsPerfTest.java
@@ -159,7 +159,7 @@ public class BatteryUsageStatsPerfTest {
 
     private static BatteryUsageStats buildBatteryUsageStats() {
         final BatteryUsageStats.Builder builder =
-                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, false, 0)
+                new BatteryUsageStats.Builder(new String[]{"FOO"}, true, false, false, false, 0)
                         .setBatteryCapacity(4000)
                         .setDischargePercentage(20)
                         .setDischargedPowerRange(1000, 2000)