State management solution for easily and safely handling complex async state like Promises or eventListeners.

createStore.ts

import { useCallback, useEffect, useReducer, useRef } from "react";

type Comparator = (value1: any, value2: any) => boolean;

/** Core */

type Listener = () => void;

type Destroy = () => void;

type SetState<State> = (change: Partial<State> | ((curr: State) => Partial<State>)) => void;
type GetState<State> = () => State;

interface StoreDetails<State, Actions> {
  state: State;
  init?: () => Destroy | void;
  actions: Actions;
}

type StoreCreator<State, Actions> = (setState: SetState<State>, getState: GetState<State>) => StoreDetails<State, Actions>;

interface Store<State, Actions> {
  getState: () => State;
  destroy: Destroy;
  subscribe: (listener: Listener) => () => void;
  notify: () => void;
  setState: SetState<State>;
  actions: Actions;
}

export const createStore = <S, A>(creator: StoreCreator<S, A>): Store<S, A> => {
  let listeners = new Set<Listener>();
  let state: S;

  const unsubscribe = (listener: Listener) => {
    listeners.delete(listener);
  }
  
  const subscribe = (listener: Listener) => {
    listeners.add(listener);
    return () => unsubscribe(listener);
  }
  
  const notify = () => {
    listeners.forEach(l => l());
  }
  
  const setState = (changed: Partial<S> | ((s: S) => Partial<S>)) => {
    const partial = typeof changed === 'function' ? changed(state) : changed;
    state = { ...state, ...partial };
    notify();
  };

  const getState = () => state;

  const store = creator(setState, getState);

  state = store.state;

  const initDestory = store.init?.();

  const destroy = () => {
    listeners.clear();
    initDestory?.();
  }

  return {
    subscribe,
    destroy,
    notify,
    setState,
    getState,
    actions: store.actions
  }
}

/** React bindings */

export const useStore = <S, A, U>(store: Store<S, A>, selector: (state: S) => U, equalityFn: Comparator = Object.is): [U, A] => {
  const selectedStateRef = useRef<ReturnType<typeof selector>>();
  if (selectedStateRef.current === undefined) {
    selectedStateRef.current = selector(store.getState());
  }

  const [, forceUpdate] = useReducer(s => s + 1, 0);

  const onNotify = useCallback(() => {
    const oldState = selectedStateRef.current;
    const newState = selector(store.getState());
    if (!equalityFn(oldState, newState)) {
      selectedStateRef.current = newState;
      forceUpdate();
    }
  }, []);

  useEffect(() => {
    const unsub = store.subscribe(onNotify);
    const cleanup = () => {
      unsub();
    }
    return cleanup;
  });

  return [selectedStateRef.current, store.actions];
};

type Selector<S, U> = (state: S) => U;

type CreateStoreHookResult<S, A> = <U>(selector: Selector<S, U>) => [U, A];

export const createStoreHook = <S, A>(sc: StoreCreator<S, A>): CreateStoreHookResult<S, A> => {
  const store: Store<S, A> = createStore(sc);
  return <U>(selector: Selector<S, U>, equalityFn?: Comparator) => useStore(store, selector, equalityFn);
}

/** Usage example */

type ChatState = {
  messages: string[];
}

type ChatActions = {
  addMessage: (body: string) => void;
}

const chat = createStore<ChatState, ChatActions>((set) => ({
  state: { messages: ['hi'] },
  init: () => {
    let i = 0;
    const interval = setInterval(() => {
      set(s => ({ messages: [ ...s.messages, `${i++}` ]}))
    }, 1000);
    return () => clearInterval(interval);
  },
  actions: {
    addMessage: (body: string) => set(s => ({ messages: [...s.messages, body ] }))
  }
}));

chat.subscribe(() => {
  console.log(chat.getState().messages);
});

chat.actions.addMessage('hello');
chat.actions.addMessage('bye')

setTimeout(() => {
  chat.destroy();
}, 10000);