Consume SuspenseNodes that were skipped when we're bailing out of a subtree

Author: sebmarkbage

packages/react-devtools-shared/src/backend/fiber/renderer.js

@@ -2700,6 +2700,76 @@ export function attach(
     }
   }
 
+  function isChildOf(
+    parentInstance: DevToolsInstance,
+    childInstance: DevToolsInstance,
+    grandParent: DevToolsInstance,
+  ): boolean {
+    let instance = childInstance.parent;
+    while (instance !== null) {
+      if (parentInstance === instance) {
+        return true;
+      }
+      if (instance === parentInstance.parent || instance === grandParent) {
+        // This was a sibling but not inside the FiberInstance. We can bail out.
+        break;
+      }
+      instance = instance.parent;
+    }
+    return false;
+  }
+
+  function consumeSuspenseNodesOfExistingInstance(
+    instance: DevToolsInstance,
+  ): void {
+    // We need to also consume any unchanged Suspense boundaries.
+    let suspenseNode = remainingReconcilingChildrenSuspenseNodes;
+    if (suspenseNode === null) {
+      return;
+    }
+    const parentSuspenseNode = reconcilingParentSuspenseNode;
+    if (parentSuspenseNode === null) {
+      throw new Error(
+        'The should not be any remaining suspense node children if there is no parent.',
+      );
+    }
+    let foundOne = false;
+    let previousSkippedSibling = null;
+    while (suspenseNode !== null) {
+      // Check if this SuspenseNode was a child of the bailed out FiberInstance.
+      if (
+        isChildOf(instance, suspenseNode.instance, parentSuspenseNode.instance)
+      ) {
+        foundOne = true;
+        // The suspenseNode was child of the bailed out Fiber.
+        // First, remove it from the remaining children set.
+        const nextRemainingSibling = suspenseNode.nextSibling;
+        if (previousSkippedSibling === null) {
+          remainingReconcilingChildrenSuspenseNodes = nextRemainingSibling;
+        } else {
+          previousSkippedSibling.nextSibling = nextRemainingSibling;
+        }
+        suspenseNode.nextSibling = null;
+        // Then, re-insert it into the newly reconciled set.
+        if (previouslyReconciledSiblingSuspenseNode === null) {
+          parentSuspenseNode.firstChild = suspenseNode;
+        } else {
+          previouslyReconciledSiblingSuspenseNode.nextSibling = suspenseNode;
+        }
+        previouslyReconciledSiblingSuspenseNode = suspenseNode;
+        // Continue
+        suspenseNode = nextRemainingSibling;
+      } else if (foundOne) {
+        // If we found one and then hit a miss, we assume that we're passed the sequence because
+        // they should've all been consecutive.
+        break;
+      } else {
+        previousSkippedSibling = suspenseNode;
+        suspenseNode = suspenseNode.nextSibling;
+      }
+    }
+  }
+
   function mountVirtualInstanceRecursively(
     virtualInstance: VirtualInstance,
     firstChild: Fiber,
@@ -3848,6 +3918,8 @@ export function attach(
             fiberInstance.firstChild = remainingReconcilingChildren;
             remainingReconcilingChildren = null;
 
+            consumeSuspenseNodesOfExistingInstance(fiberInstance);
+
             if (traceUpdatesEnabled) {
               // If we're tracing updates and we've bailed out before reaching a host node,
               // we should fall back to recursively marking the nearest host descendants for highlight.