Fix bug where references deep inside client elements couldn't be deduped.

Author: sebmarkbage

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js

@@ -224,14 +224,12 @@ describe('ReactFlightDOMEdge', () => {
       this: {is: 'a large objected'},
       with: {many: 'properties in it'},
     };
-    const props = {
-      items: new Array(30).fill(obj),
-    };
+    const props = {root: <div>{new Array(30).fill(obj)}</div>};
     const stream = ReactServerDOMServer.renderToReadableStream(props);
     const [stream1, stream2] = passThrough(stream).tee();
 
     const serializedContent = await readResult(stream1);
-    expect(serializedContent.length).toBeLessThan(470);
+    expect(serializedContent.length).toBeLessThan(1100);
 
     const result = await ReactServerDOMClient.createFromReadableStream(
       stream2,
@@ -242,10 +240,13 @@ describe('ReactFlightDOMEdge', () => {
         },
       },
     );
+    // TODO: Cyclic references currently cause a Lazy wrapper which is not ideal.
+    const resultElement = result.root._init(result.root._payload);
     // Should still match the result when parsed
-    expect(result).toEqual(props);
-    expect(result.items[5]).toBe(result.items[10]); // two random items are the same instance
-    // TODO: items[0] is not the same as the others in this case
+    expect(resultElement).toEqual(props.root);
+    expect(resultElement.props.children[5]).toBe(
+      resultElement.props.children[10],
+    ); // two random items are the same instance
   });
 
   it('should execute repeated server components only once', async () => {

packages/react-server/src/ReactFlightServer.js

@@ -2096,6 +2096,7 @@ function renderModelDestructive(
   if (typeof value === 'object') {
     switch ((value: any).$$typeof) {
       case REACT_ELEMENT_TYPE: {
+        let elementReference = null;
         const writtenObjects = request.writtenObjects;
         if (task.keyPath !== null || task.implicitSlot) {
           // If we're in some kind of context we can't reuse the result of this render or
@@ -2124,10 +2125,8 @@ function renderModelDestructive(
             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 + ':' + parentPropertyName,
-              );
+              elementReference = parentReference + ':' + parentPropertyName;
+              writtenObjects.set(value, elementReference);
             }
           }
         }
@@ -2162,7 +2161,7 @@ function renderModelDestructive(
         }
 
         // Attempt to render the Server Component.
-        return renderElement(
+        const newChild = renderElement(
           request,
           task,
           element.type,
@@ -2178,6 +2177,18 @@ function renderModelDestructive(
             : null,
           __DEV__ && enableOwnerStacks ? element._store.validated : 0,
         );
+        if (
+          typeof newChild === 'object' &&
+          newChild !== null &&
+          elementReference !== null
+        ) {
+          // If this element renders another object, we can now refer to that object through
+          // the same location as this element.
+          if (!writtenObjects.has(newChild)) {
+            writtenObjects.set(newChild, elementReference);
+          }
+        }
+        return newChild;
       }
       case REACT_LAZY_TYPE: {
         // Reset the task's thenable state before continuing. If there was one, it was