Build your own Zustand!
Brag at parties by building your own state management library for React! 🚀🚀
Note: All the code is taken from Zustand's github repo
You're here, you have explored the React ecosystem, taught yourself about state management, built a few cool projects and now you're here, thinking about redux when you should be sleeping.
How does state management work? What happens under the hood?
You have found yourself lurking around the Github pages of different state management tools trying to go through the code hoping to understand it this time, you fail again. But you don't give up, you deserve a good night's sleep.
Okay, I was exaggerating, you can still sleep. But if you're curious, you still think about it from time to time. And the best way to learn anything is by implementing it. So, let's build Zustand!
Before we proceed, there are a few pre-requisites. Concepts like closures, scope, the useRef hook offered by React and some knowledge of design patterns like PubSub will help you to understand what we're about to build.
If you're not yet familiar, come back to this article when you have experimented with these concepts a little bit.
Enough chit-chat, let's begin!
Let's first build the core of our tool, the API our tool will offer to manage state.
The API will expose the following methods: set, get, subscribe and clearAll to manage our store.
const create = (initStore) => {
let state = {};
let listeners = new Set();
/** Compare the previous state with the next state */
/** If the state has changed, we run our listeners and update our state
(refer to subscribe method to see why) */
const setState = (partial, replace) => {
const nextState = typeof partial === "function" ? partial(state) : partial;
if (nextState !== state) {
state = replace ? nextState : { ...state, ...nextState };
listeners.forEach((listener) => listener(state, nextState));
}
};
/** return the entire state */
const getState = () => state;
/** Subscribe to a slice of state similar to redux */
/** If the slice changes, then and only then */
/** run the listener */
/** listenerToAdd is closing over the currentState, each time the listenerToAdd runs, */
/** It compares the new state slice with the previous slice and runs the listener*/
/** We then update currentSlice to keep it in sync */
const subscribeWithListener = (
listener,
selector,
equalityFn = Object.is
) => {
let currentSlice = selector(state);
function listenerToAdd() {
const nextSlice = selector(state);
if (!equalityFn(currentSlice, nextSlice)) {
const previousSlice = currentSlice;
listener((currentSlice = nextSlice), previousSlice);
}
}
listeners.add(listenerToAdd);
return () => listeners.delete(listenerToAdd);
};
/** Subscribe without a listener */
const subscribe = (listener, selector, equalityFn) => {
if (selector || equalityFn) {
return subscribeWithListener(listener, selector, equalityFn);
}
listeners.add(listener);
return () => listeners.delete(listener);
};
const destroy = () => listeners.clear();
const api = {
setState,
getState,
subscribe,
destroy
};
/* Run our initialState with the following arguments to build our store */
state = initStore(setState, getState, api);
return api;
};
export default create;
This is pretty straightforward, but let's walk through it anyway:
1. setState
The setState method accepts a partial which is a simple function, we can use this function to derive our new state.
Let's consider an example of a partial:
/* State object */
let state = { value: 0 }
/* Partial to derive new state */
const partial = state => ({ ...state, value: state.value + 1 })
/* Run this partial */
const newState = partial(state)
// { value: 1 }
That's all there is to it, notice that it returns a new reference to the state every time we use a partial to derive the new state, you must be familiar with the reference change principle React uses to determine if our state has changed. We are using the same thing.
2. getState
The getState method just returns the entire state, that's it.
3. subscribe
This is a two-parter. We can subscribe to a slice of state (similar to redux). Our API exposes two ways to do so, the first-way subscribeWithListener allows you to add a listener to the state change along with an equality function.
Why though? Because this adds compatibility for vanilla JS! We'll explore this later.
The basic principle is, we compare the previousSlice with the nextSlice every time the listenerToAdd is run to determine if we should run our listener. This should remind you of how with redux, a component can subscribe to a slice of state and our component re-renders only when the slice changes.
The function returns a function that can be used to unsubscribe remove a listener listening to a particular slice of the state.
4. Run the initStore passed with set, get and api
The following are passed as arguments to our initStore to build our store. Let's build one now.
Congratulations! You have made a state manager for vanilla JS 🎉🎉, but we're not done yet.
We can use this to power a react application as well, bonus points if you figure out how! 👀👀
Now, we need to build an API for React, to create a store, consume it via hooks and re-render the component whenever the state slice changes. Ready? Let's go!
The createStore API
The createStore API is responsible for returning a hook that we can use in our React Component.
const createStore = (initState) => {
// create a new store by passing the initial state
// to the create function we wrote above
// returns API which we can use to manage the state of this store
const api = typeof initState === "function" ? create(initState) : initState;
const useStore = (selector, equalityFn) => {
// useStore implementation
}
return useStore;
}
Nothing much here, we pass the initial state to our createStore and pass it to the create function we wrote above. This gives us back the api that we can use to manage the state.
Next up, let's implement the useStore hook. Make sure to go through the comments in the code to get a clear understanding of how everything is working.
import { useRef, useReducer, useEffect } from "react";
import create from "./store";
const createStore = (initState) => {
// create a new store by passing the initial state
// to the create function we created above
const api = typeof initState === "function" ? create(initState) : initState;
const useStore = (selector, equalityFn = Object.is) => {
// useStore accepts a selector and an equality Fn
// using ref to store to prevent unnecessary re-renders
const state = api.getState();
const stateRef = useRef(state);
const selectorRef = useRef(selector);
const errorRef = useRef(false);
const equalityFnRef = useRef(equalityFn);
// When the slice of state changes, the component
// using the hook must re-render to reflect the latest state.
// Using a force-render mechanism using useReducer is
// one way to re-render the component
const [, forceRender] = useReducer((c) => c + 1, 0);
const currentStateSlice = useRef();
if (!currentStateSlice.current) {
currentStateSlice.current = selector(state);
}
// using local variables to avoid mutations
// in the render phase
let hasNewStateSlice;
let newStateSlice;
// If the state, the selector, the equality Fn changes
// it means, the state *might* have changed
if (
selectorRef.current !== selector ||
equalityFnRef.current !== equalityFnRef ||
stateRef.current !== state ||
errorRef.current
) {
newStateSlice = selector(state);
hasNewStateSlice = !equalityFn(newStateSlice, currentStateSlice.current);
}
// At *each* render, Updating everything to the latest version
useEffect(() => {
if (hasNewStateSlice) {
currentStateSlice.current = newStateSlice;
}
stateRef.current = state;
selectorRef.current = selectorRef;
equalityFnRef.current = equalityFn;
errorRef.current = false;
});
// recording the state just before our listener is initialized
const stateBeforeListenerRef = useRef(state);
// *the heart of the hook* //
useEffect(() => {
// When the hook is first called, (the first render),
// a listener is initalized which checks if the state has changed
// Compares the previous slice of state with the new slice of state
// using the equality function
// If anything changes, re-render the component to reflect the latest state
// otherwise, if there is an error, store that error in a ref to retry
// at the next render.
// This listener is stored in the set defined with the
// *create* API above
function listener() {
try {
const nextState = api.getState();
const nextStateSlice = selectorRef(nextState);
if (
!equalityFnRef.current(nextStateSlice, currentStateSlice.current)
) {
// update the state
stateRef.current = nextState;
currentStateSlice.current = nextStateSlice;
forceRender();
}
} catch (e) {
errorRef.current = true;
forceRender();
}
}
// an unsubscribe function to be called
// when the component unmounts, *see useEffect return value*
const unsubscribe = api.subscribe(listener);
// Call the listener immediately if the state changed before the
// listener was initialised
if (api.getState() !== stateBeforeListenerRef.current) {
listener();
}
return unsubscribe;
}, []);
// return the new slice of state if the slice of state has changed
return hasNewStateSlice ? newStateSlice : currentStateSlice.current;
};
Object.assign(useStore, api);
return useStore;
};
export default createStore;
That was a lot of code Let's go through the code step by step.
Store the state, the selector, the equality function and an error ref as
refsto avoid unnecessary re-renders and persist these values across component re-renders.Anytime any of the values change, i.e, the selector or the equality function, this means the state might have possibly changed.
Calculate the new state (calculated by running the selector), and store it in a variable for future use.
We update our selector, the equality function, the new state, and the state derived from the selector at each render.
We store the current state before initialising the listener. Remember the
createAPI above, we use listeners which run when the state of slice derived from the selector changesInside
useEffect, create a listener which compares the current state with the previous state to determine if the component needs to re-render.This is the core of the hook and a very powerful approach for state management. The component only re-renders if the state defined by the selector changes.
This listener is then stored, using the subscribe API.
When the state changes, the component re-renders using a force-render mechanism to re-render the component and reflect the latest state.
The hook then returns a new state if the state has changed. (As calculated in step 3)
Let's see a working example now.
Some points to remember, define your selectors outside the component if they do not require a state.
If the selectors do require a state value, memorize the selectors to prevent unnecessary calculations.
const selector = useCallback((state) => state.decrement, [])
const decrement = useStore(selector);
And that's it! 🔥 🔥
You can now brag to your friends at a party about how you built a popular state management library for React!
All credit goes to Daishi Kato
Follow him on: twitter
I really hope you learned a lot from it. Show some love to the article if you liked it! :)
I am on Twitter at: https://twitter.com/rohitpotato
