[Fizz] Fix root segment IDs (#27371)

Typically we assign IDs lazily when flushing to minimize the ids we have
to assign and we try to maximize inlining.

When we prerender we could always flush into a buffer before returning
but that's not actually what we do right now. We complete rendering
before returning but we don't actually run the flush path until someone
reads the resulting stream. We assign IDs eagerly when something
postpone so that we can refer to those ids in the holes without flushing
first.

This leads to some interesting conditions that needs to consider this.

This PR also deals with rootSegmentId which is the ID of the segment
that contains the body of a Suspense boundary. The boundary needs to
know this in addition to the Suspense Boundary's ID to know how to
inject the root segment into the boundary.

Why is the suspense boundary ID and the rootSegmentID just not the same
ID? Why don't segments and suspense boundary not share ID namespace?
That's a good question and I don't really remember. It might not be a
good reason and maybe they should just be the same.

packages/react-server/src/ReactFizzServer.js

@@ -176,6 +176,7 @@ type ReplaySuspenseBoundary = [
   string | number /* key */,
   Array<ResumableNode> /* children */,
   SuspenseBoundaryID /* id */,
+  number /* rootSegmentID */,
 ];
 
 type ReplayNode =
@@ -1957,6 +1958,7 @@ function trackPostpone(
       request.renderState,
       request.resumableState,
     );
+    boundary.rootSegmentID = request.nextSegmentId++;
 
     const boundaryKeyPath = boundary.keyPath;
     if (boundaryKeyPath === null) {
@@ -1971,6 +1973,7 @@ function trackPostpone(
       boundaryKeyPath[2],
       children,
       boundary.id,
+      boundary.rootSegmentID,
     ];
     trackedPostpones.workingMap.set(boundaryKeyPath, boundaryNode);
     addToReplayParent(boundaryNode, boundaryKeyPath[0], trackedPostpones);
@@ -2021,7 +2024,7 @@ function injectPostponedHole(
   segment.children.push(newSegment);
   // Reset lastPushedText for current Segment since the new Segment "consumed" it
   segment.lastPushedText = false;
-  return segment;
+  return newSegment;
 }
 
 function spawnNewSuspendedTask(
@@ -2316,7 +2319,9 @@ function queueCompletedSegment(
   if (
     segment.chunks.length === 0 &&
     segment.children.length === 1 &&
-    segment.children[0].boundary === null
+    segment.children[0].boundary === null &&
+    // Typically the id would not be assigned yet but if it's a postponed segment it might be.
+    segment.children[0].id === -1
   ) {
     // This is an empty segment. There's nothing to write, so we can instead transfer the ID
     // to the child. That way any existing references point to the child.
@@ -2668,20 +2673,22 @@ function flushSegment(
     );
   } else if (boundary.status !== COMPLETED) {
     if (boundary.status === PENDING) {
+      // For pending boundaries we lazily assign an ID to the boundary
+      // and root segment.
       boundary.id = assignSuspenseBoundaryID(
         request.renderState,
         request.resumableState,
       );
+      boundary.rootSegmentID = request.nextSegmentId++;
     }
-    // This boundary is still loading. Emit a pending suspense boundary wrapper.
 
-    // Assign an ID to refer to the future content by.
-    boundary.rootSegmentID = request.nextSegmentId++;
     if (boundary.completedSegments.length > 0) {
       // If this is at least partially complete, we can queue it to be partially emitted early.
       request.partialBoundaries.push(boundary);
     }
 
+    // This boundary is still loading. Emit a pending suspense boundary wrapper.
+
     /// This is the first time we should have referenced this ID.
     const id = boundary.id;
 
@@ -2868,6 +2875,10 @@ function flushPartiallyCompletedSegment(
       );
     }
 
+    return flushSegmentContainer(request, destination, segment);
+  } else if (segmentID === boundary.rootSegmentID) {
+    // When we emit postponed boundaries, we might have assigned the ID already
+    // but it's still the root segment so we can't inject it into the parent yet.
     return flushSegmentContainer(request, destination, segment);
   } else {
     flushSegmentContainer(request, destination, segment);