Allow fragment refs to attempt focus/focusLast on nested host children

Author: jackpope

fixtures/dom/src/components/fixtures/fragment-refs/FocusCase.js

@@ -43,11 +43,18 @@ export default function FocusCase() {
         </Fixture.Controls>
         <div className="highlight-focused-children" style={{display: 'flex'}}>
           <Fragment ref={fragmentRef}>
-            <div style={{outline: '1px solid black'}}>Unfocusable div</div>
-            <button>Button 1</button>
+            <div style={{outline: '1px solid black'}}>
+              <p>Unfocusable div</p>
+            </div>
+            <div style={{outline: '1px solid black'}}>
+              <p>Unfocusable div with nested focusable button</p>
+              <button>Button 1</button>
+            </div>
             <button>Button 2</button>
             <input type="text" placeholder="Input field" />
-            <div style={{outline: '1px solid black'}}>Unfocusable div</div>
+            <div style={{outline: '1px solid black'}}>
+              <p>Unfocusable div</p>
+            </div>
           </Fragment>
         </div>
       </Fixture>

fixtures/dom/src/style.css

@@ -365,6 +365,10 @@ tbody tr:nth-child(even) {
   background-color: green;
 }
 
+.highlight-focused-children * {
+  margin-left: 10px;
+}
+
 .highlight-focused-children *:focus {
   outline: 2px solid green;
 }

packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js

@@ -60,6 +60,7 @@ import {
 } from './ReactDOMComponentTree';
 import {
   traverseFragmentInstance,
+  traverseFragmentInstanceDeeply,
   getFragmentParentHostInstance,
 } from 'react-reconciler/src/ReactFiberTreeReflection';
 
@@ -2610,7 +2611,7 @@ FragmentInstance.prototype.focus = function (
   this: FragmentInstanceType,
   focusOptions?: FocusOptions,
 ): void {
-  traverseFragmentInstance(
+  traverseFragmentInstanceDeeply(
     this._fragmentFiber,
     setFocusIfFocusable,
     focusOptions,
@@ -2622,7 +2623,11 @@ FragmentInstance.prototype.focusLast = function (
   focusOptions?: FocusOptions,
 ): void {
   const children: Array<Instance> = [];
-  traverseFragmentInstance(this._fragmentFiber, collectChildren, children);
+  traverseFragmentInstanceDeeply(
+    this._fragmentFiber,
+    collectChildren,
+    children,
+  );
   for (let i = children.length - 1; i >= 0; i--) {
     const child = children[i];
     if (setFocusIfFocusable(child, focusOptions)) {

packages/react-dom/src/__tests__/ReactDOMFragmentRefs-test.js

@@ -140,6 +140,33 @@ describe('FragmentRefs', () => {
         document.activeElement.blur();
       });
 
+      // @gate enableFragmentRefs
+      it('focuses deeply nested focusable children, depth first', async () => {
+        const fragmentRef = React.createRef();
+        const root = ReactDOMClient.createRoot(container);
+        const childRef = React.createRef();
+
+        function Test() {
+          return (
+            <Fragment ref={fragmentRef}>
+              <div id="child-a">
+                <div tabIndex={0} id="grandchild-a">
+                  <a id="greatgrandchild-a" href="/"></a>
+                </div>
+              </div>
+              <a id="child-b" href="/"></a>
+            </Fragment>
+          );
+        }
+        await act(() => {
+          root.render(<Test />);
+        });
+        await act(() => {
+          fragmentRef.current.focus();
+        });
+        expect(document.activeElement.id).toEqual('grandchild-a');
+      });
+
       // @gate enableFragmentRefs
       it('preserves document order when adding and removing children', async () => {
         const fragmentRef = React.createRef();
@@ -223,6 +250,35 @@ describe('FragmentRefs', () => {
         expect(document.activeElement.id).toEqual('child-c');
         document.activeElement.blur();
       });
+
+      // @gate enableFragmentRefs
+      it('focuses deeply nested focusable children, depth first', async () => {
+        const fragmentRef = React.createRef();
+        const root = ReactDOMClient.createRoot(container);
+        const childRef = React.createRef();
+
+        function Test() {
+          return (
+            <Fragment ref={fragmentRef}>
+              <div id="child-a" href="/">
+                <a id="grandchild-a" href="/"></a>
+                <a id="grandchild-b" href="/"></a>
+              </div>
+              <div tabIndex={0} id="child-b">
+                <a id="grandchild-a" href="/"></a>
+                <a id="grandchild-b" href="/"></a>
+              </div>
+            </Fragment>
+          );
+        }
+        await act(() => {
+          root.render(<Test />);
+        });
+        await act(() => {
+          fragmentRef.current.focusLast();
+        });
+        expect(document.activeElement.id).toEqual('grandchild-b');
+      });
     });
 
     describe('blur()', () => {

packages/react-reconciler/src/ReactFiberTreeReflection.js

@@ -352,31 +352,53 @@ export function traverseFragmentInstance<A, B, C>(
   b: B,
   c: C,
 ): void {
-  traverseFragmentInstanceChildren(fragmentFiber.child, fn, a, b, c);
+  traverseFragmentInstanceChildren(fragmentFiber.child, false, fn, a, b, c);
+}
+
+export function traverseFragmentInstanceDeeply<A, B, C>(
+  fragmentFiber: Fiber,
+  fn: (Instance, A, B, C) => boolean,
+  a: A,
+  b: B,
+  c: C,
+): void {
+  traverseFragmentInstanceChildren(fragmentFiber.child, true, fn, a, b, c);
 }
 
 function traverseFragmentInstanceChildren<A, B, C>(
   child: Fiber | null,
+  searchWithinHosts: boolean,
   fn: (Instance, A, B, C) => boolean,
   a: A,
   b: B,
   c: C,
-): void {
+): boolean {
   while (child !== null) {
-    if (child.tag === HostComponent) {
-      if (fn(child.stateNode, a, b, c)) {
-        return;
-      }
+    if (child.tag === HostComponent && fn(child.stateNode, a, b, c)) {
+      return true;
     } else if (
       child.tag === OffscreenComponent &&
       child.memoizedState !== null
     ) {
       // Skip hidden subtrees
     } else {
-      traverseFragmentInstanceChildren(child.child, fn, a, b, c);
+      if (
+        (searchWithinHosts || child.tag !== HostComponent) &&
+        traverseFragmentInstanceChildren(
+          child.child,
+          searchWithinHosts,
+          fn,
+          a,
+          b,
+          c,
+        )
+      ) {
+        return true;
+      }
     }
     child = child.sibling;
   }
+  return false;
 }
 
 export function getFragmentParentHostInstance(fiber: Fiber): null | Instance {