Merge "Fix typo in NPMS dump" into main
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index e0c3230..88b5275 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -423,6 +423,7 @@
   public final class PictureInPictureParams implements android.os.Parcelable {
     method public float getAspectRatioFloat();
     method public float getExpandedAspectRatioFloat();
+    method public static boolean isSameAspectRatio(@NonNull android.graphics.Rect, @NonNull android.util.Rational);
   }
 
   public final class PictureInPictureUiState implements android.os.Parcelable {
diff --git a/core/java/android/app/PictureInPictureParams.java b/core/java/android/app/PictureInPictureParams.java
index 96d874e..afe915e 100644
--- a/core/java/android/app/PictureInPictureParams.java
+++ b/core/java/android/app/PictureInPictureParams.java
@@ -654,6 +654,33 @@
                 && !hasSetSubtitle() && mIsLaunchIntoPip == null;
     }
 
+    /**
+     * Compare a given {@link Rect} against the aspect ratio, with rounding error tolerance.
+     * @param bounds The {@link Rect} represents the source rect hint, this check is not needed
+     *               if app provides a null source rect hint.
+     * @param aspectRatio {@link Rational} representation of aspect ratio, this check is not needed
+     *                    if app provides a null aspect ratio.
+     * @return {@code true} if the given {@link Rect} matches the aspect ratio.
+     * @hide
+     */
+    @SuppressWarnings("UnflaggedApi")
+    @TestApi
+    public static boolean isSameAspectRatio(@NonNull Rect bounds, @NonNull Rational aspectRatio) {
+        // Validations
+        if (bounds.isEmpty() || aspectRatio.floatValue() <= 0) {
+            return false;
+        }
+        // Check against both the width and height.
+        final int exactWidth = (aspectRatio.getNumerator() * bounds.height())
+                / aspectRatio.getDenominator();
+        if (Math.abs(exactWidth - bounds.width()) <= 1) {
+            return true;
+        }
+        final int exactHeight = (aspectRatio.getDenominator() * bounds.width())
+                / aspectRatio.getNumerator();
+        return Math.abs(exactHeight - bounds.height()) <= 1;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
diff --git a/core/java/android/os/AggregateBatteryConsumer.java b/core/java/android/os/AggregateBatteryConsumer.java
index 67e2195..c7f8878 100644
--- a/core/java/android/os/AggregateBatteryConsumer.java
+++ b/core/java/android/os/AggregateBatteryConsumer.java
@@ -55,7 +55,7 @@
 
     @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 744f6a8..2447ff9 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -19,14 +19,19 @@
 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 @@
         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 @@
                 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 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 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));
             }
-            return mShortString;
+            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 sb.toString();
+        }
+
+        @Override
+        public String toString() {
+            return toString(powerComponent, processState, screenState, powerState);
         }
     }
 
@@ -335,11 +471,18 @@
     }
 
     /**
+     * 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 @@
      * 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 @@
      * @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 @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 @@
     }
 
     /**
+     * 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 @@
             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);
-                }
+        boolean hasValue(int columnIndex) {
+            if (mCursorRow == -1) {
+                return false;
             }
-            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;
-                    }
-                }
-            }
-            return null;
+            return mCursorWindow.getType(mCursorRow, columnIndex) != Cursor.FIELD_TYPE_NULL;
         }
 
         void putInt(int columnIndex, int value) {
@@ -693,91 +822,44 @@
         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> perComponentKeys = new ArrayList<>();
-            for (int componentId = 0; componentId < POWER_COMPONENT_COUNT; componentId++) {
-                perComponentKeys.clear();
-
-                // Declare the Key for the power component, ignoring other dimensions.
-                perComponentKeys.add(
-                        new Key(componentId, PROCESS_STATE_ANY,
-                                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 = 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
-                                    ));
-                        }
-                    }
+            ArrayList<Key> keyList = new ArrayList<>();
+            for (int screenState = 0; screenState < SCREEN_STATE_COUNT; screenState++) {
+                if (!includeScreenState && screenState != SCREEN_STATE_UNSPECIFIED) {
+                    continue;
                 }
-
-                keys[componentId] = perComponentKeys.toArray(KEY_ARRAY);
-            }
-
-            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) {
+                for (int powerState = 0; powerState < POWER_STATE_COUNT; powerState++) {
+                    if (!includePowerState && powerState != POWER_STATE_UNSPECIFIED) {
                         continue;
                     }
-
-                    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]);
-                            }
-                        }
+                    for (int componentId = 0; componentId < POWER_COMPONENT_COUNT; componentId++) {
+                        columnIndex = addKeys(keyList, powerModelsIncluded, includeProcessStateData,
+                                componentId, screenState, powerState, columnIndex);
                     }
-                    processStateKeys[processState] = perProcStateKeys.toArray(KEY_ARRAY);
                 }
-            } else {
-                processStateKeys = null;
             }
 
             firstCustomConsumedPowerColumn = columnIndex;
@@ -787,19 +869,111 @@
             columnIndex += customPowerComponentCount;
 
             columnCount = columnIndex;
+
+            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);
+            }
+        }
+
+        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
+                                columnIndex++,      // power
+                                columnIndex++       // usage duration
+                        ));
+                    }
+                }
+            }
+            return columnIndex;
+        }
+
+        Key getKey(@PowerComponent int componentId, @ProcessState int processState,
+                @ScreenState int screenState, @PowerState int powerState) {
+            return indexedKeys.get(keyIndex(componentId, processState, screenState, powerState));
+        }
+
+        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);
+                        }
+                    }
+                    componentKeys = out.toArray(new Key[out.size()]);
+                    mPerComponentKeys[componentId] = componentKeys;
+                }
+                return componentKeys;
+            }
+        }
+
+        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 @@
 
         @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 61cc23d..dd484f6 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -102,9 +102,13 @@
     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 @@
     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 @@
         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 @@
         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 @@
         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 @@
         dest.writeStringArray(mCustomPowerComponentNames);
         dest.writeBoolean(mIncludesPowerModels);
         dest.writeBoolean(mIncludesProcessStateData);
+        dest.writeBoolean(mIncludesScreenStateData);
+        dest.writeBoolean(mIncludesPowerStateData);
 
         mBatteryConsumersCursorWindow.writeToParcel(dest, flags);
 
@@ -598,23 +612,16 @@
 
         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 @@
                     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 @@
     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 @@
             BatteryStats.formatTimeMs(sb, durationMs);
         }
 
-        pw.println(sb.toString());
+        pw.println(sb);
     }
 
     private void dumpSortedBatteryConsumers(PrintWriter pw, String prefix,
@@ -670,9 +730,8 @@
                 continue;
             }
             pw.print(prefix);
-            pw.print("    ");
+            pw.print("  ");
             consumer.dump(pw);
-            pw.println();
         }
     }
 
@@ -686,6 +745,10 @@
         }
         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 @@
 
                 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 @@
         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 @@
         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 @@
             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 203ef47..d0ed297 100644
--- a/core/java/android/os/BatteryUsageStatsQuery.java
+++ b/core/java/android/os/BatteryUsageStatsQuery.java
@@ -73,6 +73,10 @@
 
     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 @@
         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 @@
         }
 
         /**
+         * 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 b035f12..f22e1ea 100644
--- a/core/java/android/os/PowerComponents.java
+++ b/core/java/android/os/PowerComponents.java
@@ -17,8 +17,12 @@
 
 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,24 +60,101 @@
      * 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);
-            }
-            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 totalPowerMah;
-        } else {
+        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);
+            }
+            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);
+                }
+            }
+            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 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 @@
      * @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 @@
      * @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 @@
         }
     }
 
-    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) {
+            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);
+                }
+            }
+        }
+
+        // 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(separator);
-                separator = " ";
-                sb.append(key.toShortString());
+                sb.append(getCustomPowerComponentName(customComponentId));
                 sb.append("=");
-                sb.append(BatteryStats.formatCharge(componentPower));
-
-                if (durationMs != 0) {
-                    sb.append(" (");
-                    BatteryStats.formatTimeMsNoSpace(sb, durationMs);
-                    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 @@
 
         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 @@
 
     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;
-                }
-
-                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);
+        for (BatteryConsumer.Key key : mData.layout.keys) {
+            if (!mData.hasValue(key.mPowerColumnIndex)
+                    && !mData.hasValue(key.mDurationColumnIndex)) {
+                continue;
             }
+
+            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 @@
                     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 @@
                                 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 @@
                                     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 @@
         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 @@
                                 + ", 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;
-                        }
-                    }
+            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;
+                }
 
-                    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 @@
                 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 @@
             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 @@
          */
         @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/TEST_MAPPING b/core/java/android/os/TEST_MAPPING
index b5029a6..2fde5e7 100644
--- a/core/java/android/os/TEST_MAPPING
+++ b/core/java/android/os/TEST_MAPPING
@@ -119,7 +119,7 @@
         "PowerComponents\\.java",
         "[^/]*BatteryConsumer[^/]*\\.java"
       ],
-      "name": "BatteryUsageStatsProtoTests"
+      "name": "PowerStatsTests"
     },
     {
       "file_patterns": ["SharedMemory[^/]*\\.java"],
diff --git a/core/java/android/os/UidBatteryConsumer.java b/core/java/android/os/UidBatteryConsumer.java
index 53af838..9b5a378 100644
--- a/core/java/android/os/UidBatteryConsumer.java
+++ b/core/java/android/os/UidBatteryConsumer.java
@@ -140,12 +140,50 @@
                     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 @@
                 .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 23ba0c6..ea2be7b 100644
--- a/core/java/android/os/UserBatteryConsumer.java
+++ b/core/java/android/os/UserBatteryConsumer.java
@@ -60,10 +60,10 @@
         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/core/java/com/android/internal/os/TEST_MAPPING b/core/java/com/android/internal/os/TEST_MAPPING
index d552e0b..ae43acf 100644
--- a/core/java/com/android/internal/os/TEST_MAPPING
+++ b/core/java/com/android/internal/os/TEST_MAPPING
@@ -18,7 +18,7 @@
         "Kernel[^/]*\\.java",
         "[^/]*Power[^/]*\\.java"
       ],
-      "name": "BatteryUsageStatsProtoTests"
+      "name": "PowerStatsTests"
     },
     {
       "file_patterns": [
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/CursorAnchorInfoTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/CursorAnchorInfoTest.java
index 9d7d71d..616c72e 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/CursorAnchorInfoTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/CursorAnchorInfoTest.java
@@ -20,8 +20,8 @@
 
 import android.graphics.Matrix;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/DeleteRangeGestureTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/DeleteRangeGestureTest.java
index d7b911d..d0a4141 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/DeleteRangeGestureTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/DeleteRangeGestureTest.java
@@ -22,8 +22,8 @@
 import android.graphics.RectF;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.compatibility.common.util.ApiTest;
 
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/EditorInfoTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/EditorInfoTest.java
index 4839dd2..013117e 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/EditorInfoTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/EditorInfoTest.java
@@ -45,8 +45,8 @@
 import android.view.MotionEvent;
 import android.view.autofill.AutofillId;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodInfoTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodInfoTest.java
index ce85a76..61bf137 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodInfoTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodInfoTest.java
@@ -28,9 +28,9 @@
 import android.os.Parcel;
 import android.platform.test.flag.junit.SetFlagsRule;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.frameworks.inputmethodcoretests.R;
 
@@ -135,7 +135,7 @@
 
     private InputMethodInfo buildInputMethodForTest(final @XmlRes int metaDataRes)
             throws Exception {
-        final Context context = InstrumentationRegistry.getContext();
+        final Context context = InstrumentationRegistry.getInstrumentation().getContext();
         final ServiceInfo serviceInfo = new ServiceInfo();
         serviceInfo.applicationInfo = context.getApplicationInfo();
         serviceInfo.packageName = context.getPackageName();
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodManagerTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodManagerTest.java
index d705724..812b3f5 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodManagerTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodManagerTest.java
@@ -24,9 +24,9 @@
 import android.content.Context;
 import android.hardware.display.DisplayManager;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeArrayTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeArrayTest.java
index e7b1110..73ff304 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeArrayTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeArrayTest.java
@@ -26,8 +26,8 @@
 import android.platform.test.annotations.Presubmit;
 import android.view.inputmethod.InputMethodSubtype.InputMethodSubtypeBuilder;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeTest.java
index 5095cad..4c76992 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InputMethodSubtypeTest.java
@@ -27,8 +27,8 @@
 import android.os.Parcel;
 import android.view.inputmethod.InputMethodSubtype.InputMethodSubtypeBuilder;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertGestureTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertGestureTest.java
index 47a724d..608dd4d 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertGestureTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertGestureTest.java
@@ -22,8 +22,8 @@
 import android.graphics.PointF;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.compatibility.common.util.ApiTest;
 
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertModeGestureTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertModeGestureTest.java
index a94f877..bb6a944b 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertModeGestureTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/InsertModeGestureTest.java
@@ -24,8 +24,8 @@
 import android.os.CancellationSignalBeamer;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.compatibility.common.util.ApiTest;
 
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectGestureTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectGestureTest.java
index b2eb07c..4cbd7ab 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectGestureTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectGestureTest.java
@@ -22,8 +22,8 @@
 import android.graphics.RectF;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.compatibility.common.util.ApiTest;
 
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectRangeGestureTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectRangeGestureTest.java
index df63a4a..c1e2197 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectRangeGestureTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SelectRangeGestureTest.java
@@ -22,8 +22,8 @@
 import android.graphics.RectF;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.compatibility.common.util.ApiTest;
 
diff --git a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SparseRectFArrayTest.java b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SparseRectFArrayTest.java
index f264cc6..724d729 100644
--- a/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SparseRectFArrayTest.java
+++ b/core/tests/InputMethodCoreTests/src/android/view/inputmethod/SparseRectFArrayTest.java
@@ -25,8 +25,8 @@
 import android.os.Parcel;
 import android.view.inputmethod.SparseRectFArray.SparseRectFArrayBuilder;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputConnectionWrapperTest.java b/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputConnectionWrapperTest.java
index a626294..9eb5dd1 100644
--- a/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputConnectionWrapperTest.java
+++ b/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputConnectionWrapperTest.java
@@ -22,8 +22,8 @@
 import android.view.inputmethod.InputConnection;
 import android.view.inputmethod.InputConnectionWrapper;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputMethodDebugTest.java b/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputMethodDebugTest.java
index 32bfdcb..6ac3639 100644
--- a/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputMethodDebugTest.java
+++ b/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/InputMethodDebugTest.java
@@ -20,8 +20,8 @@
 
 import android.view.WindowManager.LayoutParams;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/SubtypeLocaleUtilsTest.java b/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/SubtypeLocaleUtilsTest.java
index ba63908..7b0a5da 100644
--- a/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/SubtypeLocaleUtilsTest.java
+++ b/core/tests/InputMethodCoreTests/src/com/android/internal/inputmethod/SubtypeLocaleUtilsTest.java
@@ -18,8 +18,8 @@
 
 import static org.junit.Assert.assertEquals;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
deleted file mode 100644
index 1fb5f2c..0000000
--- a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
+++ /dev/null
@@ -1,30 +0,0 @@
-package {
-    // See: http://go/android-license-faq
-    // A large-scale-change added 'default_applicable_licenses' to import
-    // all of the 'license_kinds' from "frameworks_base_license"
-    // to get the below license kinds:
-    //   SPDX-license-identifier-Apache-2.0
-    default_applicable_licenses: ["frameworks_base_license"],
-}
-
-android_test {
-    name: "BatteryUsageStatsProtoTests",
-    srcs: ["src/**/*.java"],
-
-    static_libs: [
-        "androidx.test.rules",
-        "junit",
-        "mockito-target-minus-junit4",
-        "platform-test-annotations",
-        "platformprotosnano",
-        "statsdprotolite",
-        "truth",
-    ],
-
-    libs: ["android.test.runner"],
-
-    platform_apis: true,
-    certificate: "platform",
-
-    test_suites: ["device-tests"],
-}
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml
deleted file mode 100644
index 9128dca..0000000
--- a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml
+++ /dev/null
@@ -1,28 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-  ~ Copyright (C) 2021 The Android Open Source Project
-  ~
-  ~ Licensed under the Apache License, Version 2.0 (the "License");
-  ~ you may not use this file except in compliance with the License.
-  ~ You may obtain a copy of the License at
-  ~
-  ~      http://www.apache.org/licenses/LICENSE-2.0
-  ~
-  ~ Unless required by applicable law or agreed to in writing, software
-  ~ distributed under the License is distributed on an "AS IS" BASIS,
-  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-  ~ See the License for the specific language governing permissions and
-  ~ limitations under the License.
-  -->
-
-<manifest xmlns:android="http://schemas.android.com/apk/res/android"
-          package="com.android.frameworks.core.batteryusagestatsprototests">
-
-    <uses-permission android:name="android.permission.BATTERY_STATS"/>
-
-    <instrumentation
-        android:name="androidx.test.runner.AndroidJUnitRunner"
-        android:targetPackage="com.android.frameworks.core.batteryusagestatsprototests"
-        android:label="BatteryUsageStats Proto Tests" />
-
-</manifest>
diff --git a/core/tests/coretests/src/android/inputmethodservice/ImsConfigurationTrackerTest.java b/core/tests/coretests/src/android/inputmethodservice/ImsConfigurationTrackerTest.java
index 6998c32..f5e9cc6 100644
--- a/core/tests/coretests/src/android/inputmethodservice/ImsConfigurationTrackerTest.java
+++ b/core/tests/coretests/src/android/inputmethodservice/ImsConfigurationTrackerTest.java
@@ -26,8 +26,8 @@
 import android.content.res.Configuration;
 import android.platform.test.annotations.RequiresFlagsEnabled;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/location/java/android/location/flags/location.aconfig b/location/java/android/location/flags/location.aconfig
index acfe473..0edaaef 100644
--- a/location/java/android/location/flags/location.aconfig
+++ b/location/java/android/location/flags/location.aconfig
@@ -2,6 +2,20 @@
 container: "system"
 
 flag {
+    name: "keep_gnss_stationary_throttling"
+    namespace: "location"
+    description: "Keeps stationary throttling for the GNSS provider even if the disable_stationary_throttling flag is true."
+    bug: "354000147"
+}
+
+flag {
+    name: "disable_stationary_throttling"
+    namespace: "location"
+    description: "Disables stationary throttling for all providers"
+    bug: "354000147"
+}
+
+flag {
     name: "new_geocoder"
     namespace: "location"
     description: "Flag for new Geocoder APIs"
diff --git a/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodPreferenceTest.java b/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodPreferenceTest.java
index 1e75014..3906749 100644
--- a/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodPreferenceTest.java
+++ b/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodPreferenceTest.java
@@ -24,9 +24,9 @@
 import android.view.inputmethod.InputMethodInfo;
 import android.view.inputmethod.InputMethodSubtype;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -109,7 +109,7 @@
             final boolean systemIme,
             final String name) {
         return new InputMethodPreference(
-                InstrumentationRegistry.getTargetContext(),
+                InstrumentationRegistry.getInstrumentation().getTargetContext(),
                 createInputMethodInfo(systemIme, name),
                 title,
                 true /* isAllowedByOrganization */,
@@ -119,7 +119,8 @@
 
     private static InputMethodInfo createInputMethodInfo(
             final boolean systemIme, final String name) {
-        final Context targetContext = InstrumentationRegistry.getTargetContext();
+        final Context targetContext =
+                InstrumentationRegistry.getInstrumentation().getTargetContext();
         final Locale systemLocale = targetContext
                 .getResources()
                 .getConfiguration()
diff --git a/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodSubtypePreferenceTest.java b/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodSubtypePreferenceTest.java
index f1c0bea..2c3478d 100644
--- a/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodSubtypePreferenceTest.java
+++ b/packages/SettingsLib/tests/integ/src/com/android/settingslib/inputmethod/InputMethodSubtypePreferenceTest.java
@@ -18,9 +18,9 @@
 
 import android.text.TextUtils;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -104,7 +104,7 @@
         final Locale subtypeLocale = TextUtils.isEmpty(subtypeLanguageTag)
                 ? null : Locale.forLanguageTag(subtypeLanguageTag);
         return new InputMethodSubtypePreference(
-                InstrumentationRegistry.getTargetContext(),
+                InstrumentationRegistry.getInstrumentation().getTargetContext(),
                 key,
                 subtypeName,
                 subtypeLocale,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
index ce134e6..75ecb2c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
@@ -41,8 +41,12 @@
 import com.android.systemui.communal.shared.model.CommunalScenes;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.keyguard.shared.model.KeyguardState;
+import com.android.systemui.keyguard.shared.model.TransitionState;
+import com.android.systemui.keyguard.shared.model.TransitionStep;
 import com.android.systemui.kosmos.KosmosJavaAdapter;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
+import com.android.systemui.scene.shared.flag.SceneContainerFlag;
 import com.android.systemui.shade.data.repository.FakeShadeRepository;
 import com.android.systemui.shade.data.repository.ShadeAnimationRepository;
 import com.android.systemui.shade.data.repository.ShadeRepository;
@@ -91,6 +95,7 @@
     @Mock private HeadsUpManager mHeadsUpManager;
     @Mock private VisibilityLocationProvider mVisibilityLocationProvider;
     @Mock private VisualStabilityProvider mVisualStabilityProvider;
+    @Mock private VisualStabilityCoordinatorLogger mLogger;
 
     @Captor private ArgumentCaptor<WakefulnessLifecycle.Observer> mWakefulnessObserverCaptor;
     @Captor private ArgumentCaptor<StatusBarStateController.StateListener> mSBStateListenerCaptor;
@@ -128,7 +133,9 @@
                 mVisibilityLocationProvider,
                 mVisualStabilityProvider,
                 mWakefulnessLifecycle,
-                mKosmos.getCommunalInteractor());
+                mKosmos.getCommunalInteractor(),
+                mKosmos.getKeyguardTransitionInteractor(),
+                mLogger);
         mCoordinator.attach(mNotifPipeline);
         mTestScope.getTestScheduler().runCurrent();
 
@@ -241,6 +248,38 @@
     }
 
     @Test
+    public void testLockscreenPartlyShowing_groupAndSectionChangesNotAllowed() {
+        // GIVEN the panel true expanded and device isn't pulsing
+        setFullyDozed(false);
+        setSleepy(false);
+        setLockscreenShowing(0.5f);
+        setPulsing(false);
+
+        // THEN group changes are NOT allowed
+        assertFalse(mNotifStabilityManager.isGroupChangeAllowed(mEntry));
+        assertFalse(mNotifStabilityManager.isGroupPruneAllowed(mGroupEntry));
+
+        // THEN section changes are NOT allowed
+        assertFalse(mNotifStabilityManager.isSectionChangeAllowed(mEntry));
+    }
+
+    @Test
+    public void testLockscreenFullyShowing_groupAndSectionChangesNotAllowed() {
+        // GIVEN the panel true expanded and device isn't pulsing
+        setFullyDozed(false);
+        setSleepy(false);
+        setLockscreenShowing(1.0f);
+        setPulsing(false);
+
+        // THEN group changes are NOT allowed
+        assertFalse(mNotifStabilityManager.isGroupChangeAllowed(mEntry));
+        assertFalse(mNotifStabilityManager.isGroupPruneAllowed(mGroupEntry));
+
+        // THEN section changes are NOT allowed
+        assertFalse(mNotifStabilityManager.isSectionChangeAllowed(mEntry));
+    }
+
+    @Test
     public void testPulsing_screenOff_groupAndSectionChangesNotAllowed() {
         // GIVEN the device is pulsing and screen is off
         setFullyDozed(true);
@@ -614,7 +653,37 @@
     }
 
     private void setPanelExpanded(boolean expanded) {
-        mStatusBarStateListener.onExpandedChanged(expanded);
+        setPanelExpandedAndLockscreenShowing(expanded, /* lockscreenShowing = */ 0.0f);
     }
 
+    private void setLockscreenShowing(float lockscreenShowing) {
+        setPanelExpandedAndLockscreenShowing(/* panelExpanded = */ false, lockscreenShowing);
+    }
+
+    private void setPanelExpandedAndLockscreenShowing(boolean panelExpanded,
+            float lockscreenShowing) {
+        if (SceneContainerFlag.isEnabled()) {
+            mStatusBarStateListener.onExpandedChanged(panelExpanded);
+            mKosmos.getKeyguardTransitionRepository().sendTransitionStepJava(
+                    mTestScope,
+                    makeLockscreenTransitionStep(lockscreenShowing),
+                    /* validateStep = */ false);
+        } else {
+            mStatusBarStateListener.onExpandedChanged(panelExpanded || lockscreenShowing > 0.0f);
+        }
+    }
+
+    private TransitionStep makeLockscreenTransitionStep(float value) {
+        if (value <= 0.0f) {
+            return new TransitionStep(KeyguardState.GONE);
+        } else if (value >= 1.0f) {
+            return new TransitionStep(KeyguardState.LOCKSCREEN);
+        } else {
+            return new TransitionStep(
+                    KeyguardState.GONE,
+                    KeyguardState.LOCKSCREEN,
+                    value,
+                    TransitionState.RUNNING);
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
index d1c9b8e..b2ba0e1 100644
--- a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
@@ -169,6 +169,14 @@
         return factory.create("NotifRemoteInputLog", 50 /* maxSize */, false /* systrace */);
     }
 
+    /** Provides a logging buffer for all logs related to notification visual stability. */
+    @Provides
+    @SysUISingleton
+    @VisualStabilityLog
+    public static LogBuffer provideVisualStabilityLogBuffer(LogBufferFactory factory) {
+        return factory.create("VisualStabilityLog", 50 /* maxSize */, false /* systrace */);
+    }
+
     /** Provides a logging buffer for all logs related to keyguard media controller. */
     @Provides
     @SysUISingleton
diff --git a/packages/SystemUI/src/com/android/systemui/log/dagger/VisualStabilityLog.kt b/packages/SystemUI/src/com/android/systemui/log/dagger/VisualStabilityLog.kt
new file mode 100644
index 0000000..b45ffc1
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/VisualStabilityLog.kt
@@ -0,0 +1,25 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.log.dagger
+
+import javax.inject.Qualifier
+
+/** A [com.android.systemui.log.LogBuffer] for visual stability-related messages. */
+@Qualifier
+@MustBeDocumented
+@Retention(AnnotationRetention.RUNTIME)
+annotation class VisualStabilityLog
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
index 71c98b8..696298e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
@@ -18,8 +18,6 @@
 
 import static com.android.systemui.keyguard.WakefulnessLifecycle.WAKEFULNESS_ASLEEP;
 
-import android.util.Log;
-
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 
@@ -29,7 +27,10 @@
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
+import com.android.systemui.keyguard.shared.model.KeyguardState;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
+import com.android.systemui.scene.shared.flag.SceneContainerFlag;
 import com.android.systemui.shade.domain.interactor.ShadeAnimationInteractor;
 import com.android.systemui.statusbar.notification.VisibilityLocationProvider;
 import com.android.systemui.statusbar.notification.collection.GroupEntry;
@@ -41,7 +42,6 @@
 import com.android.systemui.statusbar.notification.domain.interactor.SeenNotificationsInteractor;
 import com.android.systemui.statusbar.notification.shared.NotificationMinimalismPrototype;
 import com.android.systemui.statusbar.policy.HeadsUpManager;
-import com.android.systemui.util.Compile;
 import com.android.systemui.util.concurrency.DelayableExecutor;
 import com.android.systemui.util.kotlin.JavaAdapter;
 
@@ -61,8 +61,6 @@
 // TODO(b/204468557): Move to @CoordinatorScope
 @SysUISingleton
 public class VisualStabilityCoordinator implements Coordinator, Dumpable {
-    public static final String TAG = "VisualStability";
-    public static final boolean DEBUG = Compile.IS_DEBUG && Log.isLoggable(TAG, Log.VERBOSE);
     private final DelayableExecutor mDelayableExecutor;
     private final HeadsUpManager mHeadsUpManager;
     private final SeenNotificationsInteractor mSeenNotificationsInteractor;
@@ -73,6 +71,8 @@
     private final VisualStabilityProvider mVisualStabilityProvider;
     private final WakefulnessLifecycle mWakefulnessLifecycle;
     private final CommunalInteractor mCommunalInteractor;
+    private final KeyguardTransitionInteractor mKeyguardTransitionInteractor;
+    private final VisualStabilityCoordinatorLogger mLogger;
 
     private boolean mSleepy = true;
     private boolean mFullyDozed;
@@ -81,6 +81,7 @@
     private boolean mNotifPanelCollapsing;
     private boolean mNotifPanelLaunchingActivity;
     private boolean mCommunalShowing = false;
+    private boolean mLockscreenShowing = false;
 
     private boolean mPipelineRunAllowed;
     private boolean mReorderingAllowed;
@@ -109,7 +110,9 @@
             VisibilityLocationProvider visibilityLocationProvider,
             VisualStabilityProvider visualStabilityProvider,
             WakefulnessLifecycle wakefulnessLifecycle,
-            CommunalInteractor communalInteractor) {
+            CommunalInteractor communalInteractor,
+            KeyguardTransitionInteractor keyguardTransitionInteractor,
+            VisualStabilityCoordinatorLogger logger) {
         mHeadsUpManager = headsUpManager;
         mShadeAnimationInteractor = shadeAnimationInteractor;
         mJavaAdapter = javaAdapter;
@@ -120,6 +123,8 @@
         mStatusBarStateController = statusBarStateController;
         mDelayableExecutor = delayableExecutor;
         mCommunalInteractor = communalInteractor;
+        mKeyguardTransitionInteractor = keyguardTransitionInteractor;
+        mLogger = logger;
 
         dumpManager.registerDumpable(this);
     }
@@ -138,6 +143,9 @@
                 this::onLaunchingActivityChanged);
         mJavaAdapter.alwaysCollectFlow(mCommunalInteractor.isIdleOnCommunal(),
                 this::onCommunalShowingChanged);
+        mJavaAdapter.alwaysCollectFlow(mKeyguardTransitionInteractor.transitionValue(
+                        KeyguardState.LOCKSCREEN),
+                this::onLockscreenKeyguardStateTransitionValueChanged);
 
         pipeline.setVisualStabilityManager(mNotifStabilityManager);
     }
@@ -221,12 +229,12 @@
         boolean wasReorderingAllowed = mReorderingAllowed;
         mPipelineRunAllowed = !isPanelCollapsingOrLaunchingActivity();
         mReorderingAllowed = isReorderingAllowed();
-        if (DEBUG && (wasPipelineRunAllowed != mPipelineRunAllowed
-                || wasReorderingAllowed != mReorderingAllowed)) {
-            Log.d(TAG, "Stability allowances changed:"
-                    + "  pipelineRunAllowed " + wasPipelineRunAllowed + "->" + mPipelineRunAllowed
-                    + "  reorderingAllowed " + wasReorderingAllowed + "->" + mReorderingAllowed
-                    + "  when setting " + field + "=" + value);
+        if (wasPipelineRunAllowed != mPipelineRunAllowed
+                || wasReorderingAllowed != mReorderingAllowed) {
+            mLogger.logAllowancesChanged(
+                    wasPipelineRunAllowed, mPipelineRunAllowed,
+                    wasReorderingAllowed, mReorderingAllowed,
+                    field, value);
         }
         if (mPipelineRunAllowed && mIsSuppressingPipelineRun) {
             mNotifStabilityManager.invalidateList("pipeline run suppression ended");
@@ -251,7 +259,9 @@
     }
 
     private boolean isReorderingAllowed() {
-        return ((mFullyDozed && mSleepy) || !mPanelExpanded || mCommunalShowing) && !mPulsing;
+        final boolean sleepyAndDozed = mFullyDozed && mSleepy;
+        final boolean stackShowing = mPanelExpanded || mLockscreenShowing;
+        return (sleepyAndDozed || !stackShowing || mCommunalShowing) && !mPulsing;
     }
 
     /**
@@ -364,4 +374,14 @@
         mCommunalShowing = isShowing;
         updateAllowedStates("communalShowing", isShowing);
     }
+
+    private void onLockscreenKeyguardStateTransitionValueChanged(float value) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
+            return;
+        }
+
+        final boolean isShowing = value > 0.0f;
+        mLockscreenShowing = isShowing;
+        updateAllowedStates("lockscreenShowing", isShowing);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorLogger.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorLogger.kt
new file mode 100644
index 0000000..fe23e4e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorLogger.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.notification.collection.coordinator
+
+import com.android.systemui.log.LogBuffer
+import com.android.systemui.log.core.LogLevel
+import com.android.systemui.log.dagger.VisualStabilityLog
+import javax.inject.Inject
+
+private const val TAG = "VisualStability"
+
+class VisualStabilityCoordinatorLogger
+@Inject
+constructor(@VisualStabilityLog private val buffer: LogBuffer) {
+    fun logAllowancesChanged(
+        wasRunAllowed: Boolean,
+        isRunAllowed: Boolean,
+        wasReorderingAllowed: Boolean,
+        isReorderingAllowed: Boolean,
+        field: String,
+        value: Boolean
+    ) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                bool1 = wasRunAllowed
+                bool2 = isRunAllowed
+                bool3 = wasReorderingAllowed
+                bool4 = isReorderingAllowed
+                str1 = field
+                str2 = value.toString()
+            },
+            {
+                "stability allowances changed:" +
+                    " pipelineRunAllowed $bool1->$bool2" +
+                    " reorderingAllowed $bool3->$bool4" +
+                    " when setting $str1=$str2"
+            }
+        )
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
index 677d1fd..6dcea14 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
@@ -449,96 +449,115 @@
     }
 
     @Test
-    fun shows_authenticated_no_errors_no_confirmation_required() = runGenericTest {
+    fun shows_error_to_unlock_or_success() {
+        // Face-only auth does not use error -> unlock or error -> success assets
+        if (testCase.isFingerprintOnly || testCase.isCoex) {
+            runGenericTest {
+                // Distinct asset for error -> success only applicable for fingerprint-only /
+                // explicit co-ex auth
+                val iconAsset by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
+                val iconContentDescriptionId by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.contentDescriptionId)
+                val shouldAnimateIconView by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconView)
+
+                var forceExplicitFlow =
+                    testCase.isCoex && testCase.confirmationRequested ||
+                        testCase.authenticatedByFingerprint
+                if (forceExplicitFlow) {
+                    kosmos.promptViewModel.ensureFingerprintHasStarted(isDelayed = true)
+                }
+                verifyIconSize(forceExplicitFlow)
+
+                kosmos.promptViewModel.ensureFingerprintHasStarted(isDelayed = true)
+                kosmos.promptViewModel.iconViewModel.setPreviousIconWasError(true)
+
+                kosmos.promptViewModel.showAuthenticated(
+                    modality = testCase.authenticatedModality,
+                    dismissAfterDelay = DELAY
+                )
+
+                // TODO(b/350121748): SFPS test cases to be added after SFPS assets update
+                if (testCase.sensorType != FingerprintSensorProperties.TYPE_POWER_BUTTON) {
+                    // Non-SFPS (UDFPS / rear-FPS) test cases
+                    // Covers (1) fingerprint-only (2) co-ex, authenticated by fingerprint
+                    if (testCase.authenticatedByFingerprint) {
+                        assertThat(iconAsset)
+                            .isEqualTo(R.raw.fingerprint_dialogue_error_to_success_lottie)
+                        assertThat(iconContentDescriptionId)
+                            .isEqualTo(R.string.fingerprint_dialog_touch_sensor)
+                        assertThat(shouldAnimateIconView).isEqualTo(true)
+                    } else { //  co-ex, authenticated by face
+                        assertThat(iconAsset)
+                            .isEqualTo(R.raw.fingerprint_dialogue_error_to_unlock_lottie)
+                        assertThat(iconContentDescriptionId)
+                            .isEqualTo(R.string.fingerprint_dialog_authenticated_confirmation)
+                        assertThat(shouldAnimateIconView).isEqualTo(true)
+
+                        // Confirm authentication
+                        kosmos.promptViewModel.confirmAuthenticated()
+
+                        assertThat(iconAsset)
+                            .isEqualTo(
+                                R.raw.fingerprint_dialogue_unlocked_to_checkmark_success_lottie
+                            )
+                        assertThat(iconContentDescriptionId)
+                            .isEqualTo(R.string.fingerprint_dialog_touch_sensor)
+                        assertThat(shouldAnimateIconView).isEqualTo(true)
+                    }
+                }
+            }
+        }
+    }
+
+    @Test
+    fun shows_authenticated_no_errors_no_confirmation_required() {
         if (!testCase.confirmationRequested) {
-            val iconAsset by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
-            val iconOverlayAsset by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.iconOverlayAsset)
-            val iconContentDescriptionId by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.contentDescriptionId)
-            val shouldAnimateIconView by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconView)
-            val shouldAnimateIconOverlay by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconOverlay)
-            verifyIconSize()
+            runGenericTest {
+                val iconAsset by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
+                val iconOverlayAsset by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.iconOverlayAsset)
+                val iconContentDescriptionId by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.contentDescriptionId)
+                val shouldAnimateIconView by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconView)
+                val shouldAnimateIconOverlay by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconOverlay)
+                verifyIconSize()
 
-            kosmos.promptViewModel.showAuthenticated(
-                modality = testCase.authenticatedModality,
-                dismissAfterDelay = DELAY
-            )
+                kosmos.promptViewModel.showAuthenticated(
+                    modality = testCase.authenticatedModality,
+                    dismissAfterDelay = DELAY
+                )
 
-            if (testCase.isFingerprintOnly) {
-                // Fingerprint icon asset assertions
-                if (testCase.sensorType == FingerprintSensorProperties.TYPE_POWER_BUTTON) {
-                    assertThat(iconAsset).isEqualTo(getSfpsBaseIconAsset())
-                    assertThat(iconOverlayAsset)
-                        .isEqualTo(R.raw.biometricprompt_symbol_fingerprint_to_success_landscape)
-                    assertThat(iconContentDescriptionId)
-                        .isEqualTo(R.string.security_settings_sfps_enroll_find_sensor_message)
-                    assertThat(shouldAnimateIconView).isEqualTo(true)
-                    assertThat(shouldAnimateIconOverlay).isEqualTo(true)
-                } else {
-                    assertThat(iconAsset)
-                        .isEqualTo(R.raw.fingerprint_dialogue_fingerprint_to_success_lottie)
+                if (testCase.isFingerprintOnly) {
+                    // Fingerprint icon asset assertions
+                    if (testCase.sensorType == FingerprintSensorProperties.TYPE_POWER_BUTTON) {
+                        assertThat(iconAsset).isEqualTo(getSfpsBaseIconAsset())
+                        assertThat(iconOverlayAsset)
+                            .isEqualTo(
+                                R.raw.biometricprompt_symbol_fingerprint_to_success_landscape
+                            )
+                        assertThat(iconContentDescriptionId)
+                            .isEqualTo(R.string.security_settings_sfps_enroll_find_sensor_message)
+                        assertThat(shouldAnimateIconView).isEqualTo(true)
+                        assertThat(shouldAnimateIconOverlay).isEqualTo(true)
+                    } else {
+                        assertThat(iconAsset)
+                            .isEqualTo(R.raw.fingerprint_dialogue_fingerprint_to_success_lottie)
+                        assertThat(iconOverlayAsset).isEqualTo(-1)
+                        assertThat(iconContentDescriptionId)
+                            .isEqualTo(R.string.fingerprint_dialog_touch_sensor)
+                        assertThat(shouldAnimateIconView).isEqualTo(true)
+                        assertThat(shouldAnimateIconOverlay).isEqualTo(false)
+                    }
+                } else if (testCase.isFaceOnly || testCase.isCoex) {
+                    // Face icon asset assertions
+                    // If co-ex, use implicit flow (explicit flow always requires confirmation)
+                    assertThat(iconAsset).isEqualTo(R.raw.face_dialog_dark_to_checkmark)
                     assertThat(iconOverlayAsset).isEqualTo(-1)
                     assertThat(iconContentDescriptionId)
-                        .isEqualTo(R.string.fingerprint_dialog_touch_sensor)
-                    assertThat(shouldAnimateIconView).isEqualTo(true)
-                    assertThat(shouldAnimateIconOverlay).isEqualTo(false)
-                }
-            } else if (testCase.isFaceOnly || testCase.isCoex) {
-                // Face icon asset assertions
-                // If co-ex, use implicit flow (explicit flow always requires confirmation)
-                assertThat(iconAsset).isEqualTo(R.raw.face_dialog_dark_to_checkmark)
-                assertThat(iconOverlayAsset).isEqualTo(-1)
-                assertThat(iconContentDescriptionId)
-                    .isEqualTo(R.string.biometric_dialog_face_icon_description_authenticated)
-                assertThat(shouldAnimateIconView).isEqualTo(true)
-                assertThat(shouldAnimateIconOverlay).isEqualTo(false)
-            }
-        }
-    }
-
-    @Test
-    fun shows_pending_confirmation() = runGenericTest {
-        if (
-            (testCase.isFaceOnly || testCase.isCoex) &&
-                testCase.authenticatedByFace &&
-                testCase.confirmationRequested
-        ) {
-            val iconAsset by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
-            val iconOverlayAsset by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.iconOverlayAsset)
-            val iconContentDescriptionId by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.contentDescriptionId)
-            val shouldAnimateIconView by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconView)
-            val shouldAnimateIconOverlay by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconOverlay)
-
-            val forceExplicitFlow = testCase.isCoex && testCase.confirmationRequested
-            verifyIconSize(forceExplicitFlow = forceExplicitFlow)
-
-            kosmos.promptViewModel.showAuthenticated(
-                modality = testCase.authenticatedModality,
-                dismissAfterDelay = DELAY
-            )
-
-            if (testCase.isFaceOnly) {
-                assertThat(iconAsset).isEqualTo(R.raw.face_dialog_wink_from_dark)
-                assertThat(iconOverlayAsset).isEqualTo(-1)
-                assertThat(iconContentDescriptionId)
-                    .isEqualTo(R.string.biometric_dialog_face_icon_description_authenticated)
-                assertThat(shouldAnimateIconView).isEqualTo(true)
-                assertThat(shouldAnimateIconOverlay).isEqualTo(false)
-            } else if (testCase.isCoex) { // explicit flow, confirmation requested
-                // TODO: Update when SFPS co-ex is implemented
-                if (testCase.sensorType != FingerprintSensorProperties.TYPE_POWER_BUTTON) {
-                    assertThat(iconAsset)
-                        .isEqualTo(R.raw.fingerprint_dialogue_fingerprint_to_unlock_lottie)
-                    assertThat(iconOverlayAsset).isEqualTo(-1)
-                    assertThat(iconContentDescriptionId)
-                        .isEqualTo(R.string.fingerprint_dialog_authenticated_confirmation)
+                        .isEqualTo(R.string.biometric_dialog_face_icon_description_authenticated)
                     assertThat(shouldAnimateIconView).isEqualTo(true)
                     assertThat(shouldAnimateIconOverlay).isEqualTo(false)
                 }
@@ -547,51 +566,96 @@
     }
 
     @Test
-    fun shows_authenticated_explicitly_confirmed_iconUpdate() = runGenericTest {
-        if (
-            (testCase.isFaceOnly || testCase.isCoex) &&
-                testCase.authenticatedByFace &&
-                testCase.confirmationRequested
-        ) {
-            val iconAsset by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
-            val iconOverlayAsset by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.iconOverlayAsset)
-            val iconContentDescriptionId by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.contentDescriptionId)
-            val shouldAnimateIconView by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconView)
-            val shouldAnimateIconOverlay by
-                collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconOverlay)
-            val forceExplicitFlow = testCase.isCoex && testCase.confirmationRequested
-            verifyIconSize(forceExplicitFlow = forceExplicitFlow)
+    fun shows_pending_confirmation() {
+        if (testCase.authenticatedByFace && testCase.confirmationRequested) {
+            runGenericTest {
+                val iconAsset by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
+                val iconOverlayAsset by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.iconOverlayAsset)
+                val iconContentDescriptionId by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.contentDescriptionId)
+                val shouldAnimateIconView by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconView)
+                val shouldAnimateIconOverlay by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconOverlay)
 
-            kosmos.promptViewModel.showAuthenticated(
-                modality = testCase.authenticatedModality,
-                dismissAfterDelay = DELAY
-            )
+                val forceExplicitFlow = testCase.isCoex && testCase.confirmationRequested
+                verifyIconSize(forceExplicitFlow = forceExplicitFlow)
 
-            kosmos.promptViewModel.confirmAuthenticated()
+                kosmos.promptViewModel.showAuthenticated(
+                    modality = testCase.authenticatedModality,
+                    dismissAfterDelay = DELAY
+                )
 
-            if (testCase.isFaceOnly) {
-                assertThat(iconAsset).isEqualTo(R.raw.face_dialog_dark_to_checkmark)
-                assertThat(iconOverlayAsset).isEqualTo(-1)
-                assertThat(iconContentDescriptionId)
-                    .isEqualTo(R.string.biometric_dialog_face_icon_description_confirmed)
-                assertThat(shouldAnimateIconView).isEqualTo(true)
-                assertThat(shouldAnimateIconOverlay).isEqualTo(false)
-            }
-
-            // explicit flow because confirmation requested
-            if (testCase.isCoex) {
-                // TODO: Update when SFPS co-ex is implemented
-                if (testCase.sensorType != FingerprintSensorProperties.TYPE_POWER_BUTTON) {
-                    assertThat(iconAsset)
-                        .isEqualTo(R.raw.fingerprint_dialogue_unlocked_to_checkmark_success_lottie)
+                if (testCase.isFaceOnly) {
+                    assertThat(iconAsset).isEqualTo(R.raw.face_dialog_wink_from_dark)
                     assertThat(iconOverlayAsset).isEqualTo(-1)
                     assertThat(iconContentDescriptionId)
-                        .isEqualTo(R.string.fingerprint_dialog_touch_sensor)
+                        .isEqualTo(R.string.biometric_dialog_face_icon_description_authenticated)
                     assertThat(shouldAnimateIconView).isEqualTo(true)
                     assertThat(shouldAnimateIconOverlay).isEqualTo(false)
+                } else if (testCase.isCoex) { // explicit flow, confirmation requested
+                    // TODO: Update when SFPS co-ex is implemented
+                    if (testCase.sensorType != FingerprintSensorProperties.TYPE_POWER_BUTTON) {
+                        assertThat(iconAsset)
+                            .isEqualTo(R.raw.fingerprint_dialogue_fingerprint_to_unlock_lottie)
+                        assertThat(iconOverlayAsset).isEqualTo(-1)
+                        assertThat(iconContentDescriptionId)
+                            .isEqualTo(R.string.fingerprint_dialog_authenticated_confirmation)
+                        assertThat(shouldAnimateIconView).isEqualTo(true)
+                        assertThat(shouldAnimateIconOverlay).isEqualTo(false)
+                    }
+                }
+            }
+        }
+    }
+
+    @Test
+    fun shows_authenticated_explicitly_confirmed() {
+        if (testCase.authenticatedByFace && testCase.confirmationRequested) {
+            runGenericTest {
+                val iconAsset by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
+                val iconOverlayAsset by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.iconOverlayAsset)
+                val iconContentDescriptionId by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.contentDescriptionId)
+                val shouldAnimateIconView by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconView)
+                val shouldAnimateIconOverlay by
+                    collectLastValue(kosmos.promptViewModel.iconViewModel.shouldAnimateIconOverlay)
+                val forceExplicitFlow = testCase.isCoex && testCase.confirmationRequested
+                verifyIconSize(forceExplicitFlow = forceExplicitFlow)
+
+                kosmos.promptViewModel.showAuthenticated(
+                    modality = testCase.authenticatedModality,
+                    dismissAfterDelay = DELAY
+                )
+
+                kosmos.promptViewModel.confirmAuthenticated()
+
+                if (testCase.isFaceOnly) {
+                    assertThat(iconAsset).isEqualTo(R.raw.face_dialog_dark_to_checkmark)
+                    assertThat(iconOverlayAsset).isEqualTo(-1)
+                    assertThat(iconContentDescriptionId)
+                        .isEqualTo(R.string.biometric_dialog_face_icon_description_confirmed)
+                    assertThat(shouldAnimateIconView).isEqualTo(true)
+                    assertThat(shouldAnimateIconOverlay).isEqualTo(false)
+                }
+
+                // explicit flow because confirmation requested
+                if (testCase.isCoex) {
+                    // TODO: Update when SFPS co-ex is implemented
+                    if (testCase.sensorType != FingerprintSensorProperties.TYPE_POWER_BUTTON) {
+                        assertThat(iconAsset)
+                            .isEqualTo(
+                                R.raw.fingerprint_dialogue_unlocked_to_checkmark_success_lottie
+                            )
+                        assertThat(iconOverlayAsset).isEqualTo(-1)
+                        assertThat(iconContentDescriptionId)
+                            .isEqualTo(R.string.fingerprint_dialog_touch_sensor)
+                        assertThat(shouldAnimateIconView).isEqualTo(true)
+                        assertThat(shouldAnimateIconOverlay).isEqualTo(false)
+                    }
                 }
             }
         }
@@ -700,58 +764,68 @@
     }
 
     @Test
-    fun sfpsIconUpdates_onFoldConfigurationChanged() = runGenericTest {
+    fun sfpsIconUpdates_onFoldConfigurationChanged() {
         if (
             testCase.sensorType == FingerprintSensorProperties.TYPE_POWER_BUTTON &&
                 !testCase.isInRearDisplayMode
         ) {
-            val currentIcon by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
+            runGenericTest {
+                val currentIcon by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
 
-            kosmos.promptViewModel.iconViewModel.onConfigurationChanged(getFoldedConfiguration())
-            val foldedIcon = currentIcon
+                kosmos.promptViewModel.iconViewModel.onConfigurationChanged(
+                    getFoldedConfiguration()
+                )
+                val foldedIcon = currentIcon
 
-            kosmos.promptViewModel.iconViewModel.onConfigurationChanged(getUnfoldedConfiguration())
-            val unfoldedIcon = currentIcon
+                kosmos.promptViewModel.iconViewModel.onConfigurationChanged(
+                    getUnfoldedConfiguration()
+                )
+                val unfoldedIcon = currentIcon
 
-            assertThat(foldedIcon).isNotEqualTo(unfoldedIcon)
+                assertThat(foldedIcon).isNotEqualTo(unfoldedIcon)
+            }
         }
     }
 
     @Test
-    fun sfpsIconUpdates_onRotation() = runGenericTest {
+    fun sfpsIconUpdates_onRotation() {
         if (testCase.sensorType == FingerprintSensorProperties.TYPE_POWER_BUTTON) {
-            val currentIcon by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
+            runGenericTest {
+                val currentIcon by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
 
-            kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_0)
-            val iconRotation0 = currentIcon
+                kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_0)
+                val iconRotation0 = currentIcon
 
-            kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_90)
-            val iconRotation90 = currentIcon
+                kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_90)
+                val iconRotation90 = currentIcon
 
-            kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_180)
-            val iconRotation180 = currentIcon
+                kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_180)
+                val iconRotation180 = currentIcon
 
-            kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_270)
-            val iconRotation270 = currentIcon
+                kosmos.displayStateRepository.setCurrentRotation(DisplayRotation.ROTATION_270)
+                val iconRotation270 = currentIcon
 
-            assertThat(iconRotation0).isEqualTo(iconRotation180)
-            assertThat(iconRotation0).isNotEqualTo(iconRotation90)
-            assertThat(iconRotation0).isNotEqualTo(iconRotation270)
+                assertThat(iconRotation0).isEqualTo(iconRotation180)
+                assertThat(iconRotation0).isNotEqualTo(iconRotation90)
+                assertThat(iconRotation0).isNotEqualTo(iconRotation270)
+            }
         }
     }
 
     @Test
-    fun sfpsIconUpdates_onRearDisplayMode() = runGenericTest {
+    fun sfpsIconUpdates_onRearDisplayMode() {
         if (testCase.sensorType == FingerprintSensorProperties.TYPE_POWER_BUTTON) {
-            val currentIcon by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
+            runGenericTest {
+                val currentIcon by collectLastValue(kosmos.promptViewModel.iconViewModel.iconAsset)
 
-            kosmos.displayStateRepository.setIsInRearDisplayMode(false)
-            val iconNotRearDisplayMode = currentIcon
+                kosmos.displayStateRepository.setIsInRearDisplayMode(false)
+                val iconNotRearDisplayMode = currentIcon
 
-            kosmos.displayStateRepository.setIsInRearDisplayMode(true)
-            val iconRearDisplayMode = currentIcon
+                kosmos.displayStateRepository.setIsInRearDisplayMode(true)
+                val iconRearDisplayMode = currentIcon
 
-            assertThat(iconNotRearDisplayMode).isNotEqualTo(iconRearDisplayMode)
+                assertThat(iconNotRearDisplayMode).isNotEqualTo(iconRearDisplayMode)
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 67985ef..d2c6e555 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -102,6 +102,7 @@
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.app.IBatteryStats;
 import com.android.internal.os.BinderCallsStats;
 import com.android.internal.os.Clock;
@@ -671,6 +672,12 @@
                 BatteryConsumer.POWER_COMPONENT_FLASHLIGHT,
                 Flags.streamlinedMiscBatteryStats());
 
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_GNSS,
+                Flags.streamlinedMiscBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_GNSS,
+                Flags.streamlinedMiscBatteryStats());
+
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CAMERA,
                 Flags.streamlinedMiscBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
@@ -1191,7 +1198,7 @@
                                     .setMinConsumedPowerThreshold(minConsumedPowerThreshold)
                                     .build();
                     bus = getBatteryUsageStats(List.of(query)).get(0);
-                    return StatsPerUidLogger.logStats(bus, data);
+                    return new StatsPerUidLogger(new FrameworkStatsLogger()).logStats(bus, data);
                 }
                 default:
                     throw new UnsupportedOperationException("Unknown tagId=" + atomTag);
@@ -1204,7 +1211,35 @@
         }
     }
 
-    private static class StatsPerUidLogger {
+    public static class FrameworkStatsLogger {
+        /**
+         * Wrapper for the FrameworkStatsLog.buildStatsEvent method that makes it easier
+         * for mocking.
+         */
+        @VisibleForTesting
+        public StatsEvent buildStatsEvent(long sessionStartTs, long sessionEndTs,
+                long sessionDuration, int sessionDischargePercentage, long sessionDischargeDuration,
+                int uid, @BatteryConsumer.ProcessState int processState, long timeInStateMillis,
+                String powerComponentName, float totalConsumedPowerMah, float powerComponentMah,
+                long powerComponentDurationMillis) {
+            return FrameworkStatsLog.buildStatsEvent(
+                    FrameworkStatsLog.BATTERY_USAGE_STATS_PER_UID,
+                    sessionStartTs,
+                    sessionEndTs,
+                    sessionDuration,
+                    sessionDischargePercentage,
+                    sessionDischargeDuration,
+                    uid,
+                    processState,
+                    timeInStateMillis,
+                    powerComponentName,
+                    totalConsumedPowerMah,
+                    powerComponentMah,
+                    powerComponentDurationMillis);
+        }
+    }
+
+    public static class StatsPerUidLogger {
 
         private static final int STATSD_METRIC_MAX_DIMENSIONS_COUNT = 3000;
 
@@ -1224,7 +1259,18 @@
                 long dischargeDuration) {}
         ;
 
-        static int logStats(BatteryUsageStats bus, List<StatsEvent> data) {
+        private final FrameworkStatsLogger mFrameworkStatsLogger;
+
+        public StatsPerUidLogger(FrameworkStatsLogger frameworkStatsLogger) {
+            mFrameworkStatsLogger = frameworkStatsLogger;
+        }
+
+        /**
+         * Generates StatsEvents for the supplied battery usage stats and adds them to
+         * the supplied list.
+         */
+        @VisibleForTesting
+        public int logStats(BatteryUsageStats bus, List<StatsEvent> data) {
             final SessionInfo sessionInfo =
                     new SessionInfo(
                             bus.getStatsStartTimestamp(),
@@ -1340,7 +1386,7 @@
             return StatsManager.PULL_SUCCESS;
         }
 
-        private static boolean addStatsForPredefinedComponent(
+        private boolean addStatsForPredefinedComponent(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1380,7 +1426,7 @@
                     powerComponentDurationMillis);
         }
 
-        private static boolean addStatsForCustomComponent(
+        private boolean addStatsForCustomComponent(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1422,7 +1468,7 @@
          * Returns true on success and false if reached max atoms capacity and no more atoms should
          * be added
          */
-        private static boolean addStatsAtom(
+        private boolean addStatsAtom(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1432,9 +1478,7 @@
                 float totalConsumedPowerMah,
                 float powerComponentMah,
                 long powerComponentDurationMillis) {
-            data.add(
-                    FrameworkStatsLog.buildStatsEvent(
-                            FrameworkStatsLog.BATTERY_USAGE_STATS_PER_UID,
+            data.add(mFrameworkStatsLogger.buildStatsEvent(
                             sessionInfo.startTs(),
                             sessionInfo.endTs(),
                             sessionInfo.duration(),
@@ -3214,6 +3258,9 @@
                 .setMaxStatsAgeMs(0)
                 .includeProcessStateData()
                 .includePowerModels();
+        if (Flags.batteryUsageStatsByPowerAndScreenState()) {
+            builder.includeScreenStateData().includePowerStateData();
+        }
         if (model == BatteryConsumer.POWER_MODEL_POWER_PROFILE) {
             builder.powerProfileModeledOnly();
         }
@@ -3232,7 +3279,7 @@
         if (proto) {
             batteryUsageStats.dumpToProto(fd);
         } else {
-            batteryUsageStats.dump(pw, "");
+            batteryUsageStats.dump(pw, "  ");
         }
     }
 
diff --git a/services/core/java/com/android/server/location/LocationManagerService.java b/services/core/java/com/android/server/location/LocationManagerService.java
index 6e991b4..2e167ef 100644
--- a/services/core/java/com/android/server/location/LocationManagerService.java
+++ b/services/core/java/com/android/server/location/LocationManagerService.java
@@ -376,6 +376,11 @@
                             mContext.getContentResolver(),
                             Settings.Global.LOCATION_ENABLE_STATIONARY_THROTTLE,
                             defaultStationaryThrottlingSetting) != 0;
+                    if (Flags.disableStationaryThrottling() && !(
+                            Flags.keepGnssStationaryThrottling() && enableStationaryThrottling
+                                    && GPS_PROVIDER.equals(manager.getName()))) {
+                        enableStationaryThrottling = false;
+                    }
                     if (enableStationaryThrottling) {
                         realProvider = new StationaryThrottlingLocationProvider(manager.getName(),
                                 mInjector, realProvider);
diff --git a/services/core/java/com/android/server/location/contexthub/ContextHubService.java b/services/core/java/com/android/server/location/contexthub/ContextHubService.java
index ed451ff..7de1045 100644
--- a/services/core/java/com/android/server/location/contexthub/ContextHubService.java
+++ b/services/core/java/com/android/server/location/contexthub/ContextHubService.java
@@ -333,8 +333,14 @@
         return new IContextHubClientCallback.Stub() {
             private void finishCallback() {
                 try {
-                    IContextHubClient client = mDefaultClientMap.get(contextHubId);
-                    client.callbackFinished();
+                    if (mDefaultClientMap != null && mDefaultClientMap.containsKey(contextHubId)) {
+                        IContextHubClient client = mDefaultClientMap.get(contextHubId);
+                        client.callbackFinished();
+                    } else {
+                        Log.e(TAG, "Default client not found for hub (ID = " + contextHubId + "): "
+                                + mDefaultClientMap == null ? "map was null"
+                                                            : "map did not contain the hub");
+                    }
                 } catch (RemoteException e) {
                     Log.e(
                             TAG,
diff --git a/services/core/java/com/android/server/notification/ZenModeHelper.java b/services/core/java/com/android/server/notification/ZenModeHelper.java
index c078409..b12a917 100644
--- a/services/core/java/com/android/server/notification/ZenModeHelper.java
+++ b/services/core/java/com/android/server/notification/ZenModeHelper.java
@@ -1159,11 +1159,18 @@
                 rule.conditionId = azr.getConditionId();
                 modified = true;
             }
-            boolean shouldPreserveCondition = Flags.modesApi() && Flags.modesUi()
-                    && !isNew && origin == UPDATE_ORIGIN_USER
-                    && rule.enabled == azr.isEnabled()
-                    && rule.conditionId != null && rule.condition != null
-                    && rule.conditionId.equals(rule.condition.id);
+            // This can be removed when {@link Flags#modesUi} is fully ramped up
+            final boolean isWatch =
+                    mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_WATCH);
+            boolean shouldPreserveCondition =
+                    Flags.modesApi()
+                            && (Flags.modesUi() || isWatch)
+                            && !isNew
+                            && origin == UPDATE_ORIGIN_USER
+                            && rule.enabled == azr.isEnabled()
+                            && rule.conditionId != null
+                            && rule.condition != null
+                            && rule.conditionId.equals(rule.condition.id);
             if (!shouldPreserveCondition) {
                 // Do not update 'modified'. If only this changes we treat it as a no-op updateAZR.
                 rule.condition = null;
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 ad146af..fb54c5d 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.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 @@
                 .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 8127b82..ac68966 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 @@
 
             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 @@
         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 @@
                             + " 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 549a97e..0f13492 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 @@
         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 @@
         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 @@
                 .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 @@
             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 @@
                 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 @@
         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 d34498a..05496cf 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 @@
     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/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java
index 63224bb..c54ff5f 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java
@@ -25,7 +25,7 @@
 import android.view.inputmethod.InputMethodInfo;
 import android.view.inputmethod.InputMethodSubtype;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import org.junit.Test;
 
@@ -54,7 +54,8 @@
 
         // Save & load.
         AtomicFile atomicFile = new AtomicFile(
-                new File(InstrumentationRegistry.getContext().getCacheDir(), "subtypes.xml"));
+                new File(InstrumentationRegistry.getInstrumentation().getContext().getCacheDir(),
+                        "subtypes.xml"));
         AdditionalSubtypeUtils.saveToFile(AdditionalSubtypeMap.of(allSubtypes),
                 InputMethodMap.of(methodMap), atomicFile);
         AdditionalSubtypeMap loadedSubtypes = AdditionalSubtypeUtils.loadFromFile(atomicFile);
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodUtilsTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodUtilsTest.java
index 2857619..3cf895e 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodUtilsTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodUtilsTest.java
@@ -40,7 +40,7 @@
 import android.view.inputmethod.InputMethodSubtype.InputMethodSubtypeBuilder;
 
 import androidx.annotation.NonNull;
-import androidx.test.InstrumentationRegistry;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.internal.inputmethod.StartInputFlags;
 
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index f2b4136..b2a5b02 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -59,6 +59,7 @@
     name: "PowerStatsTestsRavenwood",
     static_libs: [
         "services.core",
+        "platformprotosnano",
         "coretests-aidl",
         "ravenwood-junit",
         "truth",
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
similarity index 72%
rename from core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
index ac1f7d0..37d8f2f 100644
--- a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.internal.os;
+package com.android.server.power.stats;
 
 import static android.os.BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE;
 
@@ -23,39 +23,262 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.os.AggregateBatteryConsumer;
 import android.os.BatteryConsumer;
 import android.os.BatteryUsageStats;
+import android.os.Process;
 import android.os.UidBatteryConsumer;
 import android.os.nano.BatteryUsageStatsAtomsProto;
 import android.os.nano.BatteryUsageStatsAtomsProto.BatteryConsumerData.PowerComponentUsage;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.StatsEvent;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.server.am.BatteryStatsService;
+
 import com.google.protobuf.nano.InvalidProtocolBufferNanoException;
 
+import org.junit.Rule;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
-
 @SmallTest
-public class BatteryUsageStatsPulledTest {
+public class BatteryUsageStatsAtomTest {
+
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
 
     private static final int UID_0 = 1000;
     private static final int UID_1 = 2000;
     private static final int UID_2 = 3000;
     private static final int UID_3 = 4000;
-    private static final int[] UID_USAGE_TIME_PROCESS_STATES = {
-            BatteryConsumer.PROCESS_STATE_FOREGROUND,
-            BatteryConsumer.PROCESS_STATE_BACKGROUND,
-            BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE
-    };
 
     @Test
-    public void testGetStatsProto() {
+    public void testAtom_BatteryUsageStatsPerUid() {
+        final BatteryUsageStats bus = buildBatteryUsageStats();
+        BatteryStatsService.FrameworkStatsLogger statsLogger =
+                mock(BatteryStatsService.FrameworkStatsLogger.class);
+
+        List<StatsEvent> actual = new ArrayList<>();
+        new BatteryStatsService.StatsPerUidLogger(statsLogger).logStats(bus, actual);
+
+        // Device-wide totals
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                Process.INVALID_UID,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "cpu",
+                30000.0f,
+                20100.0f,
+                20300L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                Process.INVALID_UID,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "camera",
+                30000.0f,
+                20150.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                Process.INVALID_UID,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "CustomConsumer1",
+                30000.0f,
+                20200.0f,
+                20400L
+        );
+
+        // Per-proc state estimates for UID_0
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "screen",
+                1650.0f,
+                300.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "cpu",
+                1650.0f,
+                400.0f,
+                600L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                1000L,
+                "cpu",
+                1650.0f,
+                9100.0f,
+                8100L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND,
+                2000L,
+                "cpu",
+                1650.0f,
+                9200.0f,
+                8200L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE,
+                0L,
+                "cpu",
+                1650.0f,
+                9300.0f,
+                8400L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_CACHED,
+                0L,
+                "cpu",
+                1650.0f,
+                9400.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                1000L,
+                "CustomConsumer1",
+                1650.0f,
+                450.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND,
+                2000L,
+                "CustomConsumer1",
+                1650.0f,
+                450.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                1000L,
+                "CustomConsumer2",
+                1650.0f,
+                500.0f,
+                800L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND,
+                2000L,
+                "CustomConsumer2",
+                1650.0f,
+                500.0f,
+                800L
+        );
+
+        // Nothing for UID_1, because its power consumption is 0
+
+        // Only "screen" is populated for UID_2
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_2,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "screen",
+                766.0f,
+                766.0f,
+                0L
+        );
+
+        verifyNoMoreInteractions(statsLogger);
+    }
+
+    @Test
+    public void testAtom_BatteryUsageStatsAtomsProto() {
         final BatteryUsageStats bus = buildBatteryUsageStats();
         final byte[] bytes = bus.getStatsProto();
         BatteryUsageStatsAtomsProto proto;
@@ -68,9 +291,7 @@
 
         assertEquals(bus.getStatsStartTimestamp(), proto.sessionStartMillis);
         assertEquals(bus.getStatsEndTimestamp(), proto.sessionEndMillis);
-        assertEquals(
-                bus.getStatsEndTimestamp() - bus.getStatsStartTimestamp(),
-                proto.sessionDurationMillis);
+        assertEquals(10000, proto.sessionDurationMillis);
         assertEquals(bus.getDischargePercentage(), proto.sessionDischargePercentage);
         assertEquals(bus.getDischargeDurationMs(), proto.dischargeDurationMillis);
 
@@ -90,8 +311,8 @@
         final List<android.os.UidBatteryConsumer> uidConsumers = bus.getUidBatteryConsumers();
         uidConsumers.sort((a, b) -> a.getUid() - b.getUid());
 
-        final BatteryUsageStatsAtomsProto.UidBatteryConsumer[] uidConsumersProto
-                = proto.uidBatteryConsumers;
+        final BatteryUsageStatsAtomsProto.UidBatteryConsumer[] uidConsumersProto =
+                proto.uidBatteryConsumers;
         Arrays.sort(uidConsumersProto, (a, b) -> a.uid - b.uid);
 
         // UID_0 - After sorting, UID_0 should be in position 0 for both data structures
@@ -186,6 +407,12 @@
         }
     }
 
+    private static final int[] UID_USAGE_TIME_PROCESS_STATES = {
+            BatteryConsumer.PROCESS_STATE_FOREGROUND,
+            BatteryConsumer.PROCESS_STATE_BACKGROUND,
+            BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE
+    };
+
     private void assertSameUidBatteryConsumer(
             android.os.UidBatteryConsumer uidConsumer,
             BatteryUsageStatsAtomsProto.UidBatteryConsumer uidConsumerProto,
@@ -195,10 +422,10 @@
         assertEquals("Uid consumers had mismatched uids", uid, uidConsumer.getUid());
 
         assertEquals("For uid " + uid,
-                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND),
+                uidConsumer.getTimeInProcessStateMs(BatteryConsumer.PROCESS_STATE_FOREGROUND),
                 uidConsumerProto.timeInForegroundMillis);
         assertEquals("For uid " + uid,
-                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND),
+                uidConsumer.getTimeInProcessStateMs(BatteryConsumer.PROCESS_STATE_BACKGROUND),
                 uidConsumerProto.timeInBackgroundMillis);
         for (int processState : UID_USAGE_TIME_PROCESS_STATES) {
             final long timeInStateMillis = uidConsumer.getTimeInProcessStateMs(processState);
@@ -261,11 +488,15 @@
                 new BatteryUsageStats.Builder(new String[]{"CustomConsumer1", "CustomConsumer2"},
                         /* includePowerModels */ true,
                         /* includeProcessStats */ true,
+                        /* includeScreenStateData */ false,
+                        /* includePowerStateData */ false,
                         /* minConsumedPowerThreshold */ 0)
                         .setDischargePercentage(20)
                         .setDischargedPowerRange(1000, 2000)
                         .setDischargeDurationMs(1234)
-                        .setStatsStartTimestamp(1000);
+                        .setStatsStartTimestamp(1000)
+                        .setStatsEndTimestamp(20000)
+                        .setStatsDuration(10000);
         final UidBatteryConsumer.Builder uidBuilder = builder
                 .getOrCreateUidBatteryConsumerBuilder(UID_0)
                 .setPackageWithHighestDrain("myPackage0")
@@ -345,7 +576,7 @@
     @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 6edfede..624b189 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 @@
                 & 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 a3f0770..52bb5e8 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.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 @@
         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 @@
         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 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 @@
         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 @@
     @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 @@
     @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 @@
         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 @@
 
         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 @@
                 .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 @@
 
     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 @@
                         .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 @@
         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 @@
                     .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 32bfb2c..7f7967b 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 @@
 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 @@
 
     @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 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 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 @@
         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/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
index 4bbbc2b..b07940a 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
@@ -186,6 +186,9 @@
 import org.mockito.MockitoAnnotations;
 import org.xmlpull.v1.XmlPullParserException;
 
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
+import platform.test.runner.parameterized.Parameters;
+
 import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
 import java.io.ByteArrayInputStream;
@@ -205,9 +208,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
-import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
-import platform.test.runner.parameterized.Parameters;
-
 @SmallTest
 @SuppressLint("GuardedBy") // It's ok for this test to access guarded methods from the service.
 @RunWith(ParameterizedAndroidJunit4.class)
@@ -5022,6 +5022,34 @@
     }
 
     @Test
+    @EnableFlags(FLAG_MODES_API)
+    public void updateAutomaticZenRule_ruleChangedByUser_doesNotDeactivateRule_forWatch() {
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_WATCH)).thenReturn(true);
+        assertThat(mZenModeHelper.getZenMode()).isEqualTo(ZEN_MODE_OFF);
+        AutomaticZenRule rule =
+                new AutomaticZenRule.Builder("rule", CONDITION_ID)
+                        .setConfigurationActivity(new ComponentName(mPkg, "cls"))
+                        .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
+                        .build();
+        String ruleId =
+                mZenModeHelper.addAutomaticZenRule(
+                        mPkg, rule, UPDATE_ORIGIN_APP, "reason", CUSTOM_PKG_UID);
+        mZenModeHelper.setAutomaticZenRuleState(
+                ruleId, CONDITION_TRUE, UPDATE_ORIGIN_APP, CUSTOM_PKG_UID);
+        assertThat(mZenModeHelper.getZenMode()).isEqualTo(ZEN_MODE_IMPORTANT_INTERRUPTIONS);
+
+        AutomaticZenRule updateWithDiff =
+                new AutomaticZenRule.Builder(rule).setTriggerDescription("Whenever").build();
+        mZenModeHelper.updateAutomaticZenRule(ruleId, updateWithDiff, UPDATE_ORIGIN_USER, "reason",
+                CUSTOM_PKG_UID);
+
+        assertThat(mZenModeHelper.getZenMode()).isEqualTo(ZEN_MODE_IMPORTANT_INTERRUPTIONS);
+        assertThat(mZenModeHelper.mConfig.automaticRules.get(ruleId).condition).isEqualTo(
+                CONDITION_TRUE);
+    }
+
+    @Test
     @EnableFlags({FLAG_MODES_API, FLAG_MODES_UI})
     public void updateAutomaticZenRule_ruleDisabledByUser_doesNotReactivateOnReenable() {
         assertThat(mZenModeHelper.getZenMode()).isEqualTo(ZEN_MODE_OFF);
diff --git a/tests/BatteryStatsPerfTest/src/com/android/internal/os/BatteryUsageStatsPerfTest.java b/tests/BatteryStatsPerfTest/src/com/android/internal/os/BatteryUsageStatsPerfTest.java
index 08430f2..4143f59 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 @@
 
     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)