Always outline object keys in Map/Set

This won't always work but it'll work if we discover the Map/Set before the key.

The idea is that you'd pass the Map to something like a parent context and
then select out of it with a key passed to a child.

Author: sebmarkbage

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

@@ -374,12 +374,13 @@ describe('ReactFlight', () => {
   });
 
   it('can transport Map', async () => {
-    function ComponentClient({prop}) {
+    function ComponentClient({prop, selected}) {
       return `
         map: ${prop instanceof Map}
         size: ${prop.size}
         greet: ${prop.get('hi').greet}
         content: ${JSON.stringify(Array.from(prop))}
+        selected: ${prop.get(selected)}
       `;
     }
     const Component = clientReference(ComponentClient);
@@ -389,7 +390,7 @@ describe('ReactFlight', () => {
       ['hi', {greet: 'world'}],
       [objKey, 123],
     ]);
-    const model = <Component prop={map} />;
+    const model = <Component prop={map} selected={objKey} />;
 
     const transport = ReactNoopFlightServer.render(model);
 
@@ -402,23 +403,25 @@ describe('ReactFlight', () => {
         size: 2
         greet: world
         content: [["hi",{"greet":"world"}],[{"obj":"key"},123]]
+        selected: 123
       `);
   });
 
   it('can transport Set', async () => {
-    function ComponentClient({prop}) {
+    function ComponentClient({prop, selected}) {
       return `
         set: ${prop instanceof Set}
         size: ${prop.size}
         hi: ${prop.has('hi')}
         content: ${JSON.stringify(Array.from(prop))}
+        selected: ${prop.has(selected)}
       `;
     }
     const Component = clientReference(ComponentClient);
 
     const objKey = {obj: 'key'};
     const set = new Set(['hi', objKey]);
-    const model = <Component prop={set} />;
+    const model = <Component prop={set} selected={objKey} />;
 
     const transport = ReactNoopFlightServer.render(model);
 
@@ -431,6 +434,7 @@ describe('ReactFlight', () => {
         size: 2
         hi: true
         content: ["hi",{"obj":"key"}]
+        selected: true
       `);
   });
 

packages/react-server/src/ReactFlightServer.js

@@ -822,12 +822,36 @@ function serializeMap(
   request: Request,
   map: Map<ReactClientValue, ReactClientValue>,
 ): string {
-  const id = outlineModel(request, Array.from(map));
+  const entries = Array.from(map);
+  for (let i = 0; i < entries.length; i++) {
+    const key = entries[i][0];
+    if (typeof key === 'object' && key !== null) {
+      const writtenObjects = request.writtenObjects;
+      const existingId = writtenObjects.get(key);
+      if (existingId === undefined) {
+        // Mark all object keys as seen so that they're always outlined.
+        writtenObjects.set(key, -1);
+      }
+    }
+  }
+  const id = outlineModel(request, entries);
   return '$Q' + id.toString(16);
 }
 
 function serializeSet(request: Request, set: Set<ReactClientValue>): string {
-  const id = outlineModel(request, Array.from(set));
+  const entries = Array.from(set);
+  for (let i = 0; i < entries.length; i++) {
+    const key = entries[i];
+    if (typeof key === 'object' && key !== null) {
+      const writtenObjects = request.writtenObjects;
+      const existingId = writtenObjects.get(key);
+      if (existingId === undefined) {
+        // Mark all object keys as seen so that they're always outlined.
+        writtenObjects.set(key, -1);
+      }
+    }
+  }
+  const id = outlineModel(request, entries);
   return '$W' + id.toString(16);
 }