Fix size and color of priority modes icons in status bar

* Allow StatusBarIcons to specify whether they want a "fixed" size instead of the variable width they normally have. This also affects their scaling calculations. This is mostly for app-provided icons which may have nonstandard sizes and proportions.
* Create a new drawable from constant state so that theming isn't shared with the other usages.

Additionally: Add a flow to ZenModeInteractor with the main active mode name and icon exclusively. While the QS tile cares about which other modes are active, neither status bar nor Smartspace do, which means we can use a flow that doesn't emit when more modes are activated (and thus avoid wasted work).

Bug: 360399800
Test: coming
Flag: android.app.modes_ui
Change-Id: I6a17d5a49f76169f15efc89432e1ce0f993d040e
diff --git a/core/java/com/android/internal/statusbar/StatusBarIcon.java b/core/java/com/android/internal/statusbar/StatusBarIcon.java
index f8a1436..1938cdb 100644
--- a/core/java/com/android/internal/statusbar/StatusBarIcon.java
+++ b/core/java/com/android/internal/statusbar/StatusBarIcon.java
@@ -51,6 +51,19 @@
         ResourceIcon
     }
 
+    public enum Shape {
+        /**
+         * Icon view should use WRAP_CONTENT -- so that the horizontal space occupied depends on the
+         * icon's shape (skinny/fat icons take less/more). Most icons will want to use this option
+         * for a nicer-looking overall spacing in the status bar, as long as the icon is "known"
+         * (i.e. not coming from a 3P package).
+         */
+        WRAP_CONTENT,
+
+        /** Icon should always be displayed in a space as wide as the status bar is tall. */
+        FIXED_SPACE,
+    }
+
     public UserHandle user;
     public String pkg;
     public Icon icon;
@@ -59,6 +72,7 @@
     public int number;
     public CharSequence contentDescription;
     public Type type;
+    public Shape shape;
 
     /**
      * Optional {@link Drawable} corresponding to {@link #icon}. This field is not parcelable, so
@@ -68,7 +82,7 @@
     @Nullable public Drawable preloadedIcon;
 
     public StatusBarIcon(UserHandle user, String resPackage, Icon icon, int iconLevel, int number,
-            CharSequence contentDescription, Type type) {
+            CharSequence contentDescription, Type type, Shape shape) {
         if (icon.getType() == Icon.TYPE_RESOURCE
                 && TextUtils.isEmpty(icon.getResPackage())) {
             // This is an odd situation where someone's managed to hand us an icon without a
@@ -83,6 +97,13 @@
         this.number = number;
         this.contentDescription = contentDescription;
         this.type = type;
+        this.shape = shape;
+    }
+
+    public StatusBarIcon(UserHandle user, String resPackage, Icon icon, int iconLevel, int number,
+            CharSequence contentDescription, Type type) {
+        this(user, resPackage, icon, iconLevel, number, contentDescription, type,
+                Shape.WRAP_CONTENT);
     }
 
     public StatusBarIcon(String iconPackage, UserHandle user,
@@ -107,7 +128,7 @@
     @Override
     public StatusBarIcon clone() {
         StatusBarIcon that = new StatusBarIcon(this.user, this.pkg, this.icon,
-                this.iconLevel, this.number, this.contentDescription, this.type);
+                this.iconLevel, this.number, this.contentDescription, this.type, this.shape);
         that.visible = this.visible;
         that.preloadedIcon = this.preloadedIcon;
         return that;
@@ -129,6 +150,7 @@
         this.number = in.readInt();
         this.contentDescription = in.readCharSequence();
         this.type = Type.valueOf(in.readString());
+        this.shape = Shape.valueOf(in.readString());
     }
 
     public void writeToParcel(Parcel out, int flags) {
@@ -140,6 +162,7 @@
         out.writeInt(this.number);
         out.writeCharSequence(this.contentDescription);
         out.writeString(this.type.name());
+        out.writeString(this.shape.name());
     }
 
     public int describeContents() {
diff --git a/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java b/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
index b183ecb..149e132 100644
--- a/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
+++ b/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
@@ -20,6 +20,7 @@
 
 import android.graphics.Color;
 import android.graphics.drawable.ColorDrawable;
+import android.graphics.drawable.Icon;
 import android.os.Parcel;
 import android.os.UserHandle;
 
@@ -69,22 +70,22 @@
         assertThat(copy.preloadedIcon).isEqualTo(original.preloadedIcon);
     }
 
-
     private static StatusBarIcon newStatusBarIcon() {
         final UserHandle dummyUserHandle = UserHandle.of(100);
         final String dummyIconPackageName = "com.android.internal.statusbar.test";
-        final int dummyIconId = 123;
+        final Icon dummyIcon = Icon.createWithResource(dummyIconPackageName, 123);
         final int dummyIconLevel = 1;
         final int dummyIconNumber = 2;
         final CharSequence dummyIconContentDescription = "dummyIcon";
         return new StatusBarIcon(
-                dummyIconPackageName,
                 dummyUserHandle,
-                dummyIconId,
+                dummyIconPackageName,
+                dummyIcon,
                 dummyIconLevel,
                 dummyIconNumber,
                 dummyIconContentDescription,
-                StatusBarIcon.Type.SystemIcon);
+                StatusBarIcon.Type.SystemIcon,
+                StatusBarIcon.Shape.FIXED_SPACE);
     }
 
     private static void assertSerializableFieldsEqual(StatusBarIcon copy, StatusBarIcon original) {
@@ -96,6 +97,7 @@
         assertThat(copy.number).isEqualTo(original.number);
         assertThat(copy.contentDescription).isEqualTo(original.contentDescription);
         assertThat(copy.type).isEqualTo(original.type);
+        assertThat(copy.shape).isEqualTo(original.shape);
     }
 
     private static StatusBarIcon parcelAndUnparcel(StatusBarIcon original) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
index cb743bc..639d34d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
@@ -22,8 +22,10 @@
 import android.provider.Settings.Secure.ZEN_DURATION
 import android.provider.Settings.Secure.ZEN_DURATION_FOREVER
 import android.provider.Settings.Secure.ZEN_DURATION_PROMPT
+import android.service.notification.SystemZenRules
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.R
 import com.android.settingslib.notification.data.repository.updateNotificationPolicy
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.SysuiTestCase
@@ -253,4 +255,56 @@
             assertThat(activeModes?.modeNames).hasSize(0)
             assertThat(activeModes?.mainMode).isNull()
         }
+
+    @Test
+    fun mainActiveMode_flows() =
+        testScope.runTest {
+            val mainActiveMode by collectLastValue(underTest.mainActiveMode)
+
+            zenModeRepository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setId("Bedtime")
+                        .setName("Mode Bedtime")
+                        .setType(AutomaticZenRule.TYPE_BEDTIME)
+                        .setActive(false)
+                        .setPackage(mContext.packageName)
+                        .setIconResId(R.drawable.ic_zen_mode_type_bedtime)
+                        .build(),
+                    TestModeBuilder()
+                        .setId("Other")
+                        .setName("Mode Other")
+                        .setType(AutomaticZenRule.TYPE_OTHER)
+                        .setActive(false)
+                        .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                        .setIconResId(R.drawable.ic_zen_mode_type_other)
+                        .build(),
+                )
+            )
+
+            runCurrent()
+            assertThat(mainActiveMode).isNull()
+
+            zenModeRepository.activateMode("Other")
+            runCurrent()
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Other")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_other)
+
+            zenModeRepository.activateMode("Bedtime")
+            runCurrent()
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Bedtime")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_bedtime)
+
+            zenModeRepository.deactivateMode("Other")
+            runCurrent()
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Bedtime")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_bedtime)
+
+            zenModeRepository.deactivateMode("Bedtime")
+            runCurrent()
+            assertThat(mainActiveMode).isNull()
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
index 6eadd26..2b44c2f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
@@ -58,6 +58,7 @@
 import com.android.app.animation.Interpolators;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.statusbar.StatusBarIcon;
+import com.android.internal.statusbar.StatusBarIcon.Shape;
 import com.android.internal.util.ContrastColorUtil;
 import com.android.systemui.Flags;
 import com.android.systemui.res.R;
@@ -211,16 +212,19 @@
     /** Should always be preceded by {@link #reloadDimens()} */
     @VisibleForTesting
     public void maybeUpdateIconScaleDimens() {
-        // We do not resize and scale system icons (on the right), only notification icons (on the
-        // left).
-        if (isNotification()) {
-            updateIconScaleForNotifications();
+        // We scale notification icons (on the left) plus icons on the right that explicitly
+        // want FIXED_SPACE.
+        boolean useNonSystemIconScaling = isNotification()
+                || (usesModeIcons() && mIcon != null && mIcon.shape == Shape.FIXED_SPACE);
+
+        if (useNonSystemIconScaling) {
+            updateIconScaleForNonSystemIcons();
         } else {
             updateIconScaleForSystemIcons();
         }
     }
 
-    private void updateIconScaleForNotifications() {
+    private void updateIconScaleForNonSystemIcons() {
         float iconScale;
         // we need to scale the image size to be same as the original size
         // (fit mOriginalStatusBarIconSize), then we can scale it with mScaleToFitNewIconSize
@@ -411,7 +415,9 @@
         if (!levelEquals) {
             setImageLevel(icon.iconLevel);
         }
-
+        if (usesModeIcons()) {
+            setScaleType(icon.shape == Shape.FIXED_SPACE ? ScaleType.FIT_CENTER : ScaleType.CENTER);
+        }
         if (!visibilityEquals) {
             setVisibility(icon.visible && !mBlocked ? VISIBLE : GONE);
         }
@@ -501,7 +507,12 @@
     @Nullable
     private Drawable loadDrawable(Context context, StatusBarIcon statusBarIcon) {
         if (usesModeIcons() && statusBarIcon.preloadedIcon != null) {
-            return statusBarIcon.preloadedIcon.mutate();
+            Drawable.ConstantState cached = statusBarIcon.preloadedIcon.getConstantState();
+            if (cached != null) {
+                return cached.newDrawable(mContext.getResources()).mutate();
+            } else {
+                return statusBarIcon.preloadedIcon.mutate();
+            }
         } else {
             int userId = statusBarIcon.user.getIdentifier();
             if (userId == UserHandle.USER_ALL) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
index a9b886f..ba39c3b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
@@ -42,9 +42,9 @@
 import android.util.Log;
 import android.view.View;
 
-import androidx.annotation.NonNull;
 import androidx.lifecycle.Observer;
 
+import com.android.internal.statusbar.StatusBarIcon;
 import com.android.systemui.Flags;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.qualifiers.DisplayId;
@@ -80,7 +80,6 @@
 import com.android.systemui.statusbar.policy.UserInfoController;
 import com.android.systemui.statusbar.policy.ZenModeController;
 import com.android.systemui.statusbar.policy.domain.interactor.ZenModeInteractor;
-import com.android.systemui.statusbar.policy.domain.model.ActiveZenModes;
 import com.android.systemui.statusbar.policy.domain.model.ZenModeInfo;
 import com.android.systemui.util.RingerModeTracker;
 import com.android.systemui.util.kotlin.JavaAdapter;
@@ -364,8 +363,8 @@
         if (usesModeIcons()) {
             // Note that we're not fully replacing ZenModeController with ZenModeInteractor, so
             // we listen for the extra event here but still add the ZMC callback.
-            mJavaAdapter.alwaysCollectFlow(mZenModeInteractor.getActiveModes(),
-                    this::onActiveModesChanged);
+            mJavaAdapter.alwaysCollectFlow(mZenModeInteractor.getMainActiveMode(),
+                    this::onMainActiveModeChanged);
         }
         mZenController.addCallback(mZenControllerCallback);
         if (!Flags.statusBarScreenSharingChips()) {
@@ -397,21 +396,23 @@
                 () -> mResources.getString(R.string.accessibility_managed_profile));
     }
 
-    private void onActiveModesChanged(@NonNull ActiveZenModes activeModes) {
+    private void onMainActiveModeChanged(@Nullable ZenModeInfo mainActiveMode) {
         if (!usesModeIcons()) {
-            Log.wtf(TAG, "onActiveModeChanged shouldn't be called if MODES_UI_ICONS is disabled");
+            Log.wtf(TAG, "onMainActiveModeChanged shouldn't run if MODES_UI_ICONS is disabled");
             return;
         }
 
-        ZenModeInfo mainActiveMode = activeModes.getMainMode();
         boolean visible = mainActiveMode != null;
-
         if (visible) {
+            // Shape=FIXED_SPACE because mode icons can be from 3P packages and may not be square;
+            // we don't want to allow apps to set incredibly wide icons and take up too much space
+            // in the status bar.
             mIconController.setResourceIcon(mSlotZen,
                     mainActiveMode.getIcon().key().resPackage(),
                     mainActiveMode.getIcon().key().resId(),
                     mainActiveMode.getIcon().drawable(),
-                    mainActiveMode.getName());
+                    mainActiveMode.getName(),
+                    StatusBarIcon.Shape.FIXED_SPACE);
         }
         if (visible != mZenVisible) {
             mIconController.setIconVisibility(mSlotZen, visible);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
index 8871dae..6c30330 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
@@ -16,7 +16,6 @@
 
 package com.android.systemui.statusbar.phone.ui;
 
-import android.view.ViewGroup;
 import android.widget.LinearLayout;
 
 import com.android.internal.statusbar.StatusBarIcon;
@@ -64,9 +63,8 @@
     }
 
     @Override
-    protected LinearLayout.LayoutParams onCreateLayoutParams() {
-        LinearLayout.LayoutParams lp = new LinearLayout.LayoutParams(
-                ViewGroup.LayoutParams.WRAP_CONTENT, mIconSize);
+    protected LinearLayout.LayoutParams onCreateLayoutParams(StatusBarIcon.Shape shape) {
+        LinearLayout.LayoutParams lp = super.onCreateLayoutParams(shape);
         lp.setMargins(mIconHorizontalMargin, 0, mIconHorizontalMargin, 0);
         return lp;
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
index 5ad7376..91ead61 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
@@ -20,6 +20,7 @@
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_ICON;
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_MOBILE_NEW;
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_WIFI_NEW;
+import static com.android.systemui.statusbar.phone.ui.StatusBarIconControllerImpl.usesModeIcons;
 
 import android.annotation.Nullable;
 import android.content.Context;
@@ -27,9 +28,8 @@
 import android.view.ViewGroup;
 import android.widget.LinearLayout;
 
-import androidx.annotation.VisibleForTesting;
-
 import com.android.internal.statusbar.StatusBarIcon;
+import com.android.internal.statusbar.StatusBarIcon.Shape;
 import com.android.systemui.demomode.DemoModeCommandReceiver;
 import com.android.systemui.statusbar.BaseStatusBarFrameLayout;
 import com.android.systemui.statusbar.StatusBarIconView;
@@ -155,12 +155,11 @@
         };
     }
 
-    @VisibleForTesting
     protected StatusBarIconView addIcon(int index, String slot, boolean blocked,
             StatusBarIcon icon) {
         StatusBarIconView view = onCreateStatusBarIconView(slot, blocked);
         view.set(icon);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(icon.shape));
         return view;
     }
 
@@ -174,7 +173,7 @@
             int index) {
         mBindableIcons.put(holder.getSlot(), holder);
         ModernStatusBarView view = holder.getInitializer().createAndBind(mContext);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
         if (mIsInDemoMode) {
             mDemoStatusIcons.addBindableIcon(holder);
         }
@@ -183,7 +182,7 @@
 
     protected StatusIconDisplayable addNewWifiIcon(int index, String slot) {
         ModernStatusBarWifiView view = onCreateModernStatusBarWifiView(slot);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
 
         if (mIsInDemoMode) {
             mDemoStatusIcons.addModernWifiView(mWifiViewModel);
@@ -199,7 +198,7 @@
             int subId
     ) {
         BaseStatusBarFrameLayout view = onCreateModernStatusBarMobileView(slot, subId);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
 
         if (mIsInDemoMode) {
             Context mobileContext = mMobileContextProvider
@@ -233,8 +232,12 @@
                 );
     }
 
-    protected LinearLayout.LayoutParams onCreateLayoutParams() {
-        return new LinearLayout.LayoutParams(ViewGroup.LayoutParams.WRAP_CONTENT, mIconSize);
+    protected LinearLayout.LayoutParams onCreateLayoutParams(Shape shape) {
+        int width = usesModeIcons() && shape == StatusBarIcon.Shape.FIXED_SPACE
+                ? mIconSize
+                : ViewGroup.LayoutParams.WRAP_CONTENT;
+
+        return new LinearLayout.LayoutParams(width, mIconSize);
     }
 
     protected void destroy() {
@@ -256,6 +259,13 @@
     /** Called once an icon has been set. */
     public void onSetIcon(int viewIndex, StatusBarIcon icon) {
         StatusBarIconView view = (StatusBarIconView) mGroup.getChildAt(viewIndex);
+        if (usesModeIcons()) {
+            ViewGroup.LayoutParams current = view.getLayoutParams();
+            ViewGroup.LayoutParams desired = onCreateLayoutParams(icon.shape);
+            if (desired.width != current.width || desired.height != current.height) {
+                view.setLayoutParams(desired);
+            }
+        }
         view.set(icon);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
index ee528e9..0459b97 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
@@ -70,7 +70,8 @@
      * @param preloadedIcon optional drawable corresponding to {@code iconResId}, if known
      */
     void setResourceIcon(String slot, @Nullable String resPackage, @DrawableRes int iconResId,
-            @Nullable Drawable preloadedIcon, CharSequence contentDescription);
+            @Nullable Drawable preloadedIcon, CharSequence contentDescription,
+            StatusBarIcon.Shape shape);
 
     /**
      * Sets up a wifi icon using the new data pipeline. No effect if the wifi icon has already been
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
index ad3a9e3..9b6d32b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
@@ -234,13 +234,14 @@
                 Icon.createWithResource(mContext, resourceId),
                 /* preloadedIcon= */ null,
                 contentDescription,
-                StatusBarIcon.Type.SystemIcon);
+                StatusBarIcon.Type.SystemIcon,
+                StatusBarIcon.Shape.WRAP_CONTENT);
     }
 
     @Override
     public void setResourceIcon(String slot, @Nullable String resPackage,
             @DrawableRes int iconResId, @Nullable Drawable preloadedIcon,
-            CharSequence contentDescription) {
+            CharSequence contentDescription, StatusBarIcon.Shape shape) {
         if (!usesModeIcons()) {
             Log.wtf("TAG",
                     "StatusBarIconController.setResourceIcon() should not be called without "
@@ -260,12 +261,13 @@
                 icon,
                 preloadedIcon,
                 contentDescription,
-                StatusBarIcon.Type.ResourceIcon);
+                StatusBarIcon.Type.ResourceIcon,
+                shape);
     }
 
     private void setResourceIconInternal(String slot, Icon resourceIcon,
             @Nullable Drawable preloadedIcon, CharSequence contentDescription,
-            StatusBarIcon.Type type) {
+            StatusBarIcon.Type type, StatusBarIcon.Shape shape) {
         checkArgument(resourceIcon.getType() == Icon.TYPE_RESOURCE,
                 "Expected Icon of TYPE_RESOURCE, but got " + resourceIcon.getType());
         String resPackage = resourceIcon.getResPackage();
@@ -277,7 +279,7 @@
         if (holder == null) {
             StatusBarIcon icon = new StatusBarIcon(UserHandle.SYSTEM, resPackage,
                     resourceIcon, /* iconLevel= */ 0, /* number=*/ 0,
-                    contentDescription, type);
+                    contentDescription, type, shape);
             icon.preloadedIcon = preloadedIcon;
             holder = StatusBarIconHolder.fromIcon(icon);
             setIcon(slot, holder);
@@ -286,6 +288,7 @@
             holder.getIcon().icon = resourceIcon;
             holder.getIcon().contentDescription = contentDescription;
             holder.getIcon().type = type;
+            holder.getIcon().shape = shape;
             holder.getIcon().preloadedIcon = preloadedIcon;
             handleSet(slot, holder);
         }
@@ -578,7 +581,7 @@
         }
     }
 
-    private static boolean usesModeIcons() {
+    static boolean usesModeIcons() {
         return android.app.Flags.modesApi() && android.app.Flags.modesUi()
                 && android.app.Flags.modesUiIcons();
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
index 520f3f4..93c631f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
@@ -96,6 +96,9 @@
             .flowOn(bgDispatcher)
             .distinctUntilChanged()
 
+    val mainActiveMode: Flow<ZenModeInfo?> =
+        activeModes.map { a -> a.mainMode }.distinctUntilChanged()
+
     suspend fun getModeIcon(mode: ZenMode): ZenIcon {
         return iconLoader.getIcon(context, mode).await()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
index a0b0572..2ed3473 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
@@ -34,6 +34,7 @@
 import android.testing.TestableLooper.RunWithLooper
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.statusbar.StatusBarIcon
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
@@ -429,7 +430,8 @@
                     eq(mContext.packageName),
                     eq(android.R.drawable.ic_lock_lock),
                     any(), // non-null
-                    eq("Bedtime Mode")
+                    eq("Bedtime Mode"),
+                    eq(StatusBarIcon.Shape.FIXED_SPACE)
                 )
 
             zenModeRepository.deactivateMode("bedtime")
@@ -441,7 +443,8 @@
                     eq(null),
                     eq(android.R.drawable.ic_media_play),
                     any(), // non-null
-                    eq("Other Mode")
+                    eq("Other Mode"),
+                    eq(StatusBarIcon.Shape.FIXED_SPACE)
                 )
 
             zenModeRepository.deactivateMode("other")
@@ -460,7 +463,8 @@
 
         verify(iconController, never()).setIconVisibility(eq(ZEN_SLOT), any())
         verify(iconController, never()).setIcon(eq(ZEN_SLOT), anyInt(), any())
-        verify(iconController, never()).setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any())
+        verify(iconController, never())
+            .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any(), any())
     }
 
     @Test
@@ -476,7 +480,7 @@
             verify(iconController, never()).setIconVisibility(eq(ZEN_SLOT), any())
             verify(iconController, never()).setIcon(eq(ZEN_SLOT), anyInt(), any())
             verify(iconController, never())
-                .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any())
+                .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any(), any())
         }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
new file mode 100644
index 0000000..90732d01
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
@@ -0,0 +1,122 @@
+/*
+ * 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.phone.ui
+
+import android.app.Flags
+import android.graphics.drawable.Icon
+import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.view.ViewGroup
+import android.widget.ImageView
+import android.widget.LinearLayout
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.statusbar.StatusBarIcon
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.StatusBarIconView
+import com.android.systemui.statusbar.connectivity.ui.MobileContextProvider
+import com.android.systemui.statusbar.phone.StatusBarLocation
+import com.android.systemui.statusbar.pipeline.mobile.ui.MobileUiAdapter
+import com.android.systemui.statusbar.pipeline.wifi.ui.WifiUiAdapter
+import com.android.systemui.util.Assert
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.RETURNS_DEEP_STUBS
+import org.mockito.kotlin.mock
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class IconManagerTest : SysuiTestCase() {
+
+    private lateinit var underTest: IconManager
+    private lateinit var viewGroup: ViewGroup
+
+    @Before
+    fun setUp() {
+        Assert.setTestThread(Thread.currentThread())
+        viewGroup = LinearLayout(context)
+        underTest =
+            IconManager(
+                viewGroup,
+                StatusBarLocation.HOME,
+                mock<WifiUiAdapter>(defaultAnswer = RETURNS_DEEP_STUBS),
+                mock<MobileUiAdapter>(defaultAnswer = RETURNS_DEEP_STUBS),
+                mock<MobileContextProvider>(defaultAnswer = RETURNS_DEEP_STUBS),
+            )
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_shapeWrapContent_addsIconViewWithVariableWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.WRAP_CONTENT)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.CENTER)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_shapeFixedSpace_addsIconViewWithFixedWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.FIXED_SPACE)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isNotEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.layoutParams.width).isEqualTo(iconView.layoutParams.height)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.FIT_CENTER)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_iconsFlagOff_addsIconViewWithVariableWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.FIXED_SPACE)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.CENTER)
+    }
+
+    private fun newStatusBarIcon(shape: StatusBarIcon.Shape) =
+        StatusBarIcon(
+            UserHandle.CURRENT,
+            context.packageName,
+            Icon.createWithResource(context, android.R.drawable.ic_media_next),
+            0,
+            0,
+            "",
+            StatusBarIcon.Type.ResourceIcon,
+            shape,
+        )
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
index 26a57e4..50a13b9 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
@@ -424,7 +424,14 @@
     @EnableFlags(android.app.Flags.FLAG_MODES_UI, android.app.Flags.FLAG_MODES_UI_ICONS)
     fun setResourceIcon_setsIconAndPreloadedIconInHolder() {
         val drawable = ColorDrawable(1)
-        underTest.setResourceIcon("slot", "some.package", 123, drawable, "description")
+        underTest.setResourceIcon(
+            "slot",
+            "some.package",
+            123,
+            drawable,
+            "description",
+            StatusBarIcon.Shape.FIXED_SPACE
+        )
 
         val iconHolder = iconList.getIconHolder("slot", 0)
         assertThat(iconHolder).isNotNull()
@@ -432,6 +439,7 @@
         assertThat(iconHolder?.icon?.icon?.resId).isEqualTo(123)
         assertThat(iconHolder?.icon?.icon?.resPackage).isEqualTo("some.package")
         assertThat(iconHolder?.icon?.contentDescription).isEqualTo("description")
+        assertThat(iconHolder?.icon?.shape).isEqualTo(StatusBarIcon.Shape.FIXED_SPACE)
         assertThat(iconHolder?.icon?.preloadedIcon).isEqualTo(drawable)
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
index 2dbac67..0089199 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
@@ -64,7 +64,8 @@
 
     @Override
     public void setResourceIcon(String slot, @Nullable String resPackage, int iconResId,
-            @Nullable Drawable preloadedIcon, CharSequence contentDescription) {
+            @Nullable Drawable preloadedIcon, CharSequence contentDescription,
+            StatusBarIcon.Shape shape) {
     }
 
     @Override