New hook/use set state with callback

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -749,6 +749,10 @@ function useEffectEvent<Args, F: (...Array<Args>) => mixed>(callback: F): F {
   return callback;
 }
 
+function unsupportedSetStateWithCallback() {
+  throw new Error('useSetStateWithCallback is not supported in React Debug Tools.');
+}
+
 const Dispatcher: DispatcherType = {
   readContext,
 
@@ -775,6 +779,7 @@ const Dispatcher: DispatcherType = {
   useMemoCache,
   useCacheRefresh,
   useEffectEvent,
+  useSetStateWithCallback: unsupportedSetStateWithCallback,
 };
 
 // create a proxy to throw a custom error
@@ -1339,7 +1344,6 @@ export function inspectHooksOfFiber(
     currentFiber = null;
     currentHook = null;
     currentContextDependency = null;
-
     restoreContexts(contextMap);
   }
 }

packages/react-reconciler/src/ReactFiberHooks.js

@@ -3897,6 +3897,7 @@ const HooksDispatcherOnMount: Dispatcher = {
   useOptimistic: mountOptimistic,
   useMemoCache,
   useCacheRefresh: mountRefresh,
+  useSetStateWithCallback: mountSetStateWithCallback,
 };
 if (enableUseEffectEventHook) {
   (HooksDispatcherOnMount: Dispatcher).useEffectEvent = mountEvent;
@@ -3927,6 +3928,7 @@ const HooksDispatcherOnUpdate: Dispatcher = {
   useOptimistic: updateOptimistic,
   useMemoCache,
   useCacheRefresh: updateRefresh,
+  useSetStateWithCallback: updateSetStateWithCallback,
 };
 if (enableUseEffectEventHook) {
   (HooksDispatcherOnUpdate: Dispatcher).useEffectEvent = updateEvent;
@@ -3957,6 +3959,7 @@ const HooksDispatcherOnRerender: Dispatcher = {
   useOptimistic: rerenderOptimistic,
   useMemoCache,
   useCacheRefresh: updateRefresh,
+  useSetStateWithCallback: rerenderSetStateWithCallback,
 };
 if (enableUseEffectEventHook) {
   (HooksDispatcherOnRerender: Dispatcher).useEffectEvent = updateEvent;
@@ -4150,6 +4153,7 @@ if (__DEV__) {
       mountHookTypesDev();
       return mountRefresh();
     },
+    useSetStateWithCallback: mountSetStateWithCallback,
   };
   if (enableUseEffectEventHook) {
     (HooksDispatcherOnMountInDEV: Dispatcher).useEffectEvent =
@@ -4317,6 +4321,7 @@ if (__DEV__) {
       updateHookTypesDev();
       return mountRefresh();
     },
+    useSetStateWithCallback: mountSetStateWithCallback,
   };
   if (enableUseEffectEventHook) {
     (HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher).useEffectEvent =
@@ -4484,6 +4489,7 @@ if (__DEV__) {
       updateHookTypesDev();
       return updateRefresh();
     },
+    useSetStateWithCallback: updateSetStateWithCallback,
   };
   if (enableUseEffectEventHook) {
     (HooksDispatcherOnUpdateInDEV: Dispatcher).useEffectEvent =
@@ -4581,7 +4587,7 @@ if (__DEV__) {
       currentHookNameInDev = 'useState';
       updateHookTypesDev();
       const prevDispatcher = ReactSharedInternals.H;
-      ReactSharedInternals.H = InvalidNestedHooksDispatcherOnRerenderInDEV;
+      ReactSharedInternals.H = InvalidNestedHooksDispatcherOnUpdateInDEV;
       try {
         return rerenderState(initialState);
       } finally {
@@ -4651,6 +4657,7 @@ if (__DEV__) {
       updateHookTypesDev();
       return updateRefresh();
     },
+    useSetStateWithCallback: rerenderSetStateWithCallback,
   };
   if (enableUseEffectEventHook) {
     (HooksDispatcherOnRerenderInDEV: Dispatcher).useEffectEvent =
@@ -4842,6 +4849,7 @@ if (__DEV__) {
       mountHookTypesDev();
       return mountRefresh();
     },
+    useSetStateWithCallback: mountSetStateWithCallback,
   };
   if (enableUseEffectEventHook) {
     (InvalidNestedHooksDispatcherOnMountInDEV: Dispatcher).useEffectEvent =
@@ -5034,6 +5042,7 @@ if (__DEV__) {
       updateHookTypesDev();
       return updateRefresh();
     },
+    useSetStateWithCallback: updateSetStateWithCallback,
   };
   if (enableUseEffectEventHook) {
     (InvalidNestedHooksDispatcherOnUpdateInDEV: Dispatcher).useEffectEvent =
@@ -5226,6 +5235,7 @@ if (__DEV__) {
       updateHookTypesDev();
       return updateRefresh();
     },
+    useSetStateWithCallback: rerenderSetStateWithCallback,
   };
   if (enableUseEffectEventHook) {
     (InvalidNestedHooksDispatcherOnRerenderInDEV: Dispatcher).useEffectEvent =
@@ -5239,3 +5249,57 @@ if (__DEV__) {
       };
   }
 }
+
+function mountSetStateWithCallback<T>(initialValue: T): [T, (value: T, callback?: () => void) => void] {
+  const [state, setState] = mountState(initialValue);
+  const callbackRef = mountRef(null);
+  mountEffect(() => {
+    if (callbackRef.current) {
+      callbackRef.current();
+      callbackRef.current = null;
+    }
+  }, [state]);
+  const setStateWithCallback = (value, callback) => {
+    if (callback) {
+      callbackRef.current = callback;
+    }
+    setState(value);
+  };
+  return [state, setStateWithCallback];
+}
+
+function updateSetStateWithCallback<T>(initialValue: T): [T, (value: T, callback?: () => void) => void] {
+  const [state, setState] = updateState(initialValue);
+  const callbackRef = updateRef(null);
+  updateEffect(() => {
+    if (callbackRef.current) {
+      callbackRef.current();
+      callbackRef.current = null;
+    }
+  }, [state]);
+  const setStateWithCallback = (value, callback) => {
+    if (callback) {
+      callbackRef.current = callback;
+    }
+    setState(value);
+  };
+  return [state, setStateWithCallback];
+}
+
+function rerenderSetStateWithCallback<T>(initialValue: T): [T, (value: T, callback?: () => void) => void] {
+  const [state, setState] = rerenderState(initialValue);
+  const callbackRef = updateRef(null);
+  updateEffect(() => {
+    if (callbackRef.current) {
+      callbackRef.current();
+      callbackRef.current = null;
+    }
+  }, [state]);
+  const setStateWithCallback = (value, callback) => {
+    if (callback) {
+      callbackRef.current = callback;
+    }
+    setState(value);
+  };
+  return [state, setStateWithCallback];
+}

packages/react-reconciler/src/ReactInternalTypes.js

@@ -62,7 +62,8 @@ export type HookType =
   | 'useCacheRefresh'
   | 'useOptimistic'
   | 'useFormState'
-  | 'useActionState';
+  | 'useActionState'
+  | 'useSetStateWithCallback';
 
 export type ContextDependency<T> = {
   context: ReactContext<T>,
@@ -455,6 +456,7 @@ export type Dispatcher = {
     initialState: Awaited<S>,
     permalink?: string,
   ) => [Awaited<S>, (P) => void, boolean],
+  useSetStateWithCallback: <T>(initialValue: T) => [T, (value: T, callback?: () => void) => void],
 };
 
 export type AsyncDispatcher = {

packages/react-server/src/ReactFizzHooks.js

@@ -804,6 +804,10 @@ function clientHookNotSupported() {
   );
 }
 
+function unsupportedSetStateWithCallback() {
+  throw new Error('useSetStateWithCallback is not supported in Server Components.');
+}
+
 export const HooksDispatcher: Dispatcher = supportsClientAPIs
   ? {
       readContext,
@@ -833,6 +837,7 @@ export const HooksDispatcher: Dispatcher = supportsClientAPIs
       useHostTransitionStatus,
       useMemoCache,
       useCacheRefresh,
+      useSetStateWithCallback: unsupportedSetStateWithCallback,
     }
   : {
       readContext,
@@ -858,6 +863,7 @@ export const HooksDispatcher: Dispatcher = supportsClientAPIs
       useOptimistic,
       useMemoCache,
       useCacheRefresh,
+      useSetStateWithCallback: unsupportedSetStateWithCallback,
     };
 
 if (enableUseEffectEventHook) {

packages/react-server/src/ReactFlightHooks.js

@@ -101,6 +101,7 @@ export const HooksDispatcher: Dispatcher = {
   useCacheRefresh(): <T>(?() => T, ?T) => void {
     return unsupportedRefresh;
   },
+  useSetStateWithCallback: unsupportedSetStateWithCallback,
 };
 if (enableUseEffectEventHook) {
   HooksDispatcher.useEffectEvent = (unsupportedHook: any);
@@ -120,6 +121,10 @@ function unsupportedContext(): void {
   throw new Error('Cannot read a Client Context from a Server Component.');
 }
 
+function unsupportedSetStateWithCallback() {
+  throw new Error('useSetStateWithCallback is not supported in Server Components.');
+}
+
 function useId(): string {
   if (currentRequest === null) {
     throw new Error('useId can only be used while React is rendering');