refactor useRootContainers and introduce MainTreeProvider

RobinMalfait · RobinMalfait · commit 2453c374e70c · 2024-07-04T16:58:25.000+02:00
As a general recap, when an outside click happens, we need to react to it and typically use the `useOutsideClick` hook. We also require the context of "allowed root containers", this means that clicking on a 3rd party toast when a dialog is open, that we allow this even though we are technically clicking outside of the dialog. This is simply because we don't have control over these elements. We also need a reference to know what the "main tree" container is, because this is the container where your application lives and we _know_ that we are not allowed to click on anything in this container. The complex part is getting a reference to this element. ```html <html> <head> <title></title> </head> <body> <div id="app"> <-- main root container --> <div></div> <div> <Popover></Popover>  </div> <div></div> </div>  <3rd-party-toast-container></3rd-party-toast-container> </body>  <grammarly-extension></grammarly-extension> </html> ``` Some examples: - In case of a `Dialog`, the `Dialog` is rendered in a `Portal` which means that a DOM ref to the `Dialog` or anything inside will not point to the "main tree" node. - In case of a `Popover` we can use the `PopoverButton` as an element that lives in the main tree. However, if you work with nested `Popover` components, and the outer `PopoverPanel` uses the `anchor` or `portal` props, then the inner `PortalButton` will not be in the main tree either because it will live in the portalled `PopoverPanel` of the parent. This is where the `MainTreeProvider` comes in handy. This component will use the passed in `node` as the main tree node reference and pass this via the context through the React tree. This means that a nested `Popover` will still use a reference from the parent `Popover`. In case of the `Dialog`, we wrap the `Dialog` itself with this provider which means that the provider will be in the main tree and can be used inside the portalled `Dialog`. Another part of the `MainTreeProvider` is that if no node exists in the parent (reading from context), and no node is provided via props, then we will briefly render a hidden element, find the root node of the main tree (aka, the parent element that is a direct child of the body, `body > *`). Once we found it, we remove the hidden element again. This way we don't keep unnecessary DOM nodes around.
diff --git a/packages/@headlessui-react/src/components/dialog/dialog.tsx b/packages/@headlessui-react/src/components/dialog/dialog.tsx
@@ -25,12 +25,15 @@ import { useIsTouchDevice } from '../../hooks/use-is-touch-device'
 import { useOnDisappear } from '../../hooks/use-on-disappear'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 import { useOwnerDocument } from '../../hooks/use-owner'
-import { useRootContainers } from '../../hooks/use-root-containers'
+import {
+  MainTreeProvider,
+  useMainTreeNode,
+  useRootContainers,
+} from '../../hooks/use-root-containers'
 import { useScrollLock } from '../../hooks/use-scroll-lock'
 import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
 import { CloseProvider } from '../../internal/close-provider'
-import { HoistFormFields } from '../../internal/form-fields'
 import { ResetOpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { ForcePortalRoot } from '../../internal/portal-force-root'
 import type { Props } from '../../types'
@@ -180,11 +183,9 @@ let InternalDialog = forwardRefWithAs(function InternalDialog<
     },
   }
 
-  let {
-    resolveContainers: resolveRootContainers,
-    mainTreeNodeRef,
-    MainTreeNode,
-  } = useRootContainers({
+  let mainTreeNode = useMainTreeNode()
+  let { resolveContainers: resolveRootContainers } = useRootContainers({
+    mainTreeNode,
     portals,
     defaultContainers: [defaultContainer],
   })
@@ -207,8 +208,7 @@ let InternalDialog = forwardRefWithAs(function InternalDialog<
     ]),
     disallowed: useEvent(() => [
       // Disallow the "main" tree root node
-      mainTreeNodeRef.current?.closest<HTMLElement>('body > *:not(#headlessui-portal-root)') ??
-        null,
+      mainTreeNode?.closest<HTMLElement>('body > *:not(#headlessui-portal-root)') ?? null,
     ]),
   })
 
@@ -320,9 +320,6 @@ let InternalDialog = forwardRefWithAs(function InternalDialog<
           </DialogContext.Provider>
         </Portal>
       </ForcePortalRoot>
-      <HoistFormFields>
-        <MainTreeNode />
-      </HoistFormFields>
     </ResetOpenClosedProvider>
   )
 })
@@ -395,13 +392,19 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
 
   if ((open !== undefined || transition) && !rest.static) {
     return (
-      <Transition show={open} transition={transition} unmount={rest.unmount}>
-        <InternalDialog ref={ref} {...rest} />
-      </Transition>
+      <MainTreeProvider>
+        <Transition show={open} transition={transition} unmount={rest.unmount}>
+          <InternalDialog ref={ref} {...rest} />
+        </Transition>
+      </MainTreeProvider>
     )
   }
 
-  return <InternalDialog ref={ref} open={open} {...rest} />
+  return (
+    <MainTreeProvider>
+      <InternalDialog ref={ref} open={open} {...rest} />
+    </MainTreeProvider>
+  )
 }
 
 // ---
diff --git a/packages/@headlessui-react/src/components/popover/popover.tsx b/packages/@headlessui-react/src/components/popover/popover.tsx
@@ -32,7 +32,11 @@ import { useOnDisappear } from '../../hooks/use-on-disappear'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 import { useOwnerDocument } from '../../hooks/use-owner'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
-import { useMainTreeNode, useRootContainers } from '../../hooks/use-root-containers'
+import {
+  MainTreeProvider,
+  useMainTreeNode,
+  useRootContainers,
+} from '../../hooks/use-root-containers'
 import { useScrollLock } from '../../hooks/use-scroll-lock'
 import { optionalRef, useSyncRefs } from '../../hooks/use-sync-refs'
 import { Direction as TabDirection, useTabDirection } from '../../hooks/use-tab-direction'
@@ -195,7 +199,6 @@ let PopoverGroupContext = createContext<{
   unregisterPopover(registerBag: PopoverRegisterBag): void
   isFocusWithinPopoverGroup(): boolean
   closeOthers(buttonId: string): void
-  mainTreeNodeRef: MutableRefObject<HTMLElement | null>
 } | null>(null)
 PopoverGroupContext.displayName = 'PopoverGroupContext'
 
@@ -343,8 +346,9 @@ function PopoverFn<TTag extends ElementType = typeof DEFAULT_POPOVER_TAG>(
   useEffect(() => registerPopover?.(registerBag), [registerPopover, registerBag])
 
   let [portals, PortalWrapper] = useNestedPortals()
+  let mainTreeNode = useMainTreeNode(button)
   let root = useRootContainers({
-    mainTreeNodeRef: groupContext?.mainTreeNodeRef,
+    mainTreeNode,
     portals,
     defaultContainers: [button, panel],
   })
@@ -416,33 +420,34 @@ function PopoverFn<TTag extends ElementType = typeof DEFAULT_POPOVER_TAG>(
   let ourProps = { ref: popoverRef }
 
   return (
-    <FloatingProvider>
-      <PopoverPanelContext.Provider value={null}>
-        <PopoverContext.Provider value={reducerBag}>
-          <PopoverAPIContext.Provider value={api}>
-            <CloseProvider value={close}>
-              <OpenClosedProvider
-                value={match(popoverState, {
-                  [PopoverStates.Open]: State.Open,
-                  [PopoverStates.Closed]: State.Closed,
-                })}
-              >
-                <PortalWrapper>
-                  {render({
-                    ourProps,
-                    theirProps,
-                    slot,
-                    defaultTag: DEFAULT_POPOVER_TAG,
-                    name: 'Popover',
+    <MainTreeProvider node={mainTreeNode}>
+      <FloatingProvider>
+        <PopoverPanelContext.Provider value={null}>
+          <PopoverContext.Provider value={reducerBag}>
+            <PopoverAPIContext.Provider value={api}>
+              <CloseProvider value={close}>
+                <OpenClosedProvider
+                  value={match(popoverState, {
+                    [PopoverStates.Open]: State.Open,
+                    [PopoverStates.Closed]: State.Closed,
                   })}
-                  <root.MainTreeNode />
-                </PortalWrapper>
-              </OpenClosedProvider>
-            </CloseProvider>
-          </PopoverAPIContext.Provider>
-        </PopoverContext.Provider>
-      </PopoverPanelContext.Provider>
-    </FloatingProvider>
+                >
+                  <PortalWrapper>
+                    {render({
+                      ourProps,
+                      theirProps,
+                      slot,
+                      defaultTag: DEFAULT_POPOVER_TAG,
+                      name: 'Popover',
+                    })}
+                  </PortalWrapper>
+                </OpenClosedProvider>
+              </CloseProvider>
+            </PopoverAPIContext.Provider>
+          </PopoverContext.Provider>
+        </PopoverPanelContext.Provider>
+      </FloatingProvider>
+    </MainTreeProvider>
   )
 }
 
@@ -1110,7 +1115,6 @@ function GroupFn<TTag extends ElementType = typeof DEFAULT_GROUP_TAG>(
   let internalGroupRef = useRef<HTMLElement | null>(null)
   let groupRef = useSyncRefs(internalGroupRef, ref)
   let [popovers, setPopovers] = useState<PopoverRegisterBag[]>([])
-  let root = useMainTreeNode()
 
   let unregisterPopover = useEvent((registerBag: PopoverRegisterBag) => {
     setPopovers((existing) => {
@@ -1157,15 +1161,8 @@ function GroupFn<TTag extends ElementType = typeof DEFAULT_GROUP_TAG>(
       unregisterPopover: unregisterPopover,
       isFocusWithinPopoverGroup,
       closeOthers,
-      mainTreeNodeRef: root.mainTreeNodeRef,
     }),
-    [
-      registerPopover,
-      unregisterPopover,
-      isFocusWithinPopoverGroup,
-      closeOthers,
-      root.mainTreeNodeRef,
-    ]
+    [registerPopover, unregisterPopover, isFocusWithinPopoverGroup, closeOthers]
   )
 
   let slot = useMemo(() => ({}) satisfies GroupRenderPropArg, [])
@@ -1174,16 +1171,17 @@ function GroupFn<TTag extends ElementType = typeof DEFAULT_GROUP_TAG>(
   let ourProps = { ref: groupRef }
 
   return (
-    <PopoverGroupContext.Provider value={contextBag}>
-      {render({
-        ourProps,
-        theirProps,
-        slot,
-        defaultTag: DEFAULT_GROUP_TAG,
-        name: 'Popover.Group',
-      })}
-      <root.MainTreeNode />
-    </PopoverGroupContext.Provider>
+    <MainTreeProvider>
+      <PopoverGroupContext.Provider value={contextBag}>
+        {render({
+          ourProps,
+          theirProps,
+          slot,
+          defaultTag: DEFAULT_GROUP_TAG,
+          name: 'Popover.Group',
+        })}
+      </PopoverGroupContext.Provider>
+    </MainTreeProvider>
   )
 }
 
diff --git a/packages/@headlessui-react/src/hooks/use-root-containers.tsx b/packages/@headlessui-react/src/hooks/use-root-containers.tsx
@@ -1,20 +1,21 @@
-import React, { useMemo, useRef, type MutableRefObject } from 'react'
+import React, { createContext, useContext, useState, type MutableRefObject } from 'react'
 import { Hidden, HiddenFeatures } from '../internal/hidden'
+import { getOwnerDocument } from '../utils/owner'
 import { useEvent } from './use-event'
 import { useOwnerDocument } from './use-owner'
 
 export function useRootContainers({
   defaultContainers = [],
   portals,
-  mainTreeNodeRef: _mainTreeNodeRef,
+
+  // Reference to a node in the "main" tree, not in the portalled Dialog tree.
+  mainTreeNode,
 }: {
   defaultContainers?: (HTMLElement | null | MutableRefObject<HTMLElement | null>)[]
   portals?: MutableRefObject<HTMLElement[]>
-  mainTreeNodeRef?: MutableRefObject<HTMLElement | null>
+  mainTreeNode?: HTMLElement | null
 } = {}) {
-  // Reference to a node in the "main" tree, not in the portalled Dialog tree.
-  let mainTreeNodeRef = useRef<HTMLElement | null>(_mainTreeNodeRef?.current ?? null)
-  let ownerDocument = useOwnerDocument(mainTreeNodeRef)
+  let ownerDocument = useOwnerDocument(mainTreeNode)
 
   let resolveContainers = useEvent(() => {
     let containers: HTMLElement[] = []
@@ -42,8 +43,10 @@ export function useRootContainers({
       if (container === document.head) continue // Skip `<head>`
       if (!(container instanceof HTMLElement)) continue // Skip non-HTMLElements
       if (container.id === 'headlessui-portal-root') continue // Skip the Headless UI portal root
-      if (container.contains(mainTreeNodeRef.current)) continue // Skip if it is the main app
-      if (container.contains((mainTreeNodeRef.current?.getRootNode() as ShadowRoot)?.host)) continue // Skip if it is the main app (and the component is inside a shadow root)
+      if (mainTreeNode) {
+        if (container.contains(mainTreeNode)) continue // Skip if it is the main app
+        if (container.contains((mainTreeNode?.getRootNode() as ShadowRoot)?.host)) continue // Skip if it is the main app (and the component is inside a shadow root)
+      }
       if (containers.some((defaultContainer) => container.contains(defaultContainer))) continue // Skip if the current container is part of a container we've already seen (e.g.: default container / portal)
 
       containers.push(container)
@@ -57,25 +60,89 @@ export function useRootContainers({
     contains: useEvent((element: HTMLElement) =>
       resolveContainers().some((container) => container.contains(element))
     ),
-    mainTreeNodeRef,
-    MainTreeNode: useMemo(() => {
-      return function MainTreeNode() {
-        if (_mainTreeNodeRef != null) return null
-        return <Hidden features={HiddenFeatures.Hidden} ref={mainTreeNodeRef} />
-      }
-    }, [mainTreeNodeRef, _mainTreeNodeRef]),
   }
 }
 
-export function useMainTreeNode() {
-  let mainTreeNodeRef = useRef<HTMLElement | null>(null)
+let MainTreeContext = createContext<HTMLElement | null>(null)
 
-  return {
-    mainTreeNodeRef,
-    MainTreeNode: useMemo(() => {
-      return function MainTreeNode() {
-        return <Hidden features={HiddenFeatures.Hidden} ref={mainTreeNodeRef} />
-      }
-    }, [mainTreeNodeRef]),
-  }
+/**
+ * A provider for the main tree node.
+ *
+ * When a component is rendered in a `Portal`, it is no longer part of the main
+ * tree. This provider helps to find the main tree node and pass it along to the
+ * components that need it.
+ *
+ * The main tree node is used for features such as outside click behavior, where
+ * we allow clicks in 3rd party contains, but not in the parent of the "main
+ * tree".
+ *
+ * In case of a `Popover`, we can use the `PopoverButton` as a marker in the
+ * "main tree", the `PopoverPanel` can't be used because it could be rendered in
+ * a `Portal` (e.g.: when using the `anchor` props).
+ *
+ * However, we can't use the `PopoverButton` when it's nested inside of another
+ * `Popover`'s `PopoverPanel` component (if that parent `PopoverPanel` is
+ * rendered in a `Portal`).
+ *
+ * This is where the `MainTreeProvider` comes in. It will find the "main tree"
+ * node and pass it on. The top-level `PopoverButton` will be used as a marker
+ * in the "main tree" and nested `Popover` use this button as well.
+ */
+export function MainTreeProvider({
+  children,
+  node,
+}: {
+  children: React.ReactNode
+  node?: HTMLElement | null
+}) {
+  let [mainTreeNode, setMainTreeNode] = useState<HTMLElement | null>(null)
+
+  // 1. Prefer the main tree node from context
+  // 2. Prefer the provided node
+  // 3. Create a new node at this point, and find the main tree node
+  let resolvedMainTreeNode = useMainTreeNode(node ?? mainTreeNode)
+
+  return (
+    <MainTreeContext.Provider value={resolvedMainTreeNode}>
+      {children}
+
+      {/**
+       * If no main tree node is found at this point, then we briefly render an
+       * element to find the main tree node and pass it along.
+       */}
+      {resolvedMainTreeNode === null && (
+        <Hidden
+          features={HiddenFeatures.Hidden}
+          ref={(el) => {
+            if (el) {
+              // We will only render this when no `mainTreeNode` is found. This
+              // means that if we render this element and use it as the
+              // `mainTreeNode` that we will be unmounting it later again.
+              //
+              // However, we can resolve the actual root container of the main
+              // tree node and use that instead.
+              for (let container of getOwnerDocument(el)?.querySelectorAll('html > *, body > *') ??
+                []) {
+                if (container === document.body) continue // Skip `<body>`
+                if (container === document.head) continue // Skip `<head>`
+                if (!(container instanceof HTMLElement)) continue // Skip non-HTMLElements
+                if (container?.contains(el)) {
+                  setMainTreeNode(container)
+                  break
+                }
+              }
+            }
+          }}
+        />
+      )}
+    </MainTreeContext.Provider>
+  )
+}
+
+/**
+ * Get the main tree node from context or fallback to the optionally provided node.
+ */
+export function useMainTreeNode(fallbackMainTreeNode: HTMLElement | null = null) {
+  // Prefer the main tree node from context, but fallback to the provided node.
+  return useContext(MainTreeContext) ?? fallbackMainTreeNode
 }
