Turn (Async)Generator into an (Async)Iterable if it's an (Async) Generator ServerComponent

sebmarkbage · sebmarkbage · commit b580fa3a5ce6 · 2024-04-17T13:03:14.000-04:00
diff --git a/packages/react-client/src/__tests__/ReactFlight-test.js b/packages/react-client/src/__tests__/ReactFlight-test.js
@@ -295,6 +295,33 @@ describe('ReactFlight', () => {
     expect(Array.from(result)).toEqual([]);
   });
 
+  it('can render a Generator Server Component as a fragment', async () => {
+    function ItemListClient(props) {
+      return <span>{props.children}</span>;
+    }
+    const ItemList = clientReference(ItemListClient);
+
+    function* Items() {
+      yield 'A';
+      yield 'B';
+      yield 'C';
+    }
+
+    const model = (
+      <ItemList>
+        <Items />
+      </ItemList>
+    );
+
+    const transport = ReactNoopFlightServer.render(model);
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(<span>ABC</span>);
+  });
+
   it('can render undefined', async () => {
     function Undefined() {
       return undefined;
@@ -2151,16 +2178,9 @@ describe('ReactFlight', () => {
     }
     const Stateful = clientReference(StatefulClient);
 
-    function ServerComponent({item, initial}) {
-      // While the ServerComponent itself could be an async generator, single-shot iterables
-      // are not supported as React children since React might need to re-map them based on
-      // state updates. So we create an AsyncIterable instead.
-      return {
-        async *[Symbol.asyncIterator]() {
-          yield <Stateful key="a" initial={'a' + initial} />;
-          yield <Stateful key="b" initial={'b' + initial} />;
-        },
-      };
+    async function* ServerComponent({item, initial}) {
+      yield <Stateful key="a" initial={'a' + initial} />;
+      yield <Stateful key="b" initial={'b' + initial} />;
     }
 
     function ListClient({children}) {
@@ -2172,6 +2192,11 @@ describe('ReactFlight', () => {
         expect(fragment.type).toBe(React.Fragment);
         const fragmentChildren = [];
         const iterator = fragment.props.children[Symbol.asyncIterator]();
+        if (iterator === fragment.props.children) {
+          console.error(
+            'AyncIterators are not valid children of React. It must be a multi-shot AsyncIterable.',
+          );
+        }
         for (let entry; !(entry = React.use(iterator.next())).done; ) {
           fragmentChildren.push(entry.value);
         }
@@ -2316,23 +2341,21 @@ describe('ReactFlight', () => {
     let resolve;
     const iteratorPromise = new Promise(r => (resolve = r));
 
-    function ThirdPartyAsyncIterableComponent({item, initial}) {
-      // While the ServerComponent itself could be an async generator, single-shot iterables
-      // are not supported as React children since React might need to re-map them based on
-      // state updates. So we create an AsyncIterable instead.
-      return {
-        async *[Symbol.asyncIterator]() {
-          yield <span>Who</span>;
-          yield <span>dis?</span>;
-          resolve();
-        },
-      };
+    async function* ThirdPartyAsyncIterableComponent({item, initial}) {
+      yield <span>Who</span>;
+      yield <span>dis?</span>;
+      resolve();
     }
 
     function ListClient({children: fragment}) {
       // TODO: Unwrap AsyncIterables natively in React. For now we do it in this wrapper.
       const resolvedChildren = [];
       const iterator = fragment.props.children[Symbol.asyncIterator]();
+      if (iterator === fragment.props.children) {
+        console.error(
+          'AyncIterators are not valid children of React. It must be a multi-shot AsyncIterable.',
+        );
+      }
       for (let entry; !(entry = React.use(iterator.next())).done; ) {
         resolvedChildren.push(entry.value);
       }
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
@@ -866,20 +866,95 @@ function renderFunctionComponent<Props>(
   } else {
     result = Component(props, secondArg);
   }
-  if (
-    typeof result === 'object' &&
-    result !== null &&
-    typeof result.then === 'function'
-  ) {
-    // When the return value is in children position we can resolve it immediately,
-    // to its value without a wrapper if it's synchronously available.
-    const thenable: Thenable<any> = result;
-    if (thenable.status === 'fulfilled') {
-      return thenable.value;
-    }
-    // TODO: Once we accept Promises as children on the client, we can just return
-    // the thenable here.
-    result = createLazyWrapperAroundWakeable(result);
+  if (typeof result === 'object' && result !== null) {
+    if (typeof result.then === 'function') {
+      // When the return value is in children position we can resolve it immediately,
+      // to its value without a wrapper if it's synchronously available.
+      const thenable: Thenable<any> = result;
+      if (thenable.status === 'fulfilled') {
+        return thenable.value;
+      }
+      // TODO: Once we accept Promises as children on the client, we can just return
+      // the thenable here.
+      result = createLazyWrapperAroundWakeable(result);
+    }
+
+    // Normally we'd serialize an Iterator/AsyncIterator as a single-shot which is not compatible
+    // to be rendered as a React Child. However, because we have the function to recreate
+    // an iterable from rendering the element again, we can effectively treat it as multi-
+    // shot. Therefore we treat this as an Iterable/AsyncIterable, whether it was one or not, by
+    // adding a wrapper so that this component effectively renders down to an AsyncIterable.
+    const iteratorFn = getIteratorFn(result);
+    if (iteratorFn) {
+      const iterableChild = result;
+      result = {
+        [Symbol.iterator]: function () {
+          const iterator = iteratorFn.call(iterableChild);
+          if (__DEV__) {
+            // If this was an Iterator but not a GeneratorFunction we warn because
+            // it might have been a mistake. Technically you can make this mistake with
+            // GeneratorFunctions and even single-shot Iterables too but it's extra
+            // tempting to try to return the value from a generator.
+            if (iterator === iterableChild) {
+              const isGeneratorComponent =
+                // $FlowIgnore[method-unbinding]
+                Object.prototype.toString.call(Component) ===
+                  '[object GeneratorFunction]' &&
+                // $FlowIgnore[method-unbinding]
+                Object.prototype.toString.call(iterableChild) ===
+                  '[object Generator]';
+              if (!isGeneratorComponent) {
+                console.error(
+                  'Returning an Iterator from a Server Component is not supported ' +
+                    'since it cannot be looped over more than once. ',
+                );
+              }
+            }
+          }
+          return (iterator: any);
+        },
+      };
+      if (__DEV__) {
+        (result: any)._debugInfo = iterableChild._debugInfo;
+      }
+    } else if (
+      enableFlightReadableStream &&
+      typeof (result: any)[ASYNC_ITERATOR] === 'function' &&
+      (typeof ReadableStream !== 'function' ||
+        !(result instanceof ReadableStream))
+    ) {
+      const iterableChild = result;
+      result = {
+        [ASYNC_ITERATOR]: function () {
+          const iterator = (iterableChild: any)[ASYNC_ITERATOR]();
+          if (__DEV__) {
+            // If this was an AsyncIterator but not an AsyncGeneratorFunction we warn because
+            // it might have been a mistake. Technically you can make this mistake with
+            // AsyncGeneratorFunctions and even single-shot AsyncIterables too but it's extra
+            // tempting to try to return the value from a generator.
+            if (iterator === iterableChild) {
+              const isGeneratorComponent =
+                // $FlowIgnore[method-unbinding]
+                Object.prototype.toString.call(Component) ===
+                  '[object AsyncGeneratorFunction]' &&
+                // $FlowIgnore[method-unbinding]
+                Object.prototype.toString.call(iterableChild) ===
+                  '[object AsyncGenerator]';
+              if (!isGeneratorComponent) {
+                console.error(
+                  'Returning an AsyncIterator from a Server Component is not supported ' +
+                    'since it cannot be looped over more than once. ',
+                );
+              }
+            }
+          }
+          return iterator;
+        },
+      };
+      if (__DEV__) {
+        (result: any)._debugInfo = iterableChild._debugInfo;
+      }
+    }
   }
   // Track this element's key on the Server Component on the keyPath context..
   const prevKeyPath = task.keyPath;
