Fix focus set for delegated and already focused elements (facebook#36010)

I found two focus bugs when working on documentation for Fragment Refs.

1) If an element delegates focus handling, it will return false from
setFocusIfFocusable even though a focus event has occured on a different
element. The fix for this is a document level event listener rather than
only listening on the current element.

For example, if you have a form with multiple nested label>inputs.
Calling focus on the label will focus its input but not fire an event on
the label. setFocusIfFocusable returns false and you end up continuing
to attempt focus down the form tree.

2) If an element is already focused, setFocusIfFocusable will return
false. The fix for this is checking the document's activeElement with an
early return.

In the same form example, if the first input is already focused and you
call fragmentInstance.focus() at the form level, the second input would
end up getting focused since the focus event on the first is not
triggered.

Author: jackpope

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

@@ -4520,18 +4520,30 @@ export function setFocusIfFocusable(
   //
   // We could compare the node to document.activeElement after focus,
   // but this would not handle the case where application code managed focus to automatically blur.
+  const element = ((node: any): HTMLElement);
+
+  // If this element is already the active element, it's focusable and already
+  // focused. Calling .focus() on it would be a no-op (no focus event fires),
+  // so we short-circuit here.
+  if (element.ownerDocument.activeElement === element) {
+    return true;
+  }
+
   let didFocus = false;
   const handleFocus = () => {
     didFocus = true;
   };
 
-  const element = ((node: any): HTMLElement);
   try {
-    element.addEventListener('focus', handleFocus);
+    // Listen on the document in the capture phase so we detect focus even when
+    // it lands on a different element than the one we called .focus() on. This
+    // happens with <label> elements (focus delegates to the associated input)
+    // and shadow hosts with delegatesFocus.
+    element.ownerDocument.addEventListener('focus', handleFocus, true);
     // $FlowFixMe[method-unbinding]
     (element.focus || HTMLElement.prototype.focus).call(element, focusOptions);
   } finally {
-    element.removeEventListener('focus', handleFocus);
+    element.ownerDocument.removeEventListener('focus', handleFocus, true);
   }
 
   return didFocus;

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

@@ -255,6 +255,102 @@ describe('FragmentRefs', () => {
         expect(document.activeElement.id).toEqual('child-b');
         document.activeElement.blur();
       });
+
+      // @gate enableFragmentRefs
+      it('keeps focus on the first focusable child if already focused', async () => {
+        const fragmentRef = React.createRef();
+        const root = ReactDOMClient.createRoot(container);
+
+        function Test() {
+          return (
+            <Fragment ref={fragmentRef}>
+              <a id="child-a" href="/">
+                A
+              </a>
+              <a id="child-b" href="/">
+                B
+              </a>
+            </Fragment>
+          );
+        }
+
+        await act(() => {
+          root.render(<Test />);
+        });
+
+        // Focus the first child manually
+        document.getElementById('child-a').focus();
+        expect(document.activeElement.id).toEqual('child-a');
+
+        // Calling fragment.focus() should keep focus on child-a,
+        // not skip to child-b
+        await act(() => {
+          fragmentRef.current.focus();
+        });
+        expect(document.activeElement.id).toEqual('child-a');
+        document.activeElement.blur();
+      });
+
+      // @gate enableFragmentRefs
+      it('keeps focus on a nested child if already focused', async () => {
+        const fragmentRef = React.createRef();
+        const root = ReactDOMClient.createRoot(container);
+
+        function Test() {
+          return (
+            <Fragment ref={fragmentRef}>
+              <div>
+                <input id="nested-input" />
+              </div>
+              <a id="sibling-link" href="/">
+                Link
+              </a>
+            </Fragment>
+          );
+        }
+
+        await act(() => {
+          root.render(<Test />);
+        });
+
+        // Focus the nested input manually
+        document.getElementById('nested-input').focus();
+        expect(document.activeElement.id).toEqual('nested-input');
+
+        // Calling fragment.focus() should keep focus on nested-input
+        await act(() => {
+          fragmentRef.current.focus();
+        });
+        expect(document.activeElement.id).toEqual('nested-input');
+        document.activeElement.blur();
+      });
+
+      // @gate enableFragmentRefs
+      it('focuses the first focusable child in a fieldset', async () => {
+        const fragmentRef = React.createRef();
+        const root = ReactDOMClient.createRoot(container);
+
+        function Test() {
+          return (
+            <Fragment ref={fragmentRef}>
+              <fieldset>
+                <legend>Shipping</legend>
+                <input id="street" name="street" />
+                <input id="city" name="city" />
+              </fieldset>
+            </Fragment>
+          );
+        }
+
+        await act(() => {
+          root.render(<Test />);
+        });
+        await act(() => {
+          fragmentRef.current.focus();
+        });
+        expect(document.activeElement.id).toEqual('street');
+        document.activeElement.blur();
+      });
     });
 
     describe('focusLast()', () => {