allow registering a Portal component to a parent

This allows us to find all the `Portal` components that are nested in a
given component without manually adding refs to every `Portal` component
itself.

This will come in handy in the `Popover` component where we will allow
focus in the child `Portal` components otherwise a focus outside of the
`Popover` will close the it. In other components we often crawl the DOM
directly using `[data-headlessui-portal]` data attributes, however this
will fetch _all_ the `Portal` components, not the ones that started in
the current component.

packages/@headlessui-react/src/components/portal/portal.tsx

@@ -10,6 +10,8 @@ import React, {
   ElementType,
   MutableRefObject,
   Ref,
+  useMemo,
+  ContextType,
 } from 'react'
 import { createPortal } from 'react-dom'
 
@@ -22,6 +24,7 @@ import { optionalRef, useSyncRefs } from '../../hooks/use-sync-refs'
 import { useOnUnmount } from '../../hooks/use-on-unmount'
 import { useOwnerDocument } from '../../hooks/use-owner'
 import { env } from '../../utils/env'
+import { useEvent } from '../../hooks/use-event'
 
 function usePortalTarget(ref: MutableRefObject<HTMLElement | null>): HTMLElement | null {
   let forceInRoot = usePortalRoot()
@@ -87,7 +90,7 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
   let [element] = useState<HTMLDivElement | null>(() =>
     env.isServer ? null : ownerDocument?.createElement('div') ?? null
   )
-
+  let parent = useContext(PortalParentContext)
   let ready = useServerHandoffComplete()
 
   useIsoMorphicEffect(() => {
@@ -101,6 +104,13 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     }
   }, [target, element])
 
+  useIsoMorphicEffect(() => {
+    if (!element) return
+    if (!parent) return
+
+    return parent.register(element)
+  }, [parent, element])
+
   useOnUnmount(() => {
     if (!target || !element) return
 
@@ -164,6 +174,41 @@ function GroupFn<TTag extends ElementType = typeof DEFAULT_GROUP_TAG>(
 
 // ---
 
+let PortalParentContext = createContext<{
+  register: (portal: HTMLElement) => () => void
+  unregister: (portal: HTMLElement) => void
+} | null>(null)
+
+export function useNestedPortals() {
+  let portals = useRef<HTMLElement[]>([])
+
+  let register = useEvent((portal: HTMLElement) => {
+    portals.current.push(portal)
+    return () => unregister(portal)
+  })
+
+  let unregister = useEvent((portal: HTMLElement) => {
+    let idx = portals.current.indexOf(portal)
+    if (idx !== -1) portals.current.splice(idx, 1)
+  })
+
+  let api = useMemo<ContextType<typeof PortalParentContext>>(
+    () => ({ register, unregister }),
+    [register, unregister]
+  )
+
+  return [
+    portals,
+    useMemo(() => {
+      return function PortalWrapper({ children }: { children: React.ReactNode }) {
+        return <PortalParentContext.Provider value={api}>{children}</PortalParentContext.Provider>
+      }
+    }, [api]),
+  ] as const
+}
+
+// ---
+
 interface ComponentPortal extends HasDisplayName {
   <TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     props: PortalProps<TTag> & RefProp<typeof PortalFn>

packages/@headlessui-react/src/index.test.ts

@@ -6,14 +6,14 @@ import * as HeadlessUI from './index'
  */
 it('should expose the correct components', () => {
   expect(Object.keys(HeadlessUI)).toEqual([
+    'Portal',
     'Combobox',
     'Dialog',
     'Disclosure',
     'FocusTrap',
     'Listbox',
     'Menu',
     'Popover',
-    'Portal',
     'RadioGroup',
     'Switch',
     'Tab',

packages/@headlessui-react/src/index.ts

@@ -7,8 +7,8 @@ export * from './components/focus-trap/focus-trap'
 export * from './components/listbox/listbox'
 export * from './components/menu/menu'
 export * from './components/popover/popover'
-export * from './components/portal/portal'
 export * from './components/radio-group/radio-group'
 export * from './components/switch/switch'
 export * from './components/tabs/tabs'
 export * from './components/transitions/transition'
+export { Portal } from './components/portal/portal'