[TrapFocus] Capture nodeToRestore via relatedTarget (#26696)

Author: eps1lon

docs/pages/api-docs/unstable-trap-focus.json

@@ -5,10 +5,6 @@
     "disableAutoFocus": { "type": { "name": "bool" } },
     "disableEnforceFocus": { "type": { "name": "bool" } },
     "disableRestoreFocus": { "type": { "name": "bool" } },
-    "getDoc": {
-      "type": { "name": "func" },
-      "default": "function defaultGetDoc() {\n  return document;\n}"
-    },
     "getTabbable": { "type": { "name": "func" } },
     "isEnabled": {
       "type": { "name": "func" },

docs/translations/api-docs/unstable-trap-focus/unstable-trap-focus.json

@@ -5,7 +5,6 @@
     "disableAutoFocus": "If <code>true</code>, the trap focus will not automatically shift focus to itself when it opens, and replace it to the last focused element when it closes. This also works correctly with any trap focus children that have the <code>disableAutoFocus</code> prop.<br>Generally this should never be set to <code>true</code> as it makes the trap focus less accessible to assistive technologies, like screen readers.",
     "disableEnforceFocus": "If <code>true</code>, the trap focus will not prevent focus from leaving the trap focus while open.<br>Generally this should never be set to <code>true</code> as it makes the trap focus less accessible to assistive technologies, like screen readers.",
     "disableRestoreFocus": "If <code>true</code>, the trap focus will not restore focus to previously focused element once trap focus is hidden.",
-    "getDoc": "Return the document the trap focus is mounted into. Provide the prop if you need the restore focus to work between different documents.",
     "getTabbable": "Returns an array of ordered tabbable nodes (i.e. in tab order) within the root. For instance, you can provide the &quot;tabbable&quot; npm dependency.<br><br><strong>Signature:</strong><br><code>function(root: HTMLElement) =&gt; void</code><br>",
     "isEnabled": "This prop extends the <code>open</code> prop. It allows to toggle the open state without having to wait for a rerender when changing the <code>open</code> prop. This prop should be memoized. It can be used to support multiple trap focus mounted at the same time.",
     "open": "If <code>true</code>, focus is locked."

packages/material-ui-lab/src/internal/pickers/PickersPopper.tsx

@@ -253,7 +253,6 @@ const PickersPopper = (props: PickerPopperProps) => {
           disableAutoFocus
           disableEnforceFocus={role === 'tooltip'}
           isEnabled={() => true}
-          getDoc={() => paperRef.current?.ownerDocument ?? document}
           {...TrapFocusProps}
         >
           <TransitionComponent {...TransitionProps}>

packages/material-ui-unstyled/src/ModalUnstyled/ModalUnstyled.js

@@ -273,7 +273,6 @@ const ModalUnstyled = React.forwardRef(function ModalUnstyled(props, ref) {
           disableEnforceFocus={disableEnforceFocus}
           disableAutoFocus={disableAutoFocus}
           disableRestoreFocus={disableRestoreFocus}
-          getDoc={getDoc}
           isEnabled={isTopModal}
           open={open}
         >

packages/material-ui-unstyled/src/Unstable_TrapFocus/Unstable_TrapFocus.d.ts

@@ -6,14 +6,6 @@ export interface TrapFocusProps {
    * If `true`, focus is locked.
    */
   open: boolean;
-  /**
-   * Return the document the trap focus is mounted into.
-   * Provide the prop if you need the restore focus to work between different documents.
-   * @default function defaultGetDoc() {
-   *   return document;
-   * }
-   */
-  getDoc?: () => Document;
   /**
    * Returns an array of ordered tabbable nodes (i.e. in tab order) within the root.
    * For instance, you can provide the "tabbable" npm dependency.

packages/material-ui-unstyled/src/Unstable_TrapFocus/Unstable_TrapFocus.js

@@ -108,10 +108,6 @@ function defaultGetTabbable(root) {
     .concat(regularTabNodes);
 }
 
-function defaultGetDoc() {
-  return document;
-}
-
 function defaultIsEnabled() {
   return true;
 }
@@ -125,15 +121,14 @@ function Unstable_TrapFocus(props) {
     disableAutoFocus = false,
     disableEnforceFocus = false,
     disableRestoreFocus = false,
-    getDoc = defaultGetDoc,
     getTabbable = defaultGetTabbable,
     isEnabled = defaultIsEnabled,
     open,
   } = props;
   const ignoreNextEnforceFocus = React.useRef();
   const sentinelStart = React.useRef(null);
   const sentinelEnd = React.useRef(null);
-  const nodeToRestore = React.useRef();
+  const nodeToRestore = React.useRef(null);
   const reactFocusEventTarget = React.useRef(null);
   // This variable is useful when disableAutoFocus is true.
   // It waits for the active element to move into the component to activate.
@@ -143,23 +138,6 @@ function Unstable_TrapFocus(props) {
   const handleRef = useForkRef(children.ref, rootRef);
   const lastKeydown = React.useRef(null);
 
-  const prevOpenRef = React.useRef();
-  React.useEffect(() => {
-    prevOpenRef.current = open;
-  }, [open]);
-
-  if (!prevOpenRef.current && open && typeof window !== 'undefined' && !disableAutoFocus) {
-    // WARNING: Potentially unsafe in concurrent mode.
-    // The way the read on `nodeToRestore` is setup could make this actually safe.
-    // Say we render `open={false}` -> `open={true}` but never commit.
-    // We have now written a state that wasn't committed. But no committed effect
-    // will read this wrong value. We only read from `nodeToRestore` in effects
-    // that were committed on `open={true}`
-    // WARNING: Prevents the instance from being garbage collected. Should only
-    // hold a weak ref.
-    nodeToRestore.current = getDoc().activeElement;
-  }
-
   React.useEffect(() => {
     // We might render an empty child.
     if (!open || !rootRef.current) {
@@ -325,7 +303,7 @@ function Unstable_TrapFocus(props) {
   }, [disableAutoFocus, disableEnforceFocus, disableRestoreFocus, isEnabled, open, getTabbable]);
 
   const onFocus = (event) => {
-    if (!activated.current) {
+    if (nodeToRestore.current === null) {
       nodeToRestore.current = event.relatedTarget;
     }
     activated.current = true;
@@ -338,7 +316,7 @@ function Unstable_TrapFocus(props) {
   };
 
   const handleFocusSentinel = (event) => {
-    if (!activated.current) {
+    if (nodeToRestore.current === null) {
       nodeToRestore.current = event.relatedTarget;
     }
     activated.current = true;
@@ -391,14 +369,6 @@ Unstable_TrapFocus.propTypes /* remove-proptypes */ = {
    * @default false
    */
   disableRestoreFocus: PropTypes.bool,
-  /**
-   * Return the document the trap focus is mounted into.
-   * Provide the prop if you need the restore focus to work between different documents.
-   * @default function defaultGetDoc() {
-   *   return document;
-   * }
-   */
-  getDoc: PropTypes.func,
   /**
    * Returns an array of ordered tabbable nodes (i.e. in tab order) within the root.
    * For instance, you can provide the "tabbable" npm dependency.

test/e2e/fixtures/Unstable_TrapFocus/DefaultOpenLazyTrapFocus.tsx

@@ -0,0 +1,23 @@
+import * as React from 'react';
+import TrapFocus from '@material-ui/core/Unstable_TrapFocus';
+
+export default function BaseTrapFocus() {
+  const [open, close] = React.useReducer(() => false, true);
+
+  return (
+    <React.Fragment>
+      <button type="button" autoFocus data-testid="initial-focus">
+        initial focus
+      </button>
+      <TrapFocus isEnabled={() => true} open={open} disableAutoFocus>
+        <div data-testid="root">
+          <div>Title</div>
+          <button type="button" onClick={close}>
+            close
+          </button>
+          <button type="button">noop</button>
+        </div>
+      </TrapFocus>
+    </React.Fragment>
+  );
+}

test/e2e/fixtures/Unstable_TrapFocus/OpenTrapFocus.tsx

@@ -7,7 +7,7 @@ export default function BaseTrapFocus() {
       <button type="button" autoFocus data-testid="initial-focus">
         initial focus
       </button>
-      <TrapFocus getDoc={() => document} isEnabled={() => true} open>
+      <TrapFocus isEnabled={() => true} open>
         <div tabIndex={-1} data-testid="root">
           <div>Title</div>
           <button type="button">x</button>

test/e2e/index.test.ts

@@ -137,6 +137,21 @@ describe('e2e', () => {
       await expect(screen.getByText('ok')).toHaveFocus();
     });
 
+    it('should loop the tab key after activation', async () => {
+      await renderFixture('Unstable_TrapFocus/DefaultOpenLazyTrapFocus');
+
+      await expect(screen.getByTestId('initial-focus')).toHaveFocus();
+
+      await page.keyboard.press('Tab');
+      await expect(screen.getByText('close')).toHaveFocus();
+      await page.keyboard.press('Tab');
+      await expect(screen.getByText('noop')).toHaveFocus();
+      await page.keyboard.press('Tab');
+      await expect(screen.getByText('close')).toHaveFocus();
+      await page.keyboard.press('Enter');
+      await expect(screen.getByTestId('initial-focus')).toHaveFocus();
+    });
+
     it('should focus on first focus element after last has received a tab click', async () => {
       await renderFixture('Unstable_TrapFocus/OpenTrapFocus');