Context Selector Pattern

import React, {
  createContext,
  useCallback,
  useContext,
  useRef,
  useSyncExternalStore,
} from 'react';

type Listener = () => void;

function createStore<T>(initialState: T) {
  let state = initialState;
  const listeners = new Set<Listener>();

  return {
    getState: () => state,
    setState: (updater: T | ((prev: T) => T)) => {
      state = typeof updater === 'function'
        ? (updater as (prev: T) => T)(state)
        : updater;
      listeners.forEach((l) => l());
    },
    subscribe: (listener: Listener) => {
      listeners.add(listener);
      return () => listeners.delete(listener);
    },
  };
}

type Store<T> = ReturnType<typeof createStore<T>>;

function createSelectableContext<T>() {
  const Ctx = createContext<Store<T> | null>(null);

  function Provider({ initialState, children }: { initialState: T; children: React.ReactNode }) {
    const storeRef = useRef<Store<T>>();
    if (!storeRef.current) storeRef.current = createStore(initialState);
    return <Ctx.Provider value={storeRef.current}>{children}</Ctx.Provider>;
  }

  function useSelector<S>(selector: (state: T) => S): S {
    const store = useContext(Ctx);
    if (!store) throw new Error('Missing Provider');

    return useSyncExternalStore(
      store.subscribe,
      () => selector(store.getState()),
    );
  }

  function useDispatch() {
    const store = useContext(Ctx);
    if (!store) throw new Error('Missing Provider');
    return store.setState;
  }

  return { Provider, useSelector, useDispatch };
}

// Usage
interface AppState {
  count: number;
  user: { name: string };
  theme: 'light' | 'dark';
}

const { Provider, useSelector, useDispatch } = createSelectableContext<AppState>();

function Counter() {
  const count = useSelector((s) => s.count); // Only re-renders when count changes
  const dispatch = useDispatch();
  return (
    <button onClick={() => dispatch((s) => ({ ...s, count: s.count + 1 }))}>
      Count: {count}
    </button>
  );
}

function UserName() {
  const name = useSelector((s) => s.user.name); // Only re-renders when name changes
  return <span>{name}</span>;
}

function App() {
  return (
    <Provider initialState={{ count: 0, user: { name: 'Alice' }, theme: 'dark' }}>
      <Counter />
      <UserName />
    </Provider>
  );
}

export { createSelectableContext };