Encode lazy Node slots properly in key path and resumable paths (#27359)

It's possible to postpone a specific node and not using a wrapper
component. Therefore we encode the resumable slot as the index slot.
When it's a plain client component that postpones, it's encoded as the
child slot inside that component which is the one that's postponed
rather than the component itself.

Since it's possible for a child slot to suspend (e.g. React.lazy's
microtask in this case) retryTask might need to keep its index around
when it resolves.

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

@@ -493,6 +493,47 @@ describe('ReactDOMFizzStaticBrowser', () => {
     // TODO: expect(getVisibleChildren(container)).toEqual(<div>Hello</div>);
   });
 
+  // @gate enablePostpone
+  it('supports postponing in lazy in prerender and resuming later', async () => {
+    let prerendering = true;
+    const Hole = React.lazy(async () => {
+      React.unstable_postpone();
+    });
+
+    function Postpone() {
+      return 'Hello';
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            Hi
+            {prerendering ? Hole : <Postpone />}
+          </Suspense>
+        </div>
+      );
+    }
+
+    const prerendered = await ReactDOMFizzStatic.prerender(<App />);
+    expect(prerendered.postponed).not.toBe(null);
+
+    prerendering = false;
+
+    const resumed = await ReactDOMFizzServer.resume(
+      <App />,
+      prerendered.postponed,
+    );
+
+    await readIntoContainer(prerendered.prelude);
+
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    await readIntoContainer(resumed);
+
+    // TODO: expect(getVisibleChildren(container)).toEqual(<div>Hello</div>);
+  });
+
   // @gate enablePostpone
   it('only emits end tags once when resuming', async () => {
     let prerendering = true;

packages/react-server/src/ReactFizzServer.js

@@ -186,8 +186,7 @@ type ResumableNode =
   | ResumableParentNode
   | [
       2, // RESUME_SEGMENT
-      string | null /* name */,
-      string | number /* key */,
+      number /* index */,
       number /* segment id */,
     ];
 
@@ -220,6 +219,7 @@ type SuspenseBoundary = {
 
 export type Task = {
   node: ReactNodeList,
+  childIndex: number,
   ping: () => void,
   blockedBoundary: Root | SuspenseBoundary,
   blockedSegment: Segment, // the segment we'll write to
@@ -1632,6 +1632,7 @@ function renderNodeDestructiveImpl(
   // Stash the node we're working on. We'll pick up from this task in case
   // something suspends.
   task.node = node;
+  task.childIndex = childIndex;
 
   // Handle object types
   if (typeof node === 'object' && node !== null) {
@@ -1809,18 +1810,45 @@ function renderChildrenArray(
   for (let i = 0; i < totalChildren; i++) {
     const node = children[i];
     task.treeContext = pushTreeContext(prevTreeContext, totalChildren, i);
-    if (isArray(node) || getIteratorFn(node)) {
-      // Nested arrays behave like a "fragment node" which is keyed.
-      // Therefore we need to add the current index as a parent key.
+
+    // Nested arrays behave like a "fragment node" which is keyed.
+    // Therefore we need to add the current index as a parent key.
+    // We first check if the nested nodes are arrays or iterables.
+
+    if (isArray(node)) {
       const prevKeyPath = task.keyPath;
       task.keyPath = [task.keyPath, '', childIndex];
-      renderNode(request, task, node, i);
+      renderChildrenArray(request, task, node, i);
       task.keyPath = prevKeyPath;
-    } else {
-      // We need to use the non-destructive form so that we can safely pop back
-      // up and render the sibling if something suspends.
-      renderNode(request, task, node, i);
+      continue;
     }
+
+    const iteratorFn = getIteratorFn(node);
+    if (iteratorFn) {
+      if (__DEV__) {
+        validateIterable(node, iteratorFn);
+      }
+      const iterator = iteratorFn.call(node);
+      if (iterator) {
+        let step = iterator.next();
+        if (!step.done) {
+          const prevKeyPath = task.keyPath;
+          task.keyPath = [task.keyPath, '', childIndex];
+          const nestedChildren = [];
+          do {
+            nestedChildren.push(step.value);
+            step = iterator.next();
+          } while (!step.done);
+          renderChildrenArray(request, task, nestedChildren, i);
+          task.keyPath = prevKeyPath;
+        }
+        continue;
+      }
+    }
+
+    // We need to use the non-destructive form so that we can safely pop back
+    // up and render the sibling if something suspends.
+    renderNode(request, task, node, i);
   }
   // Because this context is always set right before rendering every child, we
   // only need to reset it to the previous value at the very end.
@@ -1831,6 +1859,7 @@ function trackPostpone(
   request: Request,
   trackedPostpones: PostponedHoles,
   task: Task,
+  childIndex: number,
   segment: Segment,
 ): void {
   segment.status = POSTPONED;
@@ -1862,7 +1891,7 @@ function trackPostpone(
       boundary.id,
     ];
     trackedPostpones.workingMap.set(boundaryKeyPath, boundaryNode);
-    addToResumableParent(boundaryNode, boundaryKeyPath, trackedPostpones);
+    addToResumableParent(boundaryNode, boundaryKeyPath[0], trackedPostpones);
   }
 
   const keyPath = task.keyPath;
@@ -1872,12 +1901,7 @@ function trackPostpone(
     );
   }
 
-  const segmentNode: ResumableNode = [
-    RESUME_SEGMENT,
-    keyPath[1],
-    keyPath[2],
-    segment.id,
-  ];
+  const segmentNode: ResumableNode = [RESUME_SEGMENT, childIndex, segment.id];
   addToResumableParent(segmentNode, keyPath, trackedPostpones);
 }
 
@@ -1941,6 +1965,7 @@ function spawnNewSuspendedTask(
     task.context,
     task.treeContext,
   );
+  newTask.childIndex = task.childIndex;
 
   if (__DEV__) {
     if (task.componentStack !== null) {
@@ -2035,7 +2060,13 @@ function renderNode(
           task,
           postponeInstance.message,
         );
-        trackPostpone(request, trackedPostpones, task, postponedSegment);
+        trackPostpone(
+          request,
+          trackedPostpones,
+          task,
+          childIndex,
+          postponedSegment,
+        );
 
         // Restore the context. We assume that this will be restored by the inner
         // functions in case nothing throws so we don't use "finally" here.
@@ -2328,7 +2359,13 @@ function retryTask(request: Request, task: Task): void {
     const prevThenableState = task.thenableState;
     task.thenableState = null;
 
-    renderNodeDestructive(request, task, prevThenableState, task.node, 0);
+    renderNodeDestructive(
+      request,
+      task,
+      prevThenableState,
+      task.node,
+      task.childIndex,
+    );
     pushSegmentFinale(
       segment.chunks,
       request.renderState,
@@ -2377,8 +2414,15 @@ function retryTask(request: Request, task: Task): void {
         task.abortSet.delete(task);
         const postponeInstance: Postpone = (x: any);
         logPostpone(request, postponeInstance.message);
-        trackPostpone(request, trackedPostpones, task, segment);
+        trackPostpone(
+          request,
+          trackedPostpones,
+          task,
+          task.childIndex,
+          segment,
+        );
         finishedTask(request, task.blockedBoundary, segment);
+        return;
       }
     }
     task.abortSet.delete(task);
@@ -2975,10 +3019,9 @@ export function getResumableState(request: Request): ResumableState {
 
 function addToResumableParent(
   node: ResumableNode,
-  keyPath: KeyNode,
+  parentKeyPath: Root | KeyNode,
   trackedPostpones: PostponedHoles,
 ): void {
-  const parentKeyPath = keyPath[0];
   if (parentKeyPath === null) {
     trackedPostpones.root.push(node);
   } else {
@@ -2992,7 +3035,7 @@ function addToResumableParent(
         ([]: Array<ResumableNode>),
       ]: ResumableParentNode);
       workingMap.set(parentKeyPath, parentNode);
-      addToResumableParent(parentNode, parentKeyPath, trackedPostpones);
+      addToResumableParent(parentNode, parentKeyPath[0], trackedPostpones);
     }
     parentNode[3].push(node);
   }