Extract reusable "descendant registry" pattern from ActionBar, with new support for reordering elements

packages/react/src/utils/__tests__/descendant-registry.test.tsx

@@ -0,0 +1,190 @@
+import {describe, expect, it} from 'vitest'
+import type React from 'react'
+import {Fragment, useState} from 'react'
+import {act, render} from '@testing-library/react'
+import {createDescendantRegistry} from '../descendant-registry'
+
+/**
+ * Creates a fresh registry instance with isolated helper components for each test. This ensures
+ * no state leaks between tests via a shared Context or Provider.
+ */
+function createTestRegistry() {
+  const {Provider, useRegistryState, useRegisterDescendant} = createDescendantRegistry<string>()
+
+  /**
+   * Parent component that exposes the registry values in the DOM for assertions.
+   * State is held here and passed down to the Provider.
+   */
+  function RegistryParent({children}: {children: React.ReactNode}) {
+    const [registryState, setRegistry] = useRegistryState()
+
+    return (
+      <>
+        <div data-testid="registry-values">{Array.from(registryState.values()).join(',')}</div>
+        <Provider setRegistry={setRegistry}>{children}</Provider>
+      </>
+    )
+  }
+
+  /** A leaf component that registers itself as a descendant. */
+  function Item({value}: {value: string}) {
+    useRegisterDescendant(value)
+    return null
+  }
+
+  return {RegistryParent, Item}
+}
+
+describe('createDescendantRegistry', () => {
+  it('registers descendant items inside of other components', () => {
+    const {RegistryParent, Item} = createTestRegistry()
+
+    function Wrapper({value}: {value: string}) {
+      return <Item value={value} />
+    }
+
+    const {getByTestId} = render(
+      <RegistryParent>
+        <Wrapper value="a" />
+        <Wrapper value="b" />
+        <Wrapper value="c" />
+      </RegistryParent>,
+    )
+
+    expect(getByTestId('registry-values').textContent).toBe('a,b,c')
+  })
+
+  it('registers descendant items inside of React fragments', () => {
+    const {RegistryParent, Item} = createTestRegistry()
+
+    const {getByTestId} = render(
+      <RegistryParent>
+        <Fragment>
+          <Item value="a" />
+          <Item value="b" />
+        </Fragment>
+        <Item value="c" />
+      </RegistryParent>,
+    )
+
+    expect(getByTestId('registry-values').textContent).toBe('a,b,c')
+  })
+
+  it('registers items added to the middle of children after initial render', () => {
+    const {RegistryParent, Item} = createTestRegistry()
+
+    function Test() {
+      const [showMiddle, setShowMiddle] = useState(false)
+      return (
+        <RegistryParent>
+          <Item value="a" />
+          {showMiddle && <Item value="middle" />}
+          <Item value="b" />
+          <button type="button" onClick={() => setShowMiddle(true)}>
+            Add middle
+          </button>
+        </RegistryParent>
+      )
+    }
+
+    const {getByTestId, getByRole} = render(<Test />)
+    expect(getByTestId('registry-values').textContent).toBe('a,b')
+
+    act(() => {
+      getByRole('button').click()
+    })
+
+    expect(getByTestId('registry-values').textContent).toBe('a,middle,b')
+  })
+
+  it('drops items from the registry after they unmount', () => {
+    const {RegistryParent, Item} = createTestRegistry()
+
+    function Test() {
+      const [showLast, setShowLast] = useState(true)
+      return (
+        <RegistryParent>
+          <Item value="a" />
+          <Item value="b" />
+          {showLast && <Item value="c" />}
+          <button type="button" onClick={() => setShowLast(false)}>
+            Remove last
+          </button>
+        </RegistryParent>
+      )
+    }
+
+    const {getByTestId, getByRole} = render(<Test />)
+    expect(getByTestId('registry-values').textContent).toBe('a,b,c')
+
+    act(() => {
+      getByRole('button').click()
+    })
+
+    expect(getByTestId('registry-values').textContent).toBe('a,b')
+  })
+
+  it('updates registry order when items are reordered, using key to maintain component mount', () => {
+    const {RegistryParent, Item} = createTestRegistry()
+
+    function Test() {
+      const [items, setItems] = useState(['a', 'b', 'c'])
+      return (
+        <RegistryParent>
+          {items.map(item => (
+            <Item key={item} value={item} />
+          ))}
+          <button type="button" onClick={() => setItems(['c', 'a', 'b'])}>
+            Reorder
+          </button>
+        </RegistryParent>
+      )
+    }
+
+    const {getByTestId, getByRole} = render(<Test />)
+    expect(getByTestId('registry-values').textContent).toBe('a,b,c')
+
+    act(() => {
+      getByRole('button').click()
+    })
+
+    expect(getByTestId('registry-values').textContent).toBe('c,a,b')
+  })
+
+  it('registers deep descendants added to the beginning of the tree after initial render', () => {
+    const {RegistryParent, Item} = createTestRegistry()
+
+    function DeepItem({value}: {value: string}) {
+      return (
+        <div>
+          <div>
+            <Item value={value} />
+          </div>
+        </div>
+      )
+    }
+
+    function Test() {
+      const [showFirst, setShowFirst] = useState(false)
+      return (
+        <RegistryParent>
+          {showFirst && <DeepItem value="first" />}
+          <Item value="second" />
+          <Item value="third" />
+          <button type="button" onClick={() => setShowFirst(true)}>
+            Add first
+          </button>
+        </RegistryParent>
+      )
+    }
+
+    const {getByTestId, getByRole} = render(<Test />)
+    expect(getByTestId('registry-values').textContent).toBe('second,third')
+
+    act(() => {
+      getByRole('button').click()
+    })
+
+    expect(getByTestId('registry-values').textContent).toBe('first,second,third')
+  })
+})

packages/react/src/utils/descendant-registry.tsx

@@ -77,6 +77,13 @@ export function createDescendantRegistry<T>() {
     /** State value to trigger a re-render and force all descendants to re-register. This ensures everything remains ordered. */
     const [key, setKey] = useState(0)
 
+    /**
+     * Tracks the `children` reference from the previous render to detect when the parent
+     * re-renders with new children (e.g., due to a reorder), as opposed to a Provider re-render
+     * triggered by our own `setRegistry` call (where children stay the same reference).
+     */
+    const prevChildrenRef = useRef<ReactNode>(children)
+
     // Instantiate a new map before all descendants' effects run to populate it
     useIsomorphicLayoutEffect(function instantiateNewRegistry() {
       if (workingRegistryRef.current === 'queued') {
@@ -112,11 +119,28 @@ export function createDescendantRegistry<T>() {
       }
     }, [])
 
-    // After all descendants' effects complete, commit the working registry to state
+    // After all descendants' effects complete, commit the working registry to state. When the
+    // registry is idle and the children reference changed (indicating the parent re-rendered with
+    // new children, e.g., due to a reorder), queue a rebuild so all descendants re-register in
+    // their current render order, keeping the registry accurate.
+    //
+    // This effect intentionally omits a dependency array so it runs after every render. Adding
+    // deps would prevent the registry Map from being committed after rebuild cycles (where `key`
+    // increments but `children` and `setRegistry` stay the same).
+    // eslint-disable-next-line react-hooks/exhaustive-deps
     useEffect(function commitWorkingRegistry() {
+      const childrenChanged = prevChildrenRef.current !== children
+      prevChildrenRef.current = children
+
       if (workingRegistryRef.current instanceof Map) {
         setRegistry(workingRegistryRef.current)
         workingRegistryRef.current = 'idle'
+      } else if (workingRegistryRef.current === 'idle' && childrenChanged) {
+        // The children changed (e.g., reordering) without triggering any descendant
+        // mount/unmount/update. Trigger a rebuild to capture the new render order.
+        workingRegistryRef.current = 'queued'
+
+        setKey(prev => prev + 1)
       }
     })