Simple state management in React

store.ts

import * as React from 'react';

/**
 * The subscription callback
 */
 export type Subscribe<T extends State = State> = (
  value: Readonly<T>,
  changes: Partial<T>
) => void;

/**
 * The store state
 */
export type State = Record<string | number, any>;

/**
 * The unsubscribe callback
 */
type Unsubscribe = () => void;

/**
 * The subscribe updater callback
 */
export type Updater<T extends State = State> = (
  state: Readonly<T>
) => Partial<T>;

/**
 * The store
 */
export type Store<T extends State = State> = {
  state: Readonly<T>;
  subscribe: (callback: Subscribe<T>) => Unsubscribe;
  update: (value: Partial<T> | Updater<T>) => void;
};

/**
 * Invoke the given value if it's a function, otherwise return the value
 */
const invokeOrReturn = (
  value: State | ((...args: any[]) => void),
  ...args: any[]
) => {
  return typeof value === 'function' ? value(...args) : value;
};

/**
 * Makes the store with the given initial value
 *
 * @param initialState The initial store value
 *
 * @returns The store
 */
export const makeStore = <T extends State>(initialState: T): Store<T> => {
  // Grab a shallow copy of the initial value
  let state = Object.assign({}, initialState);
  // Keep track of the subscribers
  const subcribers = new Set<Subscribe<T>>();

  return {
    get state() {
      return state;
    },
    subscribe: (callback) => {
      subcribers.add(callback);

      return () => {
        subcribers.delete(callback);
      };
    },
    update: (values) => {
      // Invoke updater if provided
      const changes = invokeOrReturn(values, state);
      
      state = Object.assign({}, state, changes);

      // Notify subscribers
      subcribers.forEach((callback) => {
        callback(state, changes);
      });
    },
  };
};

/**
 * Listens to the accessed store values and will only re-render
 * when one of the acceseed values changes
 *
 * @param store The store to listen to
 *
 * @returns The store state
 */
export const useStore = <T extends State>(store: Store<T>): T => {
  const subscribedProps = React.useRef(new Set<keyof T>());
  const [, rerender] = React.useReducer((version) => version + 1, 0);

  React.useEffect(() => {
    const unsubscribe = store.subscribe((_, changes) => {
      // Check if the changes contains a key that the client has subscribed to
      const shouldRerender = Object.keys(changes).some((key) =>
        subscribedProps.current.has(key)
      );

      if (shouldRerender) {
        rerender();
      }
    });

    // Clean-up the subscription
    return unsubscribe;
  }, [store]);

  // https://github.com/Microsoft/TypeScript/issues/20846
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  return new Proxy<any>(store, {
    get: (target: Store<T>, prop: string) => {
      // Subscribe to the required prop
      subscribedProps.current.add(prop);

      return target.state[prop];
    },
  }) as T;
};

import * as React from 'react';

import { makeStore, useStore } from '.';

const store = makeStore({ name: 'John Doe', age: 32 });

const Name = () => {
  const { name } = useStore(store);

  return (
    <div>
      <label>
        <strong>Name</strong>
      </label>
      <p>{name}</p>
    </div>
  );
};

const Age = () => {
  const { age } = useStore(store);

  return (
    <div>
      <label>
        <strong>Age</strong>
      </label>
      <p>{age}</p>
    </div>
  );
};