useFormState: Emit comment to mark whether state matches

A planned feature of useFormState is that if 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 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.

Author: acdlite

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

@@ -160,7 +160,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 +192,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 +1290,26 @@ export function registerSuspenseInstanceRetry(
   instance._reactRetry = callback;
 }
 
+export function canHydrateFormStateMarker(
+  instance: HydratableInstance,
+): null | FormStateMarkerInstance {
+  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 +1322,9 @@ function getNextHydratable(node: ?Node) {
       if (
         nodeData === SUSPENSE_START_DATA ||
         nodeData === SUSPENSE_FALLBACK_START_DATA ||
-        nodeData === SUSPENSE_PENDING_START_DATA
+        nodeData === SUSPENSE_PENDING_START_DATA ||
+        nodeData === FORM_STATE_IS_MATCHING ||
+        nodeData === FORM_STATE_IS_NOT_MATCHING
       ) {
         break;
       }

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

@@ -1525,6 +1525,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

@@ -28,6 +28,7 @@ let SuspenseList;
 let useSyncExternalStore;
 let useSyncExternalStoreWithSelector;
 let use;
+let useFormState;
 let PropTypes;
 let textCache;
 let writable;
@@ -83,6 +84,7 @@ describe('ReactDOMFizzServer', () => {
     if (gate(flags => flags.enableSuspenseList)) {
       SuspenseList = React.unstable_SuspenseList;
     }
+    useFormState = ReactDOM.experimental_useFormState;
 
     PropTypes = require('prop-types');
 
@@ -5911,6 +5913,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,31 @@ function tryToClaimNextHydratableSuspenseInstance(fiber: Fiber): void {
   }
 }
 
+export function tryToClaimNextHydratableFormMarkerInstance(
+  fiber: Fiber,
+): boolean {
+  if (!isHydrating) {
+    return false;
+  }
+  if (nextHydratableInstance) {
+    const markerInstance = canHydrateFormStateMarker(nextHydratableInstance);
+    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 =