Compare action signatures when reusing form state

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} />
}
```

Author: acdlite

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 difference 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 () => {