Fix focusscope restore logic (#5131)

* FocusScope build tree using effects in the correct order

Author: snowystinger

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

@@ -131,43 +131,42 @@ export function FocusScope(props: FocusScopeProps) {
   useRestoreFocus(scopeRef, restoreFocus, contain);
   useAutoFocus(scopeRef, autoFocus);
 
-  // this layout effect needs to run last so that focusScopeTree cleanup happens at the last moment possible
+  // This needs to be an effect so that activeScope is updated after the FocusScope tree is complete.
+  // It cannot be a useLayoutEffect because the parent of this node hasn't been attached in the tree yet.
   useEffect(() => {
-    if (scopeRef) {
-      let activeElement = document.activeElement;
-      let scope: TreeNode | null = 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 (node.scopeRef && isElementInScope(activeElement, node.scopeRef.current)) {
-            scope = node;
-          }
+    let activeElement = document.activeElement;
+    let scope: TreeNode | null = null;
+
+    if (isElementInScope(activeElement, scopeRef.current)) {
+      // 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 (node.scopeRef && isElementInScope(activeElement, node.scopeRef.current)) {
+          scope = node;
         }
+      }
 
-        if (scope === focusScopeTree.getTreeNode(scopeRef)) {
-          activeScope = scope.scopeRef;
-        }
+      if (scope === focusScopeTree.getTreeNode(scopeRef)) {
+        activeScope = scope.scopeRef;
       }
+    }
+  }, [scopeRef]);
 
-      return () => {
-        // Scope may have been re-parented.
-        let parentScope = focusScopeTree.getTreeNode(scopeRef)?.parent?.scopeRef ?? null;
+  // This layout effect cleanup is so that the tree node is removed synchronously with react before the RAF
+  // in useRestoreFocus cleanup runs.
+  useLayoutEffect(() => {
+    return () => {
+      // Scope may have been re-parented.
+      let parentScope = focusScopeTree.getTreeNode(scopeRef)?.parent?.scopeRef ?? null;
 
-        // Restore the active scope on unmount if this scope or a descendant scope is active.
-        // Parent effect cleanups run before children, so we need to check if the
-        // parent scope actually still exists before restoring the active scope to it.
-        if (
-          (scopeRef === activeScope || isAncestorScope(scopeRef, activeScope)) &&
-          (!parentScope || focusScopeTree.getTreeNode(parentScope))
-        ) {
-          activeScope = parentScope;
-        }
-        focusScopeTree.removeTreeNode(scopeRef);
-      };
-    }
+      if (
+        (scopeRef === activeScope || isAncestorScope(scopeRef, activeScope)) &&
+        (!parentScope || focusScopeTree.getTreeNode(parentScope))
+      ) {
+        activeScope = parentScope;
+      }
+      focusScopeTree.removeTreeNode(scopeRef);
+    };
   }, [scopeRef]);
 
   let focusManager = useMemo(() => createFocusManagerForScope(scopeRef), []);

packages/@react-spectrum/list/stories/ListView.stories.tsx

@@ -1,14 +1,17 @@
 import {action} from '@storybook/addon-actions';
 import {ActionBar, ActionBarContainer} from '@react-spectrum/actionbar';
-import {ActionButton} from '@react-spectrum/button';
+import {ActionButton, Button} from '@react-spectrum/button';
 import {ActionGroup} from '@react-spectrum/actiongroup';
 import {ComponentMeta, ComponentStoryObj} from '@storybook/react';
 import {Content} from '@react-spectrum/view';
 import Copy from '@spectrum-icons/workflow/Copy';
 import Delete from '@spectrum-icons/workflow/Delete';
+import {Dialog, DialogTrigger} from '@react-spectrum/dialog';
+import {Divider} from '@react-spectrum/divider';
 import Edit from '@spectrum-icons/workflow/Edit';
 import File from '@spectrum-icons/illustrations/File';
 import {Flex} from '@react-spectrum/layout';
+import {FocusScope} from '@react-aria/focus';
 import Folder from '@spectrum-icons/illustrations/Folder';
 import {Heading, Text} from '@react-spectrum/text';
 import {IllustratedMessage} from '@react-spectrum/illustratedmessage';
@@ -466,3 +469,51 @@ function EmptyTest() {
     </div>
   );
 }
+
+export const RemoveListItems = {
+  render: (args) => (
+    <Demo {...args} />
+  )
+};
+
+function Demo(props) {
+  let [items, setItems] = useState<{key: number, label: string}[]>([
+    {key: 1, label: 'utilities'}
+  ]);
+  let onDelete = (key) => {
+    setItems(prevItems => prevItems.filter((item) => item.key !== key));
+  };
+  return (
+    <ListView
+      selectionMode="multiple"
+      maxWidth="size-6000"
+      items={items}
+      height="300px"
+      width="250px"
+      aria-label="ListView example with complex items"
+      {...props}>
+      {(item: {key: number, label: string}) => {
+        return (
+          <Item key={item.key} textValue={item.label}>
+            <Text>{item.label}</Text>
+            <DialogTrigger type="popover">
+              <ActionButton>Delete</ActionButton>
+              <Dialog>
+                <Heading>Warning, cannot undo</Heading>
+                <Divider />
+                <Content>
+                  <Text>Are you sure?</Text>
+                  <FocusScope>
+                    <Button variant="accent" onPressStart={() => onDelete(item.key)}>
+                      Delete
+                    </Button>
+                  </FocusScope>
+                </Content>
+              </Dialog>
+            </DialogTrigger>
+          </Item>
+        );
+      }}
+    </ListView>
+  );
+}

packages/@react-spectrum/list/test/ListView.test.js

@@ -10,6 +10,7 @@
  * governing permissions and limitations under the License.
  */
 
+
 jest.mock('@react-aria/live-announcer');
 import {act, fireEvent, installPointerEvent, pointerMap, render as renderComponent, triggerPress, within} from '@react-spectrum/test-utils';
 import {ActionButton} from '@react-spectrum/button';