From 06bb3fb2760bcb4e5e55d08b3b895e0b72a76b23 Mon Sep 17 00:00:00 2001
From: Hendrik Liebau <mail@hendrik-liebau.de>
Date: Thu, 4 Apr 2024 17:35:45 +0200
Subject: [PATCH] [Flight] Deduplicate suspended elements

---
 .../src/__tests__/ReactFlight-test.js         | 70 +++++++++++++++++++
 .../react-server/src/ReactFlightServer.js     |  3 +
 2 files changed, 73 insertions(+)

diff --git a/packages/react-client/src/__tests__/ReactFlight-test.js b/packages/react-client/src/__tests__/ReactFlight-test.js
index 55956ddc93ecd..7d0ee2fea5009 100644
--- a/packages/react-client/src/__tests__/ReactFlight-test.js
+++ b/packages/react-client/src/__tests__/ReactFlight-test.js
@@ -2145,4 +2145,74 @@ describe('ReactFlight', () => {
     expect(loggedFn).not.toBe(foo);
     expect(loggedFn.toString()).toBe(foo.toString());
   });
+
+  // @gate __DEV__
+  it('does not emit duplicate chunks for already outlined elements in dev mode', async () => {
+    async function Bar({text}) {
+      return <div>{text.toUpperCase()}</div>;
+    }
+
+    function Foo() {
+      const bar = <Bar text="bar" />;
+
+      return (
+        <div>
+          {bar}
+          {bar}
+        </div>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<Foo />);
+    const textDecoder = new TextDecoder();
+
+    await act(async () => {
+      const chunks = transport
+        .map(chunk => textDecoder.decode(chunk).replace(/\n$/, ''))
+        .join('\n');
+
+      expect(chunks).toEqual(
+        `
+0:D{"name":"Foo","env":"Server"}
+1:D{"name":"Bar","env":"Server"}
+0:["$","div",null,{"children":["$L1","$1"]}]
+1:["$","div",null,{"children":"BAR"}]
+  `.trim(),
+      );
+    });
+  });
+
+  // @gate !__DEV__
+  it('does not emit duplicate chunks for already outlined elements in production mode', async () => {
+    async function Bar({text}) {
+      return <div>{text.toUpperCase()}</div>;
+    }
+
+    function Foo() {
+      const bar = <Bar text="bar" />;
+
+      return (
+        <div>
+          {bar}
+          {bar}
+        </div>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<Foo />);
+    const textDecoder = new TextDecoder();
+
+    await act(async () => {
+      const chunks = transport
+        .map(chunk => textDecoder.decode(chunk).replace(/\n$/, ''))
+        .join('\n');
+
+      expect(chunks).toEqual(
+        `
+0:["$","div",null,{"children":["$L1","$1"]}]
+1:["$","div",null,{"children":"BAR"}]
+`.trim(),
+      );
+    });
+  });
 });
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index ad0a28f37175c..f4e3f1ef91342 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -1237,6 +1237,9 @@ function renderModel(
           task.implicitSlot,
           request.abortableTasks,
         );
+        // The suspended element was outlined, so we're using the same ID for the
+        // original value, thus ensuring that it's deduplicated if it's referenced again.
+        request.writtenObjects.set((value: any), newTask.id);
         const ping = newTask.ping;
         (x: any).then(ping, ping);
         newTask.thenableState = getThenableStateAfterSuspending();