Infer Zustand store types with persist and immer middleware

[jeremiahfallin]

I have a pretty large store and so far I've written the types for every slice. I'm wondering if there is a way to infer the types of the store/slices?

Here is my store hook:

import { create } from "zustand";
import { persist } from "zustand/middleware";
import { immer } from "zustand/middleware/immer";
import { enableMapSet } from "immer";
import superjson from "superjson";

import createBaseSlice, { BaseSlice } from "./baseSlice";
import createBuildingSlice from "./buildingSlice";
import createCharacterSlice, { CharacterSlice } from "./characterSlice";
import createEventSlice from "./eventSlice";
import createExplorationSlice, { ExplorationSlice } from "./explorationSlice";
import createLocationSlice, { LocationSlice } from "./locationSlice";
import createResearchSlice, { ResearchSlice } from "./researchSlice";
import createResourceGatheringSlice, {
  ResourceGatheringSlice,
} from "./resourceGatheringSlice";
import createSettingsSlice, { SettingsSlice } from "./settingsSlice";
import createStorageSlice, { StorageSlice } from "./storageSlice";
import { random } from "@/utils";

enableMapSet();

export interface StoreState
  extends BaseSlice,
    CharacterSlice,
    ExplorationSlice,
    LocationSlice,
    ResearchSlice,
    ResourceGatheringSlice,
    SettingsSlice,
    StorageSlice {}

const storage = {
  getItem: (name: string): any | null => {
    const str = localStorage.getItem(name);
    if (!str) return null;
    return superjson.parse(str);
  },
  setItem: (name: string, value: any) => {
    localStorage.setItem(name, superjson.stringify(value));
  },
  removeItem: (name: string) => localStorage.removeItem(name),
};

const useStore = create<StoreState>()(
  immer(
    persist(
      (...a) => ({
        ...createBaseSlice(...a),
        ...createBuildingSlice(...a),
        ...createCharacterSlice(...a),
        ...createEventSlice(...a),
        ...createExplorationSlice(...a),
        ...createLocationSlice(...a),
        ...createResearchSlice(...a),
        ...createResourceGatheringSlice(...a),
        ...createSettingsSlice(...a),
        ...createStorageSlice(...a),
      }),
      {
        name: "store",
        storage,
        onRehydrateStorage: () => (state) => {
          state?.setHasHydrated(true); // Set hydration state once rehydrated
        },
      }
    )
  )
);

export default useStore;
export { random };

And one of the slices:

import { StateCreator } from "zustand";
import { StoreState } from "./useStore";

export interface BaseSlice {
  zeni: number;
  land: number;
  tatami: number;
  researchLevel: number;
  getLandCost: () => { lumber: number; stone: number };
  getTatamiCost: () => {
    lumber: number;
    stone: number;
    labor: number;
  };
  canIncreaseLand: () => boolean;
  canIncreaseTatami: () => boolean;
  increaseLand: () => void;
  increaseTatami: () => void;
  getThreatLevel: () => number;
}

const createBaseSlice: StateCreator<
  StoreState,
  [["zustand/immer", never]],
  [],
  BaseSlice
> = (set: any, get: any) => ({
  zeni: 0,
  land: 0,
  tatami: 0,
  researchLevel: 1,
  getLandCost: () => {
    const land = get().land;
    const lumber = Math.floor(10 + land ** 1.15);
    const stone = Math.floor(5 + land ** 1.2);
    return { lumber, stone };
  },
  getTatamiCost: () => {
    const tatami = get().tatami;
    const lumber = Math.floor(10 + tatami ** 1.15);
    const stone = Math.floor(5 + tatami ** 1.2);
    const labor = Math.floor(5 + tatami ** 1.2);
    return { lumber, stone, labor };
  },
  canIncreaseLand: () => {
    const cost = get().getLandCost();
    return (
      get().locations.base.inventory.lumber >= cost.lumber &&
      get().locations.base.inventory.stone >= cost.stone
    );
  },
  canIncreaseTatami: () => {
    const land = get().land;
    const tatami = get().tatami;
    const cost = get().getTatamiCost();
    return (
      get().locations.base.inventory.lumber >= cost.lumber &&
      get().locations.base.inventory.stone >= cost.stone &&
      tatami < land
    );
  },
  increaseLand: () => {
    if (get().canIncreaseLand()) {
      const cost = get().getLandCost();
      set((state: any) => {
        state.locations.base.inventory.lumber -= cost.lumber;
        state.locations.base.inventory.stone -= cost.stone;
        state.land++;
      });
    }
  },
  increaseTatami: () => {
    if (get().canIncreaseTatami()) {
      const cost = get().getTatamiCost();
      set((state: any) => {
        state.locations.base.inventory.lumber -= cost.lumber;
        state.locations.base.inventory.stone -= cost.stone;
        state.tatami++;
      });
    }
  },
  getThreatLevel: () => {
    const land = get().land;
    const tatami = get().tatami;
    const landThreat = Math.floor(land / 10);
    const tatamiThreat = Math.floor(tatami / 10);
    return landThreat + tatamiThreat;
  },
});

export default createBaseSlice;

[jeremiahfallin]

Going to use zustand-slices