Add a wrapper fragment if an AsyncIterable needs to be keyed

This is consistent with other Iterables.
react · sebmarkbage · Apr 16, 2024 · Jun 15, 2023 · Jun 15, 2023 · Jun 18, 2023
commit 105fa03747585cdc399b9312c882500cb734164b
diff --git a/packages/react-client/src/__tests__/ReactFlight-test.js b/packages/react-client/src/__tests__/ReactFlight-test.js
@@ -2068,6 +2068,147 @@ describe('ReactFlight', () => {
     expect(ReactNoop).toMatchRenderedOutput(<div>Ba</div>);
   });
 
+  it('shares state when moving keyed Server Components that render fragments', async () => {
+    function StatefulClient({name, initial}) {
+      const [state] = React.useState(initial);
+      return <span>{state}</span>;
+    }
+    const Stateful = clientReference(StatefulClient);
+
+    function ServerComponent({item, initial}) {
+      return [
+        <Stateful key="a" initial={'a' + initial} />,
+        <Stateful key="b" initial={'b' + initial} />,
+      ];
+    }
+
+    const transport = ReactNoopFlightServer.render(
+      <div>
+        <ServerComponent key="A" initial={1} />
+        <ServerComponent key="B" initial={2} />
+      </div>,
+    );
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(
+      <div>
+        <span>a1</span>
+        <span>b1</span>
+        <span>a2</span>
+        <span>b2</span>
+      </div>,
+    );
+
+    // We swap the Server Components and the state of each child inside each fragment should move.
+    // Really the Fragment itself moves.
+    const transport2 = ReactNoopFlightServer.render(
+      <div>
+        <ServerComponent key="B" initial={4} />
+        <ServerComponent key="A" initial={3} />
+      </div>,
+    );
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport2));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(
+      <div>
+        <span>a2</span>
+        <span>b2</span>
+        <span>a1</span>
+        <span>b1</span>
+      </div>,
+    );
+  });
+
+  it('shares state when moving keyed Server Components that render async iterables', async () => {
+    function StatefulClient({name, initial}) {
+      const [state] = React.useState(initial);
+      return <span>{state}</span>;
+    }
+    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} />;
+        },
+      };
+    }
+
+    function ListClient({children}) {
+      // TODO: Unwrap AsyncIterables natively in React. For now we do it in this wrapper.
+      let resolvedChildren = [];
+      for (let fragment of children) {
+        // We should've wrapped each child in a keyed Fragment.
+        expect(fragment.type).toBe(React.Fragment);
+        let fragmentChildren = [];
+        const iterator = fragment.props.children[Symbol.asyncIterator]();
+        for (let entry; !(entry = React.use(iterator.next())).done; ) {
+          fragmentChildren.push(entry.value);
+        }
+        resolvedChildren.push(
+          <React.Fragment key={fragment.key}>
+            {fragmentChildren}
+          </React.Fragment>,
+        );
+      }
+      return <div>{resolvedChildren}</div>;
+    }
+
+    const List = clientReference(ListClient);
+
+    const transport = ReactNoopFlightServer.render(
+      <List>
+        <ServerComponent key="A" initial={1} />
+        <ServerComponent key="B" initial={2} />
+      </List>,
+    );
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(
+      <div>
+        <span>a1</span>
+        <span>b1</span>
+        <span>a2</span>
+        <span>b2</span>
+      </div>,
+    );
+
+    // We swap the Server Components and the state of each child inside each fragment should move.
+    // Really the Fragment itself moves.
+    const transport2 = ReactNoopFlightServer.render(
+      <List>
+        <ServerComponent key="B" initial={4} />
+        <ServerComponent key="A" initial={3} />
+      </List>,
+    );
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport2));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(
+      <div>
+        <span>a2</span>
+        <span>b2</span>
+        <span>a1</span>
+        <span>b1</span>
+      </div>,
+    );
+  });
+
   it('preserves debug info for server-to-server pass through', async () => {
     function ThirdPartyLazyComponent() {
       return <span>!</span>;

diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
@@ -885,7 +885,8 @@ function renderFunctionComponent<Props>(
 function renderFragment(
   request: Request,
   task: Task,
-  children: $ReadOnlyArray<ReactClientValue>,
+  children:
+    | $ReadOnlyArray<ReactClientValue>
 ): ReactJSONValue {
   if (__DEV__) {
     const debugInfo: ?ReactDebugInfo = (children: any)._debugInfo;
@@ -899,6 +900,7 @@ function renderFragment(
         // Forward any debug info we have the first time we see it.
         // We do this after init so that we have received all the debug info
         // from the server by the time we emit it.
+        // TODO: We might see this fragment twice if it ends up wrapped below.
         forwardDebugInfo(request, debugID, debugInfo);
       }
     }
@@ -940,6 +942,68 @@ function renderFragment(
   return children;
 }
 
+function renderAsyncFragment(
+  request: Request,
+  task: Task,
+  children: $AsyncIterable<ReactClientValue, ReactClientValue, void>,
+  getAsyncIterator: () => $AsyncIterator<any, any, any>
+): ReactJSONValue {
+  if (__DEV__) {
+    const debugInfo: ?ReactDebugInfo = (children: any)._debugInfo;
+    if (debugInfo) {
+      // If this came from Flight, forward any debug info into this new row.
+      if (debugID === null) {
+        // We don't have a chunk to assign debug info. We need to outline this
+        // component to assign it an ID.
+        return outlineTask(request, task);
+      } else {
+        // Forward any debug info we have the first time we see it.
+        // We do this after init so that we have received all the debug info
+        // from the server by the time we emit it.
+        // TODO: We might see this fragment twice if it ends up wrapped below.
+        forwardDebugInfo(request, debugID, debugInfo);
+      }
+    }
+  }
+  if (!enableServerComponentKeys) {
+    const asyncIterator = getAsyncIterator.call(children);
+    return serializeAsyncIterable(request, children, asyncIterator);
+  }
+  if (task.keyPath !== null) {
+    // We have a Server Component that specifies a key but we're now splitting
+    // the tree using a fragment.
+    const fragment = [
+      REACT_ELEMENT_TYPE,
+      REACT_FRAGMENT_TYPE,
+      task.keyPath,
+      {children},
+    ];
+    if (!task.implicitSlot) {
+      // If this was keyed inside a set. I.e. the outer Server Component was keyed
+      // then we need to handle reorders of the whole set. To do this we need to wrap
+      // this array in a keyed Fragment.
+      return fragment;
+    }
+    // If the outer Server Component was implicit but then an inner one had a key
+    // we don't actually need to be able to move the whole set around. It'll always be
+    // in an implicit slot. The key only exists to be able to reset the state of the
+    // children. We could achieve the same effect by passing on the keyPath to the next
+    // set of components inside the fragment. This would also allow a keyless fragment
+    // reconcile against a single child.
+    // Unfortunately because of JSON.stringify, we can't call the recursive loop for
+    // each child within this context because we can't return a set with already resolved
+    // values. E.g. a string would get double encoded. Returning would pop the context.
+    // So instead, we wrap it with an unkeyed fragment and inner keyed fragment.
+    return [fragment];
+  }
+  // Since we're yielding here, that implicitly resets the keyPath context on the
+  // way up. Which is what we want since we've consumed it. If this changes to
+  // be recursive serialization, we need to reset the keyPath and implicitSlot,
+  // before recursing here.
+  const asyncIterator = getAsyncIterator.call(children);
+  return serializeAsyncIterable(request, children, asyncIterator);
+}
+
 function renderClientElement(
   task: Task,
   type: any,
@@ -1888,8 +1952,8 @@ function renderModelDestructive(
       const getAsyncIterator: void | (() => $AsyncIterator<any, any, any>) =
         (value: any)[ASYNC_ITERATOR];
       if (typeof getAsyncIterator === 'function') {
-        const asyncIterator = getAsyncIterator.call(value);
-        return serializeAsyncIterable(request, (value: any), asyncIterator);
+        // We treat AsyncIterables as a Fragment and as such we might need to key them.
+        return renderAsyncFragment(request, task, (value: any), getAsyncIterator);
       }
     }