create-store.tsx

import produce from 'immer'
import React, { createContext, ReactNode, useContext, useReducer, useRef } from 'react'

import { debugStore } from './debug-store'

export function createStore<TState>(initialState: TState, debugKey: string, displayName: string) {
  type TProducer = {
    (currentState: TState): TState | void
  }
  type TDispatch = {
    (type: string, state: TState): void
    (type: string, producer: (state: TState) => TState | void): void
  }

  const StateCtx = createContext(initialState)
  const DispatchCtx = createContext<TDispatch>(null!)

  // debug store is a Redux store that logs to the dev tools for easy tracking of dispatches
  debugStore.dispatch({
    type: `INIT_${debugKey.toUpperCase()}`,
    payload: {
      key: debugKey,
      state: initialState,
    },
  })

  function reducer(state: TState, { action, type }: { action: TState | TProducer; type: string }) {
    const newState =
      typeof action === 'function' ? (produce(state, action as TProducer) as TState) : action
    debugStore.dispatch({
      type: type,
      payload: {
        key: debugKey,
        state: newState,
      },
    })
    return newState
  }

  function Provider({ children }: { children: ReactNode }) {
    const [state, dispatch] = useReducer(reducer, initialState)
    const dispatchWithDebug = useRef((type: string, action: TState | TProducer) => {
      dispatch({
        action,
        type,
      })
    }).current
    return (
      // separate context boundaries for dispatch and state in case a component only needs to update
      <DispatchCtx.Provider value={dispatchWithDebug}>
        <StateCtx.Provider value={state}>{children}</StateCtx.Provider>
      </DispatchCtx.Provider>
    )
  }
  Provider.displayName = displayName

  /**
   * Subscribes to the store state and allows dispatching updates to the store state.
   * Dispatch accepts the new state or a callback that allows safe mutation of the state.
   * @returns [current state; dispatch function to update the store state]
   */
  function useHook() {
    const state = useContext(StateCtx)
    const dispatch = useContext(DispatchCtx)
    if (!dispatch) {
      let error = `${displayName} hook was used outside of the context provider.`
      if (process.env.NODE_ENV === 'test')
        error += " Did you forget to call jest.mock('stores/create-store')?"
      throw new Error(error)
    }
    return tuple(state, dispatch)
  }
  /**
   * Only bind the dispatch function to this component, useful if a component
   * does not want to subscribe to state updates but needs to update it.
   * Dispatch accepts the new state or a callback that allows safe mutation of the state.
   * @returns dispatch function to update the store state
   */
  useHook.useDispatchOnly = function useDispatchOnly() {
    const dispatch = useContext(DispatchCtx)
    if (!dispatch) {
      let error = `${displayName} hook was used outside of the context provider.`
      if (process.env.NODE_ENV === 'test')
        error += " Did you forget to call jest.mock('stores/create-store')?"
      throw new Error(error)
    }
    return dispatch
  }

  return {
    Provider,
    useHook,
  }
}