[Fiber] Fix hydration of useId in SuspenseList (#33491)

Includes #31412.

The issue is that `pushTreeFork` stores some global state when reconcile
children. This gets popped by `popTreeContext` in `completeWork`.
Normally `completeWork` returns its own `Fiber` again if it wants to do
a second pass which will call `pushTreeFork` again in the next pass.
However, `SuspenseList` doesn't return itself, it returns the next child
to work on.

The fix is to keep track of the count and push it again it when we
return the next child to attempt.

There are still some outstanding issues with hydration. Like the
backwards test still has the wrong behavior in it because it hydrates
backwards and so it picks up the DOM nodes in reverse order.
`tail="hidden"` also doesn't work correctly.

There's also another issue with `useId` and `AsyncIterable` in
SuspenseList when there's an unknown number of children. We don't
support those showing one at a time yet though so it's not an issue yet.
To fix it we need to add variable total count to the `useId` algorithm.
E.g. by falling back to varint encoding.

---------

Co-authored-by: Rick Hanlon <rickhanlonii@fb.com>
Co-authored-by: Ricky <rickhanlonii@gmail.com>

packages/react-dom/src/__tests__/ReactDOMUseId-test.js

@@ -7,7 +7,6 @@
  * @emails react-core
  * @jest-environment ./scripts/jest/ReactDOMServerIntegrationEnvironment
  */
-
 let JSDOM;
 let React;
 let ReactDOMClient;
@@ -24,6 +23,8 @@ let buffer = '';
 let hasErrored = false;
 let fatalError = undefined;
 let waitForPaint;
+let SuspenseList;
+let assertConsoleErrorDev;
 
 describe('useId', () => {
   beforeEach(() => {
@@ -32,11 +33,16 @@ describe('useId', () => {
     React = require('react');
     ReactDOMClient = require('react-dom/client');
     clientAct = require('internal-test-utils').act;
+    assertConsoleErrorDev =
+      require('internal-test-utils').assertConsoleErrorDev;
     ReactDOMFizzServer = require('react-dom/server');
     Stream = require('stream');
     Suspense = React.Suspense;
     useId = React.useId;
     useState = React.useState;
+    if (gate(flags => flags.enableSuspenseList)) {
+      SuspenseList = React.unstable_SuspenseList;
+    }
 
     const InternalTestUtils = require('internal-test-utils');
     waitForPaint = InternalTestUtils.waitForPaint;
@@ -375,6 +381,370 @@ describe('useId', () => {
     `);
   });
 
+  // @gate enableSuspenseList
+  it('Supports SuspenseList (reveal order independent)', async () => {
+    function Baz({id, children}) {
+      return <span id={id}>{children}</span>;
+    }
+
+    function Bar({children}) {
+      const id = useId();
+      return <Baz id={id}>{children}</Baz>;
+    }
+
+    function Foo() {
+      return (
+        <SuspenseList revealOrder="independent">
+          <Bar>A</Bar>
+          <Bar>B</Bar>
+        </SuspenseList>
+      );
+    }
+
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<Foo />);
+      pipe(writable);
+    });
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+
+    await clientAct(async () => {
+      ReactDOMClient.hydrateRoot(container, <Foo />);
+    });
+
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+  });
+
+  // @gate enableSuspenseList
+  it('Supports SuspenseList (reveal order "together")', async () => {
+    function Baz({id, children}) {
+      return <span id={id}>{children}</span>;
+    }
+
+    function Bar({children}) {
+      const id = useId();
+      return <Baz id={id}>{children}</Baz>;
+    }
+
+    function Foo() {
+      return (
+        <SuspenseList revealOrder="together">
+          <Bar>A</Bar>
+          <Bar>B</Bar>
+        </SuspenseList>
+      );
+    }
+
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<Foo />);
+      pipe(writable);
+    });
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+
+    await clientAct(async () => {
+      ReactDOMClient.hydrateRoot(container, <Foo />);
+    });
+
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+  });
+
+  // @gate enableSuspenseList
+  it('Supports SuspenseList (reveal order "forwards")', async () => {
+    function Baz({id, children}) {
+      return <span id={id}>{children}</span>;
+    }
+
+    function Bar({children}) {
+      const id = useId();
+      return <Baz id={id}>{children}</Baz>;
+    }
+
+    function Foo() {
+      return (
+        <SuspenseList revealOrder="forwards" tail="visible">
+          <Bar>A</Bar>
+          <Bar>B</Bar>
+        </SuspenseList>
+      );
+    }
+
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<Foo />);
+      pipe(writable);
+    });
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+
+    await clientAct(async () => {
+      ReactDOMClient.hydrateRoot(container, <Foo />);
+    });
+
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+  });
+
+  // @gate enableSuspenseList
+  it('Supports SuspenseList (reveal order "backwards") with a single child in a list of many', async () => {
+    function Baz({id, children}) {
+      return <span id={id}>{children}</span>;
+    }
+
+    function Bar({children}) {
+      const id = useId();
+      return <Baz id={id}>{children}</Baz>;
+    }
+
+    function Foo() {
+      return (
+        <SuspenseList revealOrder="unstable_legacy-backwards" tail="visible">
+          {null}
+          <Bar>A</Bar>
+          {null}
+        </SuspenseList>
+      );
+    }
+
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<Foo />);
+      pipe(writable);
+    });
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_2_"
+        >
+          A
+        </span>
+        <!-- -->
+      </div>
+    `);
+
+    await clientAct(async () => {
+      ReactDOMClient.hydrateRoot(container, <Foo />);
+    });
+
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_2_"
+        >
+          A
+        </span>
+        <!-- -->
+      </div>
+    `);
+  });
+
+  // @gate enableSuspenseList
+  it('Supports SuspenseList (reveal order "backwards")', async () => {
+    function Baz({id, children}) {
+      return <span id={id}>{children}</span>;
+    }
+
+    function Bar({children}) {
+      const id = useId();
+      return <Baz id={id}>{children}</Baz>;
+    }
+
+    function Foo() {
+      return (
+        <SuspenseList revealOrder="unstable_legacy-backwards" tail="visible">
+          <Bar>A</Bar>
+          <Bar>B</Bar>
+        </SuspenseList>
+      );
+    }
+
+    await serverAct(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<Foo />);
+      pipe(writable);
+    });
+    expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+
+    if (gate(flags => flags.favorSafetyOverHydrationPerf)) {
+      // TODO: This is a bug with revealOrder="backwards" in that it hydrates in reverse.
+      await expect(async () => {
+        await clientAct(async () => {
+          ReactDOMClient.hydrateRoot(container, <Foo />);
+        });
+      }).rejects.toThrowError(
+        `Hydration failed because the server rendered text didn't match the client. As a result this tree will be regenerated on the client.`,
+      );
+
+      expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_r_1_"
+        >
+          A
+        </span>
+        <span
+          id="_r_0_"
+        >
+          B
+        </span>
+      </div>
+    `);
+    } else {
+      await clientAct(async () => {
+        ReactDOMClient.hydrateRoot(container, <Foo />);
+      });
+
+      // TODO: This is a bug with revealOrder="backwards" in that it hydrates in reverse.
+      assertConsoleErrorDev([
+        `A tree hydrated but some attributes of the server rendered HTML didn't match the client properties. This won't be patched up. This can happen if a SSR-ed Client Component used:
+
+- A server/client branch \`if (typeof window !== 'undefined')\`.
+- Variable input such as \`Date.now()\` or \`Math.random()\` which changes each time it's called.
+- Date formatting in a user's locale which doesn't match the server.
+- External changing data without sending a snapshot of it along with the HTML.
+- Invalid HTML tag nesting.
+
+It can also happen if the client has a browser extension installed which messes with the HTML before React loaded.
+
+https://react.dev/link/hydration-mismatch
+
+  <Foo>
+    <SuspenseList revealOrder="unstable_l..." tail="visible">
+      <Bar>
+      <Bar>
+        <Baz id="_R_2_">
+          <span
++           id="_R_2_"
+-           id="_R_1_"
+          >
++           B
+-           A
+`,
+      ]);
+
+      expect(container).toMatchInlineSnapshot(`
+      <div
+        id="container"
+      >
+        <span
+          id="_R_1_"
+        >
+          A
+        </span>
+        <span
+          id="_R_2_"
+        >
+          B
+        </span>
+      </div>
+    `);
+    }
+  });
+
   it('basic incremental hydration', async () => {
     function App() {
       return (

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -3342,6 +3342,7 @@ function initSuspenseListRenderState(
   tail: null | Fiber,
   lastContentRow: null | Fiber,
   tailMode: SuspenseListTailMode,
+  treeForkCount: number,
 ): void {
   const renderState: null | SuspenseListRenderState =
     workInProgress.memoizedState;
@@ -3353,6 +3354,7 @@ function initSuspenseListRenderState(
       last: lastContentRow,
       tail: tail,
       tailMode: tailMode,
+      treeForkCount: treeForkCount,
     }: SuspenseListRenderState);
   } else {
     // We can reuse the existing object from previous renders.
@@ -3362,6 +3364,7 @@ function initSuspenseListRenderState(
     renderState.last = lastContentRow;
     renderState.tail = tail;
     renderState.tailMode = tailMode;
+    renderState.treeForkCount = treeForkCount;
   }
 }
 
@@ -3404,6 +3407,8 @@ function updateSuspenseListComponent(
   validateSuspenseListChildren(newChildren, revealOrder);
 
   reconcileChildren(current, workInProgress, newChildren, renderLanes);
+  // Read how many children forks this set pushed so we can push it every time we retry.
+  const treeForkCount = getIsHydrating() ? getForksAtLevel(workInProgress) : 0;
 
   if (!shouldForceFallback) {
     const didSuspendBefore =
@@ -3446,6 +3451,7 @@ function updateSuspenseListComponent(
           tail,
           lastContentRow,
           tailMode,
+          treeForkCount,
         );
         break;
       }
@@ -3478,6 +3484,7 @@ function updateSuspenseListComponent(
           tail,
           null, // last
           tailMode,
+          treeForkCount,
         );
         break;
       }
@@ -3488,6 +3495,7 @@ function updateSuspenseListComponent(
           null, // tail
           null, // last
           undefined,
+          treeForkCount,
         );
         break;
       }

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -184,7 +184,7 @@ import {resetChildFibers} from './ReactChildFiber';
 import {createScopeInstance} from './ReactFiberScope';
 import {transferActualDuration} from './ReactProfilerTimer';
 import {popCacheProvider} from './ReactFiberCacheComponent';
-import {popTreeContext} from './ReactFiberTreeContext';
+import {popTreeContext, pushTreeFork} from './ReactFiberTreeContext';
 import {popRootTransition, popTransition} from './ReactFiberTransition';
 import {
   popMarkerInstance,
@@ -1764,6 +1764,10 @@ function completeWork(
                     ForceSuspenseFallback,
                   ),
                 );
+                if (getIsHydrating()) {
+                  // Re-apply tree fork since we popped the tree fork context in the beginning of this function.
+                  pushTreeFork(workInProgress, renderState.treeForkCount);
+                }
                 // Don't bubble properties in this case.
                 return workInProgress.child;
               }
@@ -1890,6 +1894,10 @@ function completeWork(
         }
         pushSuspenseListContext(workInProgress, suspenseContext);
         // Do a pass over the next row.
+        if (getIsHydrating()) {
+          // Re-apply tree fork since we popped the tree fork context in the beginning of this function.
+          pushTreeFork(workInProgress, renderState.treeForkCount);
+        }
         // Don't bubble properties in this case.
         return next;
       }

packages/react-reconciler/src/ReactFiberSuspenseComponent.js

@@ -54,6 +54,8 @@ export type SuspenseListRenderState = {
   tail: null | Fiber,
   // Tail insertions setting.
   tailMode: SuspenseListTailMode,
+  // Keep track of total number of forks during multiple passes
+  treeForkCount: number,
 };
 
 export type RetryQueue = Set<Wakeable>;