Jotai under the hood

State management plays a crucial role in a React application. One might say, "Just use react-query and you're good to go..." except when you need to manage some pure client-side state. In this case, react-query starts to fade out, and you might need to reach out to a different solution. You can just use useSyncExternalStore to define a store yourself, or you can choose a built-in solution like Jotai.

That's the topic of this post — learning a bit about Jotai by trying to implement it from scratch.

Core concepts

Jotai is an atomic state management solution. It lets developers define granular pieces of state, called atoms, that can be combined into a much more complex state. These atoms are defined outside the React tree, providing a stable reference throughout the whole application, thus removing the need for manual memoization.

There are two main concepts: atoms and stores. Atoms provide a way to read/write state, and the store offers the underlying functionality for getting/setting the actual data stored in that state.

How I like to think about it is that the store is like the person who has the keys to all the locker rooms, and an atom is an actual room. Atoms have a way to read/write the state (i.e. a locker room has a locker that can be opened), but the store is responsible for providing the mechanism for getting/setting the data (i.e. opening the locker using the appropriate key).

Having this in mind, here is the code snippet for defining an atom:

function defaultRead<T>(this: Atom<T>, get: Getter) {
  return get(this)}
function defaultWrite<T>(this: Atom<T>, set: Setter, value: T) {
  set(this, value)}
export function atom<T>(initialValue: InitialValue<T>): Atom<T> {
  if (typeof initialValue === 'function') {
    return {
      read: initialValue as Read<T>,
      write: () => {
        throw new Error('atom is readonly')
      },
    }
  } else {
    return { init: initialValue, read: defaultRead, write: defaultWrite }
  }
}

Atoms act as an interface that defines methods for reading and writing the state, but they expect a getter/setter to know how to actually do that. The store is the one responsible for providing that. Ignore the sub function for now; it will come in handy when we wire up atoms in React components.

const store: Store = {
  sub<T>(atom: Atom<T>, l: () => void) {
    let atomState = atomsStateMap.get(atom)

    if (!atomState) {
      atomState = createAtomState(atom)
    }

    atomState.s.add(l)

    return () => {
      atomState.s.delete(l)
    }
  },
  set<T>(atom: Atom<T>, value: T) {
    let atomState = atomsStateMap.get(atom)

    if (!atomState) {
      atomState = createAtomState(atom)
    }

    atomState.v = value
    atomState.s.forEach((l) => l())
    atomState.d.forEach((dependent) => {
      const dependentValue = atomsStateMap.get(dependent)
      dependentValue?.s.forEach((l) => {
        l()
      })
    })
  },
  get<T>(atom: Atom<T>) {
    let atomState = atomsStateMap.get(atom)

    if (!atomState) {
      atomState = createAtomState(atom)
    }

    if (atom.init !== undefined) {
      return atomState.v
    }

    return atom.read(store.get)
  },
}

set and get functions are used alongside an atom's read and write methods to operate on the underlying data.

const store = getDefaultStore()
const countAtom = atom(0)

countAtom.read(store.get)
countAtom.write(store.set, 5)

The important thing to notice is that the store is the actual entity that keeps track of the state. Atoms are just primitives that define how the state should look and provide a way of working with it.

A cool feature that Jotai offers is that atoms can be derived from one another. Take this example:

const countAtom = atom(0)
const countDoubledAtom = atom((get) => get(countAtom) * 2)

countAtom.write(store.get, 2)
console.log(countDoubledAtom.read(store.get)) // This will print 4

In this case, the countDoubledAtom is a readonly atom. It's state is derived from countAtom. The count doubled does not store any state, but it uses it's dependency to compute it at access time. This is illustrated in the code for the store.get function below:

  get<T>(atom: Atom<T>) {
    let atomState = atomsStateMap.get(atom)

    if (!atomState) {
      atomState = createAtomState(atom)
    }

    if (atom.init !== undefined) {
      return atomState.v
    }

    return atom.read(store.get)  },

If an atom does not have an init state, the it is always accessed by reading its dependecy.

Wiring up atoms in React components

For components to be informed about any changes in the atoms, there needs to be a reactive way of subscribing to any state change associated with an atom. This is done by using the useAtomValue hook.

export function useAtomValue<T>(atom: Atom<T>) {
  const store = useStore()
  const sub = React.useCallback(
    (l: () => void) => {
      return store.sub(atom, l)
    },
    [store, atom]
  )
  const getSnapshot = React.useCallback(() => {
    return store.get(atom)
  }, [store, atom])

  return React.useSyncExternalStore(sub, getSnapshot)
}

This uses useSyncExteranlStore under the hood in order to subscribe to a particular atom in the store. Now it's the time to take a closer look at store.sub function.

sub<T>(atom: Atom<T>, l: () => void) {
  let atomState = atomsStateMap.get(atom)

  if (!atomState) {
    atomState = createAtomState(atom)
  }

  atomState.s.add(l)

  return () => {
    atomState.s.delete(l)
  }
}

When subscribing, a new record is added in the atom's listener set. To make things a bit more clear, let's define how useAtomValue implementation will look using the good old useEffect

export function useAtomValue<T>(atom: Atom<T>) {
  const store = useStore()
  const [atomValue, setAtomValue] = React.useState(() => store.get(atom))

  React.useEffect(() => {
    const unsubscribe = store.sub(atom, setSatomValue)
    return () => {
      unsubscribe()
    }
  }, [store, atom])

  return atomValue
}

What this does is essentially registering setAtomValue as a listener for any state change that occurs in the atom. Every time atom.write is invoked, the listener will be called, triggering a re-render of the component with the updated atom value.

export function useAtomSet<T>(atom: Atom<T>) {
  const store = useStore()
  const set = React.useCallback(
    (value: T) => {
      atom.write(store.set, value)
    },
    [store, atom]
  )

  return set
}

and to complete the picture, here is how an implementation for useAtom will look like

export function useAtom<T>(atom: Atom<T>) {
  return [useAtomValue(atom), useAtomSet(atom)] as const
}

Conclusion

State management is a core aspect of any React application. Jotai is a powerful solution that offers more features than just working with local state. You can learn more about it on their official page.

I hope you found this post interesting. If you want to take a look at the source code, it's available on my GitHub repo: https://github.com/radu-ux/jotai-under-the-hood.