[Flight Reply] Dedupe Objects and Support Cyclic References (#28997)

sebmarkbage · web-flow · commit 38d9f156b898 · 2024-05-09T19:24:37.000-04:00
Uses the same technique as in #28996 to encode references to already emitted objects. This now means that Reply can support cyclic objects too for parity.
diff --git a/packages/react-client/src/ReactFlightReplyClient.js b/packages/react-client/src/ReactFlightReplyClient.js
@@ -176,6 +176,8 @@ function escapeStringValue(value: string): string {
   }
 }
 
+interface Reference {}
+
 export function processReply(
   root: ReactServerValue,
   formFieldPrefix: string,
@@ -186,6 +188,8 @@ export function processReply(
   let nextPartId = 1;
   let pendingParts = 0;
   let formData: null | FormData = null;
+  const writtenObjects: WeakMap<Reference, string> = new WeakMap();
+  let modelRoot: null | ReactServerValue = root;
 
   function serializeTypedArray(
     tag: string,
@@ -427,7 +431,7 @@ export function processReply(
             // We always outline this as a separate part even though we could inline it
             // because it ensures a more deterministic encoding.
             const lazyId = nextPartId++;
-            const partJSON = JSON.stringify(resolvedModel, resolveToJSON);
+            const partJSON = serializeModel(resolvedModel, lazyId);
             // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
             const data: FormData = formData;
             // eslint-disable-next-line react-internal/safe-string-coercion
@@ -447,7 +451,7 @@ export function processReply(
                 // While the first promise resolved, its value isn't necessarily what we'll
                 // resolve into because we might suspend again.
                 try {
-                  const partJSON = JSON.stringify(value, resolveToJSON);
+                  const partJSON = serializeModel(value, lazyId);
                   // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
                   const data: FormData = formData;
                   // eslint-disable-next-line react-internal/safe-string-coercion
@@ -488,7 +492,7 @@ export function processReply(
         thenable.then(
           partValue => {
             try {
-              const partJSON = JSON.stringify(partValue, resolveToJSON);
+              const partJSON = serializeModel(partValue, promiseId);
               // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
               const data: FormData = formData;
               // eslint-disable-next-line react-internal/safe-string-coercion
@@ -507,6 +511,28 @@ export function processReply(
         );
         return serializePromiseID(promiseId);
       }
+
+      const existingReference = writtenObjects.get(value);
+      if (existingReference !== undefined) {
+        if (modelRoot === value) {
+          // This is the ID we're currently emitting so we need to write it
+          // once but if we discover it again, we refer to it by id.
+          modelRoot = null;
+        } else {
+          // We've already emitted this as an outlined object, so we can
+          // just refer to that by its existing ID.
+          return existingReference;
+        }
+      } else if (key.indexOf(':') === -1) {
+        // TODO: If the property name contains a colon, we don't dedupe. Escape instead.
+        const parentReference = writtenObjects.get(parent);
+        if (parentReference !== undefined) {
+          // If the parent has a reference, we can refer to this object indirectly
+          // through the property name inside that parent.
+          writtenObjects.set(value, parentReference + ':' + key);
+        }
+      }
+
       if (isArray(value)) {
         // $FlowFixMe[incompatible-return]
         return value;
@@ -530,20 +556,20 @@ export function processReply(
         return serializeFormDataReference(refId);
       }
       if (value instanceof Map) {
-        const partJSON = JSON.stringify(Array.from(value), resolveToJSON);
+        const mapId = nextPartId++;
+        const partJSON = serializeModel(Array.from(value), mapId);
         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);
+        const setId = nextPartId++;
+        const partJSON = serializeModel(Array.from(value), setId);
         if (formData === null) {
           formData = new FormData();
         }
-        const setId = nextPartId++;
         formData.append(formFieldPrefix + setId, partJSON);
         return serializeSetID(setId);
       }
@@ -622,14 +648,14 @@ export function processReply(
         const iterator = iteratorFn.call(value);
         if (iterator === value) {
           // Iterator, not Iterable
-          const partJSON = JSON.stringify(
+          const iteratorId = nextPartId++;
+          const partJSON = serializeModel(
             Array.from((iterator: any)),
-            resolveToJSON,
+            iteratorId,
           );
           if (formData === null) {
             formData = new FormData();
           }
-          const iteratorId = nextPartId++;
           formData.append(formFieldPrefix + iteratorId, partJSON);
           return serializeIteratorID(iteratorId);
         }
@@ -784,8 +810,17 @@ export function processReply(
     );
   }
 
-  // $FlowFixMe[incompatible-type] it's not going to be undefined because we'll encode it.
-  const json: string = JSON.stringify(root, resolveToJSON);
+  function serializeModel(model: ReactServerValue, id: number): string {
+    if (typeof model === 'object' && model !== null) {
+      writtenObjects.set(model, serializeByValueID(id));
+    }
+    modelRoot = model;
+    // $FlowFixMe[incompatible-return] it's not going to be undefined because we'll encode it.
+    return JSON.stringify(model, resolveToJSON);
+  }
+
+  const json = serializeModel(root, 0);
+
   if (formData === null) {
     // If it's a simple data structure, we just use plain JSON.
     resolve(json);
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
@@ -537,4 +537,13 @@ describe('ReactFlightDOMReply', () => {
       'Values cannot be passed to next() of AsyncIterables passed to Client Components.',
     );
   });
+
+  it('can transport cyclic objects', async () => {
+    const cyclic = {obj: null};
+    cyclic.obj = cyclic;
+
+    const body = await ReactServerDOMClient.encodeReply({prop: cyclic});
+    const root = await ReactServerDOMServer.decodeReply(body, webpackServerMap);
+    expect(root.prop.obj).toBe(root.prop);
+  });
 });
diff --git a/packages/react-server/src/ReactFlightReplyServer.js b/packages/react-server/src/ReactFlightReplyServer.js
@@ -47,6 +47,7 @@ export type JSONValue =
 
 const PENDING = 'pending';
 const BLOCKED = 'blocked';
+const CYCLIC = 'cyclic';
 const RESOLVED_MODEL = 'resolved_model';
 const INITIALIZED = 'fulfilled';
 const ERRORED = 'rejected';
@@ -65,6 +66,13 @@ type BlockedChunk<T> = {
   _response: Response,
   then(resolve: (T) => mixed, reject?: (mixed) => mixed): void,
 };
+type CyclicChunk<T> = {
+  status: 'cyclic',
+  value: null | Array<(T) => mixed>,
+  reason: null | Array<(mixed) => mixed>,
+  _response: Response,
+  then(resolve: (T) => mixed, reject?: (mixed) => mixed): void,
+};
 type ResolvedModelChunk<T> = {
   status: 'resolved_model',
   value: string,
@@ -98,6 +106,7 @@ type ErroredChunk<T> = {
 type SomeChunk<T> =
   | PendingChunk<T>
   | BlockedChunk<T>
+  | CyclicChunk<T>
   | ResolvedModelChunk<T>
   | InitializedChunk<T>
   | ErroredChunk<T>;
@@ -132,6 +141,7 @@ Chunk.prototype.then = function <T>(
       break;
     case PENDING:
     case BLOCKED:
+    case CYCLIC:
       if (resolve) {
         if (chunk.value === null) {
           chunk.value = ([]: Array<(T) => mixed>);
@@ -187,6 +197,7 @@ function wakeChunkIfInitialized<T>(
       break;
     case PENDING:
     case BLOCKED:
+    case CYCLIC:
       chunk.value = resolveListeners;
       chunk.reason = rejectListeners;
       break;
@@ -334,6 +345,7 @@ function loadServerReference<T>(
       false,
       response,
       createModel,
+      [],
     ),
     createModelReject(parentChunk),
   );
@@ -348,8 +360,19 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
   const prevBlocked = initializingChunkBlockedModel;
   initializingChunk = chunk;
   initializingChunkBlockedModel = null;
+
+  const resolvedModel = chunk.value;
+
+  // We go to the CYCLIC state until we've fully resolved this.
+  // We do this before parsing in case we try to initialize the same chunk
+  // while parsing the model. Such as in a cyclic reference.
+  const cyclicChunk: CyclicChunk<T> = (chunk: any);
+  cyclicChunk.status = CYCLIC;
+  cyclicChunk.value = null;
+  cyclicChunk.reason = null;
+
   try {
-    const value: T = JSON.parse(chunk.value, chunk._response._fromJSON);
+    const value: T = JSON.parse(resolvedModel, chunk._response._fromJSON);
     if (
       initializingChunkBlockedModel !== null &&
       initializingChunkBlockedModel.deps > 0
@@ -362,9 +385,13 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
       blockedChunk.value = null;
       blockedChunk.reason = null;
     } else {
+      const resolveListeners = cyclicChunk.value;
       const initializedChunk: InitializedChunk<T> = (chunk: any);
       initializedChunk.status = INITIALIZED;
       initializedChunk.value = value;
+      if (resolveListeners !== null) {
+        wakeChunk(resolveListeners, value);
+      }
     }
   } catch (error) {
     const erroredChunk: ErroredChunk<T> = (chunk: any);
@@ -416,6 +443,7 @@ function createModelResolver<T>(
   cyclic: boolean,
   response: Response,
   map: (response: Response, model: any) => T,
+  path: Array<string>,
 ): (value: any) => void {
   let blocked;
   if (initializingChunkBlockedModel) {
@@ -430,6 +458,9 @@ function createModelResolver<T>(
     };
   }
   return value => {
+    for (let i = 1; i < path.length; i++) {
+      value = value[path[i]];
+    }
     parentObject[key] = map(response, value);
 
     // If this is the root object for a model reference, where `blocked.value`
@@ -460,11 +491,13 @@ function createModelReject<T>(chunk: SomeChunk<T>): (error: mixed) => void {
 
 function getOutlinedModel<T>(
   response: Response,
-  id: number,
+  reference: string,
   parentObject: Object,
   key: string,
   map: (response: Response, model: any) => T,
 ): T {
+  const path = reference.split(':');
+  const id = parseInt(path[0], 16);
   const chunk = getChunk(response, id);
   switch (chunk.status) {
     case RESOLVED_MODEL:
@@ -474,18 +507,24 @@ function getOutlinedModel<T>(
   // The status might have changed after initialization.
   switch (chunk.status) {
     case INITIALIZED:
-      return map(response, chunk.value);
+      let value = chunk.value;
+      for (let i = 1; i < path.length; i++) {
+        value = value[path[i]];
+      }
+      return map(response, value);
     case PENDING:
     case BLOCKED:
+    case CYCLIC:
       const parentChunk = initializingChunk;
       chunk.then(
         createModelResolver(
           parentChunk,
           parentObject,
           key,
-          false,
+          chunk.status === CYCLIC,
           response,
           map,
+          path,
         ),
         createModelReject(parentChunk),
       );
@@ -548,6 +587,7 @@ function parseTypedArray(
       false,
       response,
       createModel,
+      [],
     ),
     createModelReject(parentChunk),
   );
@@ -789,10 +829,10 @@ function parseModelString(
       }
       case 'F': {
         // Server Reference
-        const id = parseInt(value.slice(2), 16);
+        const ref = value.slice(2);
         // TODO: Just encode this in the reference inline instead of as a model.
         const metaData: {id: ServerReferenceId, bound: Thenable<Array<any>>} =
-          getOutlinedModel(response, id, obj, key, createModel);
+          getOutlinedModel(response, ref, obj, key, createModel);
         return loadServerReference(
           response,
           metaData.id,
@@ -808,13 +848,13 @@ function parseModelString(
       }
       case 'Q': {
         // Map
-        const id = parseInt(value.slice(2), 16);
-        return getOutlinedModel(response, id, obj, key, createMap);
+        const ref = value.slice(2);
+        return getOutlinedModel(response, ref, obj, key, createMap);
       }
       case 'W': {
         // Set
-        const id = parseInt(value.slice(2), 16);
-        return getOutlinedModel(response, id, obj, key, createSet);
+        const ref = value.slice(2);
+        return getOutlinedModel(response, ref, obj, key, createSet);
       }
       case 'K': {
         // FormData
@@ -835,8 +875,8 @@ function parseModelString(
       }
       case 'i': {
         // Iterator
-        const id = parseInt(value.slice(2), 16);
-        return getOutlinedModel(response, id, obj, key, extractIterator);
+        const ref = value.slice(2);
+        return getOutlinedModel(response, ref, obj, key, extractIterator);
       }
       case 'I': {
         // $Infinity
@@ -933,8 +973,8 @@ function parseModelString(
     }
 
     // We assume that anything else is a reference ID.
-    const id = parseInt(value.slice(1), 16);
-    return getOutlinedModel(response, id, obj, key, createModel);
+    const ref = value.slice(1);
+    return getOutlinedModel(response, ref, obj, key, createModel);
   }
   return value;
 }
