Refactor Partial Hydration (#16346)

* Move dehydrated to be child of regular SuspenseComponent

We now store the comment node on SuspenseState instead and that indicates
that this SuspenseComponent is still dehydrated.

We also store a child but that is only used to represent the DOM node for
deletions and getNextHostSibling.

* Move logic from DehydratedSuspenseComponent to SuspenseComponent

Forked based on SuspenseState.dehydrated instead.

* Retry logic for dehydrated boundary

We can now simplify the logic for retrying dehydrated boundaries without
hydrating. This is becomes simply a reconciliation against the dehydrated
fragment which gets deleted, and the new children gets inserted.

* Remove dehydrated from throw

Instead we use the regular Suspense path. To save code, we attach retry
listeners in the commit phase even though technically we don't have to.

* Pop to nearest Suspense

I think this is right...?

* Popping hydration state should skip past the dehydrated instance

* Split mount from update and special case suspended second pass

The DidCapture flag isn't used consistently in the same way. We need
further refactor for this.

* Reorganize update path

If we remove the dehydration status in the first pass and then do a second
pass because we suspended, then we need to continue as if it didn't
previously suspend. Since there is no fragment child etc.

However, we must readd the deletion.

* Schedule context work on the boundary and not the child

* Warn for Suspense hydration in legacy mode

It does a two pass render that client renders the content.

* Rename DehydratedSuspenseComponent -> DehydratedFragment

This now doesn't represent a suspense boundary itself. Its parent does.

This Fiber represents the fragment around the dehydrated content.

* Refactor returns

Avoids the temporary mutable variables. I kept losing track of them.

* Add a comment explaining the type.

Placing it in the type since that's the central point as opposed to spread
out.

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -90,6 +90,77 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(span);
   });
 
+  it('warns and replaces the boundary content in legacy mode', async () => {
+    let suspend = false;
+    let resolve;
+    let promise = new Promise(resolvePromise => (resolve = resolvePromise));
+    let ref = React.createRef();
+
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return 'Hello';
+      }
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <span ref={ref}>
+              <Child />
+            </span>
+          </Suspense>
+        </div>
+      );
+    }
+
+    // Don't suspend on the server.
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    let span = container.getElementsByTagName('span')[0];
+
+    // On the client we try to hydrate.
+    suspend = true;
+    expect(() => {
+      act(() => {
+        ReactDOM.hydrate(<App />, container);
+      });
+    }).toWarnDev(
+      'Warning: Cannot hydrate Suspense in legacy mode. Switch from ' +
+        'ReactDOM.hydrate(element, container) to ' +
+        'ReactDOM.unstable_createSyncRoot(container, { hydrate: true })' +
+        '.render(element) or remove the Suspense components from the server ' +
+        'rendered components.' +
+        '\n    in Suspense (at **)' +
+        '\n    in div (at **)' +
+        '\n    in App (at **)',
+    );
+
+    // We're now in loading state.
+    expect(container.textContent).toBe('Loading...');
+
+    let span2 = container.getElementsByTagName('span')[0];
+    // This is a new node.
+    expect(span).not.toBe(span2);
+    expect(ref.current).toBe(span2);
+
+    // Resolving the promise should render the final content.
+    suspend = false;
+    resolve();
+    await promise;
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    // We should now have hydrated with a ref on the existing span.
+    expect(container.textContent).toBe('Hello');
+  });
+
   it('can insert siblings before the dehydrated boundary', () => {
     let suspend = false;
     let promise = new Promise(() => {});
@@ -135,7 +206,8 @@ describe('ReactDOMServerPartialHydration', () => {
     suspend = true;
 
     act(() => {
-      ReactDOM.hydrate(<App />, container);
+      let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+      root.render(<App />);
     });
 
     expect(container.firstChild.firstChild.tagName).not.toBe('DIV');
@@ -191,7 +263,8 @@ describe('ReactDOMServerPartialHydration', () => {
     // hydrating anyway.
     suspend = true;
     act(() => {
-      ReactDOM.hydrate(<App />, container);
+      let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+      root.render(<App />);
     });
 
     expect(container.firstChild.children[1].textContent).toBe('Middle');

packages/react-dom/src/__tests__/ReactDOMServerSuspense-test.internal.js

@@ -37,7 +37,7 @@ function initModules() {
   };
 }
 
-const {resetModules, serverRender, itRenders} = ReactDOMServerIntegrationUtils(
+const {resetModules, serverRender} = ReactDOMServerIntegrationUtils(
   initModules,
 );
 
@@ -102,23 +102,35 @@ describe('ReactDOMServerSuspense', () => {
     );
   });
 
-  itRenders('a SuspenseList component and its children', async render => {
-    const element = await render(
+  it('server renders a SuspenseList component and its children', async () => {
+    const example = (
       <React.unstable_SuspenseList>
         <React.Suspense fallback="Loading A">
           <div>A</div>
         </React.Suspense>
         <React.Suspense fallback="Loading B">
           <div>B</div>
         </React.Suspense>
-      </React.unstable_SuspenseList>,
+      </React.unstable_SuspenseList>
     );
+    const element = await serverRender(example);
     const parent = element.parentNode;
     const divA = parent.children[0];
     expect(divA.tagName).toBe('DIV');
     expect(divA.textContent).toBe('A');
     const divB = parent.children[1];
     expect(divB.tagName).toBe('DIV');
     expect(divB.textContent).toBe('B');
+
+    ReactTestUtils.act(() => {
+      const root = ReactDOM.unstable_createSyncRoot(parent, {hydrate: true});
+      root.render(example);
+    });
+
+    const parent2 = element.parentNode;
+    const divA2 = parent2.children[0];
+    const divB2 = parent2.children[1];
+    expect(divA).toBe(divA2);
+    expect(divB).toBe(divB2);
   });
 });

packages/react-reconciler/src/ReactDebugFiberPerf.js

@@ -21,7 +21,6 @@ import {
   ContextConsumer,
   Mode,
   SuspenseComponent,
-  DehydratedSuspenseComponent,
 } from 'shared/ReactWorkTags';
 
 type MeasurementPhase =
@@ -317,8 +316,7 @@ export function stopFailedWorkTimer(fiber: Fiber): void {
     }
     fiber._debugIsCurrentlyTiming = false;
     const warning =
-      fiber.tag === SuspenseComponent ||
-      fiber.tag === DehydratedSuspenseComponent
+      fiber.tag === SuspenseComponent
         ? 'Rendering was suspended'
         : 'An error was thrown inside this error boundary';
     endFiberMark(fiber, null, warning);

packages/react-reconciler/src/ReactFiber.js

@@ -24,6 +24,7 @@ import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {UpdateQueue} from './ReactUpdateQueue';
 import type {ContextDependency} from './ReactFiberNewContext';
 import type {HookType} from './ReactFiberHooks';
+import type {SuspenseInstance} from './ReactFiberHostConfig';
 
 import invariant from 'shared/invariant';
 import warningWithoutStack from 'shared/warningWithoutStack';
@@ -48,6 +49,7 @@ import {
   Profiler,
   SuspenseComponent,
   SuspenseListComponent,
+  DehydratedFragment,
   FunctionComponent,
   MemoComponent,
   SimpleMemoComponent,
@@ -843,6 +845,14 @@ export function createFiberFromHostInstanceForDeletion(): Fiber {
   return fiber;
 }
 
+export function createFiberFromDehydratedFragment(
+  dehydratedNode: SuspenseInstance,
+): Fiber {
+  const fiber = createFiber(DehydratedFragment, null, null, NoMode);
+  fiber.stateNode = dehydratedNode;
+  return fiber;
+}
+
 export function createFiberFromPortal(
   portal: ReactPortal,
   mode: TypeOfMode,