[Flight] Support (Async) Generator ServerComponent (#28849)

Stacked on #28853 and #28854.

React supports rendering `Iterable` and will soon support
`AsyncIterable`. As long as it's multi-shot since during an update we
may have to rerender with new inputs an loop over the iterable again.
Therefore the `Iterator` and `AsyncIterator` types are not supported
directly as a child of React - and really it shouldn't pass between
Hooks or components neither for this reason. For parity, that's also the
case when used in Server Components.

However, there is a special case when the component rendered itself is a
generator function. While it returns as a child an `Iterator`, the React
Element itself can act as an `Iterable` because we can re-evaluate the
function to create a new generator whenever we need to.

It's also very convenient to use generator functions over constructing
an `AsyncIterable`. So this is a proposal to special case the
`Generator`/`AsyncGenerator` returned by a (Async) Generator Function.

In Flight this means that when we render a Server Component we can
serialize this value as an `Iterable`/`AsyncIterable` since that's
effectively what rendering it on the server reduces down to. That way if
Fiber can receive the result in any position.

For SuspenseList this would also need another special case because the
children of SuspenseList represent "rows".

`<SuspenseList><Component /></SuspenseList>` currently is a single "row"
even if the component renders multiple children or is an iterator. This
is currently different if Component is a Server Component because it'll
reduce down to an array/AsyncIterable and therefore be treated as one
row per its child. This is different from `<SuspenseList><Component
/><Component /></SuspenseList>` since that has a wrapper array and so
this is always two rows.

It probably makes sense to special case a single-element child in
`SuspenseList` to represent a component that generates rows. That way
you can use an `AsyncGeneratorFunction` to do this.

Author: sebmarkbage

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);
       }

packages/react-server/src/ReactFlightServer.js

@@ -865,20 +865,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;