createGlobalState

createContainer.ts

import {
  Dispatch,
  SetStateAction,
  useCallback,
  useEffect,
  useState,
} from 'react';

const isFunction = (fn: unknown): fn is Function => (typeof fn === 'function');

const updateValue = <Value>(oldValue: Value, newValue: SetStateAction<Value>) => {
  if (isFunction(newValue)) {
    return newValue(oldValue);
  }
  return newValue;
};

export const createContainer = <State>(initialState: State) => {
  type StateKeys = keyof State;
  const keys = new Set(Object.keys(initialState));

  let globalState = initialState;

  const listeners = {} as {
    [StateKey in StateKeys]: Set<Dispatch<SetStateAction<State[StateKey]>>>;
  };
  keys.forEach((key) => { listeners[key as StateKeys] = new Set(); });

  const removeStateKey = <StateKey extends StateKeys>(stateKey: StateKey) => {
    delete listeners[stateKey];
    keys.delete(stateKey + '');
  }

  const setGlobalState = <StateKey extends StateKeys>(
      stateKey: StateKey,
      update: SetStateAction<State[StateKey]>,
  ) => {
    // Only update state / set listeners if the ref-equality changes.
    const updatedValue = updateValue(globalState[stateKey], update);
    if (globalState[stateKey] !== updatedValue) {
      // TODO(acornwall): Is it really necessary to replace global state like this? It's not like
      //  anything can get a reference to globalState anyways, so it _should_ be safe to just
      //  replace the one field that updated.
      globalState = {...globalState, [stateKey]: updatedValue};
      // globalState[stateKey] = updatedValue;
      listeners[stateKey]?.forEach((listener) => listener(updatedValue));
    }
  };

  const updateGlobalState = (
      update: Partial<{[key in StateKeys]: SetStateAction<State[key]>}>) => {
    // TODO(acornwall): Optimize this somehow so we don't cause the same component to be re-rendered
    //  multiple times if possible.
    //  Maybe look at ReactDOM.unstable_batchedUpdates(() => {for (const key in update) { ... }});
    for (const key in update) {
      setGlobalState(key, update[key]!);
    }
  };

  const useGlobalState = <StateKey extends StateKeys>(stateKey: StateKey) => {
    const [partialState, setPartialState] = useState(globalState[stateKey]);
    useEffect(() => {
      // If our key doesn't exist, then add a new one.
      if (!(stateKey in listeners)) {
        listeners[stateKey] = new Set();
        keys.add(stateKey + '');
      }

      listeners[stateKey].add(setPartialState);
      setPartialState(globalState[stateKey]); // in case it's changed before this effect is handled
      return () => {
        listeners[stateKey]?.delete(setPartialState);
        // If no more listeners, and no state for this key, then we can free up this state key.
        const hasListeners = listeners[stateKey] && listeners[stateKey].size > 0;
        if (!hasListeners && !(stateKey in globalState)) {
          removeStateKey(stateKey);
        }
      };
    }, [stateKey]);
    const updater = useCallback(
        (u: SetStateAction<State[StateKey]>) => setGlobalState(stateKey, u),
        [stateKey],
    );
    return [partialState, updater] as const;
  };

  const getGlobalState = <StateKey extends StateKeys>(stateKey: StateKey) => globalState[stateKey];

  const getWholeState = () => globalState;

  const notifyListeners = (prevState: State, nextState: State) => {
    keys.forEach(((key: StateKeys) => {
      const nextPartialState = nextState[key];
      // Notify listeners for a key if ref-equality check fails.
      if (prevState[key] !== nextPartialState) {
        listeners[key].forEach((listener) => listener(nextPartialState));
      }
      // If a key is going away, and it has no listeners, free up any memory for it.
      const hasListeners = listeners[key] && listeners[key].size > 0;
      if (!hasListeners && !(key in nextState)) {
        removeStateKey(key);
      }
    }) as (key: any) => {});
  };
  const setWholeState = (nextGlobalState: SetStateAction<State>) => {
    const prevGlobalState = globalState;
    globalState = updateValue(prevGlobalState, nextGlobalState);
    notifyListeners(prevGlobalState, globalState);
  };

  return {
    useGlobalState,
    getGlobalState,
    setGlobalState,
    updateGlobalState,
    getState: getWholeState,
    setState: setWholeState,
  };
};

createGlobalState.ts

import {createContainer} from "./createContainer";

type ExportFields =
    | 'useGlobalState'
    | 'getGlobalState'
    | 'setGlobalState'
    | 'updateGlobalState'
    | 'getState'
    | 'setState';

/**
 * Create a global state
 *
 * It returns a set of functions
 * - `useGlobalState`: a custom hook that works like React.useState
 * - `getGlobalState`: a function to get a global state by key outside React
 * - `setGlobalState`: a function to set a global state by key outside React
 * - `setState`: a function to set the entire global state outside React
 *
 * Based on https://github.com/dai-shi/react-hooks-global-state.
 *
 * @example
 * import { createGlobalState } from './createGlobalState';
 *
 * const { useGlobalState } = createGlobalState({ count: 0 });
 *
 * const Component = () => {
 *   const [count, setCount] = useGlobalState('count');
 *   ...
 * };
 */
export const createGlobalState = <State>(initialState: State) => {
  const store = createContainer(initialState);
  return store as Pick<typeof store, ExportFields>;
};

Thanks for sharing!