Regression test: startTransition in render phase

useTransition uses the state hook as part of its implementation, so we
need to fork it in the re-render dispatcher, too.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1198,6 +1198,26 @@ function updateDeferredValue<T>(
   return prevValue;
 }
 
+function rerenderDeferredValue<T>(
+  value: T,
+  config: TimeoutConfig | void | null,
+): T {
+  const [prevValue, setValue] = rerenderState(value);
+  updateEffect(
+    () => {
+      const previousConfig = ReactCurrentBatchConfig.suspense;
+      ReactCurrentBatchConfig.suspense = config === undefined ? null : config;
+      try {
+        setValue(value);
+      } finally {
+        ReactCurrentBatchConfig.suspense = previousConfig;
+      }
+    },
+    [value, config],
+  );
+  return prevValue;
+}
+
 function startTransition(setPending, config, callback) {
   const priorityLevel = getCurrentPriorityLevel();
   runWithPriority(
@@ -1243,6 +1263,17 @@ function updateTransition(
   return [start, isPending];
 }
 
+function rerenderTransition(
+  config: SuspenseConfig | void | null,
+): [(() => void) => void, boolean] {
+  const [isPending, setPending] = rerenderState(false);
+  const start = updateCallback(startTransition.bind(null, setPending, config), [
+    setPending,
+    config,
+  ]);
+  return [start, isPending];
+}
+
 function dispatchAction<S, A>(
   fiber: Fiber,
   queue: UpdateQueue<S, A>,
@@ -1420,8 +1451,8 @@ const HooksDispatcherOnRerender: Dispatcher = {
   useState: rerenderState,
   useDebugValue: updateDebugValue,
   useResponder: createDeprecatedResponderListener,
-  useDeferredValue: updateDeferredValue,
-  useTransition: updateTransition,
+  useDeferredValue: rerenderDeferredValue,
+  useTransition: rerenderTransition,
 };
 
 let HooksDispatcherOnMountInDEV: Dispatcher | null = null;
@@ -1913,14 +1944,14 @@ if (__DEV__) {
     useDeferredValue<T>(value: T, config: TimeoutConfig | void | null): T {
       currentHookNameInDev = 'useDeferredValue';
       updateHookTypesDev();
-      return updateDeferredValue(value, config);
+      return rerenderDeferredValue(value, config);
     },
     useTransition(
       config: SuspenseConfig | void | null,
     ): [(() => void) => void, boolean] {
       currentHookNameInDev = 'useTransition';
       updateHookTypesDev();
-      return updateTransition(config);
+      return rerenderTransition(config);
     },
   };
 
@@ -2305,15 +2336,15 @@ if (__DEV__) {
       currentHookNameInDev = 'useDeferredValue';
       warnInvalidHookAccess();
       updateHookTypesDev();
-      return updateDeferredValue(value, config);
+      return rerenderDeferredValue(value, config);
     },
     useTransition(
       config: SuspenseConfig | void | null,
     ): [(() => void) => void, boolean] {
       currentHookNameInDev = 'useTransition';
       warnInvalidHookAccess();
       updateHookTypesDev();
-      return updateTransition(config);
+      return rerenderTransition(config);
     },
   };
 }

packages/react-reconciler/src/__tests__/ReactHooksWithNoopRenderer-test.internal.js

@@ -645,6 +645,29 @@ describe('ReactHooksWithNoopRenderer', () => {
       expect(Scheduler).toFlushAndYield(['B:0']);
       expect(root).toMatchRenderedOutput(<span prop="B:0" />);
     });
+
+    // TODO: This should probably warn
+    it.experimental('calling startTransition inside render phase', async () => {
+      let startTransition;
+      function App() {
+        let [counter, setCounter] = useState(0);
+        let [_startTransition] = useTransition();
+        startTransition = _startTransition;
+
+        if (counter === 0) {
+          startTransition(() => {
+            setCounter(c => c + 1);
+          });
+        }
+
+        return <Text text={counter} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      root.render(<App />);
+      expect(Scheduler).toFlushAndYield([1]);
+      expect(root).toMatchRenderedOutput(<span prop={1} />);
+    });
   });
 
   describe('useReducer', () => {