Focus Management within Shadow DOM

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

@@ -12,7 +12,7 @@
 
 import {FocusableElement} from '@react-types/shared';
 import {focusSafely} from './focusSafely';
-import {getOwnerDocument, useLayoutEffect} from '@react-aria/utils';
+import {getOwnerDocument, getRootBody, getRootNode, useLayoutEffect} from '@react-aria/utils';
 import {isElementVisible} from './isElementVisible';
 import React, {ReactNode, RefObject, useContext, useEffect, useMemo, useRef} from 'react';
 
@@ -54,7 +54,7 @@ export interface FocusManager {
   focusPrevious(opts?: FocusManagerOptions): FocusableElement | null,
   /** Moves focus to the first focusable or tabbable element in the focus scope. */
   focusFirst(opts?: FocusManagerOptions): FocusableElement | null,
-    /** Moves focus to the last focusable or tabbable element in the focus scope. */
+  /** Moves focus to the last focusable or tabbable element in the focus scope. */
   focusLast(opts?: FocusManagerOptions): FocusableElement | null
 }
 
@@ -133,7 +133,8 @@ export function FocusScope(props: FocusScopeProps) {
   // 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(() => {
-    const activeElement = getOwnerDocument(scopeRef.current ? scopeRef.current[0] : undefined).activeElement;
+    // eslint-disable-next-line no-undef
+    const activeElement = getRootNode(scopeRef.current ? scopeRef.current[0] : undefined).activeElement;
     let scope: TreeNode | null = null;
 
     if (isElementInScope(activeElement, scopeRef.current)) {
@@ -197,7 +198,7 @@ function createFocusManagerForScope(scopeRef: React.RefObject<Element[]>): Focus
     focusNext(opts: FocusManagerOptions = {}) {
       let scope = scopeRef.current!;
       let {from, tabbable, wrap, accept} = opts;
-      let node = from || getOwnerDocument(scope[0]).activeElement!;
+      let node = from || getRootNode(scope[0]).activeElement!;
       let sentinel = scope[0].previousElementSibling!;
       let scopeRoot = getScopeRoot(scope);
       let walker = getFocusableTreeWalker(scopeRoot, {tabbable, accept}, scope);
@@ -215,7 +216,7 @@ function createFocusManagerForScope(scopeRef: React.RefObject<Element[]>): Focus
     focusPrevious(opts: FocusManagerOptions = {}) {
       let scope = scopeRef.current!;
       let {from, tabbable, wrap, accept} = opts;
-      let node = from || getOwnerDocument(scope[0]).activeElement!;
+      let node = from || getRootNode(scope[0]).activeElement!;
       let sentinel = scope[scope.length - 1].nextElementSibling!;
       let scopeRoot = getScopeRoot(scope);
       let walker = getFocusableTreeWalker(scopeRoot, {tabbable, accept}, scope);
@@ -314,7 +315,7 @@ function useFocusContainment(scopeRef: RefObject<Element[]>, contain?: boolean)
       return;
     }
 
-    const ownerDocument = getOwnerDocument(scope ? scope[0] : undefined);
+    const ownerDocument = getRootNode(scope ? scope[0] : undefined);
 
     // Handle the Tab key to contain focus within the scope
     let onKeyDown = (e) => {
@@ -374,7 +375,7 @@ function useFocusContainment(scopeRef: RefObject<Element[]>, contain?: boolean)
         // Use document.activeElement instead of e.relatedTarget so we can tell if user clicked into iframe
         if (ownerDocument.activeElement && shouldContainFocus(scopeRef) && !isElementInChildScope(ownerDocument.activeElement, scopeRef)) {
           activeScope = scopeRef;
-          if (ownerDocument.body.contains(e.target)) {
+          if (getRootBody(ownerDocument).contains(e.target)) {
             focusedNode.current = e.target;
             focusedNode.current?.focus();
           } else if (activeScope.current) {
@@ -511,7 +512,7 @@ function useActiveScopeTracker(scopeRef: RefObject<Element[]>, restore?: boolean
     }
 
     let scope = scopeRef.current;
-    const ownerDocument = getOwnerDocument(scope ? scope[0] : undefined);
+    const ownerDocument = getRootNode(scope ? scope[0] : undefined);
 
     let onFocus = (e) => {
       let target = e.target as Element;
@@ -553,7 +554,7 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
   // restoring-non-containing scopes should only care if they become active so they can perform the restore
   useLayoutEffect(() => {
     let scope = scopeRef.current;
-    const ownerDocument = getOwnerDocument(scope ? scope[0] : undefined);
+    const ownerDocument = getRootNode(scope ? scope[0] : undefined);
     if (!restoreFocus || contain) {
       return;
     }
@@ -562,7 +563,7 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
       // 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)) &&
-      isElementInScope(ownerDocument.activeElement, scopeRef.current)
+        isElementInScope(ownerDocument.activeElement, scopeRef.current)
       ) {
         activeScope = scopeRef;
       }
@@ -574,11 +575,11 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
       ownerDocument.removeEventListener('focusin', onFocus, false);
       scope?.forEach(element => element.removeEventListener('focusin', onFocus, false));
     };
-  // eslint-disable-next-line react-hooks/exhaustive-deps
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [scopeRef, contain]);
 
   useLayoutEffect(() => {
-    const ownerDocument = getOwnerDocument(scopeRef.current ? scopeRef.current[0] : undefined);
+    const ownerDocument = getRootNode(scopeRef.current ? scopeRef.current[0] : undefined);
 
     if (!restoreFocus) {
       return;
@@ -603,14 +604,16 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
       }
       let nodeToRestore = treeNode.nodeToRestore;
 
+      const rootBody = getRootBody(ownerDocument);
+
       // Create a DOM tree walker that matches all tabbable elements
-      let walker = getFocusableTreeWalker(ownerDocument.body, {tabbable: true});
+      let walker = getFocusableTreeWalker(rootBody, {tabbable: true});
 
       // Find the next tabbable element after the currently focused element
       walker.currentNode = focusedElement;
       let nextElement = (e.shiftKey ? walker.previousNode() : walker.nextNode()) as FocusableElement;
 
-      if (!nodeToRestore || !ownerDocument.body.contains(nodeToRestore) || nodeToRestore === ownerDocument.body) {
+      if (!nodeToRestore || !rootBody.contains(nodeToRestore) || nodeToRestore === rootBody) {
         nodeToRestore = undefined;
         treeNode.nodeToRestore = undefined;
       }
@@ -630,9 +633,9 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
         if (nextElement) {
           focusElement(nextElement, true);
         } else {
-           // 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 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)) {
             focusedElement.blur();
           } else {
@@ -643,19 +646,20 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
     };
 
     if (!contain) {
-      ownerDocument.addEventListener('keydown', onKeyDown, true);
+      ownerDocument.addEventListener('keydown', onKeyDown as EventListener, true);
     }
 
     return () => {
       if (!contain) {
-        ownerDocument.removeEventListener('keydown', onKeyDown, true);
+        ownerDocument.removeEventListener('keydown', onKeyDown as EventListener, true);
       }
     };
   }, [scopeRef, restoreFocus, contain]);
 
   // useLayoutEffect instead of useEffect so the active element is saved synchronously instead of asynchronously.
   useLayoutEffect(() => {
-    const ownerDocument = getOwnerDocument(scopeRef.current ? scopeRef.current[0] : undefined);
+    const ownerDocument = getRootNode(scopeRef.current ? scopeRef.current[0] : undefined);
+    const rootBody = getRootBody(ownerDocument);
 
     if (!restoreFocus) {
       return;
@@ -680,14 +684,14 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
         && (
           // eslint-disable-next-line react-hooks/exhaustive-deps
           isElementInScope(ownerDocument.activeElement, scopeRef.current)
-          || (ownerDocument.activeElement === ownerDocument.body && shouldRestoreFocus(scopeRef))
+          || (ownerDocument.activeElement === rootBody && shouldRestoreFocus(scopeRef))
         )
       ) {
         // freeze the focusScopeTree so it persists after the raf, otherwise during unmount nodes are removed from it
         let clonedTree = focusScopeTree.clone();
         requestAnimationFrame(() => {
           // Only restore focus if we've lost focus to the body, the alternative is that focus has been purposefully moved elsewhere
-          if (ownerDocument.activeElement === ownerDocument.body) {
+          if (ownerDocument.activeElement === rootBody) {
             // look up the tree starting with our scope to find a nodeToRestore still in the DOM
             let treeNode = clonedTree.getTreeNode(scopeRef);
             while (treeNode) {
@@ -719,10 +723,13 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
  * Create a [TreeWalker]{@link https://developer.mozilla.org/en-US/docs/Web/API/TreeWalker}
  * that matches all focusable/tabbable elements.
  */
-export function getFocusableTreeWalker(root: Element, opts?: FocusManagerOptions, scope?: Element[]) {
+export function getFocusableTreeWalker(root: Element | ShadowRoot, opts?: FocusManagerOptions, scope?: Element[]) {
   let selector = opts?.tabbable ? TABBABLE_ELEMENT_SELECTOR : FOCUSABLE_ELEMENT_SELECTOR;
-  let walker = getOwnerDocument(root).createTreeWalker(
-    root,
+  // Adjusted to directly handle root being a Document or ShadowRoot
+  let doc = root instanceof ShadowRoot ? root :  getRootNode(root);
+  let effectiveDocument = doc instanceof ShadowRoot ? doc.ownerDocument : doc;
+  let walker = effectiveDocument.createTreeWalker(
+    root || doc,
     NodeFilter.SHOW_ELEMENT,
     {
       acceptNode(node) {
@@ -762,7 +769,7 @@ export function createFocusManager(ref: RefObject<Element>, defaultOptions: Focu
         return null;
       }
       let {from, tabbable = defaultOptions.tabbable, wrap = defaultOptions.wrap, accept = defaultOptions.accept} = opts;
-      let node = from || getOwnerDocument(root).activeElement;
+      let node = from || getRootNode(root).activeElement;
       let walker = getFocusableTreeWalker(root, {tabbable, accept});
       if (root.contains(node)) {
         walker.currentNode = node!;
@@ -783,7 +790,7 @@ export function createFocusManager(ref: RefObject<Element>, defaultOptions: Focu
         return null;
       }
       let {from, tabbable = defaultOptions.tabbable, wrap = defaultOptions.wrap, accept = defaultOptions.accept} = opts;
-      let node = from || getOwnerDocument(root).activeElement;
+      let node = from || getRootNode(root).activeElement;
       let walker = getFocusableTreeWalker(root, {tabbable, accept});
       if (root.contains(node)) {
         walker.currentNode = node!;

packages/@react-aria/focus/src/focusSafely.ts

@@ -11,7 +11,7 @@
  */
 
 import {FocusableElement} from '@react-types/shared';
-import {focusWithoutScrolling, getOwnerDocument, runAfterTransition} from '@react-aria/utils';
+import {focusWithoutScrolling, getRootNode, runAfterTransition} from '@react-aria/utils';
 import {getInteractionModality} from '@react-aria/interactions';
 
 /**
@@ -24,7 +24,7 @@ export function focusSafely(element: FocusableElement) {
   // the page before shifting focus. This avoids issues with VoiceOver on iOS
   // causing the page to scroll when moving focus if the element is transitioning
   // from off the screen.
-  const ownerDocument = getOwnerDocument(element);
+  const ownerDocument = getRootNode(element);
   if (getInteractionModality() === 'virtual') {
     let lastFocusedElement = ownerDocument.activeElement;
     runAfterTransition(() => {