Focus first element in scope if nodeToRestore does not exist

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

@@ -13,7 +13,7 @@
 import {FocusableElement} from '@react-types/shared';
 import {focusSafely} from './focusSafely';
 import {isElementVisible} from './isElementVisible';
-import React, {ReactNode, RefObject, useContext, useEffect, useRef} from 'react';
+import React, {ReactNode, RefObject, useContext, useEffect, useMemo, useRef} from 'react';
 import {useLayoutEffect} from '@react-aria/utils';
 
 
@@ -85,8 +85,12 @@ export function FocusScope(props: FocusScopeProps) {
   let endRef = useRef<HTMLSpanElement>();
   let scopeRef = useRef<Element[]>([]);
   let ctx = useContext(FocusContext);
-  // if there is no scopeRef on the context, then the parent is the focusScopeTree's root, represented by null
-  let parentScope = ctx?.scopeRef ?? null;
+
+  // 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]);
 
   useLayoutEffect(() => {
     // Find all rendered nodes between the sentinels and add them to the scope.
@@ -109,14 +113,18 @@ export function FocusScope(props: FocusScopeProps) {
   let node = focusScopeTree.getTreeNode(scopeRef);
   node.contain = contain;
 
+  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(() => {
-    if (scopeRef && (parentScope || parentScope == null)) {
+    if (scopeRef) {
       return () => {
+        // Scope may have been re-parented.
+        let parentScope = focusScopeTree.getTreeNode(scopeRef).parent.scopeRef;
+
         // 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.
@@ -294,7 +302,7 @@ function useFocusContainment(scopeRef: RefObject<Element[]>, contain: boolean) {
     let onFocus = (e) => {
       // If focusing an element in a child scope of the currently active scope, the child becomes active.
       // Moving out of the active scope to an ancestor is not allowed.
-      if (!activeScope || isAncestorScope(activeScope, scopeRef)) {
+      if ((!activeScope || isAncestorScope(activeScope, scopeRef)) && isElementInScope(e.target, scopeRef.current)) {
         activeScope = scopeRef;
         focusedNode.current = e.target;
       } else if (shouldContainFocus(scopeRef) && !isElementInChildScope(e.target, scopeRef)) {
@@ -424,6 +432,47 @@ function useAutoFocus(scopeRef: RefObject<Element[]>, autoFocus: boolean) {
   }, [scopeRef]);
 }
 
+function useActiveScopeTracker(scopeRef: RefObject<Element[]>, restore: boolean, contain: boolean) {
+  // tracks the active scope, in case restore and contain are both false.
+  // if either are true, this is tracked in useRestoreFocus or useFocusContainment.
+  useLayoutEffect(() => {
+    if (restore || contain) {
+      return;
+    }
+
+    let scope = scopeRef.current;
+
+    let onFocus = (e: FocusEvent) => {
+      let target = e.target as Element;
+      if (isElementInScope(target, scopeRef.current)) {
+        activeScope = scopeRef;
+      } else if (!isElementInAnyScope(target)) {
+        activeScope = null;
+      }
+    };
+
+    document.addEventListener('focusin', onFocus, false);
+    scope.forEach(element => element.addEventListener('focusin', onFocus, false));
+    return () => {
+      document.removeEventListener('focusin', onFocus, false);
+      scope.forEach(element => element.removeEventListener('focusin', onFocus, false));
+    };
+  }, [scopeRef, restore, contain]);
+}
+
+function shouldRestoreFocus(scopeRef: ScopeRef) {
+  let scope = focusScopeTree.getTreeNode(activeScope);
+  while (scope && scope.scopeRef !== scopeRef) {
+    if (scope.nodeToRestore) {
+      return false;
+    }
+
+    scope = scope.parent;
+  }
+
+  return true;
+}
+
 function useRestoreFocus(scopeRef: RefObject<Element[]>, 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 FocusableElement : null);
@@ -454,11 +503,12 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus: boolean,
 
   // useLayoutEffect instead of useEffect so the active element is saved synchronously instead of asynchronously.
   useLayoutEffect(() => {
-    focusScopeTree.getTreeNode(scopeRef).nodeToRestore = nodeToRestoreRef.current;
     if (!restoreFocus) {
       return;
     }
 
+    focusScopeTree.getTreeNode(scopeRef).nodeToRestore = nodeToRestoreRef.current;
+
     // 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
@@ -529,7 +579,7 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus: boolean,
         && nodeToRestore
         && (
           isElementInScope(document.activeElement, scopeRef.current)
-          || (document.activeElement === document.body && activeScope === scopeRef)
+          || (document.activeElement === document.body && shouldRestoreFocus(scopeRef))
         )
       ) {
         // freeze the focusScopeTree so it persists after the raf, otherwise during unmount nodes are removed from it
@@ -546,6 +596,17 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus: boolean,
               }
               treeNode = treeNode.parent;
             }
+
+            // If no nodeToRestore was found, focus the first element in the nearest
+            // ancestor scope that is still in the tree.
+            treeNode = clonedTree.getTreeNode(scopeRef);
+            while (treeNode) {
+              if (treeNode.scopeRef && focusScopeTree.getTreeNode(treeNode.scopeRef)) {
+                focusFirstInScope(treeNode.scopeRef.current, true);
+                return;
+              }
+              treeNode = treeNode.parent;
+            }
           }
         });
       }

packages/@react-spectrum/actionbar/stories/Example.tsx

@@ -50,6 +50,7 @@ let defaultItems = [
 export function Example(props: any = {}) {
   const [selectedKeys, setSelectedKeys] = useState<Selection>(props.defaultSelectedKeys || new Set());
   let [items, setItems] = useState(defaultItems);
+
   let ref = useRef(null);
   return (
     <ActionBarContainer height={props.containerHeight || 300}>
@@ -88,28 +89,26 @@ export function Example(props: any = {}) {
               }
               setItems(newItems);
               setSelectedKeys(new Set());
-              // TODO need to solve restore focus when the row we arrived from is deleted, tableview ref is not focusable
-              // ref.current.focus();
             }
           }
         })}>
-        <Item key="edit">
+        <Item key="edit" textValue="Edit">
           <Edit />
           <Text>Edit</Text>
         </Item>
-        <Item key="copy">
+        <Item key="copy" textValue="Copy">
           <Copy />
           <Text>Copy</Text>
         </Item>
-        <Item key="delete">
+        <Item key="delete" textValue="Delete">
           <Delete />
           <Text>Delete</Text>
         </Item>
-        <Item key="move">
+        <Item key="move" textValue="Move">
           <Move />
           <Text>Move</Text>
         </Item>
-        <Item key="duplicate">
+        <Item key="duplicate" textValue="Duplicate">
           <Duplicate />
           <Text>Duplicate</Text>
         </Item>

packages/@react-spectrum/actionbar/test/ActionBar.test.js

@@ -167,6 +167,118 @@ describe('ActionBar', () => {
     expect(document.activeElement).toBe(rows[1]);
   });
 
+  it('should restore focus to the the new first row if the row we wanted to restore to was removed', async () => {
+    let tree = render(<Provider theme={theme}><Example /></Provider>);
+    act(() => {jest.runAllTimers();});
+
+    let table = tree.getByRole('grid');
+    let rows = within(table).getAllByRole('row');
+    let checkbox = within(rows[1]).getByRole('checkbox');
+
+    userEvent.tab();
+    expect(document.activeElement).toBe(rows[1]);
+    fireEvent.keyDown(document.activeElement, {key: 'Enter'});
+    fireEvent.keyUp(document.activeElement, {key: 'Enter'});
+    expect(checkbox).toBeChecked();
+    act(() => jest.runAllTimers());
+
+    let toolbar = tree.getByRole('toolbar');
+    expect(toolbar).toBeVisible();
+    expect(document.activeElement).toBe(rows[1]);
+
+    // Simulate tabbing within the table
+    fireEvent.keyDown(document.activeElement, {key: 'Tab'});
+    let walker = getFocusableTreeWalker(document.body, {tabbable: true});
+    walker.currentNode = document.activeElement;
+    act(() => {walker.nextNode().focus();});
+    fireEvent.keyUp(document.activeElement, {key: 'Tab'});
+    expect(document.activeElement).toBe(within(toolbar).getAllByRole('button')[0]);
+
+    fireEvent.keyDown(document.activeElement, {key: 'ArrowRight'});
+    fireEvent.keyUp(document.activeElement, {key: 'ArrowRight'});
+
+    fireEvent.keyDown(document.activeElement, {key: 'ArrowRight'});
+    fireEvent.keyUp(document.activeElement, {key: 'ArrowRight'});
+
+    expect(document.activeElement).toBe(within(toolbar).getAllByRole('button')[2]);
+
+    fireEvent.keyDown(document.activeElement, {key: 'Enter'});
+    fireEvent.keyUp(document.activeElement, {key: 'Enter'});
+
+    act(() => jest.runAllTimers());
+    act(() => jest.runAllTimers());
+    await act(async () => Promise.resolve());
+    expect(rows[1]).not.toBeInTheDocument();
+
+    rows = within(table).getAllByRole('row');
+    expect(toolbar).not.toBeInTheDocument();
+    expect(document.activeElement).toBe(rows[1]);
+  });
+
+  it('should restore focus to the new first row if the row we wanted to restore to was removed via actiongroup menu', async () => {
+    let tree = render(<Provider theme={theme}><Example /></Provider>);
+    act(() => {jest.runAllTimers();});
+
+    let table = tree.getByRole('grid');
+    let rows = within(table).getAllByRole('row');
+    let checkbox = within(rows[1]).getByRole('checkbox');
+
+    userEvent.tab();
+    expect(document.activeElement).toBe(rows[1]);
+    fireEvent.keyDown(document.activeElement, {key: 'Enter'});
+    fireEvent.keyUp(document.activeElement, {key: 'Enter'});
+    expect(checkbox).toBeChecked();
+    act(() => jest.runAllTimers());
+
+    let toolbar = tree.getByRole('toolbar');
+    expect(toolbar).toBeVisible();
+    expect(document.activeElement).toBe(rows[1]);
+
+    jest.spyOn(toolbar.parentNode, 'offsetWidth', 'get').mockImplementation(() => 200);
+    for (let action of toolbar.childNodes) {
+      jest.spyOn(action, 'offsetWidth', 'get').mockImplementation(() => 75);
+    }
+    fireEvent(window, new Event('resize'));
+
+    // Simulate tabbing within the table
+    fireEvent.keyDown(document.activeElement, {key: 'Tab'});
+    let walker = getFocusableTreeWalker(document.body, {tabbable: true});
+    walker.currentNode = document.activeElement;
+    act(() => {walker.nextNode().focus();});
+    fireEvent.keyUp(document.activeElement, {key: 'Tab'});
+    expect(document.activeElement).toBe(within(toolbar).getAllByRole('button')[0]);
+
+    fireEvent.keyDown(document.activeElement, {key: 'ArrowRight'});
+    fireEvent.keyUp(document.activeElement, {key: 'ArrowRight'});
+
+    fireEvent.keyDown(document.activeElement, {key: 'ArrowRight'});
+    fireEvent.keyUp(document.activeElement, {key: 'ArrowRight'});
+
+    expect(within(toolbar).getAllByRole('button')).toHaveLength(3);
+    expect(document.activeElement).toBe(within(toolbar).getAllByRole('button')[2]);
+
+    fireEvent.keyDown(document.activeElement, {key: 'Enter'});
+    fireEvent.keyUp(document.activeElement, {key: 'Enter'});
+
+    let listItems = tree.getAllByRole('menuitem');
+    expect(document.activeElement).toBe(listItems[0]);
+    expect(document.activeElement.textContent).toBe('Delete');
+
+    fireEvent.keyDown(document.activeElement, {key: 'Enter'});
+    fireEvent.keyUp(document.activeElement, {key: 'Enter'});
+
+    act(() => jest.runAllTimers());
+    act(() => jest.runAllTimers());
+    await act(async () => Promise.resolve());
+    rows = within(table).getAllByRole('row');
+    // row reused by the virtualizer, so it's still in dom, but its contents have been swapped out
+    expect(rows[1].textContent).not.toBe('Foo 1Bar 1Baz 1');
+
+    rows = within(table).getAllByRole('row');
+    expect(toolbar).not.toBeInTheDocument();
+    expect(document.activeElement).toBe(rows[1]);
+  });
+
   it('should fire onAction when clicking on an action', () => {
     let onAction = jest.fn();
     let tree = render(<Provider theme={theme}><Example onAction={onAction} /></Provider>);

packages/@react-spectrum/actiongroup/src/ActionGroup.tsx

@@ -23,6 +23,7 @@ import {
   useStyleProps
 } from '@react-spectrum/utils';
 import {filterDOMProps, mergeProps, useId, useLayoutEffect, useResizeObserver, useValueEffect} from '@react-aria/utils';
+import {FocusScope} from '@react-aria/focus';
 import {Item, Menu, MenuTrigger} from '@react-spectrum/menu';
 import {ListState, useListState} from '@react-stately/list';
 import More from '@spectrum-icons/workflow/More';
@@ -215,42 +216,44 @@ function ActionGroup<T extends object>(props: SpectrumActionGroupProps<T>, ref:
   };
 
   return (
-    <div {...styleProps} style={style} className={classNames(styles, 'flex-container', styleProps.className)} ref={wrapperRef}>
-      <div
-        {...actionGroupProps}
-        ref={domRef}
-        className={
-          classNames(
-            styles,
-            'flex-gap',
-            'spectrum-ActionGroup',
-            {
-              'spectrum-ActionGroup--quiet': isQuiet,
-              'spectrum-ActionGroup--vertical': isVertical,
-              'spectrum-ActionGroup--compact': density === 'compact',
-              'spectrum-ActionGroup--justified': isJustified && !isMeasuring,
-              'spectrum-ActionGroup--overflowCollapse': overflowMode === 'collapse'
-            },
-            otherProps.UNSAFE_className
-          )
-        }>
-        <Provider {...providerProps}>
-          {children.map((item) => (
-            <ActionGroupItem
-              key={item.key}
-              onAction={onAction}
-              isDisabled={isDisabled}
-              isEmphasized={isEmphasized}
-              staticColor={staticColor}
-              item={item}
-              state={state}
-              hideButtonText={hideButtonText}
-              orientation={orientation} />
-          ))}
-          {menuItem}
-        </Provider>
+    <FocusScope>
+      <div {...styleProps} style={style} className={classNames(styles, 'flex-container', styleProps.className)} ref={wrapperRef}>
+        <div
+          {...actionGroupProps}
+          ref={domRef}
+          className={
+            classNames(
+              styles,
+              'flex-gap',
+              'spectrum-ActionGroup',
+              {
+                'spectrum-ActionGroup--quiet': isQuiet,
+                'spectrum-ActionGroup--vertical': isVertical,
+                'spectrum-ActionGroup--compact': density === 'compact',
+                'spectrum-ActionGroup--justified': isJustified && !isMeasuring,
+                'spectrum-ActionGroup--overflowCollapse': overflowMode === 'collapse'
+              },
+              otherProps.UNSAFE_className
+            )
+          }>
+          <Provider {...providerProps}>
+            {children.map((item) => (
+              <ActionGroupItem
+                key={item.key}
+                onAction={onAction}
+                isDisabled={isDisabled}
+                isEmphasized={isEmphasized}
+                staticColor={staticColor}
+                item={item}
+                state={state}
+                hideButtonText={hideButtonText}
+                orientation={orientation} />
+            ))}
+            {menuItem}
+          </Provider>
+        </div>
       </div>
-    </div>
+    </FocusScope>
   );
 }