Jotai under the Hood: Reactive Magic by a Simple Solution

Jotai under the Hood: Reactive Magic by a Simple Solution

Jotai under the Hood: Reactive Magic by a Simple Solution Photo by Almos Bechtold on Unsplash

Javascript passed a tough way in state management and stopped on a crossline. Almost all solutions are based on a DOM-manipulated framework - signals in Solid and Angular, reactives in Vue, and many different approaches for managing external data in React. The question of reactivity in vanilla JS is still hung in the air. There is a Signal proposal to TC39, but it’s on Stage 1, and we will not receive a common solution for at least a few years. A few approaches like MobX Observable and RxJS BehaviorSubject help to reach signal-like reactivity. Still, they use a sledgehammer to crack a nut in simplest use cases. And here Jotai might come to the stage. Although it’s a React-based library, but it’s vanilla JS API has a potential to become an easy-weight reactive engine for Javascript ecosystem.

In this article, I want to shed the light on the principles of work of one of the simpliest and the same the prettiest state management libraries of Javascript.

Minimal acceptance criteria of reactivity

With all due respect to the grand concepts of functional-reactive programming, let me formulate a simplified version of “reactive” value. In this article, we will define value as reactive when we can subscribe to value change with an immediate reaction to that. “Immediate” in this context means that it happens literally on the following process tick without postponing to the subsequent phases of the event loop.

Based on the above, the reactive value should have three fundamental and, at the same time, simple methods in its public API — getter, setter, and subscription method. The simple Javascript variable satisfies the condition with the first two methods. Still, even a class-based implementation of this reactive value would take some time to implement the subscription method.

Jotai takes care of this by itself. It defines the concept of “atom.” What is an atom under the hood? Let’s define the most straightforward substitution of the actual implementation to demonstrate basic concepts.

// Simplifiled interpretation of atom, sufficient for the article
type Atom = {
initialValue: T;
}

// In Jotai, this method calls “atom”
// and has a lot of overrides for different atom derivations
const createAtom = (initialValue: T): Atom => ({ initialValue });

Wait, how can we subscribe to this?

Usually, when we bump into Jotai, we come up with a request to solve some state management challenge in React. We open Jotai documentation, and with the magic of atom hooks API, we solve our issue. We might intuitively think the feature of reactive re-rendering is the reward of “subscription” to Jotai atom values. It’s not. Hopefully, I will explain React-world updates’ tricks in the following article. But to understand the basic principle of how Jotai works in the vanilla JS world, we have to keep in mind the following:

The Jotai atom doesn’t hold your value. All your reactive data is preserved in the Jotai store**.**

We haven’t even started to work with atoms. What is the store?

In short, the store is the entity with all reactive APIs described above_._

type Store = {
get: (atom: Atom) => T;
set: (atom: Atom, value: T) => void;
sub: (atom: Atom, callback: () => any) => void;
}

Now, we can see some relation between atoms and stores. Still, it’s an API part. We must disclose the internal store responsible for data preservation for a complete picture.

The internal store is a Map-based data structure, where the Map key is an atom, and the Map value is a pair of the value and the list of callbacks for value change. It lives inside the store closure and is not disclosed to the outside world.

// Actually, the types of the atom and the value should be the same,
// it will be restricted on the API level
type InternalStore = Map value: any,
callbacks: Array<() => any>
}>;

NB: In real Jotai implementation, Jotai internal store is a WeakMap. Because of atoms are the objects and in the same time they are the keys of a store, Jotai allow garbage collection of store values, when atoms are collected.

It looks easy till this moment. And that’s all with our typing declaration for today! Let’s push some life into this with the appropriate implementation.

Stores and atoms are not “one-to-many”

Before implementing the store factory, let’s dot the i’s and cross the t’s on one potential delusion in Jotai architecture.

The store is not a singleton. You are not limited to create several stores and even use the same atom in different stores!

Atom is just a “configuration”. You can hold different values by the same atom because, as you remember, the context of a value is established by store.

Let’s cook!

The basic frame of the store factory would look like this

const createStore = (): Store => {
const internalStore: InternalStore = new Map();

// …methods implementation

return {
get,
set,
sub,
} satisfies Store;
}

Let’s recall some points from above. First, atoms hold the initial value. Second, atoms can participate in zero, one, or many stores. Even if you do not preserve any atom-related value in any store, you can get the initial value of an atom by any store.

// store’s method implementation
const get = (atom: Atom): T => {
const state = internalStore.get(atom);
return state ? state.value : atom.initialValue;
}

Naturally, we should implement a setter next. But let me squeeze in a subscription implementation here to demonstrate how we enrich the callbacks list.

// store’s method implementation
const sub = (atom: Atom, callback: () => any): void => {
const state = internalStore.get(atom);
if (state) {
state.callbacks.push(callback);
} else {
internalStore.set({
value: atom.initialValue,
callbacks: \[callback\],
});
}
}

Finally, the setter. Remember that we need to fire all reactions to the value change. All preserved callbacks will be iterated and invoked on every setter call. It’s the core of all Jotai reactivity. Simple, right?

// store’s method implementation
const set = (atom: Atom, value: T): void => {
const state = internalStore.get(atom);
if (state) {
state.value = value;
// Here all reactive magic happens!
state.callbacks.forEach(callback => callback());
} else {
internalStore.set({
value: atom.initialValue,
callbacks: \[\],
});
}
}

How to use all of this?

Now, we have some interpretation of the most uncomplicated reactive state management based on the Jotai paradigm.

// Let’s create the store
const store = createStore();
// And create some atom with initial value
const messageAtom = createAtom(“Hello from Ukraine!”);

// Initialize the subscription to atom for furher observation
// of all its changes
store.sub(messageAtom, () => {
const message = store.get(messageAtom);
console.log(message);
});

// Some iterations happen here…

store.set(messageAtom, “Have a good day!”);
// In the next iteration the console logs the message: “Have a good day!”

Thank you for reading!

If you have experience with Jotai, please comment on whether this article helped you better understand it or if you would like to start using Jotai after looking under the hood.

Useful references:

In Plain English 🚀

Thank you for being a part of the In Plain English community! Before you go:


Jotai under the Hood: Reactive Magic by a Simple Solution was originally published in JavaScript in Plain English on Medium, where people are continuing the conversation by highlighting and responding to this story. You can include dynamic values by using placeholders like: https://drewdru.local.press/articles/b78daaf2-9ccd-4d23-aeb3-daa013baf70c, Drew Dru, https://javascript.plainenglish.io/jotai-under-the-hood-simple-magic-by-a-simple-solution-54ba5b58da9c?source=rss-b714250038a5------2, drewdru, drewdru, drewdru, drewdru These will automatically be replaced with the actual data when the message is sent.

Write a comment
No comments yet.