Serialize Map and Set - Client to Server

Author: sebmarkbage

packages/react-client/src/ReactFlightReplyClient.js

@@ -119,6 +119,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 +237,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-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);