Resolve references to deduped owner objects (facebook#30549)

This is a follow-up from facebook#30528 to not only handle props (the critical
change), but also the owner ~and stack~ of a referenced element.

~Handling stacks here is rather academic because the Flight Server
currently does not deduplicate owner stacks. And if they are really
identical, we should probably just dedupe the whole element.~ EDIT:
Removed from the PR.

Handling owner objects on the other hand is an actual requirement as
reported in vercel/next.js#69545. This problem
only affects the stable release channel, as the absence of owner stacks
allows for the specific kind of shared owner deduping as demonstrated in
the unit test.

Author: unstubbable

packages/react-client/src/ReactFlightClient.js

@@ -937,26 +937,35 @@ function waitForReference<T>(
       }
       value = value[path[i]];
     }
-    parentObject[key] = map(response, value);
+    const mappedValue = map(response, value);
+    parentObject[key] = mappedValue;
 
     // If this is the root object for a model reference, where `handler.value`
     // is a stale `null`, the resolved value can be used directly.
     if (key === '' && handler.value === null) {
-      handler.value = parentObject[key];
+      handler.value = mappedValue;
     }
 
-    // If the parent object is an unparsed React element tuple and its outlined
-    // props have now been resolved, we also need to update the props of the
-    // parsed element object (i.e. handler.value).
+    // If the parent object is an unparsed React element tuple, we also need to
+    // update the props and owner of the parsed element object (i.e.
+    // handler.value).
     if (
       parentObject[0] === REACT_ELEMENT_TYPE &&
-      key === '3' &&
       typeof handler.value === 'object' &&
       handler.value !== null &&
-      handler.value.$$typeof === REACT_ELEMENT_TYPE &&
-      handler.value.props === null
+      handler.value.$$typeof === REACT_ELEMENT_TYPE
     ) {
-      handler.value.props = parentObject[key];
+      const element: any = handler.value;
+      switch (key) {
+        case '3':
+          element.props = mappedValue;
+          break;
+        case '4':
+          if (__DEV__) {
+            element._owner = mappedValue;
+          }
+          break;
+      }
     }
 
     handler.deps--;

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

@@ -630,6 +630,119 @@ describe('ReactFlightDOMBrowser', () => {
     expect(container.innerHTML).toBe('<div></div>');
   });
 
+  it('should handle references to deduped owner objects', async () => {
+    // This is replicating React components as generated by @svgr/webpack:
+    let path1a: React.ReactNode;
+    let path1b: React.ReactNode;
+    let path2: React.ReactNode;
+
+    function Svg1() {
+      return ReactServer.createElement(
+        'svg',
+        {id: '1'},
+        path1a || (path1a = ReactServer.createElement('path', {})),
+        path1b || (path1b = ReactServer.createElement('path', {})),
+      );
+    }
+
+    function Svg2() {
+      return ReactServer.createElement(
+        'svg',
+        {id: '2'},
+        path2 || (path2 = ReactServer.createElement('path', {})),
+      );
+    }
+
+    function Server() {
+      return ReactServer.createElement(
+        ReactServer.Fragment,
+        {},
+        ReactServer.createElement(Svg1),
+        ReactServer.createElement(Svg2),
+      );
+    }
+
+    let stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(<Server />, webpackMap),
+    );
+
+    function ClientRoot({response}) {
+      return use(response);
+    }
+
+    let response = ReactServerDOMClient.createFromReadableStream(stream);
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+
+    await act(() => {
+      root.render(<ClientRoot response={response} />);
+    });
+
+    const expectedHtml =
+      '<svg id="1"><path></path><path></path></svg><svg id="2"><path></path></svg>';
+
+    expect(container.innerHTML).toBe(expectedHtml);
+
+    // Render a second time:
+
+    // Assigning the path elements to variables in module scope (here simulated
+    // with the test's function scope), and rendering a second time, prevents
+    // the owner of the path elements (i.e. Svg1/Svg2) to be deduped. The owner
+    // of the path in Svg1 is fully inlined. The owner of the owner of the path
+    // in Svg2 is Server, which is deduped and replaced with a reference to the
+    // owner of the owner of the path in Svg1. This nested owner is actually
+    // Server from the previous render pass, which is kinda broken and libraries
+    // probably shouldn't generate code like this. This reference can only be
+    // resolved properly if owners are specifically handled when resolving
+    // outlined models.
+
+    stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(<Server />, webpackMap),
+    );
+
+    response = ReactServerDOMClient.createFromReadableStream(stream);
+
+    await act(() => {
+      root.render(<ClientRoot response={response} />);
+    });
+
+    expect(container.innerHTML).toBe(expectedHtml);
+
+    if (__DEV__) {
+      const resolvedPath1b = await response.value[0].props.children[1]._payload;
+
+      expect(resolvedPath1b._owner).toEqual(
+        expect.objectContaining({
+          name: 'Svg1',
+          env: 'Server',
+          key: null,
+          owner: expect.objectContaining({
+            name: 'Server',
+            env: 'Server',
+            key: null,
+            owner: null,
+          }),
+        }),
+      );
+
+      const resolvedPath2 = response.value[1].props.children;
+
+      expect(resolvedPath2._owner).toEqual(
+        expect.objectContaining({
+          name: 'Svg2',
+          env: 'Server',
+          key: null,
+          owner: expect.objectContaining({
+            name: 'Server',
+            env: 'Server',
+            key: null,
+            owner: null,
+          }),
+        }),
+      );
+    }
+  });
+
   it('should progressively reveal server components', async () => {
     let reportedErrors = [];
 

packages/react-server/src/ReactFlightServer.js

@@ -2665,6 +2665,9 @@ function renderModelDestructive(
             case '3':
               propertyName = 'props';
               break;
+            case '4':
+              propertyName = '_owner';
+              break;
           }
         }
         writtenObjects.set(value, parentReference + ':' + propertyName);