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

@@ -68,6 +68,7 @@ import {
   getFragmentParentHostFiber,
   getNextSiblingHostFiber,
   getInstanceFromHostFiber,
+  traverseFragmentInstanceDeeply,
 } from 'react-reconciler/src/ReactFiberTreeReflection';
 
 export {detachDeletedInstance};
@@ -2698,7 +2699,7 @@ FragmentInstance.prototype.focus = function (
   this: FragmentInstanceType,
   focusOptions?: FocusOptions,
 ): void {
-  traverseFragmentInstance(
+  traverseFragmentInstanceDeeply(
     this._fragmentFiber,
     setFocusOnFiberIfFocusable,
     focusOptions,
@@ -2717,7 +2718,11 @@ FragmentInstance.prototype.focusLast = function (
   focusOptions?: FocusOptions,
 ): void {
   const children: Array<Fiber> = [];
-  traverseFragmentInstance(this._fragmentFiber, collectChildren, children);
+  traverseFragmentInstanceDeeply(
+    this._fragmentFiber,
+    collectChildren,
+    children,
+  );
   for (let i = children.length - 1; i >= 0; i--) {
     const child = children[i];
     if (setFocusOnFiberIfFocusable(child, focusOptions)) {

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

@@ -145,6 +145,32 @@ 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);
+
+        function Test() {
+          return (
+            <Fragment ref={fragmentRef}>
+              <div id="child-a">
+                <div tabIndex={0} id="grandchild-a">
+                  <a id="greatgrandchild-a" href="/" />
+                </div>
+              </div>
+              <a id="child-b" href="/" />
+            </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();
@@ -228,6 +254,34 @@ 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);
+
+        function Test() {
+          return (
+            <Fragment ref={fragmentRef}>
+              <div id="child-a" href="/">
+                <a id="grandchild-a" href="/" />
+                <a id="grandchild-b" href="/" />
+              </div>
+              <div tabIndex={0} id="child-b">
+                <a id="grandchild-a" href="/" />
+                <a id="grandchild-b" href="/" />
+              </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

@@ -354,6 +354,16 @@ export function traverseFragmentInstance<A, B, C>(
   traverseVisibleHostChildren(fragmentFiber.child, false, fn, a, b, c);
 }
 
+export function traverseFragmentInstanceDeeply<A, B, C>(
+  fragmentFiber: Fiber,
+  fn: (Fiber, A, B, C) => boolean,
+  a: A,
+  b: B,
+  c: C,
+): void {
+  traverseVisibleHostChildren(fragmentFiber.child, true, fn, a, b, c);
+}
+
 function traverseVisibleHostChildren<A, B, C>(
   child: Fiber | null,
   searchWithinHosts: boolean,
@@ -363,31 +373,16 @@ function traverseVisibleHostChildren<A, B, C>(
   c: C,
 ): boolean {
   while (child !== null) {
-    if (child.tag === HostComponent) {
-      if (fn(child, a, b, c)) {
-        return true;
-      }
-      if (searchWithinHosts) {
-        if (
-          traverseVisibleHostChildren(
-            child.child,
-            searchWithinHosts,
-            fn,
-            a,
-            b,
-            c,
-          )
-        ) {
-          return true;
-        }
-      }
+    if (child.tag === HostComponent && fn(child, a, b, c)) {
+      return true;
     } else if (
       child.tag === OffscreenComponent &&
       child.memoizedState !== null
     ) {
       // Skip hidden subtrees
     } else {
       if (
+        (searchWithinHosts || child.tag !== HostComponent) &&
         traverseVisibleHostChildren(child.child, searchWithinHosts, fn, a, b, c)
       ) {
         return true;