(RSP-1596 and RSP-1592) Update onFocus and onBlur behavior of FocusScope containment

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

@@ -36,6 +36,7 @@ interface FocusManager {
 const FocusContext = React.createContext<FocusManager>(null);
 
 let activeScope: RefObject<HTMLElement[]> = null;
+let scopes: Set<RefObject<HTMLElement[]>> = new Set();
 
 // This is a hacky DOM-based implementation of a FocusScope until this RFC lands in React:
 // https://github.com/reactjs/rfcs/pull/109
@@ -57,6 +58,10 @@ export function FocusScope(props: FocusScopeProps) {
     }
 
     scopeRef.current = nodes;
+    scopes.add(scopeRef);
+    return () => {
+      scopes.delete(scopeRef);
+    };
   }, [children]);
 
   useFocusContainment(scopeRef, contain);
@@ -148,6 +153,7 @@ function useFocusContainment(scopeRef: RefObject<HTMLElement[]>, contain: boolea
   let focusedNode = useRef<HTMLElement>();
 
   useEffect(() => {
+    let scope = scopeRef.current;
     if (!contain) {
       return;
     }
@@ -159,11 +165,11 @@ function useFocusContainment(scopeRef: RefObject<HTMLElement[]>, contain: boolea
       }
 
       let focusedElement = document.activeElement as HTMLElement;
-      if (!isElementInScope(focusedElement, scopeRef.current)) {
+      if (!isElementInScope(focusedElement, scope)) {
         return;
       }
 
-      let elements = getFocusableElementsInScope(scopeRef.current, {tabbable: true});
+      let elements = getFocusableElementsInScope(scope, {tabbable: true});
       let position = elements.indexOf(focusedElement);
       let lastPosition = elements.length - 1;
       let nextElement = null;
@@ -189,38 +195,55 @@ function useFocusContainment(scopeRef: RefObject<HTMLElement[]>, contain: boolea
     };
 
     let onFocus = (e) => {
-      // If the focused element is in the current scope, and not in the active scope,
-      // update the active scope to point to this scope.
-      let isInScope = isElementInScope(e.target, scopeRef.current);
-      if (isInScope && (!activeScope || !isElementInScope(e.target, activeScope.current))) {
+      // 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 (focusedNode.current) {
+          focusedNode.current.focus();
+        } else if (activeScope) {
+          focusFirstInScope(activeScope.current);
+        } 
+      } else {
+        e.stopPropagation();
         activeScope = scopeRef;
-      }
-
-      // Save the currently focused node in this scope
-      if (isInScope) {
         focusedNode.current = e.target;
       }
+    };
 
-      // 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 if none.
-      if (activeScope === scopeRef && !isInScope) {
-        if (focusedNode.current) {
-          focusedNode.current.focus();
-        } else {
-          focusFirstInScope(scopeRef.current);
-        }
+    let onBlur = (e) => {
+      e.stopPropagation();
+      let isInAnyScope = isElementInAnyScope(e.relatedTarget, scopes);
+
+      if (!isInAnyScope) {
+        activeScope = scopeRef;
+        focusedNode.current = e.target;
+        focusedNode.current.focus();
       }
     };
 
     document.addEventListener('keydown', onKeyDown, false);
     document.addEventListener('focusin', onFocus, false);
+    scope.forEach(element => element.addEventListener('focusin', onFocus, false));
+    scope.forEach(element => element.addEventListener('focusout', onBlur, false));
     return () => {
       document.removeEventListener('keydown', onKeyDown, false);
       document.removeEventListener('focusin', onFocus, false);
+      scope.forEach(element => element.removeEventListener('focusin', onFocus, false));
+      scope.forEach(element => element.removeEventListener('focusout', onBlur, false));
     };
   }, [scopeRef, contain]);
 }
 
+function isElementInAnyScope(element: Element, scopes: Set<RefObject<HTMLElement[]>>) {
+  for (let scope of scopes.values()) {
+    if (isElementInScope(element, scope.current)) { 
+      return true; 
+    }
+  }  
+  return false;
+}
+
 function isElementInScope(element: Element, scope: HTMLElement[]) {
   return scope.some(node => node.contains(element));
 }

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

@@ -13,6 +13,7 @@
 import {cleanup, fireEvent, render} from '@testing-library/react';
 import {FocusScope, useFocusManager} from '../';
 import React from 'react';
+import ReactDOM from 'react-dom';
 
 describe('FocusScope', function () {
   afterEach(cleanup);
@@ -252,6 +253,33 @@ describe('FocusScope', function () {
       fireEvent.focusIn(outside);
       expect(document.activeElement).toBe(input2);
     });
+
+    it('should restore focus to the last focused element in the scope on focus out', function () {
+      let {getByTestId} = render(
+        <div>
+          <FocusScope contain>
+            <input data-testid="input1" />
+            <input data-testid="input2" />
+          </FocusScope>
+        </div>
+      );
+
+      let input1 = getByTestId('input1');
+      let input2 = getByTestId('input2');
+
+      input1.focus();
+      fireEvent.focusIn(input1); // jsdom doesn't fire this automatically
+      expect(document.activeElement).toBe(input1);
+
+      fireEvent.keyDown(document.activeElement, {key: 'Tab'});
+      fireEvent.focusIn(input2);
+      expect(document.activeElement).toBe(input2);
+
+      input2.blur();
+      expect(document.activeElement).toBe(document.body);
+      fireEvent.focusOut(input2);
+      expect(document.activeElement).toBe(input2);
+    });
   });
 
   describe('focus restoration', function () {
@@ -593,4 +621,43 @@ describe('FocusScope', function () {
       expect(document.activeElement).toBe(item1);
     });
   });
+  describe('nested focus scopes', function () {
+    it('should make child FocusScopes the active scope regardless of DOM structure', function () {
+      function ChildComponent(props) {
+        return ReactDOM.createPortal(props.children, document.body);
+      }
+
+      function Test({show}) {
+        return (
+          <div>
+            <input data-testid="outside" />
+            <FocusScope restoreFocus contain>
+              <input data-testid="input1" />
+              {show &&
+                <ChildComponent>
+                  <FocusScope restoreFocus contain>
+                    <input data-testid="input3" />
+                  </FocusScope>
+                </ChildComponent>
+              }
+            </FocusScope>
+          </div>
+        );
+      }
+
+      let {getByTestId, rerender} = render(<Test />);
+      // Set a focused node and make first FocusScope the active scope
+      let input1 = getByTestId('input1');
+      input1.focus();
+      fireEvent.focusIn(input1);
+      expect(document.activeElement).toBe(input1);
+
+      rerender(<Test show />);
+      expect(document.activeElement).toBe(input1);
+      let input3 = getByTestId('input3');
+      input3.focus();
+      fireEvent.focusIn(input3);
+      expect(document.activeElement).toBe(input3);
+    }); 
+  });
 });