From b0657fde6aef2e067e6d7808809deb194a857278 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <trueadm@users.noreply.github.com>
Date: Wed, 15 May 2019 01:08:30 +0100
Subject: [PATCH] Event API: ensure getFocusableElementsInScope handles
 suspended trees (#15651)

---
 .../src/events/DOMEventResponderSystem.js     | 70 ++++++++++++++++---
 .../react-dom/src/shared/assertValidProps.js  |  1 +
 .../src/__tests__/FocusScope-test.internal.js | 50 +++++++++++++
 3 files changed, 112 insertions(+), 9 deletions(-)

diff --git a/packages/react-dom/src/events/DOMEventResponderSystem.js b/packages/react-dom/src/events/DOMEventResponderSystem.js
index 251e5c6695f5..aae2ab55dbc8 100644
--- a/packages/react-dom/src/events/DOMEventResponderSystem.js
+++ b/packages/react-dom/src/events/DOMEventResponderSystem.js
@@ -16,6 +16,8 @@ import {
   EventComponent,
   EventTarget as EventTargetWorkTag,
   HostComponent,
+  SuspenseComponent,
+  Fragment,
 } from 'shared/ReactWorkTags';
 import type {
   ReactEventResponder,
@@ -352,14 +354,24 @@ const eventResponderContext: ReactResponderContext = {
     let node = ((eventComponentInstance.currentFiber: any): Fiber).child;
 
     while (node !== null) {
-      if (isFiberHostComponentFocusable(node)) {
-        focusableElements.push(node.stateNode);
+      if (node.tag === SuspenseComponent) {
+        const suspendedChild = isFiberSuspenseAndTimedOut(node)
+          ? getSuspenseFallbackChild(node)
+          : getSuspenseChild(node);
+        if (suspendedChild !== null) {
+          node = suspendedChild;
+          continue;
+        }
       } else {
-        const child = node.child;
+        if (isFiberHostComponentFocusable(node)) {
+          focusableElements.push(node.stateNode);
+        } else {
+          const child = node.child;
 
-        if (child !== null) {
-          node = child;
-          continue;
+          if (child !== null) {
+            node = child;
+            continue;
+          }
         }
       }
       const sibling = node.sibling;
@@ -368,9 +380,14 @@ const eventResponderContext: ReactResponderContext = {
         node = sibling;
         continue;
       }
-      const parent = node.return;
-      if (parent === null) {
-        break;
+      let parent;
+      if (isFiberSuspenseChild(node)) {
+        parent = getSuspenseFiberFromChild(node);
+      } else {
+        parent = node.return;
+        if (parent === null) {
+          break;
+        }
       }
       if (parent.stateNode === currentInstance) {
         break;
@@ -588,6 +605,41 @@ export function processEventQueue(): void {
   }
 }
 
+function isFiberSuspenseAndTimedOut(fiber: Fiber): boolean {
+  return fiber.tag === SuspenseComponent && fiber.memoizedState !== null;
+}
+
+function isFiberSuspenseChild(fiber: Fiber | null): boolean {
+  if (fiber === null) {
+    return false;
+  }
+  const parent = fiber.return;
+  if (parent !== null && parent.tag === Fragment) {
+    const grandParent = parent.return;
+
+    if (
+      grandParent !== null &&
+      grandParent.tag === SuspenseComponent &&
+      grandParent.stateNode !== null
+    ) {
+      return true;
+    }
+  }
+  return false;
+}
+
+function getSuspenseFiberFromChild(fiber: Fiber): Fiber {
+  return ((((fiber.return: any): Fiber).return: any): Fiber);
+}
+
+function getSuspenseFallbackChild(fiber: Fiber): Fiber | null {
+  return ((((fiber.child: any): Fiber).sibling: any): Fiber).child;
+}
+
+function getSuspenseChild(fiber: Fiber): Fiber | null {
+  return (((fiber.child: any): Fiber): Fiber).child;
+}
+
 function getTargetEventTypesSet(
   eventTypes: Array<ReactEventResponderEventType>,
 ): Set<string> {
diff --git a/packages/react-dom/src/shared/assertValidProps.js b/packages/react-dom/src/shared/assertValidProps.js
index 4c6d9eb2f04d..30fc5c6ce5ff 100644
--- a/packages/react-dom/src/shared/assertValidProps.js
+++ b/packages/react-dom/src/shared/assertValidProps.js
@@ -31,6 +31,7 @@ function assertValidProps(tag: string, props: ?Object) {
     invariant(
       (props.children == null ||
         (enableEventAPI &&
+          props.children.type &&
           props.children.type.$$typeof === REACT_EVENT_TARGET_TYPE)) &&
         props.dangerouslySetInnerHTML == null,
       '%s is a void element tag and must neither have `children` nor ' +
diff --git a/packages/react-events/src/__tests__/FocusScope-test.internal.js b/packages/react-events/src/__tests__/FocusScope-test.internal.js
index 22c72418b4a7..3930b25e1fa3 100644
--- a/packages/react-events/src/__tests__/FocusScope-test.internal.js
+++ b/packages/react-events/src/__tests__/FocusScope-test.internal.js
@@ -191,4 +191,54 @@ describe('FocusScope event responder', () => {
     document.activeElement.dispatchEvent(createTabBackward());
     expect(document.activeElement).toBe(button2Ref.current);
   });
+
+  it('should work as expected with suspense fallbacks', () => {
+    const buttonRef = React.createRef();
+    const button2Ref = React.createRef();
+    const button3Ref = React.createRef();
+    const button4Ref = React.createRef();
+    const button5Ref = React.createRef();
+
+    function SuspendedComponent() {
+      throw new Promise(() => {
+        // Never resolve
+      });
+    }
+
+    function Component() {
+      return (
+        <React.Fragment>
+          <button ref={button5Ref} id={5} />
+          <SuspendedComponent />
+        </React.Fragment>
+      );
+    }
+
+    const SimpleFocusScope = () => (
+      <div>
+        <FocusScope>
+          <button ref={buttonRef} id={1} />
+          <button ref={button2Ref} id={2} />
+          <React.Suspense fallback={<button ref={button3Ref} id={3} />}>
+            <Component />
+          </React.Suspense>
+          <button ref={button4Ref} id={4} />
+        </FocusScope>
+      </div>
+    );
+
+    ReactDOM.render(<SimpleFocusScope />, container);
+    buttonRef.current.focus();
+    expect(document.activeElement).toBe(buttonRef.current);
+    document.activeElement.dispatchEvent(createTabForward());
+    expect(document.activeElement).toBe(button2Ref.current);
+    document.activeElement.dispatchEvent(createTabForward());
+    expect(document.activeElement).toBe(button3Ref.current);
+    document.activeElement.dispatchEvent(createTabForward());
+    expect(document.activeElement).toBe(button4Ref.current);
+    document.activeElement.dispatchEvent(createTabBackward());
+    expect(document.activeElement).toBe(button3Ref.current);
+    document.activeElement.dispatchEvent(createTabBackward());
+    expect(document.activeElement).toBe(button2Ref.current);
+  });
 });