React useSharedState hook

useSharedState.ts

import { useCallback, useEffect, useMemo, useState } from "react";

const store = new Map<symbol, unknown>();
const listeners = new Map<symbol, Set<(value: any) => void>>();

function useSharedState<T>(key: symbol, initialValue?: T) {
  const [state, setState] = useState(() =>
    store.has(key) ? (store.get(key) as T) : initialValue,
  );

  const currentListeners = useMemo(
    () => listeners.get(key) ?? new Set([setState]),
    [],
  );

  const handleUpdate = useCallback<typeof setState>((next) => {
    const nextState =
      next instanceof Function ? next(store.get(key) as T) : next;

    for (const listener of currentListeners) {
      listener(nextState);
    }

    store.set(key, nextState);
  }, []);

  useEffect(() => {
    if (!store.has(key)) {
      store.set(key, initialValue);
    }

    if (!listeners.has(key)) {
      listeners.set(key, currentListeners);
    }

    return () => {
      currentListeners.delete(setState);

      if (currentListeners.size === 0) {
        store.delete(key);
        listeners.delete(key);
      }
    };
  }, [currentListeners, initialValue, key]);

  return [state, handleUpdate] as const;
}

export function createSharedState<T>(defaultValue?: T) {
  const key = Symbol("useSharedState");
  return (initialValue = defaultValue) => useSharedState(key, initialValue);
}