[Flight / Flight Reply] Encode Iterator separately from Iterable (#28854)

For [`AsyncIterable`](https://github.com/facebook/react/pull/28847) we
encode `AsyncIterator` as a separate tag.

Previously we encoded `Iterator` as just an Array. This adds a special
encoding for this. Technically this is a breaking change.

This is kind of an edge case that you'd care about the difference but it
becomes more important to treat these correctly for the warnings here
#28853.

packages/react-client/src/ReactFlightClient.js

@@ -825,6 +825,11 @@ function createFormData(
   return formData;
 }
 
+function extractIterator(response: Response, model: Array<any>): Iterator<any> {
+  // $FlowFixMe[incompatible-use]: This uses raw Symbols because we're extracting from a native array.
+  return model[Symbol.iterator]();
+}
+
 function createModel(response: Response, model: any): any {
   return model;
 }
@@ -919,6 +924,17 @@ function parseModelString(
           createFormData,
         );
       }
+      case 'i': {
+        // Iterator
+        const id = parseInt(value.slice(2), 16);
+        return getOutlinedModel(
+          response,
+          id,
+          parentObject,
+          key,
+          extractIterator,
+        );
+      }
       case 'I': {
         // $Infinity
         return Infinity;

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-client/src/__tests__/ReactFlight-test.js

@@ -277,6 +277,24 @@ describe('ReactFlight', () => {
     expect(ReactNoop).toMatchRenderedOutput(<span>ABC</span>);
   });
 
+  it('can render an iterator as a single shot iterator', async () => {
+    const iterator = (function* () {
+      yield 'A';
+      yield 'B';
+      yield 'C';
+    })();
+
+    const transport = ReactNoopFlightServer.render(iterator);
+    const result = await ReactNoopFlightClient.read(transport);
+
+    // 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 render undefined', async () => {
     function Undefined() {
       return undefined;

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;

packages/react-server/src/ReactFlightServer.js

@@ -241,7 +241,9 @@ export type ReactClientValue =
   | bigint
   | ReadableStream
   | $AsyncIterable<ReactClientValue, ReactClientValue, void>
+  | $AsyncIterator<ReactClientValue, ReactClientValue, void>
   | Iterable<ReactClientValue>
+  | Iterator<ReactClientValue>
   | Array<ReactClientValue>
   | Map<ReactClientValue, ReactClientValue>
   | Set<ReactClientValue>
@@ -1457,6 +1459,14 @@ function serializeSet(request: Request, set: Set<ReactClientValue>): string {
   return '$W' + id.toString(16);
 }
 
+function serializeIterator(
+  request: Request,
+  iterator: Iterator<ReactClientValue>,
+): string {
+  const id = outlineModel(request, Array.from(iterator));
+  return '$i' + id.toString(16);
+}
+
 function serializeTypedArray(
   request: Request,
   tag: string,
@@ -1910,7 +1920,13 @@ function renderModelDestructive(
 
     const iteratorFn = getIteratorFn(value);
     if (iteratorFn) {
-      return renderFragment(request, task, Array.from((value: any)));
+      // TODO: Should we serialize the return value as well like we do for AsyncIterables?
+      const iterator = iteratorFn.call(value);
+      if (iterator === value) {
+        // Iterator, not Iterable
+        return serializeIterator(request, (iterator: any));
+      }
+      return renderFragment(request, task, Array.from((iterator: any)));
     }
 
     if (enableFlightReadableStream) {