React descendant registry

descendant-registry.tsx

function createDescendantRegistry<T>() {
  const Context = createContext<{
    register: (id: string, value: T) => () => void
    key: number
  }>({
    register: () => () => {},
    key: -1,
  })

  function useRegistryState() {
    return useState(() => new Map<string, T>())
  }

  function useRegisterDescendant(value: T) {
    const {register, key} = useContext(Context)

    const id = useId()

    useEffect(() => register(id, value), [register, id, value, key])
  }

  function Provider({
    children,
    setRegistry,
  }: {
    children: ReactNode
    setRegistry: Dispatch<React.SetStateAction<Map<string, T>>>
  }) {
    const workingRegistryRef = useRef<Map<string, T> | 'queued' | 'idle'>('queued')

    /** State value to trigger a re-render and force all descendants to re-register. This ensures everything remains ordered. */
    const [key, setKey] = useState(0)

    // Instantiate a new map before all descendants' effects run to populate it
    useLayoutEffect(function instantiateNewRegistry() {
      if (workingRegistryRef.current === 'queued') {
        workingRegistryRef.current = new Map<string, T>()
      }
    })

    // Each descendant will register itself. Initially this will just update the map. After the initial render and
    // commit, any descendant updates (mounts, unmounts, or value changes) will trigger a full recalculation of the
    // registry. This is necessary because a descendant might mount into the middle of the tree, and the only way for
    // the provider to know that is to re-render the tree and listen to every descendant effect in order.
    const register = useCallback((id: string, value: T) => {
      function queueRebuild() {
        // If we instantiate a new map now, we'd commit it in the ensuing provider effect, which would trigger a
        // new state update and cause an infinite render loop. So we wait to instantiate it until the next layout effect,
        // which we know will happen because we update state to queue a new render.
        workingRegistryRef.current = 'queued'
        setKey(prev => prev + 1)
      }

      if (workingRegistryRef.current instanceof Map) {
        workingRegistryRef.current.set(id, value)
      } else if (workingRegistryRef.current === 'idle') {
        queueRebuild()
      }

      return () => {
        if (workingRegistryRef.current instanceof Map) {
          workingRegistryRef.current.delete(id)
        } else if (workingRegistryRef.current === 'idle') {
          queueRebuild()
        }
      }
    }, [])

    // After all descendants' effects complete, commit the working registry to state
    useEffect(function commitWorkingRegistry() {
      if (workingRegistryRef.current instanceof Map) {
        setRegistry(workingRegistryRef.current)
        workingRegistryRef.current = 'idle'
      }
    })

    const contextValue = useMemo(() => ({register, key}), [register, key])

    return <Context.Provider value={contextValue}>{children}</Context.Provider>
  }

  return {Provider, useRegistryState, useRegisterDescendant}
}