useImmerStore.ts

import produce, { Draft } from "immer";
import { useCallback, useRef, useState } from "react";

export type GetStateFunc<T> = () => Readonly<T>;
export type SetStateFunc<T> = (callback: T | ((draft: Draft<T>) => void)) => void;

/**
 *
 * @param initialValue
 */
export function useImmerStore<T>(initialValue: T): [GetStateFunc<T>, SetStateFunc<T>] {
  const [innerState, setInnerState] = useState<Readonly<T>>(initialValue);
  const stateRef = useRef<Readonly<T>>(innerState);
  const getState = useCallback((): Readonly<T> => stateRef.current, []);
  const setState = useCallback(
    (callback: T | ((draft: Draft<T>) => void)) => {
      const nextState =
        typeof callback === "function" ? produce(getState(), callback as (draft: Draft<T>) => void) : callback;
      stateRef.current = nextState;
      setInnerState(nextState);
    },
    [getState],
  );
  return [getState, setState];
}

export default useImmerStore;

Example:

export interface CustomStoreProps {
    firstName: string;
    lastName: string;
    age: number;
}

export const useCustomStore = () => {
  const [getState, setState] = useImmerStore<CustomStoreProps>({
    firstName: "John",
    lastName: "Doe",
    age: 21,
  });

  const actions = useMemo(() => ({
    setName: (firstName: string, lastName: string) => setState((state) => {
      state.firstName = firstName;
      state.lastName = lastName;
    }),
    incrementAge: (years: number) => setState((state) => state.age += years),
  }), [getState, setState]);

  return [getState(), actions];
}