Zustand context

gistfile1.txt

import {
  createContext,
  PropsWithChildren,
  ReducerAction,
  useContext,
  useMemo,
  useRef,
} from "react";
import { createStore, useStore } from "zustand";
import { devtools, redux } from "zustand/middleware";


export declare type PayloadAction<
  P = void,
  T extends string = string,
  M = never,
  E = never,
> = {
  payload: P;
  type: T;
} & ([M] extends [never]
  ? Record<string, unknown>
  : {
      meta: M;
    }) &
  ([E] extends [never]
    ? Record<string, unknown>
    : {
        error: E;
      });

export type TState = unknown;

export type TAction<P = unknown, T extends string = string> = {
  payload: P;
  type: T;
};
export type TReducer<S extends TState, A extends TAction> = (
  state: S,
  action: A
) => S;

type TActions<S extends TState, A extends TAction> = Record<
  string,
  (payload?: any) => ReducerAction<TReducer<S, A>>
>;

declare type StoreRedux<Action> = {
  dispatch: (action: Action) => Action;
  dispatchFromDevtools: true;
};

type ReduxState<Action> = {
  dispatch: StoreRedux<Action>["dispatch"];
};

declare type Write<T, U> = Omit<T, keyof U> & U;

export function createReducerStoreContext<
  Name extends string,
  Reducer extends TReducer<State, TAction>,
  State extends TState,
  Actions extends TActions<State, TAction>,
>(
  displayName: Name,
  reducer: Reducer,
  initialState: State,
  actions: Actions,
  devTools = process.env.NODE_ENV !== "test" &&
    process.env.NODE_ENV !== "production"
) {
  const createCustomStore = (reducer: Reducer, initialState: State) => {

    return createStore(
      devtools(redux(reducer, initialState), {
        name: displayName,
        enabled: devTools,
      })
    );
  };

  type Store = ReturnType<typeof createCustomStore>;

  const CustomContext = createContext<Store | null>(null);
  CustomContext.displayName = displayName;

  type ProviderProps = PropsWithChildren<{
    initialState?: State;
    reducer?: Reducer;
  }>;

  function Provider({ children, ...props }: ProviderProps) {
    const storeRef = useRef<Store>();
    if (!storeRef.current) {
      //TODO: This allows to pass a custom reducer and initial state at Provider level
      storeRef.current = createCustomStore(
        props.reducer || reducer,
        props.initialState || initialState
      );
    }
    return (
      <CustomContext.Provider value={storeRef.current}>
        {children}
      </CustomContext.Provider>
    );
  }

  type WriteableState = Write<State, ReduxState<TAction>>;

  const useCustomReducerContext = <TSelected,>(
    selector: (state: WriteableState) => TSelected,
    equalityFn?: (left: TSelected, right: TSelected) => boolean
  ) => {
    const store = useContext(CustomContext);
    if (!store) {
      throw new Error(
        `use${displayName} must be called inside a ${displayName}Provider`
      );
    }

    return [useStore(store, selector, equalityFn), store.dispatch] as const;
  };

  const useCustomReducerStore = () => {
    const [state, dispatch] = useCustomReducerContext((state) => state);

    return [state, dispatch] as const;
  };

  const useCustomReducerStoreSelector = <TSelected,>(
    selector: (state: WriteableState) => TSelected,
    equalityFn?: (left: TSelected, right: TSelected) => boolean
  ) => {
    const [state, dispatch] = useCustomReducerContext(selector, equalityFn);

    return state;
  };

  const useCustomReducerStoreDispatch = () => {
    const [state, dispatch] = useCustomReducerContext((state) => state);

    return dispatch;
  };

  const useCustomReducerStoreActions = <TState,>() => {
    const [state, dispatch] = useCustomReducerContext((state) => state);

    const memoizedHandlers = useMemo(
      () =>
        Object.entries(actions).reduce(
          (
            acc: Record<string, (payload?: unknown) => void>,
            [actionKey, action]
          ) => {
            // eslint-disable-next-line no-param-reassign
            acc[actionKey] = (payload) => {
              dispatch(action(payload));
            };

            return acc;
          },
          {}
        ),
      [dispatch]
    );
    return memoizedHandlers as unknown as Actions;
  };

  return [
    Provider,
    useCustomReducerStore,
    useCustomReducerStoreSelector,
    useCustomReducerStoreDispatch,
    useCustomReducerStoreActions,
  ] as const;
}

export default createReducerStoreContext;

//Usage
type TestState = {
  grumpiness: number;
};

type Payload = { by: number };

const initialState: TestState = { grumpiness: 1 };

export const slice = createSlice({
  name: "screen-split",
  initialState: initialState,
  reducers: {
    increase: (state, action: PayloadAction<Payload>) => {
      state.grumpiness = state.grumpiness + action.payload.by;
    },
    decrease: (state, action: PayloadAction<Payload>) => {
      state.grumpiness = state.grumpiness - action.payload.by;
    },
  },
});

const [
  Provider,
  useCustomReducerStore,
  useCustomReducerStoreSelector,
  useCustomReducerStoreDispatch,
  useCustomReducerStoreActions,
] = createReducerStoreContext(
  "Positions",
  slice.reducer,
  { grumpiness: 1 },
  slice.actions
);