[Flight] Allow values to be encoded by "reference" to a value rather than the value itself

packages/react-client/src/ReactFlightClient.js

@@ -318,6 +318,17 @@ function createLazyBlock<Props, Data>(
   return lazyType;
 }
 
+function createLazyChunkWrapper<T>(
+  chunk: SomeChunk<T>,
+): LazyComponent<T, SomeChunk<T>> {
+  const lazyType: LazyComponent<T, SomeChunk<T>> = {
+    $$typeof: REACT_LAZY_TYPE,
+    _payload: chunk,
+    _init: readChunk,
+  };
+  return lazyType;
+}
+
 function getChunk(response: Response, id: number): SomeChunk<any> {
   const chunks = response._chunks;
   let chunk = chunks.get(id);
@@ -333,26 +344,36 @@ export function parseModelString(
   parentObject: Object,
   value: string,
 ): any {
-  if (value[0] === '$') {
-    if (value === '$') {
-      return REACT_ELEMENT_TYPE;
-    } else if (value[1] === '$' || value[1] === '@') {
-      // This was an escaped string value.
-      return value.substring(1);
-    } else {
-      const id = parseInt(value.substring(1), 16);
-      const chunk = getChunk(response, id);
-      if (parentObject[0] === REACT_BLOCK_TYPE) {
-        // Block types know how to deal with lazy values.
-        return chunk;
+  switch (value[0]) {
+    case '$': {
+      if (value === '$') {
+        return REACT_ELEMENT_TYPE;
+      } else if (value[1] === '$' || value[1] === '@') {
+        // This was an escaped string value.
+        return value.substring(1);
+      } else {
+        const id = parseInt(value.substring(1), 16);
+        const chunk = getChunk(response, id);
+        if (parentObject[0] === REACT_BLOCK_TYPE) {
+          // Block types know how to deal with lazy values.
+          return chunk;
+        }
+        // For anything else we must Suspend this block if
+        // we don't yet have the value.
+        return readChunk(chunk);
+      }
+    }
+    case '@': {
+      if (value === '@') {
+        return REACT_BLOCK_TYPE;
+      } else {
+        const id = parseInt(value.substring(1), 16);
+        const chunk = getChunk(response, id);
+        // We create a React.lazy wrapper around any lazy values.
+        // When passed into React, we'll know how to suspend on this.
+        return createLazyChunkWrapper(chunk);
       }
-      // For anything else we must Suspend this block if
-      // we don't yet have the value.
-      return readChunk(chunk);
     }
-  }
-  if (value === '@') {
-    return REACT_BLOCK_TYPE;
   }
   return value;
 }

packages/react-server/src/ReactFlightServer.js

@@ -195,10 +195,14 @@ function createSegment(request: Request, query: () => ReactModel): Segment {
   return segment;
 }
 
-function serializeIDRef(id: number): string {
+function serializeByValueID(id: number): string {
   return '$' + id.toString(16);
 }
 
+function serializeByRefID(id: number): string {
+  return '@' + id.toString(16);
+}
+
 function escapeStringValue(value: string): string {
   if (value[0] === '$' || value[0] === '@') {
     // We need to escape $ or @ prefixed strings since we use those to encode
@@ -419,13 +423,13 @@ export function resolveModelToJSON(
             const newSegment = createSegment(request, load);
             const ping = newSegment.ping;
             x.then(ping, ping);
-            return serializeIDRef(newSegment.id);
+            return serializeByValueID(newSegment.id);
           } else {
             // This load failed, encode the error as a separate row and reference that.
             request.pendingChunks++;
             const errorId = request.nextChunkId++;
             emitErrorChunk(request, errorId, x);
-            return serializeIDRef(errorId);
+            return serializeByValueID(errorId);
           }
         }
       }
@@ -456,7 +460,7 @@ export function resolveModelToJSON(
         const newSegment = createSegment(request, () => value);
         const ping = newSegment.ping;
         x.then(ping, ping);
-        return serializeIDRef(newSegment.id);
+        return serializeByValueID(newSegment.id);
       } else {
         // Something errored. Don't bother encoding anything up to here.
         throw x;
@@ -479,12 +483,20 @@ export function resolveModelToJSON(
         request.pendingChunks++;
         const moduleId = request.nextChunkId++;
         emitModuleChunk(request, moduleId, moduleMetaData);
-        return serializeIDRef(moduleId);
+        if (parent[0] === REACT_ELEMENT_TYPE && key === '1') {
+          // If we're encoding the "type" of an element, we can refer
+          // to that by a lazy reference instead of directly since React
+          // knows how to deal with lazy values. This lets us suspend
+          // on this component rather than its parent until the code has
+          // loaded.
+          return serializeByRefID(moduleId);
+        }
+        return serializeByValueID(moduleId);
       } catch (x) {
         request.pendingChunks++;
         const errorId = request.nextChunkId++;
         emitErrorChunk(request, errorId, x);
-        return serializeIDRef(errorId);
+        return serializeByValueID(errorId);
       }
     }