Simple reactive store for state management

dev's store.ts

export type QueryStatus = "loading" | "error" | "success";

export type ReadonlyStore<T> = {
  getState: () => T;
  subscribe: (
    listenerOrStore: ((state: T) => void) | Store<any>
  ) => () => void;
};
export type Store<T> = ReadonlyStore<T> & {
  setState: (newStateOrReducer: T | ((state: T) => T)) => void;
};

export type AsyncStoreData<D, E> =
  | {
      data: undefined;
      error: undefined;
      isLoading: true;
      isError: false;
      isSuccess: false;
      status: "loading";
    }
  | {
      data: D;
      error: undefined;
      isLoading: false;
      isError: false;
      isSuccess: true;
      status: "success";
    }
  | {
      data: undefined;
      error: E;
      isLoading: false;
      isError: true;
      isSuccess: false;
      status: "error";
    };

export type AsyncStore<D, E> = ReadonlyStore<AsyncStoreData<D, E>> & {
  refetch: () => void;
};

const _makeStore = <T>(initialState: T): Store<T> => {
  let state = initialState;
  const listeners: Array<(state: T) => void> = [];
  const cleanups: Array<() => void> = [];

  const getState = () => state;

  const setState = (newStateOrReducer: T | ((state: T) => T)) => {
    let newState: T;
    if (typeof newStateOrReducer === "function") {
      newState = (newStateOrReducer as (state: T) => T)(state);
    } else {
      newState = newStateOrReducer;
    }
    if (state !== newState) {
      state = newState;
      listeners.forEach((listener) => listener(state));
    }
  };

  const subscribe = (
    listenerOrStore: ((state: T) => void) | Store<any>
  ): (() => void) => {
    if (typeof listenerOrStore === "function") {
      const listener = listenerOrStore;
      listeners.push(listener);

      return () => {
        const index = listeners.indexOf(listener);
        listeners.splice(index, 1);

        if (listeners.length === 0) {
          cleanups.forEach((fn) => fn());
        }
      };
    } else {
      return subscribe(listenerOrStore.setState);
    }
  };

  return {
    getState,
    setState,
    subscribe,
  };
};

const _makeComputedStore = <O>(
  sources: ReadonlyStore<any>[],
  fn: (...states: any[]) => O
): ReadonlyStore<O> | AsyncStore<O, unknown> => {
  const getStates = () => sources.map((source) => source.getState());
  const initial = fn(...getStates());

  if (initial instanceof Promise) {
    const store = _makeAsyncStore(
      async () => fn(...getStates()),
      initial
    );
    sources.forEach((source) => source.subscribe(store.refetch));
    return store;
  } else {
    const store = _makeStore(initial);
    sources.forEach((source) =>
      source.subscribe(() => store.setState(fn))
    );
    return store;
  }
};

const _makeAsyncStore = <D, E = unknown>(
  fetcher: () => Promise<D>,
  initialPromise?: Promise<D>
): AsyncStore<D, E> => {
  const store = _makeStore<AsyncStoreData<D, E>>({
    data: undefined,
    error: undefined,
    isLoading: true,
    isError: false,
    isSuccess: false,
    status: "loading",
  });

  const _setData = (data: D) =>
    store.setState({
      data,
      error: undefined,
      isLoading: false,
      isError: false,
      isSuccess: true,
      status: "success",
    });

  const _setError = (error: E) =>
    store.setState({
      data: undefined,
      error,
      isError: true,
      isLoading: false,
      isSuccess: false,
      status: "error",
    });

  const _reset = () =>
    store.setState({
      data: undefined,
      error: undefined,
      isError: false,
      isLoading: true,
      isSuccess: false,
      status: "loading",
    });

  const refetch = async () => {
    _reset();
    fetcher().then(_setData).catch(_setError);
  };

  initialPromise?.then(_setData).catch(_setError);

  return { ...store, refetch };
};

export function store<I1, I2, I3, O>(
  sources: [ReadonlyStore<I1>, ReadonlyStore<I2>, ReadonlyStore<I3>],
  fn: (state1: I1, state2: I2, state3: I3) => O
): ReadonlyStore<O>;

export function store<I1, I2, O>(
  sources: [ReadonlyStore<I1>, ReadonlyStore<I2>],
  fn: (state1: I1, state2: I2) => O
): ReadonlyStore<O>;

export function store<I, O, E = Error>(
  source: ReadonlyStore<I>,
  fetcher: (state: I) => Promise<O> | null
): AsyncStore<O, E>;

export function store<I, O>(
  source: ReadonlyStore<I>,
  fn: (state: I) => O
): ReadonlyStore<O>;

export function store<O, E = Error>(
  fetcher: () => Promise<O>
): AsyncStore<O, E>;

export function store<O>(initialState: O): Store<O>;
export function store<O>(): Store<O | undefined>;

export function store(sourceOrFetcher?: any, fnOrNothing?: any) {
  if (typeof sourceOrFetcher === "function") {
    return _makeAsyncStore(sourceOrFetcher);
  } else if (typeof fnOrNothing === "function") {
    const sources = Array.isArray(sourceOrFetcher)
      ? sourceOrFetcher
      : [sourceOrFetcher];
    return _makeComputedStore(sources, fnOrNothing);
  } else {
    return _makeStore(sourceOrFetcher);
  }
}