Fix FocusScope StrictMode issues

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

@@ -61,8 +61,8 @@ export interface FocusManager {
 
 type ScopeRef = RefObject<Element[]>;
 interface IFocusContext {
-  scopeRef: ScopeRef,
-  focusManager: FocusManager
+  focusManager: FocusManager,
+  parentNode: TreeNode | null
 }
 
 const FocusContext = React.createContext<IFocusContext>(null);
@@ -84,13 +84,33 @@ export function FocusScope(props: FocusScopeProps) {
   let startRef = useRef<HTMLSpanElement>();
   let endRef = useRef<HTMLSpanElement>();
   let scopeRef = useRef<Element[]>([]);
-  let ctx = useContext(FocusContext);
+  let {parentNode} = useContext(FocusContext) || {};
 
-  // The parent scope is based on the JSX tree, using context.
-  // However, if a new scope mounts outside the active scope (e.g. DialogContainer launched from a menu),
-  // we want the parent scope to be the active scope instead.
-  let ctxParent = ctx?.scopeRef ?? null;
-  let parentScope = useMemo(() => activeScope && focusScopeTree.getTreeNode(activeScope) && !isAncestorScope(activeScope, ctxParent) ? activeScope : ctxParent, [ctxParent]);
+  // Create a tree node here so we can add children to it even before it is added to the tree.
+  let node = useMemo(() => new TreeNode({scopeRef}), [scopeRef]);
+
+  useLayoutEffect(() => {
+    // If a new scope mounts outside the active scope, (e.g. DialogContainer launched from a menu),
+    // use the active scope as the parent instead of the parent from context. Layout effects run bottom
+    // up, so if the parent is not yet added to the tree, don't do this. Only the outer-most FocusScope
+    // that is being added should get the activeScope as its parent.
+    let parent = parentNode || focusScopeTree.root;
+    if (focusScopeTree.getTreeNode(parent.scopeRef) && activeScope && !isAncestorScope(activeScope, parent.scopeRef)) {
+      let activeNode = focusScopeTree.getTreeNode(activeScope);
+      if (activeNode) {
+        parent = activeNode;
+      }
+    }
+
+    // Add the node to the parent, and to the tree.
+    parent.addChild(node);
+    focusScopeTree.addNode(node);
+  }, [node, parentNode]);
+
+  useLayoutEffect(() => {
+    let node = focusScopeTree.getTreeNode(scopeRef);
+    node.contain = contain;
+  }, [contain]);
 
   useLayoutEffect(() => {
     // Find all rendered nodes between the sentinels and add them to the scope.
@@ -102,25 +122,34 @@ export function FocusScope(props: FocusScopeProps) {
     }
 
     scopeRef.current = nodes;
-  }, [children, parentScope]);
-
-  // add to the focus scope tree in render order because useEffects/useLayoutEffects run children first whereas render runs parent first
-  // which matters when constructing a tree
-  if (focusScopeTree.getTreeNode(parentScope) && !focusScopeTree.getTreeNode(scopeRef)) {
-    focusScopeTree.addTreeNode(scopeRef, parentScope);
-  }
-
-  let node = focusScopeTree.getTreeNode(scopeRef);
-  node.contain = contain;
+  }, [children]);
 
   useActiveScopeTracker(scopeRef, restoreFocus, contain);
   useFocusContainment(scopeRef, contain);
   useRestoreFocus(scopeRef, restoreFocus, contain);
   useAutoFocus(scopeRef, autoFocus);
 
   // this layout effect needs to run last so that focusScopeTree cleanup happens at the last moment possible
-  useLayoutEffect(() => {
+  useEffect(() => {
     if (scopeRef) {
+      let activeElement = document.activeElement;
+      let scope = null;
+      // In strict mode, active scope is incorrectly updated since cleanup will run even though scope hasn't unmounted.
+      // To fix this, we need to update the actual activeScope here
+      if (isElementInScope(activeElement, scopeRef.current)) {
+        // Since useLayoutEffect runs for children first, we need to traverse the focusScope tree and find the bottom most scope that
+        // contains the active element and set that as the activeScope
+        for (let node of focusScopeTree.traverse()) {
+          if (isElementInScope(activeElement, node.scopeRef.current)) {
+            scope = node;
+          }
+        }
+
+        if (scope === focusScopeTree.getTreeNode(scopeRef)) {
+          activeScope = scope.scopeRef;
+        }
+      }
+
       return () => {
         // Scope may have been re-parented.
         let parentScope = focusScopeTree.getTreeNode(scopeRef).parent.scopeRef;
@@ -137,12 +166,16 @@ export function FocusScope(props: FocusScopeProps) {
         focusScopeTree.removeTreeNode(scopeRef);
       };
     }
-  }, [scopeRef, parentScope]);
+  }, [scopeRef]);
 
-  let focusManager = createFocusManagerForScope(scopeRef);
+  let focusManager = useMemo(() => createFocusManagerForScope(scopeRef), []);
+  let value = useMemo(() => ({
+    focusManager,
+    parentNode: node
+  }), [node, focusManager]);
 
   return (
-    <FocusContext.Provider value={{scopeRef, focusManager}}>
+    <FocusContext.Provider value={value}>
       <span data-focus-scope-start hidden ref={startRef} />
       {children}
       <span data-focus-scope-end hidden ref={endRef} />
@@ -753,7 +786,7 @@ function last(walker: TreeWalker) {
 
 
 class Tree {
-  private root: TreeNode;
+  root: TreeNode;
   private fastMap = new Map<ScopeRef, TreeNode>();
 
   constructor() {
@@ -780,6 +813,10 @@ class Tree {
     }
   }
 
+  addNode(node: TreeNode) {
+    this.fastMap.set(node.scopeRef, node);
+  }
+
   removeTreeNode(scopeRef: ScopeRef) {
     // never remove the root
     if (scopeRef === null) {
@@ -802,9 +839,10 @@ class Tree {
     }
     let children = node.children;
     parentNode.removeChild(node);
-    if (children.length > 0) {
+    if (children.size > 0) {
       children.forEach(child => parentNode.addChild(child));
     }
+
     this.fastMap.delete(node.scopeRef);
   }
 
@@ -813,7 +851,7 @@ class Tree {
     if (node.scopeRef != null) {
       yield node;
     }
-    if (node.children.length > 0) {
+    if (node.children.size > 0) {
       for (let child of node.children) {
         yield* this.traverse(child);
       }
@@ -833,18 +871,18 @@ class TreeNode {
   public scopeRef: ScopeRef;
   public nodeToRestore: FocusableElement;
   public parent: TreeNode;
-  public children: TreeNode[] = [];
+  public children: Set<TreeNode> = new Set();
   public contain = false;
 
   constructor(props: {scopeRef: ScopeRef}) {
     this.scopeRef = props.scopeRef;
   }
   addChild(node: TreeNode) {
-    this.children.push(node);
+    this.children.add(node);
     node.parent = this;
   }
   removeChild(node: TreeNode) {
-    this.children.splice(this.children.indexOf(node), 1);
+    this.children.delete(node);
     node.parent = undefined;
   }
 }

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

@@ -1521,6 +1521,12 @@ describe('FocusScope', function () {
 
       fireEvent.keyDown(document.activeElement, {key: 'Enter'});
       fireEvent.keyUp(document.activeElement, {key: 'Enter'});
+
+      // Needed for onBlur raf in useFocusContainment
+      act(() => {
+        jest.runAllTimers();
+      });
+      // Needed for useRestoreFocus layout cleanup raf
       act(() => {
         jest.runAllTimers();
       });

packages/@react-spectrum/buttongroup/test/ButtonGroup.test.js

@@ -125,7 +125,6 @@ describe('ButtonGroup', function () {
   describe('resizing', () => {
     it('goes vertical if there is not enough room after buttongroup gets a new size', () => {
       let setUp = ({buttonGroup, button1, button2, button3}) => {
-        // can't do anything about first render, so this starts with the resize
         jest.spyOn(buttonGroup, 'offsetWidth', 'get').mockImplementationOnce(() => 88).mockImplementation(() => 90);
         jest.spyOn(button1, 'offsetLeft', 'get').mockImplementation(() => 0);
         jest.spyOn(button1, 'offsetWidth', 'get').mockImplementation(() => 30);
@@ -137,7 +136,6 @@ describe('ButtonGroup', function () {
       };
       let {getByTestId} = render(<ButtonGroupWithRefs setUp={setUp} />);
       let buttonGroup = getByTestId(buttonGroupId);
-      expect(buttonGroup).not.toHaveAttribute('class', expect.stringContaining('spectrum-ButtonGroup--vertical'));
 
       // ResizeObserver not actually implemented in jsdom, so rely on the fallback to window resize listener
       act(() => {window.dispatchEvent(new Event('resize'));});

packages/@react-spectrum/dialog/src/DialogTrigger.tsx

@@ -50,7 +50,10 @@ function DialogTrigger(props: SpectrumDialogTriggerProps) {
 
   let state = useOverlayTriggerState(props);
   let wasOpen = useRef(false);
-  wasOpen.current = state.isOpen;
+  useEffect(() => {
+    wasOpen.current = state.isOpen;
+  }, [state.isOpen]);
+
   let isExiting = useRef(false);
   let onExiting = () => isExiting.current = true;
   let onExited = () => isExiting.current = false;

packages/@react-spectrum/dialog/test/DialogTrigger.test.js

@@ -27,6 +27,7 @@ import userEvent from '@testing-library/user-event';
 
 describe('DialogTrigger', function () {
   let matchMedia;
+  let warnMock;
   beforeAll(() => {
     jest.spyOn(window.screen, 'width', 'get').mockImplementation(() => 1024);
     jest.useFakeTimers('legacy');
@@ -40,6 +41,9 @@ describe('DialogTrigger', function () {
     matchMedia = new MatchMediaMock();
     // this needs to be a setTimeout so that the dialog can be removed from the dom before the callback is invoked
     jest.spyOn(window, 'requestAnimationFrame').mockImplementation(cb => setTimeout(() => cb(), 0));
+    if (process.env.STRICT_MODE) {
+      warnMock = jest.spyOn(global.console, 'warn').mockImplementation();
+    }
   });
 
   afterEach(() => {
@@ -55,6 +59,12 @@ describe('DialogTrigger', function () {
 
     matchMedia.clear();
     window.requestAnimationFrame.mockRestore();
+
+    if (process.env.STRICT_MODE && warnMock.mock.calls.length > 0) {
+      expect(warnMock).toHaveBeenCalledTimes(1);
+      expect(warnMock).toHaveBeenCalledWith('A DialogTrigger unmounted while open. This is likely due to being placed within a trigger that unmounts or inside a conditional. Consider using a DialogContainer instead.');
+      warnMock.mockRestore();
+    }
   });
 
   it('should trigger a modal by default', function () {