[select][menu] Better handle dynamic and non-string items (#1861)

Author: atomiks

packages/react/src/menu/popup/MenuPopup.tsx

@@ -103,6 +103,7 @@ export const MenuPopup = React.forwardRef(function MenuPopup(
         (parent.type === 'menubar' && lastOpenChangeReason !== 'outside-press')
       }
       initialFocus={parent.type === 'menu' ? -1 : 0}
+      restoreFocus
     >
       {renderElement()}
     </FloatingFocusManager>

packages/react/src/menubar/Menubar.test.tsx

@@ -316,7 +316,9 @@ describe('<Menubar />', () => {
 
       // First submenu item should be focused
       const submenuItem = screen.getByTestId('share-item-1');
-      expect(submenuItem).toHaveFocus();
+      await waitFor(() => {
+        expect(submenuItem).toHaveFocus();
+      });
     });
 
     it.skipIf(isJSDOM)('should close the menu with Escape key', async () => {

packages/react/src/select/item-text/SelectItemText.tsx

@@ -24,7 +24,9 @@ const InnerSelectItemText = React.memo(
       (node: HTMLElement | null) => {
         // Wait for the DOM indices to be set.
         queueMicrotask(() => {
-          if (selected || (selectedItemTextRef.current === null && indexRef.current === 0)) {
+          const hasNoSelectedItemText =
+            selectedItemTextRef.current === null || !selectedItemTextRef.current.isConnected;
+          if (selected || (hasNoSelectedItemText && indexRef.current === 0)) {
             selectedItemTextRef.current = node;
           }
         });

packages/react/src/select/popup/SelectPopup.tsx

@@ -85,7 +85,12 @@ export const SelectPopup = React.forwardRef(function SelectPopup(
           dangerouslySetInnerHTML={html}
         />
       )}
-      <FloatingFocusManager context={positioner.context} modal={false} disabled={!mounted}>
+      <FloatingFocusManager
+        context={positioner.context}
+        modal={false}
+        disabled={!mounted}
+        restoreFocus
+      >
         {element}
       </FloatingFocusManager>
     </React.Fragment>

packages/react/src/select/popup/useSelectPopup.ts

@@ -20,15 +20,18 @@ export function useSelectPopup(): useSelectPopup.ReturnValue {
     valueRef,
     selectedItemTextRef,
     popupRef,
-    scrollUpArrowVisible,
-    scrollDownArrowVisible,
-    setScrollUpArrowVisible,
-    setScrollDownArrowVisible,
     keyboardActiveRef,
     floatingRootContext,
   } = useSelectRootContext();
   const { setActiveIndex } = useSelectIndexContext();
-  const { alignItemWithTriggerActive, setControlledItemAnchor } = useSelectPositionerContext();
+  const {
+    alignItemWithTriggerActive,
+    setControlledItemAnchor,
+    scrollUpArrowVisible,
+    scrollDownArrowVisible,
+    setScrollUpArrowVisible,
+    setScrollDownArrowVisible,
+  } = useSelectPositionerContext();
 
   const initialHeightRef = React.useRef(0);
   const reachedMaxHeightRef = React.useRef(false);

packages/react/src/select/positioner/SelectPositioner.tsx

@@ -10,6 +10,9 @@ import { SelectPositionerContext } from './SelectPositionerContext';
 import { InternalBackdrop } from '../../utils/InternalBackdrop';
 import { inertValue } from '../../utils/inertValue';
 import { useRenderElement } from '../../utils/useRenderElement';
+import { clearPositionerStyles } from '../popup/utils';
+import { useSelectIndexContext } from '../root/SelectIndexContext';
+import { useEventCallback } from '../../utils/useEventCallback';
 
 /**
  * Positions the select menu popup.
@@ -42,20 +45,24 @@ export const SelectPositioner = React.forwardRef(function SelectPositioner(
   const {
     open,
     mounted,
+    positionerElement,
     setPositionerElement,
     listRef,
     labelsRef,
     floatingRootContext,
     modal,
     touchModality,
-    scrollUpArrowVisible,
-    setScrollUpArrowVisible,
-    scrollDownArrowVisible,
-    setScrollDownArrowVisible,
     alignItemWithTriggerActiveRef,
+    valuesRef,
+    value,
+    setLabel,
   } = useSelectRootContext();
+  const { setSelectedIndex } = useSelectIndexContext();
 
+  const [scrollUpArrowVisible, setScrollUpArrowVisible] = React.useState(false);
+  const [scrollDownArrowVisible, setScrollDownArrowVisible] = React.useState(false);
   const [controlledItemAnchor, setControlledItemAnchor] = React.useState(alignItemWithTrigger);
+
   const alignItemWithTriggerActive = mounted && controlledItemAnchor && !touchModality;
 
   React.useImperativeHandle(alignItemWithTriggerActiveRef, () => alignItemWithTriggerActive);
@@ -114,12 +121,58 @@ export const SelectPositioner = React.forwardRef(function SelectPositioner(
       alignItemWithTriggerActive,
       controlledItemAnchor,
       setControlledItemAnchor,
+      scrollUpArrowVisible,
+      setScrollUpArrowVisible,
+      scrollDownArrowVisible,
+      setScrollDownArrowVisible,
     }),
-    [positioner, alignItemWithTriggerActive, controlledItemAnchor],
+    [
+      positioner,
+      alignItemWithTriggerActive,
+      controlledItemAnchor,
+      scrollUpArrowVisible,
+      scrollDownArrowVisible,
+    ],
   );
 
+  const prevMapSizeRef = React.useRef(0);
+
+  const onMapChange = useEventCallback((map: Map<Element, { index?: number | null } | null>) => {
+    if (map.size === 0 && prevMapSizeRef.current === 0) {
+      return;
+    }
+
+    if (valuesRef.current.length === 0) {
+      return;
+    }
+
+    const prevSize = prevMapSizeRef.current;
+    prevMapSizeRef.current = map.size;
+
+    if (map.size === prevSize) {
+      return;
+    }
+
+    if (value !== null) {
+      const valueIndex = valuesRef.current.indexOf(value);
+      if (valueIndex === -1) {
+        setSelectedIndex(null);
+        setLabel('');
+      }
+    }
+
+    if (open && alignItemWithTriggerActive) {
+      setScrollDownArrowVisible(false);
+      setScrollUpArrowVisible(false);
+
+      if (positionerElement) {
+        clearPositionerStyles(positionerElement, { height: '' });
+      }
+    }
+  });
+
   return (
-    <CompositeList elementsRef={listRef} labelsRef={labelsRef}>
+    <CompositeList elementsRef={listRef} labelsRef={labelsRef} onMapChange={onMapChange}>
       <SelectPositionerContext.Provider value={contextValue}>
         {mounted && modal && <InternalBackdrop inert={inertValue(!open)} />}
         {element}

packages/react/src/select/positioner/SelectPositionerContext.ts

@@ -5,6 +5,10 @@ export type SelectPositionerContext = useSelectPositioner.ReturnValue & {
   alignItemWithTriggerActive: boolean;
   controlledItemAnchor: boolean;
   setControlledItemAnchor: React.Dispatch<React.SetStateAction<boolean>>;
+  scrollUpArrowVisible: boolean;
+  setScrollUpArrowVisible: React.Dispatch<React.SetStateAction<boolean>>;
+  scrollDownArrowVisible: boolean;
+  setScrollDownArrowVisible: React.Dispatch<React.SetStateAction<boolean>>;
 };
 
 export const SelectPositionerContext = React.createContext<SelectPositionerContext | null>(null);

packages/react/src/select/root/SelectRoot.test.tsx

@@ -131,6 +131,38 @@ describe('<Select.Root />', () => {
         '',
       );
     });
+
+    it('should not update the internal value if the controlled value prop does not change', async () => {
+      const onValueChange = spy();
+      await render(
+        <Select.Root value="a" onValueChange={onValueChange}>
+          <Select.Trigger data-testid="trigger">
+            <Select.Value>a</Select.Value>
+          </Select.Trigger>
+          <Select.Portal>
+            <Select.Positioner>
+              <Select.Popup>
+                <Select.Item value="a">a</Select.Item>
+                <Select.Item value="b">b</Select.Item>
+              </Select.Popup>
+            </Select.Positioner>
+          </Select.Portal>
+        </Select.Root>,
+      );
+
+      const trigger = screen.getByTestId('trigger');
+      expect(trigger).to.have.text('a');
+
+      fireEvent.click(trigger);
+      await flushMicrotasks();
+
+      const optionB = screen.getByRole('option', { name: 'b' });
+      fireEvent.click(optionB);
+      await flushMicrotasks();
+
+      expect(onValueChange.callCount).to.equal(0);
+      expect(trigger).to.have.text('a');
+    });
   });
 
   describe('prop: onValueChange', () => {
@@ -1272,5 +1304,68 @@ describe('<Select.Root />', () => {
 
       expect(screen.queryByRole('option', { name: 'Share' })).toHaveFocus();
     });
+
+    it('unselects the selected item if removed', async () => {
+      function DynamicMenu() {
+        const [items, setItems] = React.useState(['a', 'b', 'c']);
+        const [selectedItem, setSelectedItem] = React.useState('a');
+
+        return (
+          <div>
+            <button
+              onClick={() => {
+                setItems((prev) => prev.filter((item) => item !== 'a'));
+              }}
+            >
+              Remove
+            </button>
+
+            <button
+              onClick={() => {
+                setItems(['a', 'b', 'c']);
+              }}
+            >
+              Add
+            </button>
+            <div data-testid="value">{selectedItem}</div>
+
+            <Select.Root value={selectedItem} onValueChange={setSelectedItem}>
+              <Select.Trigger>Toggle</Select.Trigger>
+              <Select.Portal>
+                <Select.Positioner>
+                  <Select.Popup>
+                    {items.map((item) => (
+                      <Select.Item key={item} value={item}>
+                        {item}
+                      </Select.Item>
+                    ))}
+                  </Select.Popup>
+                </Select.Positioner>
+              </Select.Portal>
+            </Select.Root>
+          </div>
+        );
+      }
+
+      const { user } = await renderFakeTimers(<DynamicMenu />);
+
+      const trigger = screen.getByText('Toggle');
+
+      await act(async () => {
+        trigger.focus();
+      });
+      await user.keyboard('{ArrowDown}');
+
+      expect(screen.queryByRole('option', { name: 'a' })).to.have.attribute('data-selected');
+      expect(screen.getByTestId('value')).to.have.text('a');
+
+      fireEvent.click(screen.getByText('Remove'));
+
+      expect(screen.queryByRole('option', { name: 'b' })).not.to.have.attribute('data-selected');
+
+      fireEvent.click(screen.getByText('Add'));
+
+      expect(screen.queryByRole('option', { name: 'a' })).to.have.attribute('data-selected');
+    });
   });
 });

packages/react/src/select/root/SelectRootContext.ts

@@ -25,10 +25,6 @@ export interface SelectRootContext {
   setTriggerElement: React.Dispatch<React.SetStateAction<HTMLElement | null>>;
   positionerElement: HTMLElement | null;
   setPositionerElement: React.Dispatch<React.SetStateAction<HTMLElement | null>>;
-  scrollUpArrowVisible: boolean;
-  setScrollUpArrowVisible: React.Dispatch<React.SetStateAction<boolean>>;
-  scrollDownArrowVisible: boolean;
-  setScrollDownArrowVisible: React.Dispatch<React.SetStateAction<boolean>>;
   listRef: React.MutableRefObject<Array<HTMLElement | null>>;
   popupRef: React.MutableRefObject<HTMLDivElement | null>;
   triggerProps: HTMLProps;