Fix mount-or-update check in rerenderOptimistic (#27277)

I noticed this was wrong because it should call updateWorkInProgressHook
before it checks if currentHook is null.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1793,7 +1793,20 @@ function updateOptimistic<S, A>(
   reducer: ?(S, A) => S,
 ): [S, (A) => void] {
   const hook = updateWorkInProgressHook();
+  return updateOptimisticImpl(
+    hook,
+    ((currentHook: any): Hook),
+    passthrough,
+    reducer,
+  );
+}
 
+function updateOptimisticImpl<S, A>(
+  hook: Hook,
+  current: Hook | null,
+  passthrough: S,
+  reducer: ?(S, A) => S,
+): [S, (A) => void] {
   // Optimistic updates are always rebased on top of the latest value passed in
   // as an argument. It's called a passthrough because if there are no pending
   // updates, it will be returned as-is.
@@ -1820,14 +1833,20 @@ function rerenderOptimistic<S, A>(
   // So instead of a forked re-render implementation that knows how to handle
   // render phase udpates, we can use the same implementation as during a
   // regular mount or update.
+  const hook = updateWorkInProgressHook();
 
   if (currentHook !== null) {
     // This is an update. Process the update queue.
-    return updateOptimistic(passthrough, reducer);
+    return updateOptimisticImpl(
+      hook,
+      ((currentHook: any): Hook),
+      passthrough,
+      reducer,
+    );
   }
 
   // This is a mount. No updates to process.
-  const hook = updateWorkInProgressHook();
+
   // Reset the base state and memoized state to the passthrough. Future
   // updates will be applied on top of this.
   hook.baseState = hook.memoizedState = passthrough;

packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js

@@ -818,6 +818,44 @@ describe('ReactAsyncActions', () => {
     );
   });
 
+  // @gate enableAsyncActions
+  test('regression: useOptimistic during setState-in-render', async () => {
+    // This is a regression test for a very specific case where useOptimistic is
+    // the first hook in the component, it has a pending update, and a later
+    // hook schedules a local (setState-in-render) update. Don't sweat about
+    // deleting this test if the implementation details change.
+
+    let setOptimisticState;
+    let startTransition;
+    function App() {
+      const [optimisticState, _setOptimisticState] = useOptimistic(0);
+      setOptimisticState = _setOptimisticState;
+      const [, _startTransition] = useTransition();
+      startTransition = _startTransition;
+
+      const [derivedState, setDerivedState] = useState(0);
+      if (derivedState !== optimisticState) {
+        setDerivedState(optimisticState);
+      }
+
+      return <Text text={optimisticState} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(() => root.render(<App />));
+    assertLog([0]);
+    expect(root).toMatchRenderedOutput('0');
+
+    await act(() => {
+      startTransition(async () => {
+        setOptimisticState(1);
+        await getText('Wait');
+      });
+    });
+    assertLog([1]);
+    expect(root).toMatchRenderedOutput('1');
+  });
+
   // @gate enableAsyncActions
   test('useOptimistic accepts a custom reducer', async () => {
     let serverCart = ['A'];