React global state!

atom-hooks.tsx

import { useSyncExternalStore } from "react";
import type { ConsumerProps, FC, ReactNode } from "react";

export interface ReactAtom<T> {
  readonly get: () => T;
  readonly sub: (listener: () => void) => () => void;
}

/**
 * Hook: Use an atom value in a react component.
 * @param atom The atom instance to use, created with atom or derivedAtom
 */
export const useAtom = <T>({ sub, get }: ReactAtom<T>): T => useSyncExternalStore(sub, get);

/**
 * Hook: Use an atom value in a react component, with a selector function.
 * @param atom The atom instance to use, created with atom or derivedAtom
 * @param selector A function that takes the atom value and returns a derived value
 * @returns The derived value
 * @example
 * const myAtom = atom({ x:1, y:2});
 *
 * // this component will be re-rendered only when myAtom.x changes
 * export const MyComponent: FC = () => {
 *  const x = useAtomSelector(myAtom, (value) => value.x);
 *  return <div>x is: {x}</div>;
 * };
 */
export const useAtomSelector = <T, U>(atom: ReactAtom<T>, selector: (value: T) => U): U =>
  useSyncExternalStore(atom.sub, () => selector(atom.get()));

export interface AtomSelectorHook<U> {
  (): U;
  get: () => U;
}

/**
 * Higher order hook: Creates a react selector hook from an atom and a selector function.
 * @param atom The atom instance to use, created with atom or derivedAtom
 * @param selector A function that takes the atom value and returns a derived value
 * @returns A react selector hook
 * @example
 * const myAtom = atom({ x:1, y:2});
 *
 * const useMyAtomX = newAtomSelectorHook(myAtom, (value) => value.x);
 *
 * // this component will be re-rendered only when myAtom.x changes
 * export const MyComponent: FC = () => {
 *   const x = useMyAtomX();
 *   return <div>x is: {x}</div>;
 * };
 */
export const newAtomSelectorHook = <T, U>(
  { sub, get: getter }: ReactAtom<T>,
  selector: (value: T) => U,
): AtomSelectorHook<U> => {
  const get = () => selector(getter());
  const self = () => useSyncExternalStore(sub, get);
  self.get = get;
  return self;
};

export interface AtomConsumerProps<T> extends ConsumerProps<T> {
  atom: ReactAtom<T>;
}

export type AtomConsumer<T> = FC<AtomConsumerProps<T>>;

/**
 * Component: Use an atom value in a react component.
 * @param atom The atom instance to use, created with atom or derivedAtom
 *
 * @example
 *
 * const myCounterAtom = atom(123);
 *
 * const MyComponent: FC = () => {
 *   return (
 *     <AtomConsumer atom={myCounterAtom}>
 *       {(count) => <div>counter is: {count}</div>}
 *     </AtomConsumer>
 *   );
 * };
 *
 * const MyWrapper: FC = () => {
 *   return (
 *     <div>
 *       <MyComponent />
 *       <button onClick={() => ++myCounterAtom.value}>count/button>
 *     </div>
 *   );
 * };
 */
export const AtomConsumer = <T>({ atom, children }: AtomConsumerProps<T>): ReactNode => children(useAtom(atom));

atom.ts

/** The type of a function that is called when the atom's value changes */
export type AtomListenerFn = () => void;

/** The type of a function that subscribes a listener to an atom */
export type AtomSubscribeFn = (listener: AtomListenerFn) => AtomUnsubsribeFn;

/** The type of a function that unsubscribes a listener from an atom */
export type AtomUnsubsribeFn = () => boolean;

/** The type of a function that initializes the atom's value */
export type AtomGetterFn<T> = () => T;

export interface AtomSubscribeable {
  sub: AtomSubscribeFn;
}

/** The type of an object that can be subscribed to and has a value */
export interface ReadonlyAtom<T> extends AtomSubscribeable {
  /** The current value of the atom */
  get value(): T;

  /** The current value of the atom */
  readonly get: AtomGetterFn<T>;

  /** Adds a listener to the atom */
  readonly sub: AtomSubscribeFn;
}

export interface Atom<T> extends ReadonlyAtom<T> {
  /** The current value of the atom */
  get value(): T;

  /** The current value of the atom */
  set value(value: T);

  /** The current value of the atom */
  readonly set: (newState: T) => boolean;

  /** Resets the current value of the atom to the initial value */
  readonly reset: () => void;
}

const _atomInitializerSym = Symbol.for("AtomInitializer");

const _notInitialized = {};

export interface AtomInitializer<T> {
  $$typeof: typeof _atomInitializerSym;
  fn: (atom: Atom<T>) => T;
}

export const atomInitializer = <T>(initializer: (atom: Atom<T>) => T): AtomInitializer<T> => ({
  $$typeof: _atomInitializerSym,
  fn: initializer,
});

/**
 * An atom is a simple object that has a value and a list of listeners.
 * When the value is changed, all listeners are notified.
 *
 * @example
 * const myAtom0 = atom(0);
 * const myAtom1 = atom(atomInitializer(() => { console.log('atom initialized'); return 0; }));
 *
 * const unsub0 = myAtom0.sub(() => console.log('myAtom0 changed'));
 * const unsub1 = myAtom1.sub(() => console.log('myAtom1 changed'));
 *
 * console.log(myAtom0.get()); // expects 0
 * console.log(myAtom1.get()); // expects 'atom initialized' and 0
 *
 * myAtom0.set(1); // expects 'myAtom0 changed'
 * myAtom1.set(1); // expects 'myAtom1 changed'
 *
 * unsub0();
 * unsub1();
 *
 * myAtom0.set(2); // expects nothing
 * myAtom1.set(2); // expects nothing
 *
 * console.log(myAtom0.get()); // expects 2
 * console.log(myAtom1.get()); // expects 2
 *
 * // This is useful especially for testing
 * myAtom0.reset();
 *
 */
export const atom = <T>(initial: T | AtomInitializer<T>): Atom<T> => {
  let state: T | typeof _notInitialized;
  let self: Atom<T>;

  // This implementation uses a doubly linked list for performance and memory efficiency
  // sub and unsub are O(1) and performs better than using an array from the benchmarks.
  // pub is O(n) obviously, but the performance of it is comparable to the array implementation (for 100000 subscribers, the difference is negligible).

  interface PubSubNode extends AtomUnsubsribeFn {
    f?: AtomListenerFn | null;
    p?: PubSubNode | null | undefined;
    n?: PubSubNode | null | undefined;
  }

  let head: PubSubNode | null | undefined;
  let tail: PubSubNode | null | undefined;

  let get: () => T;
  let reset: () => void;

  const sub: AtomSubscribeFn = (listener: AtomListenerFn | null | undefined | false): AtomUnsubsribeFn => {
    if (!listener) {
      return () => false;
    }

    const unsub: PubSubNode = (): boolean => {
      const { f, p, n } = unsub;
      if (!f) {
        return false;
      }

      // Remove from the linked list

      unsub.f = null;

      if (p) {
        p.n = n;
        unsub.p = null;
      } else {
        head = n;
      }

      if (n) {
        n.p = p;
        unsub.n = null;
      } else {
        tail = p;
      }

      return true;
    };

    // Add to the linked list

    unsub.f = listener;
    unsub.p = tail;
    unsub.n = null;

    if (tail) {
      tail.n = unsub;
    } else {
      head = unsub;
    }
    tail = unsub;

    return unsub;
  };

  const set = (value: T): boolean => {
    if (state === value) {
      return false;
    }
    state = value;

    // Loop through the linked list and call all listeners
    let node = head;
    while (node) {
      const { f, n } = node;
      if (f) {
        node = n;
        f();
      } else {
        // List was modified while iterating, so we need to restart from the beginning
        node = head;
      }
    }

    return true;
  };

  if (initial && (initial as AtomInitializer<T>).$$typeof === _atomInitializerSym) {
    const { fn } = initial as AtomInitializer<T>;

    get = (): T => {
      if (state === _notInitialized) {
        state = fn(self);
      }
      return state as T;
    };

    reset = () => {
      state = _notInitialized;
    };
  } else {
    get = (): T => state as T;

    reset = () => {
      set(initial as T);
    };
  }

  self = { sub, get, set, reset } satisfies Omit<Atom<T>, "value"> as Atom<T>;

  Reflect.defineProperty(self, "value", { get, set });

  return self;
};

/**
 * Resets the value of the given atoms to their initial value.
 * This is useful mostly for testing.
 * @param atoms The atoms to reset
 */
export const resetAtoms = (
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  ...atoms: (readonly (Atom<any> | null | undefined | false)[] | Atom<any> | null | undefined | false | 0 | "")[]
) => {
  for (const a of atoms) {
    if (a) {
      if (Array.isArray(a)) {
        resetAtoms(...a);
      } else {
        (a as Atom<unknown>).reset();
      }
    }
  }
};

export interface DerivedAtom<T> extends Atom<T> {
  /** Disposes this derived atoms, unregistering it from all dependencies. */
  dispose: () => boolean;
}

/**
 * A derived atom is an atom that is derived from other atoms.
 * When any of the atoms it depends on changes, it is recomputed.
 *
 * @example
 * const myAtom0 = atom(0);
 * const myAtom1 = atom(1);
 *
 * const myDerivedAtom = derivedAtom((atom) => myAtom0.value + myAtom1.value, [myAtom0, myAtom1]);
 *
 * const unsub = myDerivedAtom.sub(() => console.log('myDerivedAtom changed'));
 *
 * console.log(myDerivedAtom.get()); // expects 1
 *
 * myAtom0.set(2); // expects 'myDerivedAtom changed'
 *
 * console.log(myDerivedAtom.get()); // expects 3
 *
 * unsub();
 *
 * // This unregister the derived atom from all dependencies
 * myDerivedAtom.dispose();
 *
 */
export const derivedAtom = <T>(derive: (atom: Atom<T>) => T, deps: Iterable<AtomSubscribeable>): DerivedAtom<T> => {
  const self = atom(atomInitializer(derive)) as DerivedAtom<T>;

  let unsubs: AtomUnsubsribeFn[] | null = [];

  const dispose = (): boolean => {
    if (unsubs) {
      const array = unsubs;
      unsubs = null;
      for (let i = 0; i < array.length; ++i) {
        array[i]!();
      }
      return true;
    }
    return false;
  };

  const update = () => self.set(derive(self));

  for (const dep of deps) {
    if (dep) {
      unsubs.push(dep.sub(update));
    }
  }

  self.dispose = dispose;

  return self;
};

example.tsx

// EXAMPLE USAGE

import { atom, atomInitializer } from './atom'
import { useAtom } from './atom-hooks'

export const atomMyCounter = atom(123);

export const atomIsMobile = atom(atomInitializer(({ set }) => {
  const _isMobile = () => window.innerWidth < 768;
  window.addEventListener("resize", () => set(_isMobile()));
  return _isMobile();
}));

const HookComponentExample: FC = () => {
  const state = useAtom(atomMyCounter);
  return <div>state is: {state}</div>;
};

const ConsumerComponentExample: FC = () => {
  return (
    <AtomConsumer atom={atomMyCounter}>{(state) => <div>state is: {state}</div>}</AtomConsumer>
  );
};

const IsMobileComponentExample: FC = () => {
  const mobile = useAtom(atomIsMobile);
  return <div>isMobile is: {mobile ? "true" : "false"}</div>;
};

export const MyComponent: FC = () => {
  return (
    <div>
      <HookComponentExample />
      <ConsumerComponentExample />
      <IsMobileComponentExample />
      <button onClick={() => ++myCounterGlobalState.value}>click me</button>
    </div>
  );
};