From d29a0dcfab73b1fc33b9f2821ce3e7420ed829d3 Mon Sep 17 00:00:00 2001
From: Hendrik Liebau <mail@hendrik-liebau.de>
Date: Sat, 23 Mar 2024 17:29:39 +0100
Subject: [PATCH] Choose lazy or value ID based on whether model chunk was
 emitted

fixes #28595
---
 .../react-server/src/ReactFlightServer.js     | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index d253db44c4a15..c1eab782a4ab7 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -281,6 +281,7 @@ export type Request = {
   writtenClientReferences: Map<ClientReferenceKey, number>,
   writtenServerReferences: Map<ServerReference<any>, number>,
   writtenObjects: WeakMap<Reference, number>, // -1 means "seen" but not outlined.
+  emittedModelChunkIds: Set<number>,
   identifierPrefix: string,
   identifierCount: number,
   taintCleanupQueue: Array<string | bigint>,
@@ -379,6 +380,7 @@ export function createRequest(
     writtenClientReferences: new Map(),
     writtenServerReferences: new Map(),
     writtenObjects: new WeakMap(),
+    emittedModelChunkIds: new Set(),
     identifierPrefix: identifierPrefix || '',
     identifierCount: 1,
     taintCleanupQueue: cleanupQueue,
@@ -1334,16 +1336,14 @@ function renderModelDestructive(
             // but that is able to reuse the same task if we're already in one but then that
             // will be a lazy future value rather than guaranteed to exist but maybe that's good.
             const newId = outlineModel(request, (value: any));
-            return serializeByValueID(newId);
+            return request.emittedModelChunkIds.has(newId)
+              ? serializeByValueID(newId)
+              : serializeLazyID(newId);
           } else {
-            // We've already emitted this as an outlined object, so we can refer to that by its
-            // existing ID. TODO: We should use a lazy reference since, unlike plain objects,
-            // elements might suspend so it might not have emitted yet even if we have the ID for
-            // it. However, this creates an extra wrapper when it's not needed. We should really
-            // detect whether this already was emitted and synchronously available. In that
-            // case we can refer to it synchronously and only make it lazy otherwise.
-            // We currently don't have a data structure that lets us see that though.
-            return serializeByValueID(existingId);
+            // We've already outlined this model, so we can refer to that by its existing ID.
+            return request.emittedModelChunkIds.has(existingId)
+              ? serializeByValueID(existingId)
+              : serializeLazyID(existingId);
           }
         } else {
           // This is the first time we've seen this object. We may never see it again
@@ -1884,6 +1884,7 @@ function emitSymbolChunk(request: Request, id: number, name: string): void {
 function emitModelChunk(request: Request, id: number, json: string): void {
   const row = id.toString(16) + ':' + json + '\n';
   const processedChunk = stringToChunk(row);
+  request.emittedModelChunkIds.add(id);
   request.completedRegularChunks.push(processedChunk);
 }