useStateCallback

consumeUseCallback.js

const App = () => {
  const [state, setState] = useStateCallback(0); // same API as useState + setState with cb

  const handleClick = () => {
    setState(
      prev => prev + 1,
      // 2nd argument is callback , `s` is *updated* state
      // Hey that state is set successfully what do you
      // want to do, I'm callback your friend 😄
      s => console.log("I am called after setState, state:", s)
    );
  };

  return <button onClick={handleClick}>Increment</button>;
}

useStateCallback.js

function useStateCallback(initialState) {
  const [state, setState] = useState(initialState);
  const cbRef = useRef(null); 

  const setStateCallback = (state, cb) => {
    cbRef.current = cb; // store passed callback to ref
    setState(state);
  };

  useEffect(() => {
    // cb.current is `null` on initial render, so we only execute cb on state *updates*
    if (cbRef.current) {
      cbRef.current(state);
      cbRef.current = null; // reset callback after execution
    }
  }, [state]);

  return [state, setStateCallback];
}