Fix focus scope with changing children

packages/@react-aria/focus/src/FocusScope.tsx

@@ -387,15 +387,16 @@ function useAutoFocus(scopeRef: RefObject<HTMLElement[]>, autoFocus: boolean) {
 }
 
 function useRestoreFocus(scopeRef: RefObject<HTMLElement[]>, restoreFocus: boolean, contain: boolean) {
+  // create a ref during render instead of useLayoutEffect so the active element is saved before a child with autoFocus=true mounts.
+  const nodeToRestoreRef = useRef(typeof document !== 'undefined' ? document.activeElement as HTMLElement : null);
+
   // useLayoutEffect instead of useEffect so the active element is saved synchronously instead of asynchronously.
   useLayoutEffect(() => {
+    let nodeToRestore = nodeToRestoreRef.current;
     if (!restoreFocus) {
       return;
     }
 
-    let scope = scopeRef.current;
-    let nodeToRestore = document.activeElement as HTMLElement;
-
     // Handle the Tab key so that tabbing out of the scope goes to the next element
     // after the node that had focus when the scope mounted. This is important when
     // using portals for overlays, so that focus goes to the expected element when
@@ -406,7 +407,7 @@ function useRestoreFocus(scopeRef: RefObject<HTMLElement[]>, restoreFocus: boole
       }
 
       let focusedElement = document.activeElement as HTMLElement;
-      if (!isElementInScope(focusedElement, scope)) {
+      if (!isElementInScope(focusedElement, scopeRef.current)) {
         return;
       }
 
@@ -423,13 +424,13 @@ function useRestoreFocus(scopeRef: RefObject<HTMLElement[]>, restoreFocus: boole
 
       // If there is no next element, or it is outside the current scope, move focus to the
       // next element after the node to restore to instead.
-      if ((!nextElement || !isElementInScope(nextElement, scope)) && nodeToRestore) {
+      if ((!nextElement || !isElementInScope(nextElement, scopeRef.current)) && nodeToRestore) {
         walker.currentNode = nodeToRestore;
 
         // Skip over elements within the scope, in case the scope immediately follows the node to restore.
         do {
           nextElement = (e.shiftKey ? walker.previousNode() : walker.nextNode()) as HTMLElement;
-        } while (isElementInScope(nextElement, scope));
+        } while (isElementInScope(nextElement, scopeRef.current));
 
         e.preventDefault();
         e.stopPropagation();
@@ -457,7 +458,7 @@ function useRestoreFocus(scopeRef: RefObject<HTMLElement[]>, restoreFocus: boole
         document.removeEventListener('keydown', onKeyDown, true);
       }
 
-      if (restoreFocus && nodeToRestore && isElementInScope(document.activeElement, scope)) {
+      if (restoreFocus && nodeToRestore && isElementInScope(document.activeElement, scopeRef.current)) {
         requestAnimationFrame(() => {
           if (document.body.contains(nodeToRestore)) {
             focusElement(nodeToRestore);

packages/@react-aria/focus/stories/FocusScope.stories.tsx

@@ -18,7 +18,8 @@ import ReactDOM from 'react-dom';
 const dialogsRoot = 'dialogsRoot';
 
 interface StoryProps {
-  usePortal: boolean
+  usePortal: boolean,
+  contain: boolean
 }
 
 const meta: Meta<StoryProps> = {
@@ -33,7 +34,7 @@ const meta: Meta<StoryProps> = {
 
 export default meta;
 
-const Template = (): Story<StoryProps> => ({usePortal}) => <Example usePortal={usePortal} />;
+const Template = (): Story<StoryProps> => ({usePortal, contain = true}) => <Example usePortal={usePortal} contain={contain} />;
 
 function MaybePortal({children, usePortal}: { children: ReactNode, usePortal: boolean}) {
   if (!usePortal) {
@@ -46,35 +47,47 @@ function MaybePortal({children, usePortal}: { children: ReactNode, usePortal: bo
   );
 }
 
-function NestedDialog({onClose, usePortal}: {onClose: VoidFunction, usePortal: boolean}) {
+function NestedDialog({onClose, usePortal, contain}: {onClose: VoidFunction, usePortal: boolean, contain: boolean}) {
   let [open, setOpen] = useState(false);
+  let [showNew, setShowNew] = useState(false);
+  let onKeyDown = (e) => {
+    if (e.key === 'Escape') {
+      e.stopPropagation();
+      onClose();
+    }
+  };
 
   return (
     <MaybePortal usePortal={usePortal}>
-      <FocusScope contain restoreFocus autoFocus>
-        <div>
-          <input />
-
-          <input />
-
-          <input />
-
-          <button type="button" onClick={() => setOpen(true)}>
-            Open dialog
-          </button>
-
-          <button type="button" onClick={onClose}>
-            close
-          </button>
-
-          {open && <NestedDialog onClose={() => setOpen(false)} usePortal={usePortal} />}
-        </div>
+      <FocusScope contain={contain} restoreFocus autoFocus>
+        {!showNew && (
+          <div role="dialog" onKeyDown={onKeyDown}>
+            <input />
+            <input />
+            <input />
+            <button onClick={() => setShowNew(true)}>replace focusscope children</button>
+            <button type="button" onClick={() => setOpen(true)}>
+              Open dialog
+            </button>
+            <button type="button" onClick={onClose}>
+              close
+            </button>
+            {open && <NestedDialog contain={contain} onClose={() => setOpen(false)} usePortal={usePortal} />}
+          </div>
+        )}
+        {showNew && (
+          <div role="dialog" onKeyDown={onKeyDown}>
+            <input />
+            <input autoFocus />
+            <input />
+          </div>
+        )}
       </FocusScope>
     </MaybePortal>
   );
 }
 
-function Example({usePortal}: StoryProps) {
+function Example({usePortal, contain}: StoryProps) {
   let [open, setOpen] = useState(false);
 
   return (
@@ -84,8 +97,8 @@ function Example({usePortal}: StoryProps) {
       <button type="button" onClick={() => setOpen(true)}>
         Open dialog
       </button>
-
-      {open && <NestedDialog onClose={() => setOpen(false)} usePortal={usePortal} />}
+      <input />
+      {open && <NestedDialog onClose={() => setOpen(false)} usePortal={usePortal} contain={contain} />}
 
       <div id={dialogsRoot} />
     </div>
@@ -97,3 +110,9 @@ KeyboardNavigation.args = {usePortal: false};
 
 export const KeyboardNavigationInsidePortal = Template().bind({});
 KeyboardNavigationInsidePortal.args = {usePortal: true};
+
+export const KeyboardNavigationNoContain = Template().bind({});
+KeyboardNavigationNoContain.args = {usePortal: false, contain: false};
+
+export const KeyboardNavigationInsidePortalNoContain = Template().bind({});
+KeyboardNavigationInsidePortalNoContain.args = {usePortal: true, contain: false};

packages/@react-aria/focus/test/FocusScope.test.js

@@ -331,7 +331,134 @@ describe('FocusScope', function () {
       expect(document.activeElement).toBe(outside);
     });
 
-    it('should move focus to the next element after the previously focused node on Tab', function () {
+    it('should restore focus to the previously focused node after a child with autoFocus unmounts', function () {
+      function Test({show}) {
+        return (
+          <div>
+            <input data-testid="outside" />
+            {show &&
+              <FocusScope restoreFocus>
+                <input data-testid="input1" />
+                <input data-testid="input2" autoFocus />
+                <input data-testid="input3" />
+              </FocusScope>
+            }
+          </div>
+        );
+      }
+
+      let {getByTestId, rerender} = render(<Test />);
+
+      let outside = getByTestId('outside');
+      act(() => {outside.focus();});
+
+      rerender(<Test show />);
+
+      let input2 = getByTestId('input2');
+      expect(document.activeElement).toBe(input2);
+
+      rerender(<Test />);
+
+      expect(document.activeElement).toBe(outside);
+    });
+
+    it('should move focus after the previously focused node when tabbing away from a scope with autoFocus', function () {
+      function Test({show}) {
+        return (
+          <div>
+            <input data-testid="before" />
+            <input data-testid="outside" />
+            <input data-testid="after" />
+            {show &&
+              <FocusScope restoreFocus>
+                <input data-testid="input1" />
+                <input data-testid="input2" />
+                <input data-testid="input3" autoFocus />
+              </FocusScope>
+            }
+          </div>
+        );
+      }
+
+      let {getByTestId, rerender} = render(<Test />);
+
+      let outside = getByTestId('outside');
+      act(() => {outside.focus();});
+
+      rerender(<Test show />);
+
+      let input3 = getByTestId('input3');
+      expect(document.activeElement).toBe(input3);
+
+      userEvent.tab();
+      expect(document.activeElement).toBe(getByTestId('after'));
+    });
+
+    it('should move focus before the previously focused node when tabbing away from a scope with Shift+Tab', function () {
+      function Test({show}) {
+        return (
+          <div>
+            <input data-testid="before" />
+            <input data-testid="outside" />
+            <input data-testid="after" />
+            {show &&
+              <FocusScope restoreFocus>
+                <input data-testid="input1" autoFocus />
+                <input data-testid="input2" />
+                <input data-testid="input3" />
+              </FocusScope>
+            }
+          </div>
+        );
+      }
+
+      let {getByTestId, rerender} = render(<Test />);
+
+      let outside = getByTestId('outside');
+      act(() => {outside.focus();});
+
+      rerender(<Test show />);
+
+      let input1 = getByTestId('input1');
+      expect(document.activeElement).toBe(input1);
+
+      userEvent.tab({shift: true});
+      expect(document.activeElement).toBe(getByTestId('before'));
+    });
+
+    it('should restore focus to the previously focused node after children change', function () {
+      function Test({show, showChild}) {
+        return (
+          <div>
+            <input data-testid="outside" />
+            {show &&
+              <FocusScope restoreFocus autoFocus>
+                <input data-testid="input1" />
+                {showChild && <input data-testid="dynamic" />}
+              </FocusScope>
+            }
+          </div>
+        );
+      }
+
+      let {getByTestId, rerender} = render(<Test />);
+
+      let outside = getByTestId('outside');
+      act(() => {outside.focus();});
+
+      rerender(<Test show />);
+      rerender(<Test show showChild />);
+
+      let dynamic = getByTestId('dynamic');
+      act(() => {dynamic.focus();});
+      expect(document.activeElement).toBe(dynamic);
+
+      rerender(<Test />);
+
+      expect(document.activeElement).toBe(outside);
+    });
+
+    it('should move focus to the element after the previously focused node on Tab', function () {
       function Test({show}) {
         return (
           <div>