Fix form state reset when component state is updated

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

@@ -2281,4 +2281,73 @@ describe('ReactDOMForm', () => {
     await submit(formRef.current);
     assertLog(['stringified action']);
   });
+
+  it('form actions should retain status when nested state changes', async () => {
+    const formRef = React.createRef();
+    const textRef = React.createRef();
+    const btnRef = React.createRef();
+
+    let rerenderUnrelatedStatus;
+    function UnrelatedStatus() {
+      const {pending} = useFormStatus();
+      const [counter, setCounter] = useState(0);
+      rerenderUnrelatedStatus = () => setCounter(n => n + 1);
+      Scheduler.log(`[unrelated form] pending: ${pending}, state: ${counter}`);
+    }
+
+    function Status() {
+      const {pending} = useFormStatus();
+      const [counter, setCounter] = useState(0);
+      return (
+        <div>
+          <p ref={textRef}>
+            {pending ? `Pending` : null} with state: {counter}
+          </p>
+          <button
+            ref={btnRef}
+            onClick={() => setCounter(n => n + 1)}
+            type="button">
+            Increment
+          </button>
+        </div>
+      );
+    }
+
+    function App() {
+      async function action() {
+        return new Promise(resolve => {
+          // never resolves
+        });
+      }
+
+      return (
+        <>
+          <form action={action} ref={formRef}>
+            <input type="submit" />
+            <Status />
+          </form>
+          <form>
+            <UnrelatedStatus />
+          </form>
+        </>
+      );
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => root.render(<App />));
+
+    assertLog(['[unrelated form] pending: false, state: 0']);
+
+    await submit(formRef.current);
+
+    expect(textRef.current.textContent).toBe('Pending with state: 0');
+
+    await act(() => btnRef.current.click());
+
+    expect(textRef.current.textContent).toBe('Pending with state: 1');
+
+    await act(() => rerenderUnrelatedStatus());
+
+    assertLog(['[unrelated form] pending: false, state: 1']);
+  });
 });

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -181,7 +181,6 @@ import {
   isPrimaryRenderer,
   getResource,
   createHoistableInstance,
-  HostTransitionContext,
 } from './ReactFiberConfig';
 import type {ActivityInstance, SuspenseInstance} from './ReactFiberConfig';
 import {shouldError, shouldSuspend} from './ReactFiberReconciler';
@@ -1870,8 +1869,6 @@ function updateHostComponent(
     tryToClaimNextHydratableInstance(workInProgress);
   }
 
-  pushHostContext(workInProgress);
-
   const type = workInProgress.type;
   const nextProps = workInProgress.pendingProps;
   const prevProps = current !== null ? current.memoizedProps : null;
@@ -1899,31 +1896,15 @@ function updateHostComponent(
     //
     // Once a fiber is upgraded to be stateful, it remains stateful for the
     // rest of its lifetime.
-    const newState = renderTransitionAwareHostComponentWithHooks(
+    renderTransitionAwareHostComponentWithHooks(
       current,
       workInProgress,
       renderLanes,
     );
-
-    // If the transition state changed, propagate the change to all the
-    // descendents. We use Context as an implementation detail for this.
-    //
-    // This is intentionally set here instead of pushHostContext because
-    // pushHostContext gets called before we process the state hook, to avoid
-    // a state mismatch in the event that something suspends.
-    //
-    // NOTE: This assumes that there cannot be nested transition providers,
-    // because the only renderer that implements this feature is React DOM,
-    // and forms cannot be nested. If we did support nested providers, then
-    // we would need to push a context value even for host fibers that
-    // haven't been upgraded yet.
-    if (isPrimaryRenderer) {
-      HostTransitionContext._currentValue = newState;
-    } else {
-      HostTransitionContext._currentValue2 = newState;
-    }
   }
 
+  pushHostContext(workInProgress);
+
   markRef(current, workInProgress);
   reconcileChildren(current, workInProgress, nextChildren, renderLanes);
   return workInProgress.child;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -853,8 +853,8 @@ export function renderTransitionAwareHostComponentWithHooks(
   current: Fiber | null,
   workInProgress: Fiber,
   lanes: Lanes,
-): TransitionStatus {
-  return renderWithHooks(
+): void {
+  renderWithHooks(
     current,
     workInProgress,
     TransitionAwareHostComponent,

packages/react-reconciler/src/ReactFiberHostContext.js

@@ -9,7 +9,11 @@
 
 import type {Fiber} from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack';
-import type {Container, HostContext} from './ReactFiberConfig';
+import type {
+  Container,
+  HostContext,
+  TransitionStatus,
+} from './ReactFiberConfig';
 import type {Hook} from './ReactFiberHooks';
 
 import {
@@ -92,6 +96,21 @@ function getHostContext(): HostContext {
 function pushHostContext(fiber: Fiber): void {
   const stateHook: Hook | null = fiber.memoizedState;
   if (stateHook !== null) {
+    // Propagate the state to all the descendents.
+    // We use Context as an implementation detail for this.
+    //
+    // NOTE: This assumes that there cannot be nested transition providers,
+    // because the only renderer that implements this feature is React DOM,
+    // and forms cannot be nested. If we did support nested providers, then
+    // we would need to push a context value even for host fibers that
+    // haven't been upgraded yet.
+    const transitionStatus: TransitionStatus = stateHook.memoizedState;
+    if (isPrimaryRenderer) {
+      HostTransitionContext._currentValue = transitionStatus;
+    } else {
+      HostTransitionContext._currentValue2 = transitionStatus;
+    }
+
     // Only provide context if this fiber has been upgraded by a host
     // transition. We use the same optimization for regular host context below.
     push(hostTransitionProviderCursor, fiber, fiber);