useFormState: Emit comment to mark whether state matches

packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js

@@ -98,6 +98,7 @@ import {
   enableTrustedTypesIntegration,
   diffInCommitPhase,
   enableFormActions,
+  enableAsyncActions,
 } from 'shared/ReactFeatureFlags';
 import {
   HostComponent,
@@ -160,7 +161,12 @@ export type TextInstance = Text;
 export interface SuspenseInstance extends Comment {
   _reactRetry?: () => void;
 }
-export type HydratableInstance = Instance | TextInstance | SuspenseInstance;
+type FormStateMarkerInstance = Comment;
+export type HydratableInstance =
+  | Instance
+  | TextInstance
+  | SuspenseInstance
+  | FormStateMarkerInstance;
 export type PublicInstance = Element | Text;
 export type HostContextDev = {
   context: HostContextProd,
@@ -187,6 +193,8 @@ const SUSPENSE_START_DATA = '$';
 const SUSPENSE_END_DATA = '/$';
 const SUSPENSE_PENDING_START_DATA = '$?';
 const SUSPENSE_FALLBACK_START_DATA = '$!';
+const FORM_STATE_IS_MATCHING = 'F!';
+const FORM_STATE_IS_NOT_MATCHING = 'F';
 
 const STYLE = 'style';
 
@@ -1283,6 +1291,37 @@ export function registerSuspenseInstanceRetry(
   instance._reactRetry = callback;
 }
 
+export function canHydrateFormStateMarker(
+  instance: HydratableInstance,
+  inRootOrSingleton: boolean,
+): null | FormStateMarkerInstance {
+  while (instance.nodeType !== COMMENT_NODE) {
+    if (!inRootOrSingleton || !enableHostSingletons) {
+      return null;
+    }
+    const nextInstance = getNextHydratableSibling(instance);
+    if (nextInstance === null) {
+      return null;
+    }
+    instance = nextInstance;
+  }
+  const nodeData = (instance: any).data;
+  if (
+    nodeData === FORM_STATE_IS_MATCHING ||
+    nodeData === FORM_STATE_IS_NOT_MATCHING
+  ) {
+    const markerInstance: FormStateMarkerInstance = (instance: any);
+    return markerInstance;
+  }
+  return null;
+}
+
+export function isFormStateMarkerMatching(
+  markerInstance: FormStateMarkerInstance,
+): boolean {
+  return markerInstance.data === FORM_STATE_IS_MATCHING;
+}
+
 function getNextHydratable(node: ?Node) {
   // Skip non-hydratable nodes.
   for (; node != null; node = ((node: any): Node).nextSibling) {
@@ -1295,7 +1334,11 @@ function getNextHydratable(node: ?Node) {
       if (
         nodeData === SUSPENSE_START_DATA ||
         nodeData === SUSPENSE_FALLBACK_START_DATA ||
-        nodeData === SUSPENSE_PENDING_START_DATA
+        nodeData === SUSPENSE_PENDING_START_DATA ||
+        (enableFormActions &&
+          enableAsyncActions &&
+          (nodeData === FORM_STATE_IS_MATCHING ||
+            nodeData === FORM_STATE_IS_NOT_MATCHING))
       ) {
         break;
       }

packages/react-dom-bindings/src/server/ReactFizzConfigDOM.js

@@ -1519,6 +1519,21 @@ function injectFormReplayingRuntime(
   }
 }
 
+const formStateMarkerIsMatching = stringToPrecomputedChunk('<!--F!-->');
+const formStateMarkerIsNotMatching = stringToPrecomputedChunk('<!--F-->');
+
+export function pushFormStateMarkerIsMatching(
+  target: Array<Chunk | PrecomputedChunk>,
+) {
+  target.push(formStateMarkerIsMatching);
+}
+
+export function pushFormStateMarkerIsNotMatching(
+  target: Array<Chunk | PrecomputedChunk>,
+) {
+  target.push(formStateMarkerIsNotMatching);
+}
+
 function pushStartForm(
   target: Array<Chunk | PrecomputedChunk>,
   props: Object,

packages/react-dom-bindings/src/server/ReactFizzConfigDOMLegacy.js

@@ -101,6 +101,8 @@ export {
   pushEndInstance,
   pushStartCompletedSuspenseBoundary,
   pushEndCompletedSuspenseBoundary,
+  pushFormStateMarkerIsMatching,
+  pushFormStateMarkerIsNotMatching,
   writeStartSegment,
   writeEndSegment,
   writeCompletedSegmentInstruction,

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

@@ -30,6 +30,7 @@ let SuspenseList;
 let useSyncExternalStore;
 let useSyncExternalStoreWithSelector;
 let use;
+let useFormState;
 let PropTypes;
 let textCache;
 let writable;
@@ -88,6 +89,7 @@ describe('ReactDOMFizzServer', () => {
     if (gate(flags => flags.enableSuspenseList)) {
       SuspenseList = React.unstable_SuspenseList;
     }
+    useFormState = ReactDOM.experimental_useFormState;
 
     PropTypes = require('prop-types');
 
@@ -5876,6 +5878,123 @@ describe('ReactDOMFizzServer', () => {
     expect(getVisibleChildren(container)).toEqual('Hi');
   });
 
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  it('useFormState hydrates without a mismatch', async () => {
+    // This is testing an implementation detail: useFormState emits comment
+    // nodes into the SSR stream, so this checks that they are handled correctly
+    // during hydration.
+
+    async function action(state) {
+      return state;
+    }
+
+    const childRef = React.createRef(null);
+    function Form() {
+      const [state] = useFormState(action, 0);
+      const text = `Child: ${state}`;
+      return (
+        <div id="child" ref={childRef}>
+          {text}
+        </div>
+      );
+    }
+
+    function App() {
+      return (
+        <div>
+          <div>
+            <Form />
+          </div>
+          <span>Sibling</span>
+        </div>
+      );
+    }
+
+    await act(() => {
+      const {pipe} = renderToPipeableStream(<App />);
+      pipe(writable);
+    });
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <div>
+          <div id="child">Child: 0</div>
+        </div>
+        <span>Sibling</span>
+      </div>,
+    );
+    const child = document.getElementById('child');
+
+    // Confirm that it hydrates correctly
+    await clientAct(() => {
+      ReactDOMClient.hydrateRoot(container, <App />);
+    });
+    expect(childRef.current).toBe(child);
+  });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  it("useFormState hydrates without a mismatch if there's a render phase update", async () => {
+    async function action(state) {
+      return state;
+    }
+
+    const childRef = React.createRef(null);
+    function Form() {
+      const [localState, setLocalState] = React.useState(0);
+      if (localState < 3) {
+        setLocalState(localState + 1);
+      }
+
+      // Because of the render phase update above, this component is evaluated
+      // multiple times (even during SSR), but it should only emit a single
+      // marker per useFormState instance.
+      const [formState] = useFormState(action, 0);
+      const text = `${readText('Child')}:${formState}:${localState}`;
+      return (
+        <div id="child" ref={childRef}>
+          {text}
+        </div>
+      );
+    }
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Form />
+          </Suspense>
+          <span>Sibling</span>
+        </div>
+      );
+    }
+
+    await act(() => {
+      const {pipe} = renderToPipeableStream(<App />);
+      pipe(writable);
+    });
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        Loading...<span>Sibling</span>
+      </div>,
+    );
+
+    await act(() => resolveText('Child'));
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <div id="child">Child:0:3</div>
+        <span>Sibling</span>
+      </div>,
+    );
+    const child = document.getElementById('child');
+
+    // Confirm that it hydrates correctly
+    await clientAct(() => {
+      ReactDOMClient.hydrateRoot(container, <App />);
+    });
+    expect(childRef.current).toBe(child);
+  });
+
   describe('useEffectEvent', () => {
     // @gate enableUseEffectEventHook
     it('can server render a component with useEffectEvent', async () => {

packages/react-reconciler/src/ReactFiberConfigWithNoHydration.js

@@ -26,6 +26,8 @@ export const isSuspenseInstancePending = shim;
 export const isSuspenseInstanceFallback = shim;
 export const getSuspenseInstanceFallbackErrorDetails = shim;
 export const registerSuspenseInstanceRetry = shim;
+export const canHydrateFormStateMarker = shim;
+export const isFormStateMarkerMatching = shim;
 export const getNextHydratableSibling = shim;
 export const getFirstHydratableChild = shim;
 export const getFirstHydratableChildWithinContainer = shim;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -111,7 +111,10 @@ import {
   markWorkInProgressReceivedUpdate,
   checkIfWorkInProgressReceivedUpdate,
 } from './ReactFiberBeginWork';
-import {getIsHydrating} from './ReactFiberHydrationContext';
+import {
+  getIsHydrating,
+  tryToClaimNextHydratableFormMarkerInstance,
+} from './ReactFiberHydrationContext';
 import {logStateUpdateScheduled} from './DebugTracing';
 import {
   markStateUpdateScheduled,
@@ -2010,6 +2013,12 @@ function mountFormState<S, P>(
   initialState: S,
   permalink?: string,
 ): [S, (P) => void] {
+  if (getIsHydrating()) {
+    // TODO: If this function returns true, it means we should use the form
+    // state passed to hydrateRoot instead of initialState.
+    tryToClaimNextHydratableFormMarkerInstance(currentlyRenderingFiber);
+  }
+
   // State hook. The state is stored in a thenable which is then unwrapped by
   // the `use` algorithm during render.
   const stateHook = mountWorkInProgressHook();
@@ -2145,7 +2154,8 @@ function rerenderFormState<S, P>(
   }
 
   // This is a mount. No updates to process.
-  const state = stateHook.memoizedState;
+  const thenable: Thenable<S> = stateHook.memoizedState;
+  const state = useThenable(thenable);
 
   const actionQueueHook = updateWorkInProgressHook();
   const actionQueue = actionQueueHook.queue;

packages/react-reconciler/src/ReactFiberHydrationContext.js

@@ -76,6 +76,8 @@ import {
   canHydrateInstance,
   canHydrateTextInstance,
   canHydrateSuspenseInstance,
+  canHydrateFormStateMarker,
+  isFormStateMarkerMatching,
   isHydratableText,
 } from './ReactFiberConfig';
 import {OffscreenLane} from './ReactFiberLane';
@@ -595,6 +597,34 @@ function tryToClaimNextHydratableSuspenseInstance(fiber: Fiber): void {
   }
 }
 
+export function tryToClaimNextHydratableFormMarkerInstance(
+  fiber: Fiber,
+): boolean {
+  if (!isHydrating) {
+    return false;
+  }
+  if (nextHydratableInstance) {
+    const markerInstance = canHydrateFormStateMarker(
+      nextHydratableInstance,
+      rootOrSingletonContext,
+    );
+    if (markerInstance) {
+      // Found the marker instance.
+      nextHydratableInstance = getNextHydratableSibling(markerInstance);
+      // Return true if this marker instance should use the state passed
+      // to hydrateRoot.
+      // TODO: As an optimization, Fizz should only emit these markers if form
+      // state is passed at the root.
+      return isFormStateMarkerMatching(markerInstance);
+    }
+  }
+  // Should have found a marker instance. Throw an error to trigger client
+  // rendering. We don't bother to check if we're in a concurrent root because
+  // useFormState is a new API, so backwards compat is not an issue.
+  throwOnHydrationMismatch(fiber);
+  return false;
+}
+
 function prepareToHydrateHostInstance(
   fiber: Fiber,
   hostContext: HostContext,

packages/react-reconciler/src/forks/ReactFiberConfig.custom.js

@@ -142,6 +142,8 @@ export const getSuspenseInstanceFallbackErrorDetails =
   $$$config.getSuspenseInstanceFallbackErrorDetails;
 export const registerSuspenseInstanceRetry =
   $$$config.registerSuspenseInstanceRetry;
+export const canHydrateFormStateMarker = $$$config.canHydrateFormStateMarker;
+export const isFormStateMarkerMatching = $$$config.isFormStateMarkerMatching;
 export const getNextHydratableSibling = $$$config.getNextHydratableSibling;
 export const getFirstHydratableChild = $$$config.getFirstHydratableChild;
 export const getFirstHydratableChildWithinContainer =

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

@@ -344,7 +344,7 @@ describe('ReactFlightDOMForm', () => {
     const ssrStream = await ReactDOMServer.renderToReadableStream(response);
     await readIntoContainer(ssrStream);
 
-    const form = container.firstChild;
+    const form = container.getElementsByTagName('form')[0];
     const span = container.getElementsByTagName('span')[0];
     expect(span.textContent).toBe('Count: 1');
 
@@ -382,7 +382,7 @@ describe('ReactFlightDOMForm', () => {
     const ssrStream = await ReactDOMServer.renderToReadableStream(response);
     await readIntoContainer(ssrStream);
 
-    const form = container.firstChild;
+    const form = container.getElementsByTagName('form')[0];
     const span = container.getElementsByTagName('span')[0];
     expect(span.textContent).toBe('Count: 1');
 
@@ -423,7 +423,7 @@ describe('ReactFlightDOMForm', () => {
     const ssrStream = await ReactDOMServer.renderToReadableStream(response);
     await readIntoContainer(ssrStream);
 
-    const form = container.firstChild;
+    const form = container.getElementsByTagName('form')[0];
     const span = container.getElementsByTagName('span')[0];
     expect(span.textContent).toBe('Count: 1');