ReactJS hook for passing state and a setter function with ability of getting callback for each setter call.

useStateCallback.ts

import { useRef, useCallback, useEffect, useState } from 'react';
import type { Dispatch, SetStateAction } from 'react';

type StateFunctionType<S> = Dispatch<SetStateAction<S | undefined>>;
type SetStateCallbackGeneric<S> = (
  x: S | StateFunctionType<S>,
  cb?: Function
) => void;

const useStateCallback = <T>(
  initialState: T
): [T, SetStateCallbackGeneric<T>] => {
  const [state, setState] = useState<T>(initialState);
  const cbRef = useRef(null);

  const setStateCallback = useCallback((newState, cb) => {
    cbRef.current = cb;
    setState(newState);
  }, []);

  useEffect(() => {
    if (typeof cbRef?.current === 'function') {
      cbRef?.current?.(state);
      cbRef.current = null;
    }
  }, [state]);

  return [state, setStateCallback];
};

export default useStateCallback;