Bind startTransition to queue, not setPending

Small refactor to startTransition so that it calls setState directly
instead of a bound setState method passed to it from the caller.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1880,9 +1880,7 @@ function forceStoreRerender(fiber: Fiber) {
   }
 }
 
-function mountState<S>(
-  initialState: (() => S) | S,
-): [S, Dispatch<BasicStateAction<S>>] {
+function mountStateImpl<S>(initialState: (() => S) | S): Hook {
   const hook = mountWorkInProgressHook();
   if (typeof initialState === 'function') {
     // $FlowFixMe[incompatible-use]: Flow doesn't like mixed types
@@ -1897,9 +1895,20 @@ function mountState<S>(
     lastRenderedState: (initialState: any),
   };
   hook.queue = queue;
-  const dispatch: Dispatch<BasicStateAction<S>> = (queue.dispatch =
-    (dispatchSetState.bind(null, currentlyRenderingFiber, queue): any));
-  return [hook.memoizedState, dispatch];
+  const dispatch: Dispatch<BasicStateAction<S>> = (dispatchSetState.bind(
+    null,
+    currentlyRenderingFiber,
+    queue,
+  ): any);
+  queue.dispatch = dispatch;
+  return hook;
+}
+
+function mountState<S>(
+  initialState: (() => S) | S,
+): [S, Dispatch<BasicStateAction<S>>] {
+  const hook = mountStateImpl(initialState);
+  return [hook.memoizedState, hook.queue.dispatch];
 }
 
 function updateState<S>(
@@ -2465,9 +2474,10 @@ function updateDeferredValueImpl<T>(hook: Hook, prevValue: T, value: T): T {
 }
 
 function startTransition<S>(
+  fiber: Fiber,
+  queue: UpdateQueue<S | Thenable<S>, BasicStateAction<S | Thenable<S>>>,
   pendingState: S,
   finishedState: S,
-  setPending: (Thenable<S> | S) => void,
   callback: () => mixed,
   options?: StartTransitionOptions,
 ): void {
@@ -2478,7 +2488,7 @@ function startTransition<S>(
 
   const prevTransition = ReactCurrentBatchConfig.transition;
   ReactCurrentBatchConfig.transition = null;
-  setPending(pendingState);
+  dispatchSetState(fiber, queue, pendingState);
   const currentTransition = (ReactCurrentBatchConfig.transition =
     ({}: BatchConfigTransition));
 
@@ -2505,10 +2515,10 @@ function startTransition<S>(
         returnValue,
         finishedState,
       );
-      setPending(maybeThenable);
+      dispatchSetState(fiber, queue, maybeThenable);
     } else {
       // Async actions are not enabled.
-      setPending(finishedState);
+      dispatchSetState(fiber, queue, finishedState);
       callback();
     }
   } catch (error) {
@@ -2521,7 +2531,7 @@ function startTransition<S>(
         status: 'rejected',
         reason: error,
       };
-      setPending(rejectedThenable);
+      dispatchSetState(fiber, queue, rejectedThenable);
     } else {
       // The error rethrowing behavior is only enabled when the async actions
       // feature is on, even for sync actions.
@@ -2573,36 +2583,39 @@ export function startHostTransition<F>(
     );
   }
 
-  let setPending;
+  let queue: UpdateQueue<
+    Thenable<TransitionStatus> | TransitionStatus,
+    BasicStateAction<Thenable<TransitionStatus> | TransitionStatus>,
+  >;
   if (formFiber.memoizedState === null) {
     // Upgrade this host component fiber to be stateful. We're going to pretend
     // it was stateful all along so we can reuse most of the implementation
     // for function components and useTransition.
     //
     // Create the state hook used by TransitionAwareHostComponent. This is
     // essentially an inlined version of mountState.
-    const queue: UpdateQueue<
-      Thenable<TransitionStatus> | TransitionStatus,
+    const newQueue: UpdateQueue<
       Thenable<TransitionStatus> | TransitionStatus,
+      BasicStateAction<Thenable<TransitionStatus> | TransitionStatus>,
     > = {
       pending: null,
       lanes: NoLanes,
-      dispatch: null,
+      // We're going to cheat and intentionally not create a bound dispatch
+      // method, because we can call it directly in startTransition.
+      dispatch: (null: any),
       lastRenderedReducer: basicStateReducer,
       lastRenderedState: NoPendingHostTransition,
     };
+    queue = newQueue;
+
     const stateHook: Hook = {
       memoizedState: NoPendingHostTransition,
       baseState: NoPendingHostTransition,
       baseQueue: null,
-      queue: queue,
+      queue: newQueue,
       next: null,
     };
 
-    const dispatch: (Thenable<TransitionStatus> | TransitionStatus) => void =
-      (dispatchSetState.bind(null, formFiber, queue): any);
-    setPending = queue.dispatch = dispatch;
-
     // Add the state hook to both fiber alternates. The idea is that the fiber
     // had this hook all along.
     formFiber.memoizedState = stateHook;
@@ -2613,15 +2626,14 @@ export function startHostTransition<F>(
   } else {
     // This fiber was already upgraded to be stateful.
     const stateHook: Hook = formFiber.memoizedState;
-    const dispatch: (Thenable<TransitionStatus> | TransitionStatus) => void =
-      stateHook.queue.dispatch;
-    setPending = dispatch;
+    queue = stateHook.queue;
   }
 
   startTransition(
+    formFiber,
+    queue,
     pendingState,
     NoPendingHostTransition,
-    setPending,
     // TODO: We can avoid this extra wrapper, somehow. Figure out layering
     // once more of this function is implemented.
     () => callback(formData),
@@ -2632,9 +2644,15 @@ function mountTransition(): [
   boolean,
   (callback: () => void, options?: StartTransitionOptions) => void,
 ] {
-  const [, setPending] = mountState((false: Thenable<boolean> | boolean));
+  const stateHook = mountStateImpl((false: Thenable<boolean> | boolean));
   // The `start` method never changes.
-  const start = startTransition.bind(null, true, false, setPending);
+  const start = startTransition.bind(
+    null,
+    currentlyRenderingFiber,
+    stateHook.queue,
+    true,
+    false,
+  );
   const hook = mountWorkInProgressHook();
   hook.memoizedState = start;
   return [false, start];