Better Zustand slice functionality with added bonus of better tracing in redux devtools

01-slice.ts

import { type StoreApi } from 'zustand';
import { type NamedSet } from 'zustand/middleware';

type FirstSetStateArg<T> = Parameters<StoreApi<T>['setState']>[0];
export type SetState<T> = (partial: FirstSetStateArg<T>, name?: string) => void;
export type GetState<T> = StoreApi<T>['getState'];

type StateObject = Record<string, any>;

type SliceKey<S extends StateObject> = {
    [K in keyof S]: S[K] extends S ? K : never;
}[keyof StateObject];

export const sliceGet = <S extends StateObject, K extends SliceKey<S>>(
    get: GetState<S>,
    key: K
): GetState<S[K]> => {
    return () => get()[key];
};

/**
 * differs from Zustand set in that the `replace` option is removed in
 * favor of a `name` optional arg, which makes Redux Dev Tools output
 * more readable.
 */
export const sliceSet = <S extends StateObject, K extends SliceKey<S>>(
    set: NamedSet<S>,
    key: K
): SetState<S[K]> => {
    return (partial: FirstSetStateArg<S[K]>, name?: string) => {
        const fullName = `${key}.${name ?? 'set'}`;
        set(
            (state) => {
                const nextState =
                    typeof partial === 'function'
                        ? (partial as (s: S[K]) => S[K])(state[key])
                        : partial;

                return {
                    [key]: {
                        ...state[key],
                        ...nextState
                    }
                } as Partial<S>;
            },
            false,
            fullName
        );
    };
};

export type CreateSlice<Slice extends StateObject, AppState extends StateObject = StateObject> = (
    set: SetState<Slice>,
    get: GetState<Slice>,
    store: StoreApi<AppState>,
) => Slice;

const slice =
    <S extends StateObject>(set: NamedSet<S>, get: GetState<S>, store: StoreApi<S>) =>
        <K extends SliceKey<S>>(k: K, init: CreateSlice<S[K], S>) =>
            init(sliceSet(set, k), sliceGet(get, k), store);

export default slice;

02-state.ts

mport { create } from 'zustand';
import { devtools } from 'zustand/middleware';

import slice from './slice';

import { type EditorState, createEditorSlice } from './editor';
// ... more slices

export interface AppState {
  editor: EditorState;
  // ... more slices
}

export const useAppStore = create<AppState>()(
  devtools(
    (set, get, store) => {
      const createSlice = slice(set, get, store);
      return {
        editor: createSlice<'editor'>('editor', createEditorSlice),
        // ...
      };
    }
  )
);

03-exampleSlice.ts

import { type CreateSlice } from './slice';
import { type Lesson } from '@lib/entities/Lesson';

interface EditorProperties {
  startPrompt?: string;
  buttonCompletions: string[];
  lesson?: Lesson;
}

interface EditorFunctions {
  setStartPrompt: (startPrompt?: string) => void;
  setButtonCompletions: (buttonCompletions: string[]) => void;
  setLesson: (lesson?: Lesson) => void;
}

export interface EditorState extends EditorProperties, EditorFunctions {}

const DefaultState: EditorProperties = {
  startPrompt: undefined,
  buttonCompletions: ["bullying", "anxiety", "impulse control"],
  lesson: undefined,
}

export const createEditorSlice: CreateSlice<EditorState> = (set, get) => {
  return {
    ...DefaultState,
    setStartPrompt: (startPrompt) => set({ startPrompt }, 'setStartPrompt'),
    setButtonCompletions: (buttonCompletions) => set({ buttonCompletions }, 'setButtonCompletions'),
    setLesson: (lesson) => set({ lesson }, 'setLesson'),
  };
}

@weeksie is it possibile to use this with immer?

@pinkynrg hi! Sorry, I missed this comment. I'm bad at checking my github notifications. I haven't tested this with Immer middleware. The logic is pretty simple but you'd have to give it a shot to know for sure.