useFormState: Allow sync actions

Updates useFormState to allow a sync function to be passed as an action.

A form action is almost always async, because it needs to talk to the
server. But since we support client-side actions, too, there's no reason
we can't allow sync actions, too.

I originally chose not to allow them to keep the implementation simpler
but it's not really that much more complicated because we already
support this for actions passed to startTransition. So now it's
consistent: anywhere an action is accepted, a sync client function is
a valid input.

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

@@ -1113,29 +1113,179 @@ describe('ReactDOMForm', () => {
 
   // @gate enableFormActions
   // @gate enableAsyncActions
-  test('useFormState: warns if action is not async', async () => {
-    let dispatch;
+  test('useFormState: works if action is sync', async () => {
+    let increment;
+    function App({stepSize}) {
+      const [state, dispatch] = useFormState(prevState => {
+        return prevState + stepSize;
+      }, 0);
+      increment = dispatch;
+      return <Text text={state} />;
+    }
+
+    // Initial render
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => root.render(<App stepSize={1} />));
+    assertLog([0]);
+
+    // Perform an action. This will increase the state by 1, as defined by the
+    // stepSize prop.
+    await act(() => increment());
+    assertLog([1]);
+
+    // Now increase the stepSize prop to 10. Subsequent steps will increase
+    // by this amount.
+    await act(() => root.render(<App stepSize={10} />));
+    assertLog([1]);
+
+    // Increment again. The state should increase by 10.
+    await act(() => increment());
+    assertLog([11]);
+  });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  test('useFormState: can mix sync and async actions', async () => {
+    let action;
     function App() {
-      const [state, _dispatch] = useFormState(() => {}, 0);
-      dispatch = _dispatch;
+      const [state, dispatch] = useFormState((s, a) => a, 'A');
+      action = dispatch;
       return <Text text={state} />;
     }
 
     const root = ReactDOMClient.createRoot(container);
-    await act(async () => {
-      root.render(<App />);
+    await act(() => root.render(<App />));
+    assertLog(['A']);
+
+    await act(() => action(getText('B')));
+    await act(() => action('C'));
+    await act(() => action(getText('D')));
+    await act(() => action('E'));
+
+    await act(() => resolveText('B'));
+    await act(() => resolveText('D'));
+    assertLog(['E']);
+    expect(container.textContent).toBe('E');
+  });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  test('useFormState: error handling (sync action)', async () => {
+    let resetErrorBoundary;
+    class ErrorBoundary extends React.Component {
+      state = {error: null};
+      static getDerivedStateFromError(error) {
+        return {error};
+      }
+      render() {
+        resetErrorBoundary = () => this.setState({error: null});
+        if (this.state.error !== null) {
+          return <Text text={'Caught an error: ' + this.state.error.message} />;
+        }
+        return this.props.children;
+      }
+    }
+
+    let action;
+    function App() {
+      const [state, dispatch] = useFormState((s, a) => {
+        if (a.endsWith('!')) {
+          throw new Error(a);
+        }
+        return a;
+      }, 'A');
+      action = dispatch;
+      return <Text text={state} />;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(() =>
+      root.render(
+        <ErrorBoundary>
+          <App />
+        </ErrorBoundary>,
+      ),
+    );
+    assertLog(['A']);
+
+    await act(() => action('Oops!'));
+    assertLog(['Caught an error: Oops!', 'Caught an error: Oops!']);
+    expect(container.textContent).toBe('Caught an error: Oops!');
+
+    // Reset the error boundary
+    await act(() => resetErrorBoundary());
+    assertLog(['A']);
+
+    // Trigger an error again, but this time, perform another action that
+    // overrides the first one and fixes the error
+    await act(() => {
+      action('Oops!');
+      action('B');
     });
-    assertLog([0]);
+    assertLog(['B']);
+    expect(container.textContent).toBe('B');
+  });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  test('useFormState: error handling (async action)', async () => {
+    let resetErrorBoundary;
+    class ErrorBoundary extends React.Component {
+      state = {error: null};
+      static getDerivedStateFromError(error) {
+        return {error};
+      }
+      render() {
+        resetErrorBoundary = () => this.setState({error: null});
+        if (this.state.error !== null) {
+          return <Text text={'Caught an error: ' + this.state.error.message} />;
+        }
+        return this.props.children;
+      }
+    }
+
+    let action;
+    function App() {
+      const [state, dispatch] = useFormState(async (s, a) => {
+        const text = await getText(a);
+        if (text.endsWith('!')) {
+          throw new Error(text);
+        }
+        return text;
+      }, 'A');
+      action = dispatch;
+      return <Text text={state} />;
+    }
 
-    expect(() => {
-      // This throws because React expects the action to return a promise.
-      expect(() => dispatch()).toThrow('Cannot read properties of undefined');
-    }).toErrorDev(
-      [
-        // In dev we also log a warning.
-        'The action passed to useFormState must be an async function',
-      ],
-      {withoutStack: true},
+    const root = ReactDOMClient.createRoot(container);
+    await act(() =>
+      root.render(
+        <ErrorBoundary>
+          <App />
+        </ErrorBoundary>,
+      ),
     );
+    assertLog(['A']);
+
+    await act(() => action('Oops!'));
+    assertLog([]);
+    await act(() => resolveText('Oops!'));
+    assertLog(['Caught an error: Oops!', 'Caught an error: Oops!']);
+    expect(container.textContent).toBe('Caught an error: Oops!');
+
+    // Reset the error boundary
+    await act(() => resetErrorBoundary());
+    assertLog(['A']);
+
+    // Trigger an error again, but this time, perform another action that
+    // overrides the first one and fixes the error
+    await act(() => {
+      action('Oops!');
+      action('B');
+    });
+    assertLog([]);
+    await act(() => resolveText('B'));
+    assertLog(['B']);
+    expect(container.textContent).toBe('B');
   });
 });

packages/react-reconciler/src/ReactFiberAsyncAction.js

@@ -34,7 +34,7 @@ let currentEntangledPendingCount: number = 0;
 let currentEntangledLane: Lane = NoLane;
 
 export function requestAsyncActionContext<S>(
-  actionReturnValue: Thenable<mixed>,
+  actionReturnValue: Thenable<any>,
   // If this is provided, this resulting thenable resolves to this value instead
   // of the return value of the action. This is a perf trick to avoid composing
   // an extra async function.
@@ -112,7 +112,7 @@ export function requestAsyncActionContext<S>(
 }
 
 export function requestSyncActionContext<S>(
-  actionReturnValue: mixed,
+  actionReturnValue: any,
   // If this is provided, this resulting thenable resolves to this value instead
   // of the return value of the action. This is a perf trick to avoid composing
   // an extra async function.

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1891,7 +1891,7 @@ type FormStateActionQueueNode<P> = {
 function dispatchFormState<S, P>(
   fiber: Fiber,
   actionQueue: FormStateActionQueue<S, P>,
-  setState: Dispatch<Thenable<S>>,
+  setState: Dispatch<S | Thenable<S>>,
   payload: P,
 ): void {
   if (isRenderPhaseUpdate(fiber)) {
@@ -1921,7 +1921,7 @@ function dispatchFormState<S, P>(
 
 function runFormStateAction<S, P>(
   actionQueue: FormStateActionQueue<S, P>,
-  setState: Dispatch<Thenable<S>>,
+  setState: Dispatch<S | Thenable<S>>,
   payload: P,
 ) {
   const action = actionQueue.action;
@@ -1935,39 +1935,49 @@ function runFormStateAction<S, P>(
     ReactCurrentBatchConfig.transition._updatedFibers = new Set();
   }
   try {
-    const promise = action(prevState, payload);
+    const returnValue = action(prevState, payload);
+    if (
+      returnValue !== null &&
+      typeof returnValue === 'object' &&
+      // $FlowFixMe[method-unbinding]
+      typeof returnValue.then === 'function'
+    ) {
+      const thenable = ((returnValue: any): Thenable<S>);
+
+      // Attach a listener to read the return state of the action. As soon as this
+      // resolves, we can run the next action in the sequence.
+      thenable.then(
+        (nextState: S) => {
+          actionQueue.state = nextState;
+          finishRunningFormStateAction(actionQueue, setState);
+        },
+        () => finishRunningFormStateAction(actionQueue, setState),
+      );
 
-    if (__DEV__) {
-      if (
-        promise === null ||
-        typeof promise !== 'object' ||
-        typeof (promise: any).then !== 'function'
-      ) {
-        console.error(
-          'The action passed to useFormState must be an async function.',
-        );
-      }
+      const entangledResult = requestAsyncActionContext<S>(thenable, null);
+      setState(entangledResult);
+    } else {
+      // This is either `finishedState` or a thenable that resolves to
+      // `finishedState`, depending on whether we're inside an async
+      // action scope.
+      const entangledResult = requestSyncActionContext<S>(returnValue, null);
+      setState(entangledResult);
+
+      const nextState = ((returnValue: any): S);
+      actionQueue.state = nextState;
+      finishRunningFormStateAction(actionQueue, setState);
     }
-
-    // Attach a listener to read the return state of the action. As soon as this
-    // resolves, we can run the next action in the sequence.
-    promise.then(
-      (nextState: S) => {
-        actionQueue.state = nextState;
-        finishRunningFormStateAction(actionQueue, setState);
-      },
-      () => finishRunningFormStateAction(actionQueue, setState),
-    );
-
-    // Create a thenable that resolves once the current async action scope has
-    // finished. Then stash that thenable in state. We'll unwrap it with the
-    // `use` algorithm during render. This is the same logic used
-    // by startTransition.
-    const entangledThenable: Thenable<S> = requestAsyncActionContext(
-      promise,
-      null,
-    );
-    setState(entangledThenable);
+  } catch (error) {
+    // This is a trick to get the `useFormState` hook to rethrow the error.
+    // When it unwraps the thenable with the `use` algorithm, the error
+    // will be thrown.
+    const rejectedThenable: RejectedThenable<S> = {
+      then() {},
+      status: 'rejected',
+      reason: error,
+    };
+    setState(rejectedThenable);
+    finishRunningFormStateAction(actionQueue, setState);
   } finally {
     ReactCurrentBatchConfig.transition = prevTransition;
 
@@ -1989,7 +1999,7 @@ function runFormStateAction<S, P>(
 
 function finishRunningFormStateAction<S, P>(
   actionQueue: FormStateActionQueue<S, P>,
-  setState: Dispatch<Thenable<S>>,
+  setState: Dispatch<S | Thenable<S>>,
 ) {
   // The action finished running. Pop it from the queue and run the next pending
   // action, if there are any.
@@ -2035,25 +2045,20 @@ function mountFormState<S, P>(
       }
     }
   }
-  const initialStateThenable: Thenable<S> = {
-    status: 'fulfilled',
-    value: initialState,
-    then() {},
-  };
 
   // State hook. The state is stored in a thenable which is then unwrapped by
   // the `use` algorithm during render.
   const stateHook = mountWorkInProgressHook();
-  stateHook.memoizedState = stateHook.baseState = initialStateThenable;
-  const stateQueue: UpdateQueue<Thenable<S>, Thenable<S>> = {
+  stateHook.memoizedState = stateHook.baseState = initialState;
+  const stateQueue: UpdateQueue<S | Thenable<S>, S | Thenable<S>> = {
     pending: null,
     lanes: NoLanes,
     dispatch: null,
     lastRenderedReducer: formStateReducer,
-    lastRenderedState: initialStateThenable,
+    lastRenderedState: initialState,
   };
   stateHook.queue = stateQueue;
-  const setState: Dispatch<Thenable<S>> = (dispatchSetState.bind(
+  const setState: Dispatch<S | Thenable<S>> = (dispatchSetState.bind(
     null,
     currentlyRenderingFiber,
     stateQueue,
@@ -2111,14 +2116,20 @@ function updateFormStateImpl<S, P>(
   initialState: S,
   permalink?: string,
 ): [S, (P) => void] {
-  const [thenable] = updateReducerImpl<Thenable<S>, Thenable<S>>(
+  const [actionResult] = updateReducerImpl<S | Thenable<S>, S | Thenable<S>>(
     stateHook,
     currentStateHook,
     formStateReducer,
   );
 
   // This will suspend until the action finishes.
-  const state = useThenable(thenable);
+  const state: S =
+    typeof actionResult === 'object' &&
+    actionResult !== null &&
+    // $FlowFixMe[method-unbinding]
+    typeof actionResult.then === 'function'
+      ? useThenable(((actionResult: any): Thenable<S>))
+      : (actionResult: any);
 
   const actionQueueHook = updateWorkInProgressHook();
   const actionQueue = actionQueueHook.queue;
@@ -2173,8 +2184,7 @@ function rerenderFormState<S, P>(
   }
 
   // This is a mount. No updates to process.
-  const thenable: Thenable<S> = stateHook.memoizedState;
-  const state = useThenable(thenable);
+  const state: S = stateHook.memoizedState;
 
   const actionQueueHook = updateWorkInProgressHook();
   const actionQueue = actionQueueHook.queue;