[Flight] Add Support for Map and Set (#26933)

We already support these in the sense that they're Iterable so they just
get serialized as arrays. However, these are part of the Structured
Clone algorithm [and should be
supported](https://github.com/facebook/react/issues/25687).

The encoding is simply the same form as the Iterable, which is
conveniently the same as the constructor argument. The difference is
that now there's a separate reference to it.

It's a bit awkward because for multiple reference to the same value,
it'd be a new Map/Set instance for each reference. So to encode sharing,
it needs one level of indirection with its own ID. That's not really a
big deal for other types since they're inline anyway - but since this
needs to be outlined it creates possibly two ids where there only needs
to be one or zero.

One variant would be to encode this in the row type. Another variant
would be something like what we do for React Elements where they're
arrays but tagged with a symbol. For simplicity I stick with the simple
outlining for now.

packages/react-client/src/ReactFlightClient.js

@@ -535,6 +535,24 @@ function createServerReferenceProxy<A: Iterable<any>, T>(
   return proxy;
 }
 
+function getOutlinedModel(response: Response, id: number): any {
+  const chunk = getChunk(response, id);
+  switch (chunk.status) {
+    case RESOLVED_MODEL:
+      initializeModelChunk(chunk);
+      break;
+  }
+  // The status might have changed after initialization.
+  switch (chunk.status) {
+    case INITIALIZED: {
+      return chunk.value;
+    }
+    // We always encode it first in the stream so it won't be pending.
+    default:
+      throw chunk.reason;
+  }
+}
+
 function parseModelString(
   response: Response,
   parentObject: Object,
@@ -576,22 +594,20 @@ function parseModelString(
       case 'F': {
         // Server Reference
         const id = parseInt(value.slice(2), 16);
-        const chunk = getChunk(response, id);
-        switch (chunk.status) {
-          case RESOLVED_MODEL:
-            initializeModelChunk(chunk);
-            break;
-        }
-        // The status might have changed after initialization.
-        switch (chunk.status) {
-          case INITIALIZED: {
-            const metadata = chunk.value;
-            return createServerReferenceProxy(response, metadata);
-          }
-          // We always encode it first in the stream so it won't be pending.
-          default:
-            throw chunk.reason;
-        }
+        const metadata = getOutlinedModel(response, id);
+        return createServerReferenceProxy(response, metadata);
+      }
+      case 'Q': {
+        // Map
+        const id = parseInt(value.slice(2), 16);
+        const data = getOutlinedModel(response, id);
+        return new Map(data);
+      }
+      case 'W': {
+        // Set
+        const id = parseInt(value.slice(2), 16);
+        const data = getOutlinedModel(response, id);
+        return new Set(data);
       }
       case 'I': {
         // $Infinity

packages/react-client/src/ReactFlightReplyClient.js

@@ -59,8 +59,12 @@ export type ReactServerValue =
   | symbol
   | null
   | void
+  | bigint
   | Iterable<ReactServerValue>
   | Array<ReactServerValue>
+  | Map<ReactServerValue, ReactServerValue>
+  | Set<ReactServerValue>
+  | Date
   | ReactServerObject
   | Promise<ReactServerValue>; // Thenable<ReactServerValue>
 
@@ -119,6 +123,14 @@ function serializeBigInt(n: bigint): string {
   return '$n' + n.toString(10);
 }
 
+function serializeMapID(id: number): string {
+  return '$Q' + id.toString(16);
+}
+
+function serializeSetID(id: number): string {
+  return '$W' + id.toString(16);
+}
+
 function escapeStringValue(value: string): string {
   if (value[0] === '$') {
     // We need to escape $ prefixed strings since we use those to encode
@@ -229,6 +241,24 @@ export function processReply(
         });
         return serializeFormDataReference(refId);
       }
+      if (value instanceof Map) {
+        const partJSON = JSON.stringify(Array.from(value), resolveToJSON);
+        if (formData === null) {
+          formData = new FormData();
+        }
+        const mapId = nextPartId++;
+        formData.append(formFieldPrefix + mapId, partJSON);
+        return serializeMapID(mapId);
+      }
+      if (value instanceof Set) {
+        const partJSON = JSON.stringify(Array.from(value), resolveToJSON);
+        if (formData === null) {
+          formData = new FormData();
+        }
+        const setId = nextPartId++;
+        formData.append(formFieldPrefix + setId, partJSON);
+        return serializeSetID(setId);
+      }
       if (!isArray(value)) {
         const iteratorFn = getIteratorFn(value);
         if (iteratorFn) {

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -323,6 +323,67 @@ describe('ReactFlight', () => {
     expect(ReactNoop).toMatchRenderedOutput('prop: 2009-02-13T23:31:30.123Z');
   });
 
+  it('can transport Map', async () => {
+    function ComponentClient({prop}) {
+      return `
+        map: ${prop instanceof Map}
+        size: ${prop.size}
+        greet: ${prop.get('hi').greet}
+        content: ${JSON.stringify(Array.from(prop))}
+      `;
+    }
+    const Component = clientReference(ComponentClient);
+
+    const objKey = {obj: 'key'};
+    const map = new Map([
+      ['hi', {greet: 'world'}],
+      [objKey, 123],
+    ]);
+    const model = <Component prop={map} />;
+
+    const transport = ReactNoopFlightServer.render(model);
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(`
+        map: true
+        size: 2
+        greet: world
+        content: [["hi",{"greet":"world"}],[{"obj":"key"},123]]
+      `);
+  });
+
+  it('can transport Set', async () => {
+    function ComponentClient({prop}) {
+      return `
+        set: ${prop instanceof Set}
+        size: ${prop.size}
+        hi: ${prop.has('hi')}
+        content: ${JSON.stringify(Array.from(prop))}
+      `;
+    }
+    const Component = clientReference(ComponentClient);
+
+    const objKey = {obj: 'key'};
+    const set = new Set(['hi', objKey]);
+    const model = <Component prop={set} />;
+
+    const transport = ReactNoopFlightServer.render(model);
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(`
+        set: true
+        size: 2
+        hi: true
+        content: ["hi",{"obj":"key"}]
+      `);
+  });
+
   it('can render a lazy component as a shared component on the server', async () => {
     function SharedComponent({text}) {
       return (

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

@@ -197,4 +197,32 @@ describe('ReactFlightDOMReply', () => {
     expect(d).toEqual(d2);
     expect(d % 1000).toEqual(123); // double-check the milliseconds made it through
   });
+
+  it('can pass a Map as a reply', async () => {
+    const objKey = {obj: 'key'};
+    const m = new Map([
+      ['hi', {greet: 'world'}],
+      [objKey, 123],
+    ]);
+    const body = await ReactServerDOMClient.encodeReply(m);
+    const m2 = await ReactServerDOMServer.decodeReply(body, webpackServerMap);
+
+    expect(m2 instanceof Map).toBe(true);
+    expect(m2.size).toBe(2);
+    expect(m2.get('hi').greet).toBe('world');
+    expect(m2).toEqual(m);
+  });
+
+  it('can pass a Set as a reply', async () => {
+    const objKey = {obj: 'key'};
+    const s = new Set(['hi', objKey]);
+
+    const body = await ReactServerDOMClient.encodeReply(s);
+    const s2 = await ReactServerDOMServer.decodeReply(body, webpackServerMap);
+
+    expect(s2 instanceof Set).toBe(true);
+    expect(s2.size).toBe(2);
+    expect(s2.has('hi')).toBe(true);
+    expect(s2).toEqual(s);
+  });
 });

packages/react-server/src/ReactFlightReplyServer.js

@@ -364,6 +364,18 @@ function createModelReject<T>(chunk: SomeChunk<T>): (error: mixed) => void {
   return (error: mixed) => triggerErrorOnChunk(chunk, error);
 }
 
+function getOutlinedModel(response: Response, id: number): any {
+  const chunk = getChunk(response, id);
+  if (chunk.status === RESOLVED_MODEL) {
+    initializeModelChunk(chunk);
+  }
+  if (chunk.status !== INITIALIZED) {
+    // We know that this is emitted earlier so otherwise it's an error.
+    throw chunk.reason;
+  }
+  return chunk.value;
+}
+
 function parseModelString(
   response: Response,
   parentObject: Object,
@@ -389,17 +401,9 @@ function parseModelString(
       case 'F': {
         // Server Reference
         const id = parseInt(value.slice(2), 16);
-        const chunk = getChunk(response, id);
-        if (chunk.status === RESOLVED_MODEL) {
-          initializeModelChunk(chunk);
-        }
-        if (chunk.status !== INITIALIZED) {
-          // We know that this is emitted earlier so otherwise it's an error.
-          throw chunk.reason;
-        }
         // TODO: Just encode this in the reference inline instead of as a model.
         const metaData: {id: ServerReferenceId, bound: Thenable<Array<any>>} =
-          chunk.value;
+          getOutlinedModel(response, id);
         return loadServerReference(
           response,
           metaData.id,
@@ -409,6 +413,18 @@ function parseModelString(
           key,
         );
       }
+      case 'Q': {
+        // Map
+        const id = parseInt(value.slice(2), 16);
+        const data = getOutlinedModel(response, id);
+        return new Map(data);
+      }
+      case 'W': {
+        // Set
+        const id = parseInt(value.slice(2), 16);
+        const data = getOutlinedModel(response, id);
+        return new Set(data);
+      }
       case 'K': {
         // FormData
         const stringId = value.slice(2);

packages/react-server/src/ReactFlightServer.js

@@ -137,8 +137,12 @@ export type ReactClientValue =
   | symbol
   | null
   | void
+  | bigint
   | Iterable<ReactClientValue>
   | Array<ReactClientValue>
+  | Map<ReactClientValue, ReactClientValue>
+  | Set<ReactClientValue>
+  | Date
   | ReactClientObject
   | Promise<ReactClientValue>; // Thenable<ReactClientValue>
 
@@ -683,6 +687,15 @@ function serializeClientReference(
   }
 }
 
+function outlineModel(request: Request, value: any): number {
+  request.pendingChunks++;
+  const outlinedId = request.nextChunkId++;
+  // We assume that this object doesn't suspend, but a child might.
+  const processedChunk = processModelChunk(request, outlinedId, value);
+  request.completedRegularChunks.push(processedChunk);
+  return outlinedId;
+}
+
 function serializeServerReference(
   request: Request,
   parent:
@@ -708,15 +721,7 @@ function serializeServerReference(
     id: getServerReferenceId(request.bundlerConfig, serverReference),
     bound: bound ? Promise.resolve(bound) : null,
   };
-  request.pendingChunks++;
-  const metadataId = request.nextChunkId++;
-  // We assume that this object doesn't suspend.
-  const processedChunk = processModelChunk(
-    request,
-    metadataId,
-    serverReferenceMetadata,
-  );
-  request.completedRegularChunks.push(processedChunk);
+  const metadataId = outlineModel(request, serverReferenceMetadata);
   writtenServerReferences.set(serverReference, metadataId);
   return serializeServerReferenceID(metadataId);
 }
@@ -735,6 +740,19 @@ function serializeLargeTextString(request: Request, text: string): string {
   return serializeByValueID(textId);
 }
 
+function serializeMap(
+  request: Request,
+  map: Map<ReactClientValue, ReactClientValue>,
+): string {
+  const id = outlineModel(request, Array.from(map));
+  return '$Q' + id.toString(16);
+}
+
+function serializeSet(request: Request, set: Set<ReactClientValue>): string {
+  const id = outlineModel(request, Array.from(set));
+  return '$W' + id.toString(16);
+}
+
 function escapeStringValue(value: string): string {
   if (value[0] === '$') {
     // We need to escape $ prefixed strings since we use those to encode
@@ -924,6 +942,12 @@ function resolveModelToJSON(
       }
       return (undefined: any);
     }
+    if (value instanceof Map) {
+      return serializeMap(request, value);
+    }
+    if (value instanceof Set) {
+      return serializeSet(request, value);
+    }
     if (!isArray(value)) {
       const iteratorFn = getIteratorFn(value);
       if (iteratorFn) {