useFormState: Emit comment to mark whether state matches (#27307)

A planned feature of useFormState is that if the page load is the result
of an MPA-style form submission — i.e. a form was submitted before it
was hydrated, using Server Actions — the state of the hook should
transfer to the next page.

I haven't implemented that part yet, but as a prerequisite, we need some
way for Fizz to indicate whether a useFormState hook was rendered using
the "postback" state. That way we can do all state matching logic on the
server without having to replicate it on the client, too.

The approach here is to emit a comment node for each useFormState hook.
We use one of two comment types: `<!--F-->` for a normal useFormState
hook, and `<!--F!-->` for a hook that was rendered using the postback
state. React will read these markers during hydration. This is similar
to how we encode Suspense boundaries.

Again, the actual matching algorithm is not yet implemented — for now,
the "not matching" marker is always emitted.

We can optimize this further by not emitting any markers for a render
that is not the result of a form postback, which I'll do in subsequent
PRs.

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');
 

packages/react-server/src/ReactFizzHooks.js

@@ -72,6 +72,13 @@ let isReRender: boolean = false;
 let didScheduleRenderPhaseUpdate: boolean = false;
 // Counts the number of useId hooks in this component
 let localIdCounter: number = 0;
+// Chunks that should be pushed to the stream once the component
+// finishes rendering.
+// Counts the number of useFormState calls in this component
+let formStateCounter: number = 0;
+// The index of the useFormState hook that matches the one passed in at the
+// root during an MPA navigation, if any.
+let formStateMatchingIndex: number = -1;
 // Counts the number of use(thenable) calls in this component
 let thenableIndexCounter: number = 0;
 let thenableState: ThenableState | null = null;
@@ -208,6 +215,8 @@ export function prepareToUseHooks(
   // workInProgressHook = null;
 
   localIdCounter = 0;
+  formStateCounter = 0;
+  formStateMatchingIndex = -1;
   thenableIndexCounter = 0;
   thenableState = prevThenableState;
 }
@@ -228,6 +237,8 @@ export function finishHooks(
     // restarting until no more updates are scheduled.
     didScheduleRenderPhaseUpdate = false;
     localIdCounter = 0;
+    formStateCounter = 0;
+    formStateMatchingIndex = -1;
     thenableIndexCounter = 0;
     numberOfReRenders += 1;
 
@@ -236,6 +247,7 @@ export function finishHooks(
 
     children = Component(props, refOrContext);
   }
+
   resetHooksState();
   return children;
 }
@@ -254,6 +266,19 @@ export function checkDidRenderIdHook(): boolean {
   return didRenderIdHook;
 }
 
+export function getFormStateCount(): number {
+  // This should be called immediately after every finishHooks call.
+  // Conceptually, it's part of the return value of finishHooks; it's only a
+  // separate function to avoid using an array tuple.
+  return formStateCounter;
+}
+export function getFormStateMatchingIndex(): number {
+  // This should be called immediately after every finishHooks call.
+  // Conceptually, it's part of the return value of finishHooks; it's only a
+  // separate function to avoid using an array tuple.
+  return formStateMatchingIndex;
+}
+
 // Reset the internal hooks state if an error occurs while rendering a component
 export function resetHooksState(): void {
   if (__DEV__) {
@@ -559,6 +584,11 @@ function useFormState<S, P>(
 ): [S, (P) => void] {
   resolveCurrentlyRenderingComponent();
 
+  // Count the number of useFormState hooks per component.
+  // TODO: We should also track which hook matches the form state passed at
+  // the root, if any. Matching is not yet implemented.
+  formStateCounter++;
+
   // Bind the initial state to the first argument of the action.
   // TODO: Use the keypath (or permalink) to check if there's matching state
   // from the previous page.

packages/react-server/src/ReactFizzServer.js

@@ -76,6 +76,8 @@ import {
   prepareHostDispatcher,
   supportsRequestStorage,
   requestStorage,
+  pushFormStateMarkerIsMatching,
+  pushFormStateMarkerIsNotMatching,
 } from './ReactFizzConfig';
 import {
   constructClassInstance,
@@ -104,6 +106,8 @@ import {
   setCurrentResumableState,
   getThenableStateAfterSuspending,
   unwrapThenable,
+  getFormStateCount,
+  getFormStateMatchingIndex,
 } from './ReactFizzHooks';
 import {DefaultCacheDispatcher} from './ReactFizzCache';
 import {getStackByComponentStackNode} from './ReactFizzComponentStack';
@@ -1044,6 +1048,8 @@ function renderIndeterminateComponent(
     legacyContext,
   );
   const hasId = checkDidRenderIdHook();
+  const formStateCount = getFormStateCount();
+  const formStateMatchingIndex = getFormStateMatchingIndex();
 
   if (__DEV__) {
     // Support for module components is deprecated and is removed behind a flag.
@@ -1113,28 +1119,74 @@ function renderIndeterminateComponent(
     if (__DEV__) {
       validateFunctionComponentInDev(Component);
     }
-    // We're now successfully past this task, and we don't have to pop back to
-    // the previous task every again, so we can use the destructive recursive form.
-    if (hasId) {
-      // This component materialized an id. We treat this as its own level, with
-      // a single "child" slot.
-      const prevTreeContext = task.treeContext;
-      const totalChildren = 1;
-      const index = 0;
-      // Modify the id context. Because we'll need to reset this if something
-      // suspends or errors, we'll use the non-destructive render path.
-      task.treeContext = pushTreeContext(prevTreeContext, totalChildren, index);
-      renderNode(request, task, value, 0);
-      // Like the other contexts, this does not need to be in a finally block
-      // because renderNode takes care of unwinding the stack.
-      task.treeContext = prevTreeContext;
-    } else {
-      renderNodeDestructive(request, task, null, value, 0);
-    }
+    finishFunctionComponent(
+      request,
+      task,
+      value,
+      hasId,
+      formStateCount,
+      formStateMatchingIndex,
+    );
   }
   popComponentStackInDEV(task);
 }
 
+function finishFunctionComponent(
+  request: Request,
+  task: Task,
+  children: ReactNodeList,
+  hasId: boolean,
+  formStateCount: number,
+  formStateMatchingIndex: number,
+) {
+  let didEmitFormStateMarkers = false;
+  if (formStateCount !== 0) {
+    // For each useFormState hook, emit a marker that indicates whether we
+    // rendered using the form state passed at the root.
+    // TODO: As an optimization, Fizz should only emit these markers if form
+    // state is passed at the root.
+    const segment = task.blockedSegment;
+    if (segment === null) {
+      // Implies we're in reumable mode.
+    } else {
+      didEmitFormStateMarkers = true;
+      const target = segment.chunks;
+      for (let i = 0; i < formStateCount; i++) {
+        if (i === formStateMatchingIndex) {
+          pushFormStateMarkerIsMatching(target);
+        } else {
+          pushFormStateMarkerIsNotMatching(target);
+        }
+      }
+    }
+  }
+
+  if (hasId) {
+    // This component materialized an id. We treat this as its own level, with
+    // a single "child" slot.
+    const prevTreeContext = task.treeContext;
+    const totalChildren = 1;
+    const index = 0;
+    // Modify the id context. Because we'll need to reset this if something
+    // suspends or errors, we'll use the non-destructive render path.
+    task.treeContext = pushTreeContext(prevTreeContext, totalChildren, index);
+    renderNode(request, task, children, 0);
+    // Like the other contexts, this does not need to be in a finally block
+    // because renderNode takes care of unwinding the stack.
+    task.treeContext = prevTreeContext;
+  } else if (didEmitFormStateMarkers) {
+    // If there were formState hooks, we must use the non-destructive path
+    // because this component is not a pure indirection; we emitted markers
+    // to the stream.
+    renderNode(request, task, children, 0);
+  } else {
+    // We're now successfully past this task, and we haven't modified the
+    // context stack. We don't have to pop back to the previous task every
+    // again, so we can use the destructive recursive form.
+    renderNodeDestructive(request, task, null, children, 0);
+  }
+}
+
 function validateFunctionComponentInDev(Component: any): void {
   if (__DEV__) {
     if (Component) {
@@ -1221,21 +1273,16 @@ function renderForwardRef(
     ref,
   );
   const hasId = checkDidRenderIdHook();
-  if (hasId) {
-    // This component materialized an id. We treat this as its own level, with
-    // a single "child" slot.
-    const prevTreeContext = task.treeContext;
-    const totalChildren = 1;
-    const index = 0;
-    // Modify the id context. Because we'll need to reset this if something
-    // suspends or errors, we'll use the non-destructive render path.
-    task.treeContext = pushTreeContext(prevTreeContext, totalChildren, index);
-    renderNode(request, task, children, 0);
-    // Like the other contexts, this does not need to be in a finally block
-    // because renderNode takes care of unwinding the stack.
-  } else {
-    renderNodeDestructive(request, task, null, children, 0);
-  }
+  const formStateCount = getFormStateCount();
+  const formStateMatchingIndex = getFormStateMatchingIndex();
+  finishFunctionComponent(
+    request,
+    task,
+    children,
+    hasId,
+    formStateCount,
+    formStateMatchingIndex,
+  );
   popComponentStackInDEV(task);
 }
 

packages/react-server/src/forks/ReactFizzConfig.custom.js

@@ -53,6 +53,10 @@ export const pushStartCompletedSuspenseBoundary =
 export const pushEndCompletedSuspenseBoundary =
   $$$config.pushEndCompletedSuspenseBoundary;
 export const pushSegmentFinale = $$$config.pushSegmentFinale;
+export const pushFormStateMarkerIsMatching =
+  $$$config.pushFormStateMarkerIsMatching;
+export const pushFormStateMarkerIsNotMatching =
+  $$$config.pushFormStateMarkerIsNotMatching;
 export const writeCompletedRoot = $$$config.writeCompletedRoot;
 export const writePlaceholder = $$$config.writePlaceholder;
 export const writeStartCompletedSuspenseBoundary =