How the Jotai Store API Is Inspired by the WeakMap API

Hi,

The Jotai store API was introduced in Jotai v2. Before v2, it was an internal API. Jotai conceptually functions like a WeakMap, and the store API is designed to reflect that. By the way, Jotai’s implementation internally uses multiple WeakMaps.

Recap: The WeakMap API

Let's first review the usage of the WeakMap API:

const key = {};
const map = new WeakMap();
map.set(key, 1);
console.log(map.get(key)); // ---> 1

Jotai Store API: Similar to WeakMap

The Jotai store API is almost identical in its usage:

import { atom, createStore } from 'jotai/vanilla';

const countAtom = atom(0);
const store = createStore();
store.set(countAtom, 1);
console.log(store.get(countAtom)); // ---> 1

How Is Jotai Store Different from WeakMap?

There are several differences, but one of the biggest is that you can subscribe to atom value changes.

Instead of naming this method subscribe, Jotai shortens it to sub for consistency. The three-character method aligns with other methods in the API.

import { atom, createStore } from 'jotai/vanilla';

const countAtom = atom(0);
const store = createStore();
store.sub(countAtom, () => {
  console.log('countAtom changed');
});

store.set(countAtom, 1); // will show 'countAtom changed'

The Full TypeScript Definition

Jotai’s store is essentially a subscribable WeakMap. Here is its full TypeScript definition:

type Store = {
  get: <Value>(atom: Atom<Value>) => Value
  set: <Value, Args extends unknown[], Result>(
    atom: WritableAtom<Value, Args, Result>,
    ...args: Args
  ) => Result
  sub: (atom: AnyAtom, listener: () => void) => () => void
}

Reference: Jotai Store Type Definition

You may also notice that the set method can return a Result, which is another key difference from WeakMap. The Result type is customizable based on the atom definition.

Are you interested in how the Jotai store is implemented?

Happy coding.