Regression test: startTransition in render phase

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

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1187,6 +1187,23 @@ 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(
@@ -1232,6 +1249,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>,
@@ -1409,8 +1437,8 @@ const HooksDispatcherOnRerender: Dispatcher = {
   useState: rerenderState,
   useDebugValue: updateDebugValue,
   useResponder: createDeprecatedResponderListener,
-  useDeferredValue: updateDeferredValue,
-  useTransition: updateTransition,
+  useDeferredValue: rerenderDeferredValue,
+  useTransition: rerenderTransition,
 };
 
 let HooksDispatcherOnMountInDEV: Dispatcher | null = null;
@@ -1902,14 +1930,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);
     },
   };
 
@@ -2294,15 +2322,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

@@ -648,6 +648,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', () => {