Fix Suspense fragment edge cases

shells/dev/app/SuspenseTree/index.js

@@ -0,0 +1,64 @@
+// @flow
+
+import React, { Suspense, useState } from 'react';
+
+function SuspenseTree() {
+  return (
+    <>
+      <h1>Suspense</h1>
+      <Suspense fallback={<Fallback>Loading outer</Fallback>}>
+        <Parent />
+      </Suspense>
+    </>
+  );
+}
+
+function Parent() {
+  return (
+    <div>
+      <Suspense fallback={<Fallback>Loading inner 1</Fallback>}>
+        <Child>Hello</Child>
+      </Suspense>
+      <Suspense fallback={<Fallback>Loading inner 2</Fallback>}>
+        <Child>World</Child>
+      </Suspense>
+      <Suspense fallback={<Fallback>This will never load</Fallback>}>
+        <Never />
+      </Suspense>
+      <LoadLater />
+    </div>
+  );
+}
+
+function LoadLater() {
+  const [loadChild, setLoadChild] = useState(0);
+  return (
+    <Suspense
+      fallback={
+        <Fallback onClick={() => setLoadChild(true)}>Click to load</Fallback>
+      }
+    >
+      {loadChild ? (
+        <Child onClick={() => setLoadChild(false)}>
+          Loaded! Click to suspend again.
+        </Child>
+      ) : (
+        <Never />
+      )}
+    </Suspense>
+  );
+}
+
+function Child(props) {
+  return <p {...props} />;
+}
+
+function Fallback(props) {
+  return <h3 {...props}>{props.children}</h3>;
+}
+
+function Never() {
+  throw new Promise(resolve => {});
+}
+
+export default SuspenseTree;

shells/dev/app/index.js

@@ -11,6 +11,7 @@ import InspectableElements from './InspectableElements';
 import InteractionTracing from './InteractionTracing';
 import ToDoList from './ToDoList';
 import Toggle from './Toggle';
+import SuspenseTree from './SuspenseTree';
 
 import './styles.css';
 
@@ -33,6 +34,7 @@ function mountTestApp() {
   mountHelper(ElementTypes);
   mountHelper(EditableProps);
   mountHelper(Toggle);
+  mountHelper(SuspenseTree);
   mountHelper(DeeplyNestedComponents);
 }
 

src/backend/renderer.js

@@ -716,6 +716,16 @@ export function attach(
   function enqueueUnmount(fiber) {
     const isRoot = fiber.tag === HostRoot;
     const primaryFiber = getPrimaryFiber(fiber);
+    if (!fiberToIDMap.has(primaryFiber)) {
+      // If we've never seen this Fiber, it might be because
+      // it is inside a non-current Suspense fragment tree,
+      // and so the store is not even aware of it.
+      // In that case we can just ignore it, or otherwise
+      // there will be errors later on.
+      primaryFibers.delete(primaryFiber);
+      // TODO: this is fragile and can obscure actual bugs.
+      return;
+    }
     const id = getFiberID(primaryFiber);
     if (isRoot) {
       const operation = new Uint32Array(2);
@@ -757,8 +767,24 @@ export function attach(
       enqueueMount(fiber, parentFiber);
     }
 
-    if (fiber.child !== null) {
-      mountFiber(fiber.child, shouldEnqueueMount ? fiber : parentFiber, true);
+    const isTimedOutSuspense =
+      fiber.tag === ReactTypeOfWork.SuspenseComponent &&
+      fiber.memoizedState !== null;
+
+    if (isTimedOutSuspense) {
+      // Special case: if Suspense mounts in a timed-out state,
+      // get the fallback child from the inner fragment and mount
+      // it as if it was our own child. Updates handle this too.
+      const primaryChildFragment = fiber.child;
+      const fallbackChildFragment = primaryChildFragment.sibling;
+      const fallbackChild = fallbackChildFragment.child;
+      if (fallbackChild !== null) {
+        mountFiber(fallbackChild, shouldEnqueueMount ? fiber : parentFiber, true);
+      }
+    } else {
+      if (fiber.child !== null) {
+        mountFiber(fiber.child, shouldEnqueueMount ? fiber : parentFiber, true);
+      }
     }
 
     if (traverseSiblings && fiber.sibling !== null) {