Track the key path difference between right before the first array and after

There's a subtle difference if you suspend before the first array or after.
In Fiber, we don't deal with this because we just suspend the parent and
replay it if lazy() or Usable are used in its child slots. In Fizz we try
to optimize this a bit more and enable resuming inside the component.

Semantically, it's different if you suspend/postpone before the first child
array or inside that child array. Because when you resume the inner result
might be another array and either that's part of the parent path or part
of the inner slot.

There might be more clever way of structuring this but I just use -1 to
indicate that we're not yet inside the array and is in the root child
position. If that renders an element, then that's just the same as the 0
slot.

We need to also encode this in the resuming.

Author: sebmarkbage

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

@@ -461,7 +461,7 @@ describe('ReactDOMFizzStaticBrowser', () => {
       if (prerendering) {
         React.unstable_postpone();
       }
-      return 'Hello';
+      return ['Hello', 'World'];
     }
 
     function App() {

packages/react-server/src/ReactFizzServer.js

@@ -166,7 +166,8 @@ type KeyNode = [
 
 const REPLAY_NODE = 0;
 const REPLAY_SUSPENSE_BOUNDARY = 1;
-const RESUME_SEGMENT = 2;
+const RESUME_ELEMENT = 2;
+const RESUME_SLOT = 3;
 
 type ResumableParentNode =
   | [
@@ -185,10 +186,16 @@ type ResumableParentNode =
 type ResumableNode =
   | ResumableParentNode
   | [
-      2, // RESUME_SEGMENT
-      number /* index */,
+      2, // RESUME_ELEMENT
+      string | null /* name */,
+      string | number /* key */,
       number /* segment id */,
-    ];
+    ]
+  | [
+    3, // RESUME_SLOT
+    number /* index */,
+    number /* segment id */,
+  ];
 
 type PostponedHoles = {
   workingMap: Map<KeyNode, ResumableParentNode>,
@@ -784,7 +791,7 @@ function renderSuspenseBoundary(
   }
   try {
     // We use the safe form because we don't handle suspending here. Only error handling.
-    renderNode(request, task, content, 0);
+    renderNode(request, task, content, -1);
     pushSegmentFinale(
       contentRootSegment.chunks,
       request.renderState,
@@ -873,7 +880,7 @@ function renderBackupSuspenseBoundary(
   const segment = task.blockedSegment;
 
   pushStartCompletedSuspenseBoundary(segment.chunks);
-  renderNode(request, task, content, 0);
+  renderNode(request, task, content, -1);
   pushEndCompletedSuspenseBoundary(segment.chunks);
 
   popComponentStackInDEV(task);
@@ -903,7 +910,7 @@ function renderHostElement(
 
   // We use the non-destructive form because if something suspends, we still
   // need to pop back up and finish this subtree of HTML.
-  renderNode(request, task, children, 0);
+  renderNode(request, task, children, -1);
 
   // We expect that errors will fatal the whole task and that we don't need
   // the correct context. Therefore this is not in a finally.
@@ -970,13 +977,13 @@ function finishClassComponent(
         childContextTypes,
       );
       task.legacyContext = mergedContext;
-      renderNodeDestructive(request, task, null, nextChildren, 0);
+      renderNodeDestructive(request, task, null, nextChildren, -1);
       task.legacyContext = previousContext;
       return;
     }
   }
 
-  renderNodeDestructive(request, task, null, nextChildren, 0);
+  renderNodeDestructive(request, task, null, nextChildren, -1);
 }
 
 function renderClassComponent(
@@ -1170,20 +1177,20 @@ function finishFunctionComponent(
     // Modify the id context. Because we'll need to reset this if something
     // suspends or errors, we'll use the non-destructive render path.
     task.treeContext = pushTreeContext(prevTreeContext, totalChildren, index);
-    renderNode(request, task, children, 0);
+    renderNode(request, task, children, -1);
     // Like the other contexts, this does not need to be in a finally block
     // because renderNode takes care of unwinding the stack.
     task.treeContext = prevTreeContext;
   } else if (didEmitFormStateMarkers) {
     // If there were formState hooks, we must use the non-destructive path
     // because this component is not a pure indirection; we emitted markers
     // to the stream.
-    renderNode(request, task, children, 0);
+    renderNode(request, task, children, -1);
   } else {
     // We're now successfully past this task, and we haven't modified the
     // context stack. We don't have to pop back to the previous task every
     // again, so we can use the destructive recursive form.
-    renderNodeDestructive(request, task, null, children, 0);
+    renderNodeDestructive(request, task, null, children, -1);
   }
 }
 
@@ -1353,7 +1360,7 @@ function renderContextConsumer(
   const newValue = readContext(context);
   const newChildren = render(newValue);
 
-  renderNodeDestructive(request, task, null, newChildren, 0);
+  renderNodeDestructive(request, task, null, newChildren, -1);
 }
 
 function renderContextProvider(
@@ -1370,7 +1377,7 @@ function renderContextProvider(
     prevSnapshot = task.context;
   }
   task.context = pushProvider(context, value);
-  renderNodeDestructive(request, task, null, children, 0);
+  renderNodeDestructive(request, task, null, children, -1);
   task.context = popProvider(context);
   if (__DEV__) {
     if (prevSnapshot !== task.context) {
@@ -1413,7 +1420,7 @@ function renderOffscreen(request: Request, task: Task, props: Object): void {
   } else {
     // A visible Offscreen boundary is treated exactly like a fragment: a
     // pure indirection.
-    renderNodeDestructive(request, task, null, props.children, 0);
+    renderNodeDestructive(request, task, null, props.children, -1);
   }
 }
 
@@ -1460,7 +1467,7 @@ function renderElement(
     case REACT_STRICT_MODE_TYPE:
     case REACT_PROFILER_TYPE:
     case REACT_FRAGMENT_TYPE: {
-      renderNodeDestructive(request, task, null, props.children, 0);
+      renderNodeDestructive(request, task, null, props.children, -1);
       return;
     }
     case REACT_OFFSCREEN_TYPE: {
@@ -1470,13 +1477,13 @@ function renderElement(
     case REACT_SUSPENSE_LIST_TYPE: {
       pushBuiltInComponentStackInDEV(task, 'SuspenseList');
       // TODO: SuspenseList should control the boundaries.
-      renderNodeDestructive(request, task, null, props.children, 0);
+      renderNodeDestructive(request, task, null, props.children, -1);
       popComponentStackInDEV(task);
       return;
     }
     case REACT_SCOPE_TYPE: {
       if (enableScopeAPI) {
-        renderNodeDestructive(request, task, null, props.children, 0);
+        renderNodeDestructive(request, task, null, props.children, -1);
         return;
       }
       throw new Error('ReactDOMServer does not yet support scope components.');
@@ -1645,7 +1652,11 @@ function renderNodeDestructiveImpl(
         const ref = element.ref;
         const name = getComponentNameFromType(type);
         const prevKeyPath = task.keyPath;
-        task.keyPath = [task.keyPath, name, key == null ? childIndex : key];
+        task.keyPath = [
+          task.keyPath,
+          name,
+          key == null ? (childIndex === -1 ? 0 : childIndex) : key,
+        ];
         renderElement(request, task, prevThenableState, type, props, ref);
         task.keyPath = prevKeyPath;
         return;
@@ -1805,61 +1816,29 @@ function renderChildrenArray(
   children: Array<any>,
   childIndex: number,
 ) {
+  const prevKeyPath = task.keyPath;
+  if (childIndex !== -1) {
+    task.keyPath = [task.keyPath, '', childIndex];
+  }
   const prevTreeContext = task.treeContext;
   const totalChildren = children.length;
   for (let i = 0; i < totalChildren; i++) {
     const node = children[i];
     task.treeContext = pushTreeContext(prevTreeContext, totalChildren, i);
-
-    // 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];
-      renderChildrenArray(request, task, node, i);
-      task.keyPath = prevKeyPath;
-      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.
   task.treeContext = prevTreeContext;
+  task.keyPath = prevKeyPath;
 }
 
 function trackPostpone(
   request: Request,
   trackedPostpones: PostponedHoles,
   task: Task,
-  childIndex: number,
   segment: Segment,
 ): void {
   segment.status = POSTPONED;
@@ -1901,8 +1880,20 @@ function trackPostpone(
     );
   }
 
-  const segmentNode: ResumableNode = [RESUME_SEGMENT, childIndex, segment.id];
-  addToResumableParent(segmentNode, keyPath, trackedPostpones);
+  if (task.childIndex === -1) {
+    // Resume at the position before the first array
+    const resumableElement = [
+      RESUME_ELEMENT,
+      keyPath[1],
+      keyPath[2],
+      segment.id,
+    ];
+    addToResumableParent(resumableElement, keyPath[0], trackedPostpones);
+  } else {
+    // Resume at the slot within the array
+    const resumableNode = [RESUME_SLOT, task.childIndex, segment.id];
+    addToResumableParent(resumableNode, keyPath, trackedPostpones);
+  }
 }
 
 function injectPostponedHole(
@@ -2060,13 +2051,7 @@ function renderNode(
           task,
           postponeInstance.message,
         );
-        trackPostpone(
-          request,
-          trackedPostpones,
-          task,
-          childIndex,
-          postponedSegment,
-        );
+        trackPostpone(request, trackedPostpones, task, 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.
@@ -2414,13 +2399,7 @@ function retryTask(request: Request, task: Task): void {
         task.abortSet.delete(task);
         const postponeInstance: Postpone = (x: any);
         logPostpone(request, postponeInstance.message);
-        trackPostpone(
-          request,
-          trackedPostpones,
-          task,
-          task.childIndex,
-          segment,
-        );
+        trackPostpone(request, trackedPostpones, task, segment);
         finishedTask(request, task.blockedBoundary, segment);
         return;
       }