useFormState: Compare action signatures when reusing form state (#27370)

During an MPA form submission, useFormState should only reuse the form
state if same action is passed both times. (We also compare the key
paths.)

We compare the identity of the inner closure function, disregarding the
value of the bound arguments. That way you can pass an inline Server
Action closure:

```js
function FormContainer({maxLength}) {
  function submitAction(prevState, formData) {
    'use server'
    if (formData.get('field').length > maxLength) {
      return { errorMsg: 'Too many characters' };
    }
    // ...
  }
  return <Form submitAction={submitAction} />
}
```

packages/react-client/src/ReactFlightReplyClient.js

@@ -9,6 +9,7 @@
 
 import type {
   Thenable,
+  PendingThenable,
   FulfilledThenable,
   RejectedThenable,
   ReactCustomFormAction,
@@ -489,6 +490,69 @@ export function encodeFormAction(
   };
 }
 
+function isSignatureEqual(
+  this: any => Promise<any>,
+  referenceId: ServerReferenceId,
+  numberOfBoundArgs: number,
+): boolean {
+  const reference = knownServerReferences.get(this);
+  if (!reference) {
+    throw new Error(
+      'Tried to encode a Server Action from a different instance than the encoder is from. ' +
+        'This is a bug in React.',
+    );
+  }
+  if (reference.id !== referenceId) {
+    // These are different functions.
+    return false;
+  }
+  // Now check if the number of bound arguments is the same.
+  const boundPromise = reference.bound;
+  if (boundPromise === null) {
+    // No bound arguments.
+    return numberOfBoundArgs === 0;
+  }
+  // Unwrap the bound arguments array by suspending, if necessary. As with
+  // encodeFormData, this means isSignatureEqual can only be called while React
+  // is rendering.
+  switch (boundPromise.status) {
+    case 'fulfilled': {
+      const boundArgs = boundPromise.value;
+      return boundArgs.length === numberOfBoundArgs;
+    }
+    case 'pending': {
+      throw boundPromise;
+    }
+    case 'rejected': {
+      throw boundPromise.reason;
+    }
+    default: {
+      if (typeof boundPromise.status === 'string') {
+        // Only instrument the thenable if the status if not defined.
+      } else {
+        const pendingThenable: PendingThenable<Array<any>> =
+          (boundPromise: any);
+        pendingThenable.status = 'pending';
+        pendingThenable.then(
+          (boundArgs: Array<any>) => {
+            const fulfilledThenable: FulfilledThenable<Array<any>> =
+              (boundPromise: any);
+            fulfilledThenable.status = 'fulfilled';
+            fulfilledThenable.value = boundArgs;
+          },
+          (error: mixed) => {
+            const rejectedThenable: RejectedThenable<number> =
+              (boundPromise: any);
+            rejectedThenable.status = 'rejected';
+            rejectedThenable.reason = error;
+          },
+        );
+      }
+      throw boundPromise;
+    }
+  }
+}
+
 export function registerServerReference(
   proxy: any,
   reference: {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
@@ -499,6 +563,7 @@ export function registerServerReference(
     // Only expose this in builds that would actually use it. Not needed on the client.
     Object.defineProperties((proxy: any), {
       $$FORM_ACTION: {value: encodeFormAction},
+      $$IS_SIGNATURE_EQUAL: {value: isSignatureEqual},
       bind: {value: bind},
     });
   }

packages/react-dom/src/client/ReactDOMRoot.js

@@ -57,7 +57,7 @@ export type HydrateRootOptions = {
   unstable_transitionCallbacks?: TransitionTracingCallbacks,
   identifierPrefix?: string,
   onRecoverableError?: (error: mixed) => void,
-  experimental_formState?: ReactFormState<any> | null,
+  experimental_formState?: ReactFormState<any, any> | null,
   ...
 };
 

packages/react-dom/src/server/ReactDOMFizzServerBrowser.js

@@ -41,7 +41,7 @@ type Options = {
   onPostpone?: (reason: string) => void,
   unstable_externalRuntimeSrc?: string | BootstrapScriptDescriptor,
   importMap?: ImportMap,
-  experimental_formState?: ReactFormState<any> | null,
+  experimental_formState?: ReactFormState<any, any> | null,
 };
 
 type ResumeOptions = {

packages/react-dom/src/server/ReactDOMFizzServerBun.js

@@ -39,7 +39,7 @@ type Options = {
   onPostpone?: (reason: string) => void,
   unstable_externalRuntimeSrc?: string | BootstrapScriptDescriptor,
   importMap?: ImportMap,
-  experimental_formState?: ReactFormState<any> | null,
+  experimental_formState?: ReactFormState<any, any> | null,
 };
 
 // TODO: Move to sub-classing ReadableStream.

packages/react-dom/src/server/ReactDOMFizzServerEdge.js

@@ -41,7 +41,7 @@ type Options = {
   onPostpone?: (reason: string) => void,
   unstable_externalRuntimeSrc?: string | BootstrapScriptDescriptor,
   importMap?: ImportMap,
-  experimental_formState?: ReactFormState<any> | null,
+  experimental_formState?: ReactFormState<any, any> | null,
 };
 
 type ResumeOptions = {

packages/react-dom/src/server/ReactDOMFizzServerNode.js

@@ -54,7 +54,7 @@ type Options = {
   onPostpone?: (reason: string) => void,
   unstable_externalRuntimeSrc?: string | BootstrapScriptDescriptor,
   importMap?: ImportMap,
-  experimental_formState?: ReactFormState<any> | null,
+  experimental_formState?: ReactFormState<any, any> | null,
 };
 
 type ResumeOptions = {

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -281,7 +281,7 @@ export function createHydrationContainer(
   identifierPrefix: string,
   onRecoverableError: (error: mixed) => void,
   transitionCallbacks: null | TransitionTracingCallbacks,
-  formState: ReactFormState<any> | null,
+  formState: ReactFormState<any, any> | null,
 ): OpaqueRoot {
   const hydrate = true;
   const root = createFiberRoot(

packages/react-reconciler/src/ReactFiberRoot.js

@@ -52,7 +52,7 @@ function FiberRootNode(
   hydrate: any,
   identifierPrefix: any,
   onRecoverableError: any,
-  formState: ReactFormState<any> | null,
+  formState: ReactFormState<any, any> | null,
 ) {
   this.tag = tag;
   this.containerInfo = containerInfo;
@@ -145,7 +145,7 @@ export function createFiberRoot(
   identifierPrefix: string,
   onRecoverableError: null | ((error: mixed) => void),
   transitionCallbacks: null | TransitionTracingCallbacks,
-  formState: ReactFormState<any> | null,
+  formState: ReactFormState<any, any> | null,
 ): FiberRoot {
   // $FlowFixMe[invalid-constructor] Flow no longer supports calling new on functions
   const root: FiberRoot = (new FiberRootNode(

packages/react-reconciler/src/ReactInternalTypes.js

@@ -272,7 +272,7 @@ type BaseFiberRootProperties = {
     errorInfo: {digest?: ?string, componentStack?: ?string},
   ) => void,
 
-  formState: ReactFormState<any> | null,
+  formState: ReactFormState<any, any> | null,
 };
 
 // The following attributes are only used by DevTools and are only present in DEV builds.

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

@@ -67,7 +67,7 @@ describe('ReactFlightDOMForm', () => {
       webpackServerMap,
     );
     const returnValue = boundAction();
-    const formState = ReactServerDOMServer.decodeFormState(
+    const formState = await ReactServerDOMServer.decodeFormState(
       await returnValue,
       formData,
       webpackServerMap,
@@ -435,6 +435,174 @@ describe('ReactFlightDOMForm', () => {
     }
   });
 
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  it(
+    'useFormState preserves state if arity is the same, but different ' +
+      'arguments are bound (i.e. inline closure)',
+    async () => {
+      const serverAction = serverExports(async function action(
+        stepSize,
+        prevState,
+        formData,
+      ) {
+        return prevState + stepSize;
+      });
+
+      function Form({action}) {
+        const [count, dispatch] = useFormState(action, 1);
+        return <form action={dispatch}>{count}</form>;
+      }
+
+      function Client({action}) {
+        return (
+          <div>
+            <Form action={action} />
+            <Form action={action} />
+            <Form action={action} />
+          </div>
+        );
+      }
+
+      const ClientRef = await clientExports(Client);
+
+      const rscStream = ReactServerDOMServer.renderToReadableStream(
+        // Note: `.bind` is the same as an inline closure with 'use server'
+        <ClientRef action={serverAction.bind(null, 1)} />,
+        webpackMap,
+      );
+      const response = ReactServerDOMClient.createFromReadableStream(rscStream);
+      const ssrStream = await ReactDOMServer.renderToReadableStream(response);
+      await readIntoContainer(ssrStream);
+
+      expect(container.textContent).toBe('111');
+
+      // There are three identical forms. We're going to submit the second one.
+      const form = container.getElementsByTagName('form')[1];
+      const {formState} = await submit(form);
+
+      // Simulate an MPA form submission by resetting the container and
+      // rendering again.
+      container.innerHTML = '';
+
+      // On the next page, the same server action is rendered again, but with
+      // a different bound stepSize argument. We should treat this as the same
+      // action signature.
+      const postbackRscStream = ReactServerDOMServer.renderToReadableStream(
+        // Note: `.bind` is the same as an inline closure with 'use server'
+        <ClientRef action={serverAction.bind(null, 5)} />,
+        webpackMap,
+      );
+      const postbackResponse =
+        ReactServerDOMClient.createFromReadableStream(postbackRscStream);
+      const postbackSsrStream = await ReactDOMServer.renderToReadableStream(
+        postbackResponse,
+        {experimental_formState: formState},
+      );
+      await readIntoContainer(postbackSsrStream);
+
+      // The state should have been preserved because the action signatures are
+      // the same. (Note that the amount increased by 1, because that was the
+      // value of stepSize at the time the form was submitted)
+      expect(container.textContent).toBe('121');
+
+      // Now submit the form again. This time, the state should increase by 5
+      // because the stepSize argument has changed.
+      const form2 = container.getElementsByTagName('form')[1];
+      const {formState: formState2} = await submit(form2);
+
+      container.innerHTML = '';
+
+      const postbackRscStream2 = ReactServerDOMServer.renderToReadableStream(
+        // Note: `.bind` is the same as an inline closure with 'use server'
+        <ClientRef action={serverAction.bind(null, 5)} />,
+        webpackMap,
+      );
+      const postbackResponse2 =
+        ReactServerDOMClient.createFromReadableStream(postbackRscStream2);
+      const postbackSsrStream2 = await ReactDOMServer.renderToReadableStream(
+        postbackResponse2,
+        {experimental_formState: formState2},
+      );
+      await readIntoContainer(postbackSsrStream2);
+
+      expect(container.textContent).toBe('171');
+    },
+  );
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  it('useFormState does not reuse state if action signatures are different', async () => {
+    // This is the same as the previous test, except instead of using bind to
+    // configure the server action (i.e. a closure), it swaps the action.
+    const increaseBy1 = serverExports(async function action(
+      prevState,
+      formData,
+    ) {
+      return prevState + 1;
+    });
+
+    const increaseBy5 = serverExports(async function action(
+      prevState,
+      formData,
+    ) {
+      return prevState + 5;
+    });
+
+    function Form({action}) {
+      const [count, dispatch] = useFormState(action, 1);
+      return <form action={dispatch}>{count}</form>;
+    }
+
+    function Client({action}) {
+      return (
+        <div>
+          <Form action={action} />
+          <Form action={action} />
+          <Form action={action} />
+        </div>
+      );
+    }
+
+    const ClientRef = await clientExports(Client);
+
+    const rscStream = ReactServerDOMServer.renderToReadableStream(
+      <ClientRef action={increaseBy1} />,
+      webpackMap,
+    );
+    const response = ReactServerDOMClient.createFromReadableStream(rscStream);
+    const ssrStream = await ReactDOMServer.renderToReadableStream(response);
+    await readIntoContainer(ssrStream);
+
+    expect(container.textContent).toBe('111');
+
+    // There are three identical forms. We're going to submit the second one.
+    const form = container.getElementsByTagName('form')[1];
+    const {formState} = await submit(form);
+
+    // Simulate an MPA form submission by resetting the container and
+    // rendering again.
+    container.innerHTML = '';
+
+    // On the next page, a different server action is rendered. It should not
+    // reuse the state from the previous page.
+    const postbackRscStream = ReactServerDOMServer.renderToReadableStream(
+      <ClientRef action={increaseBy5} />,
+      webpackMap,
+    );
+    const postbackResponse =
+      ReactServerDOMClient.createFromReadableStream(postbackRscStream);
+    const postbackSsrStream = await ReactDOMServer.renderToReadableStream(
+      postbackResponse,
+      {experimental_formState: formState},
+    );
+    await readIntoContainer(postbackSsrStream);
+
+    // The state should not have been preserved because the action signatures
+    // are not the same.
+    expect(container.textContent).toBe('111');
+  });
+
   // @gate enableFormActions
   // @gate enableAsyncActions
   it('useFormState can change the action URL with the `permalink` argument', async () => {

packages/react-server/src/ReactFizzHooks.js

@@ -599,70 +599,83 @@ function useFormState<S, P>(
   const formStateHookIndex = formStateCounter++;
   const request: Request = (currentlyRenderingRequest: any);
 
-  // Append a node to the key path that represents the form state hook.
-  const componentKey: KeyNode | null = (currentlyRenderingKeyPath: any);
-  const key: KeyNode = [componentKey, null, formStateHookIndex];
-  const keyJSON = JSON.stringify(key);
-
-  // Get the form state. If we received form state from a previous page, then
-  // we should reuse that, if the action identity matches. Otherwise we'll use
-  // the initial state argument. We emit a comment marker into the stream
-  // that indicates whether the state was reused.
-  let state;
-  const postbackFormState = getFormState(request);
-  if (postbackFormState !== null) {
-    const postbackKey = postbackFormState[1];
-    // TODO: Compare the action identity, too
-    // TODO: If a permalink is used, disregard the key and compare that instead.
-    if (keyJSON === postbackKey) {
-      // This was a match.
-      formStateMatchingIndex = formStateHookIndex;
-      // Reuse the state that was submitted by the form.
-      state = postbackFormState[0];
-    } else {
-      state = initialState;
-    }
-  } else {
-    // TODO: As an optimization, Fizz should only emit these markers if form
-    // state is passed at the root.
-    state = initialState;
-  }
-
-  // Bind the state to the first argument of the action.
-  const boundAction = action.bind(null, state);
-
-  // Wrap the action so the return value is void.
-  const dispatch = (payload: P): void => {
-    boundAction(payload);
-  };
-
   // $FlowIgnore[prop-missing]
-  if (typeof boundAction.$$FORM_ACTION === 'function') {
+  const formAction = action.$$FORM_ACTION;
+  if (typeof formAction === 'function') {
+    // This is a server action. These have additional features to enable
+    // MPA-style form submissions with progressive enhancement.
+
+    // Determine the current form state. If we received state during an MPA form
+    // submission, then we will reuse that, if the action identity matches.
+    // Otherwise we'll use the initial state argument. We will emit a comment
+    // marker into the stream that indicates whether the state was reused.
+    let state = initialState;
+
+    // Append a node to the key path that represents the form state hook.
+    const componentKey: KeyNode | null = (currentlyRenderingKeyPath: any);
+    const key: KeyNode = [componentKey, null, formStateHookIndex];
+    const keyJSON = JSON.stringify(key);
+
+    const postbackFormState = getFormState(request);
     // $FlowIgnore[prop-missing]
-    dispatch.$$FORM_ACTION = (prefix: string) => {
-      // $FlowIgnore[prop-missing]
-      const metadata: ReactCustomFormAction = boundAction.$$FORM_ACTION(prefix);
-
-      const formData = metadata.data;
-      if (formData) {
-        formData.append('$ACTION_KEY', keyJSON);
+    const isSignatureEqual = action.$$IS_SIGNATURE_EQUAL;
+    if (postbackFormState !== null && typeof isSignatureEqual === 'function') {
+      const postbackKeyJSON = postbackFormState[1];
+      const postbackReferenceId = postbackFormState[2];
+      const postbackBoundArity = postbackFormState[3];
+      if (
+        postbackKeyJSON === keyJSON &&
+        isSignatureEqual.call(action, postbackReferenceId, postbackBoundArity)
+      ) {
+        // This was a match
+        formStateMatchingIndex = formStateHookIndex;
+        // Reuse the state that was submitted by the form.
+        state = postbackFormState[0];
       }
+    }
 
-      // Override the action URL
-      if (permalink !== undefined) {
-        if (__DEV__) {
-          checkAttributeStringCoercion(permalink, 'target');
-        }
-        metadata.action = permalink + '';
-      }
-      return metadata;
+    // Bind the state to the first argument of the action.
+    const boundAction = action.bind(null, state);
+
+    // Wrap the action so the return value is void.
+    const dispatch = (payload: P): void => {
+      boundAction(payload);
     };
-  } else {
-    // This is not a server action, so the permalink argument has
-    // no effect. The form will have to be hydrated before it's submitted.
-  }
 
-  return [state, dispatch];
+    // $FlowIgnore[prop-missing]
+    if (typeof boundAction.$$FORM_ACTION === 'function') {
+      // $FlowIgnore[prop-missing]
+      dispatch.$$FORM_ACTION = (prefix: string) => {
+        const metadata: ReactCustomFormAction =
+          boundAction.$$FORM_ACTION(prefix);
+        const formData = metadata.data;
+        if (formData) {
+          formData.append('$ACTION_KEY', keyJSON);
+        }
+
+        // Override the action URL
+        if (permalink !== undefined) {
+          if (__DEV__) {
+            checkAttributeStringCoercion(permalink, 'target');
+          }
+          metadata.action = permalink + '';
+        }
+        return metadata;
+      };
+    }
+
+    return [state, dispatch];
+  } else {
+    // This is not a server action, so the implementation is much simpler.
+
+    // Bind the state to the first argument of the action.
+    const boundAction = action.bind(null, initialState);
+    // Wrap the action so the return value is void.
+    const dispatch = (payload: P): void => {
+      boundAction(payload);
+    };
+    return [initialState, dispatch];
+  }
 }
 
 function useId(): string {

packages/react-server/src/ReactFizzServer.js

@@ -313,7 +313,7 @@ export opaque type Request = {
   // rendering - e.g. to the client. This is considered intentional and not an error.
   onPostpone: (reason: string) => void,
   // Form state that was the result of an MPA submission, if it was provided.
-  formState: null | ReactFormState<any>,
+  formState: null | ReactFormState<any, any>,
 };
 
 // This is a default heuristic for how to split up the HTML content into progressive
@@ -352,7 +352,7 @@ export function createRequest(
   onShellError: void | ((error: mixed) => void),
   onFatalError: void | ((error: mixed) => void),
   onPostpone: void | ((reason: string) => void),
-  formState: void | null | ReactFormState<any>,
+  formState: void | null | ReactFormState<any, any>,
 ): Request {
   prepareHostDispatcher();
   const pingedTasks: Array<Task> = [];
@@ -3095,7 +3095,9 @@ export function flushResources(request: Request): void {
   enqueueFlush(request);
 }
 
-export function getFormState(request: Request): ReactFormState<any> | null {
+export function getFormState(
+  request: Request,
+): ReactFormState<any, any> | null {
   return request.formState;
 }
 

packages/react-server/src/ReactFlightActionServer.js

@@ -53,6 +53,28 @@ function loadServerReference<T>(
   }
 }
 
+function decodeBoundActionMetaData(
+  body: FormData,
+  serverManifest: ServerManifest,
+  formFieldPrefix: string,
+): {id: ServerReferenceId, bound: null | Promise<Array<any>>} {
+  // The data for this reference is encoded in multiple fields under this prefix.
+  const actionResponse = createResponse(serverManifest, formFieldPrefix, body);
+  close(actionResponse);
+  const refPromise = getRoot<{
+    id: ServerReferenceId,
+    bound: null | Promise<Array<any>>,
+  }>(actionResponse);
+  // Force it to initialize
+  // $FlowFixMe
+  refPromise.then(() => {});
+  if (refPromise.status !== 'fulfilled') {
+    // $FlowFixMe
+    throw refPromise.reason;
+  }
+  return refPromise.value;
+}
+
 export function decodeAction<T>(
   body: FormData,
   serverManifest: ServerManifest,
@@ -73,25 +95,11 @@ export function decodeAction<T>(
     // form action.
     if (key.startsWith('$ACTION_REF_')) {
       const formFieldPrefix = '$ACTION_' + key.slice(12) + ':';
-      // The data for this reference is encoded in multiple fields under this prefix.
-      const actionResponse = createResponse(
+      const metaData = decodeBoundActionMetaData(
+        body,
         serverManifest,
         formFieldPrefix,
-        body,
       );
-      close(actionResponse);
-      const refPromise = getRoot<{
-        id: ServerReferenceId,
-        bound: null | Promise<Array<any>>,
-      }>(actionResponse);
-      // Force it to initialize
-      // $FlowFixMe
-      refPromise.then(() => {});
-      if (refPromise.status !== 'fulfilled') {
-        // $FlowFixMe
-        throw refPromise.reason;
-      }
-      const metaData = refPromise.value;
       action = loadServerReference(serverManifest, metaData.id, metaData.bound);
       return;
     }
@@ -109,17 +117,47 @@ export function decodeAction<T>(
   return action.then(fn => fn.bind(null, formData));
 }
 
-// TODO: Should this be an async function to preserve the option in the future
-// to do async stuff in here? Would also make it consistent with decodeAction
 export function decodeFormState<S>(
   actionResult: S,
   body: FormData,
   serverManifest: ServerManifest,
-): ReactFormState<S> | null {
+): Promise<ReactFormState<S, ServerReferenceId> | null> {
   const keyPath = body.get('$ACTION_KEY');
   if (typeof keyPath !== 'string') {
     // This form submission did not include any form state.
-    return null;
+    return Promise.resolve(null);
   }
-  return [actionResult, keyPath];
+  // Search through the form data object to get the reference id and the number
+  // of bound arguments. This repeats some of the work done in decodeAction.
+  let metaData = null;
+  // $FlowFixMe[prop-missing]
+  body.forEach((value: string | File, key: string) => {
+    if (key.startsWith('$ACTION_REF_')) {
+      const formFieldPrefix = '$ACTION_' + key.slice(12) + ':';
+      metaData = decodeBoundActionMetaData(
+        body,
+        serverManifest,
+        formFieldPrefix,
+      );
+    }
+    // We don't check for the simple $ACTION_ID_ case because form state actions
+    // are always bound to the state argument.
+  });
+  if (metaData === null) {
+    // Should be unreachable.
+    return Promise.resolve(null);
+  }
+  const referenceId = metaData.id;
+  return Promise.resolve(metaData.bound).then(bound => {
+    if (bound === null) {
+      // Should be unreachable because form state actions are always bound to the
+      // state argument.
+      return null;
+    }
+    // The form action dispatch method is always bound to the initial state.
+    // But when comparing signatures, we compare to the original unbound action.
+    // Subtract one from the arity to account for this.
+    const boundArity = bound.length - 1;
+    return [actionResult, keyPath, referenceId, boundArity];
+  });
 }

packages/shared/ReactTypes.js

@@ -178,7 +178,9 @@ export type ReactCustomFormAction = {
 // This is an opaque type returned by decodeFormState on the server, but it's
 // defined in this shared file because the same type is used by React on
 // the client.
-export type ReactFormState<S> = [
+export type ReactFormState<S, ReferenceId> = [
   S /* actual state value */,
   string /* key path */,
+  ReferenceId /* Server Reference ID */,
+  number /* number of bound arguments */,
 ];