[Flight] Add Support for Map and Set

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) {