Serialize Map and Set - Server to Client

Author: sebmarkbage

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/__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/src/ReactFlightServer.js

@@ -139,6 +139,8 @@ export type ReactClientValue =
   | void
   | Iterable<ReactClientValue>
   | Array<ReactClientValue>
+  | Map<ReactClientValue, ReactClientValue>
+  | Set<ReactClientValue>
   | ReactClientObject
   | Promise<ReactClientValue>; // Thenable<ReactClientValue>
 
@@ -683,6 +685,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 +719,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 +738,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 +940,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) {