Fix FocusScope issues (#2223)

* Handle nested focus scopes without contain

* Fix tabbing without restoreFocus

* Fixes

* Restore active scope to the correct scope on unmount

* Fix FocusScope in portal

Co-authored-by: Daniel Lu <[email protected]>

Author: devongovett

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

@@ -54,10 +54,16 @@ interface FocusManager {
   focusPrevious(opts?: FocusManagerOptions): HTMLElement
 }
 
-const FocusContext = React.createContext<FocusManager>(null);
+type ScopeRef = RefObject<HTMLElement[]>;
+interface IFocusContext {
+  scopeRef: ScopeRef,
+  focusManager: FocusManager
+}
+
+const FocusContext = React.createContext<IFocusContext>(null);
 
-let activeScope: RefObject<HTMLElement[]> = null;
-let scopes: Set<RefObject<HTMLElement[]>> = new Set();
+let activeScope: ScopeRef = null;
+let scopes: Map<ScopeRef, ScopeRef | null> = new Map();
 
 // This is a hacky DOM-based implementation of a FocusScope until this RFC lands in React:
 // https://github.com/reactjs/rfcs/pull/109
@@ -76,6 +82,8 @@ export function FocusScope(props: FocusScopeProps) {
   let startRef = useRef<HTMLSpanElement>();
   let endRef = useRef<HTMLSpanElement>();
   let scopeRef = useRef<HTMLElement[]>([]);
+  let ctx = useContext(FocusContext);
+  let parentScope = ctx?.scopeRef;
 
   useLayoutEffect(() => {
     // Find all rendered nodes between the sentinels and add them to the scope.
@@ -87,11 +95,23 @@ export function FocusScope(props: FocusScopeProps) {
     }
 
     scopeRef.current = nodes;
-    scopes.add(scopeRef);
+  }, [children, parentScope]);
+
+  useLayoutEffect(() => {
+    scopes.set(scopeRef, parentScope);
     return () => {
+      // 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 || scopes.has(parentScope))
+      ) {
+        activeScope = parentScope;
+      }
       scopes.delete(scopeRef);
     };
-  }, [children]);
+  }, [scopeRef, parentScope]);
 
   useFocusContainment(scopeRef, contain);
   useRestoreFocus(scopeRef, restoreFocus, contain);
@@ -100,7 +120,7 @@ export function FocusScope(props: FocusScopeProps) {
   let focusManager = createFocusManagerForScope(scopeRef);
 
   return (
-    <FocusContext.Provider value={focusManager}>
+    <FocusContext.Provider value={{scopeRef, focusManager}}>
       <span data-focus-scope-start hidden ref={startRef} />
       {children}
       <span data-focus-scope-end hidden ref={endRef} />
@@ -114,7 +134,7 @@ export function FocusScope(props: FocusScopeProps) {
  * a FocusScope, e.g. in response to user events like keyboard navigation.
  */
 export function useFocusManager(): FocusManager {
-  return useContext(FocusContext);
+  return useContext(FocusContext)?.focusManager;
 }
 
 function createFocusManagerForScope(scopeRef: React.RefObject<HTMLElement[]>): FocusManager {
@@ -185,19 +205,20 @@ function useFocusContainment(scopeRef: RefObject<HTMLElement[]>, contain: boolea
   let focusedNode = useRef<HTMLElement>();
 
   let raf = useRef(null);
-  useEffect(() => {
+  useLayoutEffect(() => {
     let scope = scopeRef.current;
     if (!contain) {
       return;
     }
 
     // Handle the Tab key to contain focus within the scope
     let onKeyDown = (e) => {
-      if (e.key !== 'Tab' || e.altKey || e.ctrlKey || e.metaKey) {
+      if (e.key !== 'Tab' || e.altKey || e.ctrlKey || e.metaKey || scopeRef !== activeScope) {
         return;
       }
 
       let focusedElement = document.activeElement as HTMLElement;
+      let scope = scopeRef.current;
       if (!isElementInScope(focusedElement, scope)) {
         return;
       }
@@ -217,17 +238,20 @@ function useFocusContainment(scopeRef: RefObject<HTMLElement[]>, contain: boolea
     };
 
     let onFocus = (e) => {
-      // If a focus event occurs outside the active scope (e.g. user tabs from browser location bar),
-      // restore focus to the previously focused node or the first tabbable element in the active scope.
-      let isInAnyScope = isElementInAnyScope(e.target, scopes);
-      if (!isInAnyScope) {
+      // 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)) {
+        activeScope = scopeRef;
+        focusedNode.current = e.target;
+      } else if (scopeRef === activeScope && !isElementInChildScope(e.target, scopeRef)) {
+        // If a focus event occurs outside the active scope (e.g. user tabs from browser location bar),
+        // restore focus to the previously focused node or the first tabbable element in the active scope.
         if (focusedNode.current) {
           focusedNode.current.focus();
         } else if (activeScope) {
           focusFirstInScope(activeScope.current);
         }
-      } else {
-        activeScope = scopeRef;
+      } else if (scopeRef === activeScope) {
         focusedNode.current = e.target;
       }
     };
@@ -236,9 +260,7 @@ function useFocusContainment(scopeRef: RefObject<HTMLElement[]>, contain: boolea
       // Firefox doesn't shift focus back to the Dialog properly without this
       raf.current = requestAnimationFrame(() => {
         // Use document.activeElement instead of e.relatedTarget so we can tell if user clicked into iframe
-        let isInAnyScope = isElementInAnyScope(document.activeElement, scopes);
-
-        if (!isInAnyScope) {
+        if (scopeRef === activeScope && !isElementInChildScope(document.activeElement, scopeRef)) {
           activeScope = scopeRef;
           focusedNode.current = e.target;
           focusedNode.current.focus();
@@ -264,34 +286,42 @@ function useFocusContainment(scopeRef: RefObject<HTMLElement[]>, contain: boolea
   }, [raf]);
 }
 
-function isElementInAnyScope(element: Element, scopes: Set<RefObject<HTMLElement[]>>) {
-  for (let scope of scopes.values()) {
+function isElementInAnyScope(element: Element) {
+  for (let scope of scopes.keys()) {
     if (isElementInScope(element, scope.current)) {
       return true;
     }
   }
   return false;
 }
 
-const focusScopeDataAttrNames = [
-  'data-focus-scope-start',
-  'data-focus-scope-end'
-];
-
-function isFocusScopeDirectChild(scope: HTMLElement) {
-  return focusScopeDataAttrNames.some(name => scope.getAttribute(name) !== null);
+function isElementInScope(element: Element, scope: HTMLElement[]) {
+  return scope.some(node => node.contains(element));
 }
 
-function isFocusScopeNestedChild(scope: HTMLElement) {
-  return focusScopeDataAttrNames.some(name => scope.querySelector(`[${name}]`));
-}
+function isElementInChildScope(element: Element, scope: ScopeRef) {
+  // node.contains in isElementInScope covers child scopes that are also DOM children,
+  // but does not cover child scopes in portals.
+  for (let s of scopes.keys()) {
+    if ((s === scope || isAncestorScope(scope, s)) && isElementInScope(element, s.current)) {
+      return true;
+    }
+  }
 
-function isFocusScopeInScope(scopes: HTMLElement[]) {
-  return scopes.some(scope => isFocusScopeDirectChild(scope) || isFocusScopeNestedChild(scope));
+  return false;
 }
 
-function isElementInScope(element: Element, scope: HTMLElement[]) {
-  return !isFocusScopeInScope(scope) && scope.some(node => node.contains(element));
+function isAncestorScope(ancestor: ScopeRef, scope: ScopeRef) {
+  let parent = scopes.get(scope);
+  if (!parent) {
+    return false;
+  }
+
+  if (parent === ancestor) {
+    return true;
+  }
+
+  return isAncestorScope(ancestor, parent);
 }
 
 function focusElement(element: HTMLElement | null, scroll = false) {
@@ -333,6 +363,10 @@ function useAutoFocus(scopeRef: RefObject<HTMLElement[]>, autoFocus: boolean) {
 function useRestoreFocus(scopeRef: RefObject<HTMLElement[]>, restoreFocus: boolean, contain: boolean) {
   // useLayoutEffect instead of useEffect so the active element is saved synchronously instead of asynchronously.
   useLayoutEffect(() => {
+    if (!restoreFocus) {
+      return;
+    }
+
     let scope = scopeRef.current;
     let nodeToRestore = document.activeElement as HTMLElement;
 
@@ -379,7 +413,7 @@ function useRestoreFocus(scopeRef: RefObject<HTMLElement[]>, restoreFocus: boole
            // If there is no next element and the nodeToRestore isn't within a FocusScope (i.e. we are leaving the top level focus scope)
            // then move focus to the body.
            // Otherwise restore focus to the nodeToRestore (e.g menu within a popover -> tabbing to close the menu should move focus to menu trigger)
-          if (!isElementInAnyScope(nodeToRestore, scopes)) {
+          if (!isElementInAnyScope(nodeToRestore)) {
             focusedElement.blur();
           } else {
             focusElement(nodeToRestore, true);