Allow fragment refs to attempt focus/focusLast on nested host children (facebook#33058)

This enables `focus` and `focusLast` methods on FragmentInstances to
search nested host components, depth first. Attempts focus on each child
and bails if one is successful. Previously, only the first level of host
children would attempt focus.

Now if we have an example like

```
component MenuItem() {
  return (<div><a>{...}</a></div>)
}

component Menu() {
  return <Fragment>{items.map(i => <MenuItem i={i} />)}</Fragment>
}
```
We can target focus on the first or last a tag, rather than checking
each wrapping div and then noop.

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;