fix: update action button to use grid area to position notification badge (#9303)

* Revert "Revert "fix: update action button to use grid area to position notification badge (#8409)" (#8552)"

This reverts commit 0cdf701.

* add chromatic story

* fix lint

Author: yihuiliao

packages/@react-spectrum/s2/chromatic/ActionButton.stories.tsx

@@ -10,7 +10,9 @@
  * governing permissions and limitations under the License.
  */
 
-import {ActionButton, ActionButtonProps, Avatar, Text} from '../src';
+import {ActionButton, ActionButtonProps, Avatar, NotificationBadge, Text} from '../src';
+import BellIcon from '../s2wf-icons/S2_Icon_Bell_20_N.svg';
+import CommentIcon from '../s2wf-icons/S2_Icon_Comment_20_N.svg';
 import {Fonts, NotificationBadges, UnsafeClassName} from '../stories/ActionButton.stories';
 import {generatePowerset} from '@react-spectrum/story-utils';
 import type {Meta, StoryObj} from '@storybook/react';
@@ -100,3 +102,13 @@ export const AvatarOnly: ActionButtonStory = {
 };
 
 export {Fonts, UnsafeClassName, NotificationBadges};
+
+export const NotificationBadgesCustomWidth: ActionButtonStory = {
+  render: (args) => (
+    <div className={style({display: 'flex', flexDirection: 'column', gap: 8})}>
+      <ActionButton aria-label="Messages has new activity" styles={style({width: 200})} {...args}><CommentIcon /><NotificationBadge /></ActionButton>
+      <ActionButton styles={style({width: 200})} {...args}><BellIcon /><NotificationBadge value={10} /></ActionButton>
+      <ActionButton styles={style({width: 200})} {...args}><CommentIcon /><Text>Messages</Text><NotificationBadge value={5} /></ActionButton>
+      <ActionButton styles={style({width: 200})} {...args}><Text>Notifications</Text><NotificationBadge value={105} /></ActionButton>
+    </div>)
+};

packages/@react-spectrum/s2/src/ActionButton.tsx

@@ -73,7 +73,7 @@ export const btnStyles = style<ButtonRenderProps & ActionButtonStyleProps & Togg
   ...focusRing(),
   ...staticColor(),
   ...controlStyle,
-  position: 'relative',
+  display: 'grid',
   justifyContent: 'center',
   flexShrink: {
     default: 1,
@@ -86,9 +86,21 @@ export const btnStyles = style<ButtonRenderProps & ActionButtonStyleProps & Togg
     isJustified: 0
   },
   fontWeight: 'medium',
+  width: 'fit',
   userSelect: 'none',
   transition: 'default',
   forcedColorAdjust: 'none',
+  position: 'relative',
+  gridTemplateAreas: {
+    default: ['icon text'],
+    [iconOnly]: ['icon'],
+    [textOnly]: ['text']
+  },
+  gridTemplateColumns: {
+    default: ['auto', 'auto'],
+    [iconOnly]: ['auto'],
+    [textOnly]: ['auto']
+  },
   backgroundColor: {
     default: {
       ...baseColor('gray-100'),
@@ -230,8 +242,7 @@ export const btnStyles = style<ButtonRenderProps & ActionButtonStyleProps & Togg
   '--badgePosition': {
     type: 'width',
     value: {
-      default: 'calc(self(paddingStart) + var(--iconWidth))',
-      [iconOnly]: 'calc(self(minWidth)/2 + var(--iconWidth)/2)',
+      default: '--iconWidth',
       [textOnly]: 'full'
     }
   },
@@ -311,19 +322,18 @@ export const ActionButton = forwardRef(function ActionButton(props: ActionButton
               [SkeletonContext, null],
               [TextContext, {styles:
                 style({
-                  order: 1,
+                  gridArea: 'text',
                   truncate: true,
                   visibility: {
                     isProgressVisible: 'hidden'
                   }
                 })({isProgressVisible})
               }],
               [IconContext, {
-                render: centerBaseline({slot: 'icon', styles: style({order: 0})}),
+                render: centerBaseline({slot: 'icon', styles: style({gridArea: 'icon'})}),
                 styles: style({
                   size: fontRelative(20),
                   marginStart: '--iconMargin',
-                  flexShrink: 0,
                   visibility: {
                     isProgressVisible: 'hidden'
                   }
@@ -333,8 +343,7 @@ export const ActionButton = forwardRef(function ActionButton(props: ActionButton
                 size: avatarSize[size],
                 styles: style({
                   marginStart: '--iconMargin',
-                  flexShrink: 0,
-                  order: 0
+                  gridArea: 'icon'
                 })
               }],
               [ImageContext, {
@@ -351,9 +360,10 @@ export const ActionButton = forwardRef(function ActionButton(props: ActionButton
                 styles: style({
                   position: 'absolute',
                   top: '--badgeTop',
-                  insetStart: '--badgePosition',
                   marginTop: 'calc((self(height) * -1)/2)',
-                  marginStart: 'calc((self(height) * -1)/2)',
+                  marginStart: 'calc(var(--iconMargin) * 2 + (self(height) * -1)/4)',
+                  gridColumnStart: 1,
+                  insetStart: '--badgePosition',
                   visibility: {
                     isProgressVisible: 'hidden'
                   }