[portals] Ensure container is reactive from undefined

Author: atomiks

packages/react/src/floating-ui-react/components/FloatingPortal.test.tsx

@@ -1,13 +1,14 @@
-import { fireEvent, flushMicrotasks, render, screen } from '@mui/internal-test-utils';
 import * as React from 'react';
-
+import { fireEvent, flushMicrotasks, render, screen } from '@mui/internal-test-utils';
 import { isJSDOM } from '@base-ui-components/utils/detectBrowser';
 import { FloatingPortal, useFloating } from '../index';
+import type { UseFloatingPortalNodeProps } from './FloatingPortal';
+
+interface AppProps {
+  root?: UseFloatingPortalNodeProps['root'];
+}
 
-function App(props: {
-  root?: HTMLElement | null | React.RefObject<HTMLElement | null>;
-  id?: string;
-}) {
+function App(props: AppProps) {
   const [open, setOpen] = React.useState(false);
   const { refs } = useFloating({
     open,
@@ -25,30 +26,6 @@ function App(props: {
 }
 
 describe.skipIf(!isJSDOM)('FloatingPortal', () => {
-  test('creates a custom id node', async () => {
-    render(<App id="custom-id" />);
-    await flushMicrotasks();
-    expect(document.querySelector('#custom-id')).toBeInTheDocument();
-  });
-
-  test('uses a custom id node as the root', async () => {
-    const customRoot = document.createElement('div');
-    customRoot.id = 'custom-root';
-    document.body.appendChild(customRoot);
-    render(<App id="custom-root" />);
-    fireEvent.click(screen.getByTestId('reference'));
-    await flushMicrotasks();
-    expect(screen.getByTestId('floating').parentElement?.parentElement).toBe(customRoot);
-    customRoot.remove();
-  });
-
-  test('creates a custom id node as the root', async () => {
-    render(<App id="custom-id" />);
-    fireEvent.click(screen.getByTestId('reference'));
-    await flushMicrotasks();
-    expect(screen.getByTestId('floating').parentElement?.parentElement?.id).toBe('custom-id');
-  });
-
   test('allows custom roots', async () => {
     const customRoot = document.createElement('div');
     customRoot.id = 'custom-root';
@@ -89,7 +66,7 @@ describe.skipIf(!isJSDOM)('FloatingPortal', () => {
       return (
         <React.Fragment>
           {renderRoot && <div ref={setRoot} data-testid="root" />}
-          <App root={root} />;
+          <App root={root} />
         </React.Fragment>
       );
     }
@@ -103,4 +80,43 @@ describe.skipIf(!isJSDOM)('FloatingPortal', () => {
     const root = screen.getByTestId('root');
     expect(root).toBe(subRoot?.parentElement);
   });
+
+  test('reattaches the portal when the root changes', async () => {
+    const customRoot = document.createElement('div');
+    document.body.appendChild(customRoot);
+
+    try {
+      function RootSwitcher() {
+        const [root, setRoot] = React.useState<UseFloatingPortalNodeProps['root']>(undefined);
+
+        return (
+          <React.Fragment>
+            <App root={root} />
+            <button onClick={() => setRoot(undefined)} data-testid="use-undefined" />
+            <button onClick={() => setRoot(customRoot)} data-testid="use-element" />
+          </React.Fragment>
+        );
+      }
+
+      render(<RootSwitcher />);
+
+      fireEvent.click(screen.getByTestId('reference'));
+
+      expect((await screen.findByTestId('floating')).parentElement?.parentElement).toBe(
+        document.body,
+      );
+
+      fireEvent.click(screen.getByTestId('use-element'));
+
+      expect((await screen.findByTestId('floating')).parentElement?.parentElement).toBe(customRoot);
+
+      fireEvent.click(screen.getByTestId('use-undefined'));
+
+      const floatingInBodyAgain = await screen.findByTestId('floating');
+      expect(floatingInBodyAgain.parentElement?.parentElement).toBe(document.body);
+      expect(customRoot.contains(floatingInBodyAgain)).toBe(false);
+    } finally {
+      customRoot.remove();
+    }
+  });
 });

packages/react/src/floating-ui-react/components/FloatingPortal.tsx

@@ -38,23 +38,24 @@ export const usePortalContext = () => React.useContext(PortalContext);
 const attr = createAttribute('portal');
 
 export interface UseFloatingPortalNodeProps {
-  id?: string;
   root?: HTMLElement | ShadowRoot | null | React.RefObject<HTMLElement | ShadowRoot | null>;
 }
 
 /**
  * @see https://floating-ui.com/docs/FloatingPortal#usefloatingportalnode
  */
 export function useFloatingPortalNode(props: UseFloatingPortalNodeProps = {}) {
-  const { id, root } = props;
+  const { root } = props;
 
   const uniqueId = useId();
   const portalContext = usePortalContext();
 
   const [portalNode, setPortalNode] = React.useState<HTMLElement | null>(null);
 
   const portalNodeRef = React.useRef<HTMLDivElement | null>(null);
+  const prevContainerRef = React.useRef<HTMLElement | ShadowRoot | null>(null);
 
+  // Cleanup when the portal node instance changes or unmounts.
   useIsoLayoutEffect(() => {
     return () => {
       portalNode?.remove();
@@ -63,81 +64,60 @@ export function useFloatingPortalNode(props: UseFloatingPortalNodeProps = {}) {
       // https://github.com/floating-ui/floating-ui/issues/2454
       queueMicrotask(() => {
         portalNodeRef.current = null;
+        prevContainerRef.current = null;
       });
     };
   }, [portalNode]);
 
+  // Handle reactive `root` changes (including undefined <-> container changes)
   useIsoLayoutEffect(() => {
-    // Wait for the uniqueId to be generated before creating the portal node in
-    // React <18 (using `useFloatingId` instead of the native `useId`).
-    // https://github.com/floating-ui/floating-ui/issues/2778
-    if (!uniqueId) {
-      return;
-    }
-    if (portalNodeRef.current) {
-      return;
-    }
-    const existingIdRoot = id ? document.getElementById(id) : null;
-    if (!existingIdRoot) {
-      return;
-    }
-
-    const subRoot = document.createElement('div');
-    subRoot.id = uniqueId;
-    subRoot.setAttribute(attr, '');
-    existingIdRoot.appendChild(subRoot);
-    portalNodeRef.current = subRoot;
-    setPortalNode(subRoot);
-  }, [id, uniqueId]);
-
-  useIsoLayoutEffect(() => {
-    // Wait for the root to exist before creating the portal node. The root must
-    // be stored in state, not a ref, for this to work reactively.
+    // "Wait" mode: remove any existing node and pause until root changes.
     if (root === null) {
+      if (portalNodeRef.current) {
+        portalNodeRef.current.remove();
+        portalNodeRef.current = null;
+        setPortalNode(null);
+      }
+      prevContainerRef.current = null;
       return;
     }
+
+    // For React 17, as the id is generated in an effect instead of React.useId().
     if (!uniqueId) {
       return;
     }
-    if (portalNodeRef.current) {
-      return;
-    }
 
-    let container = root || portalContext?.portalNode;
-    if (container && !isNode(container)) {
-      container = container.current;
-    }
-    container = container || document.body;
+    const resolvedContainer =
+      (root && (isNode(root) ? root : root.current)) || portalContext?.portalNode || document.body;
+
+    const containerChanged = resolvedContainer !== prevContainerRef.current;
 
-    let idWrapper: HTMLDivElement | null = null;
-    if (id) {
-      idWrapper = document.createElement('div');
-      idWrapper.id = id;
-      container.appendChild(idWrapper);
+    if (portalNodeRef.current && containerChanged) {
+      portalNodeRef.current.remove();
+      portalNodeRef.current = null;
+      setPortalNode(null);
     }
 
-    const subRoot = document.createElement('div');
+    if (portalNodeRef.current) {
+      return;
+    }
 
-    subRoot.id = uniqueId;
-    subRoot.setAttribute(attr, '');
+    const portalElement = document.createElement('div');
+    portalElement.id = uniqueId;
+    portalElement.setAttribute(attr, '');
+    resolvedContainer.appendChild(portalElement);
 
-    container = idWrapper || container;
-    container.appendChild(subRoot);
+    portalNodeRef.current = portalElement;
+    prevContainerRef.current = resolvedContainer;
 
-    portalNodeRef.current = subRoot;
-    setPortalNode(subRoot);
-  }, [id, root, uniqueId, portalContext]);
+    setPortalNode(portalElement);
+  }, [root, uniqueId, portalContext]);
 
   return portalNode;
 }
 
 export interface FloatingPortalProps {
   children?: React.ReactNode;
-  /**
-   * Optionally selects the node with the id if it exists, or create it and
-   * append it to the specified `root` (by default `document.body`).
-   */
-  id?: string;
   /**
    * Specifies the root node the portal container will be appended to.
    */
@@ -160,9 +140,9 @@ export interface FloatingPortalProps {
  * @internal
  */
 export function FloatingPortal(props: FloatingPortalProps): React.JSX.Element {
-  const { children, id, root, preserveTabOrder = true } = props;
+  const { children, root, preserveTabOrder = true } = props;
 
-  const portalNode = useFloatingPortalNode({ id, root });
+  const portalNode = useFloatingPortalNode({ root });
   const [focusManagerState, setFocusManagerState] = React.useState<FocusManagerState>(null);
 
   const beforeOutsideRef = React.useRef<HTMLSpanElement>(null);
@@ -211,10 +191,7 @@ export function FloatingPortal(props: FloatingPortalProps): React.JSX.Element {
   }, [portalNode, preserveTabOrder, modal]);
 
   React.useEffect(() => {
-    if (!portalNode) {
-      return;
-    }
-    if (open) {
+    if (!portalNode || open) {
       return;
     }
     enableFocusInside(portalNode);