Encode Iterator separately from Iterable in Flight Reply

Author: sebmarkbage

packages/react-client/src/ReactFlightReplyClient.js

@@ -81,7 +81,10 @@ export type ReactServerValue =
   | null
   | void
   | bigint
+  | $AsyncIterable<ReactServerValue, ReactServerValue, void>
+  | $AsyncIterator<ReactServerValue, ReactServerValue, void>
   | Iterable<ReactServerValue>
+  | Iterator<ReactServerValue>
   | Array<ReactServerValue>
   | Map<ReactServerValue, ReactServerValue>
   | Set<ReactServerValue>
@@ -157,6 +160,10 @@ function serializeBlobID(id: number): string {
   return '$B' + id.toString(16);
 }
 
+function serializeIteratorID(id: number): string {
+  return '$i' + id.toString(16);
+}
+
 function escapeStringValue(value: string): string {
   if (value[0] === '$') {
     // We need to escape $ prefixed strings since we use those to encode
@@ -448,7 +455,21 @@ export function processReply(
 
       const iteratorFn = getIteratorFn(value);
       if (iteratorFn) {
-        return Array.from((value: any));
+        const iterator = iteratorFn.call(value);
+        if (iterator === value) {
+          // Iterator, not Iterable
+          const partJSON = JSON.stringify(
+            Array.from((iterator: any)),
+            resolveToJSON,
+          );
+          if (formData === null) {
+            formData = new FormData();
+          }
+          const iteratorId = nextPartId++;
+          formData.append(formFieldPrefix + iteratorId, partJSON);
+          return serializeIteratorID(iteratorId);
+        }
+        return Array.from((iterator: any));
       }
 
       // Verify that this is a simple plain object.

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js

@@ -97,6 +97,35 @@ describe('ReactFlightDOMReply', () => {
       items.push(item);
     }
     expect(items).toEqual(['A', 'B', 'C']);
+
+    // Multipass
+    const items2 = [];
+    // eslint-disable-next-line no-for-of-loops/no-for-of-loops
+    for (const item of iterable) {
+      items2.push(item);
+    }
+    expect(items2).toEqual(['A', 'B', 'C']);
+  });
+
+  it('can pass an iterator as a reply', async () => {
+    const iterator = (function* () {
+      yield 'A';
+      yield 'B';
+      yield 'C';
+    })();
+
+    const body = await ReactServerDOMClient.encodeReply(iterator);
+    const result = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+    );
+
+    // The iterator should be the same as itself.
+    expect(result[Symbol.iterator]()).toBe(result);
+
+    expect(Array.from(result)).toEqual(['A', 'B', 'C']);
+    // We've already consumed this iterator.
+    expect(Array.from(result)).toEqual([]);
   });
 
   it('can pass weird numbers as a reply', async () => {

packages/react-server/src/ReactFlightReplyServer.js

@@ -477,6 +477,12 @@ function parseModelString(
         });
         return data;
       }
+      case 'i': {
+        // Iterator
+        const id = parseInt(value.slice(2), 16);
+        const data = getOutlinedModel(response, id);
+        return data[Symbol.iterator]();
+      }
       case 'I': {
         // $Infinity
         return Infinity;