Building zustand from scratch | Akashdeep Samantra

Hello 👋! It has been some time since my last blog.

In this blog, we will try to build a toy zustand project from scratch with no external dependencies. For sake of types (I have become a fan of types since i learned rust. types are super helpful), we will be choosing typescript as our programming language. We will follow the zustand specification.

Requirement

1. Build Core Library

First of all, our toy Zustand-like library should be framework-agnostic, meaning we should be able to use it with plain old JavaScript and HTML. Let’s see how that could work.

<!doctype html>
<html lang="en">
    <head>
        <meta charset="UTF-8" />
        <link rel="icon" type="image/svg+xml" href="/favicon.svg" />
        <meta name="viewport" content="width=device-width, initial-scale=1.0" />
        <title>vanilla-app</title>
    </head>
    <body>
        <div id="app">
            <h1 id="label-age">Your Age: 0</h1>
            <button id="btn-increment-age">Increment</button>
        </div>
        <script type="module" src="/src/main.js"></script>
    </body>
</html>

The corresponding script would look something like this. Whenever Increment button is clicked, it would increment age and update it in the UI.

import { createStore } from "toy-zustand";

// create our toy-zustand store
const ageStore = createStore()((set) => ({
    age: 42,
    setAge: (nextAge) =>
        set((state) => ({
            age: typeof nextAge === "function" ? nextAge(state.age) : nextAge,
        })),
}));

// DOM elements
const ageLabel = document.getElementById("label-age");
const incrementButton = document.getElementById("btn-increment-age");

// our increment event listener
// On click of the incrementButton, we are updating the age value in our toy-zustand store
function increment() {
    ageStore.getState().setAge((currentAge) => currentAge + 1);
}
incrementButton.addEventListener("click", increment);

// render function to update age
// As argument, it would recieve the entire state
function render(state) {
    ageLabel.innerText = `Your Age: ${state.age}`;
}

// rendering it with initial state i.e. age will be shown as 42
render(ageStore.getInitialState(), ageStore.getInitialState());

// subscribe to our ageStore so that any update to the store will trigger the render function
ageStore.subscribe(render);

2. Add React Support

As for the second requirement, our store should work with react state. Check following code for more information:

import { createStore } from "toy-zustand/react";

const useAgeStore = create()((set) => ({
    age: 42,
    setAge: (nextAge) => {
        set((state) => ({
            age: typeof nextAge === "function" ? nextAge(state.age) : nextAge,
        }));
    },
}));

export default function App() {
    const age = useAgeStore((state) => state.age);
    const setAge = useAgeStore((state) => state.setAge);

    function increment() {
        setAge((currentAge) => currentAge + 1);
    }

    return (
        <>
            <h1>Your age: {age}</h1>
            <button onClick={() => increment()}>+1</button>
        </>
    );
}

3. Add Vue Support

Similarly, for vue:

<script setup>
import { createStore } from "toy-zustand/vue";

const useAgeStore = createStore((set) => ({
    age: 42,
    setAge: (nextAge) => {
        set((state) => ({
            age: typeof nextAge === "function" ? nextAge(state.age) : nextAge,
        }));
    },
}));

const age = useAgeStore((state) => state.age);
const setAge = useAgeStore((state) => state.setAge);

function increment() {
    setAge.value((currentAge) => currentAge + 1);
}
</script>

<template>
    <h1>Your age: {{ age }}</h1>
    <button @click="increment">+1</button>
</template>

Let’s build it

I have created a monorepo with following folder structure. You can find the code for it here

toy-zustand
├── apps/
| |-- vanilla-app/
│ ├── react-app/
│ └── vue-app/
├── packages/
│ └── store/
├── package.json
└── package-lock.json

Since we are trying to build framewok-agnostic, we can’t use react or vue’s internal state for it. We have to create and maintain the store in javascript and somehow hook our state to frameworks’ internal state so that whenever our store changes, it triggers a side effect which would rerender the UI.

Building the core library

Let’s refer the zustand docs for the API reference. We are interested in the createStore function.

Signature

createStore<T>()(stateCreatorFn: StateCreator<T, [], []>): StoreApi<T>

Parameters

stateCreatorFn: A function that takes set function, get function and store as arguments. Usually, you will return an object with the methods you want to expose.

Returns

createStore returns a vanilla store that exposes API utilities, setState, getState, getInitialState and subscribe.

// returns object
{
    setState,
    getState,
    getInitialState,
    subscribe
}

Sample Usage

// store.ts
import { createStore } from "zustand";

// Define types for state & actions
interface BearState {
    bears: number;
    food: string;
    feed: (food: string) => void;
}

// Create store using the curried form of `create`
export const useBearStore = createStore<BearState>()((set) => ({
    bears: 2,
    food: "honey",
    feed: (food) => set(() => ({ food })),
}));

Looking at the createStore docs, we notice few important details:

This method is triggered with createStore()() i.e. a Currying.
The inner function takes a callback as argument (set) => ({ .. }).
The inner function returns the json object we defined before.

Our skeleton of the core program would look something like:

// vanilla.ts
type StateCreator<T> = Function;

type StoreApi<T> = {
    setState: Function;
    getState: Function;
    getInitialState: Function;
    subscribe: Function;
};

export function createStore<T>() {
    let initialState: T;
    let state: T;

    function setState() {
        // TODO
    }

    function getState() {
        // TODO
    }

    function getInitialState() {
        // TODO
    }

    function subscribe() {
        // TODO
    }

    return function (stateCreatorFn: StateCreator<T>): StoreApi<T> {
        return {
            setState,
            getState,
            getInitialState,
            subscribe,
        };
    };
}

Now let’s ensure the store is initialized correctly. To implement our initialization, we need to do few things:

stateCreatorFn should get executed and the return value is updated to the initialState as well as state.
stateCreateFn is a callback with signature (set, get, api) => T, T being the state type.
Our getState function to return state.

You can check my tests here.

// vanilla.ts

// setState can receive value directly or a callback to update the state
type setStateArg<T> = Partial<T> | ((state: T) => Partial<T>);

type setFn<T> = (callback: setStateArg<T>) => T;

// This is the type for the callback for createStore
// (set, get, api) => T
type StateCreator<T> = (set: setFn<T>, get: () => T, api: StoreApi<T>) => T;

// return value of createStore
// These methods are available for us to use
type StoreApi<T> = {
    // update store's internal state
    setState: (partial: setStateArg<T>, replace?: boolean) => void;

    // get store's internal state
    getState: () => T;

    // get store's initial state even after updates
    getInitialState: () => T;

    // subscribe to a particular store
    // TODO: whenever any state changes, this method will get triggered by our toy-zustand
    subscribe: Function;
};

export function createStore<T>() {
    let initialState: T;
    let state: T;

    function setState(partial: setStateArg<T>, replace: boolean = false) {
        let nextState: Partial<T>;

        // partial can be a function or object.
        if (typeof partial === "function") {
            nextState = partial(state);
        } else {
            nextState = partial;
        }

        // if nextState and state are same, no need to update state
        // This will avoid unnecessary re-render
        if (Object.is(nextState, state)) return;

        // if replace is true, then we will replace the entire state object
        // const store = createStore<{ a: number; b: number }>()(() => ({
        //   a: 1,
        //   b: 2,
        // }));
        // store.setState({ a: 5 }, true);
        // state will be { a: 5 }
        if (replace) {
            state = nextState as T;
        } else {
            state = Object.assign({}, state, nextState);
        }
    }

    function getState() {
        return state;
    }

    function getInitialState() {
        return initialState;
    }

    function subscribe() {
        // TODO
    }

    return function (stateCreatorFn: StateCreator<T>): StoreApi<T> {
        function set(callback: setStateArg<T>): T {
            setState(callback);

            return Object.assign({}, state);
        }

        const api = {
            setState,
            getState,
            getInitialState,
            subscribe,
        };

        initialState = stateCreatorFn(set, getState, api);
        state = Object.assign({}, initialState);

        return api;
    };
}

This code works when we are working in plain javascript. We have achieved the first step in our vanilla-app requirement. Second and more importantly, the html page has to be re-rendered whenever increment button is clicked.

To handle these effects, zustand expose a function i.e subscribe. We need to trigger this method with a listener callback as argument. It’s signature looks something liket this:

type Listener<T> = (state: T, prevState: T) => void;

Each listener callback would recieve state and prevState as arguments. More importantly, it would only get triggered only when there is a state change. If there is change in state value even after update, this listener won’t get triggered.

const store = createStore<{ a: number; b: number }>()(() => ({
    a: 1,
    b: 2,
}));
store.setState({ a: 1, b: 2 }); // listener wont get triggered
store.setState({ a: 10, b: 2 }); // listener will get triggered

Here is the subscriber implementaion:

// vanilla.ts

// ---snip---

type StoreApi<T> = {
    setState: (partial: setStateArg<T>, replace?: boolean) => void;
    getState: () => T;
    getInitialState: () => T;
    subscribe: (listener: Listener<T>) => () => void;
};

type Listener<T> = (state: T, prevState: T) => void;

export function createStore<T extends object>() {
    let initialState: T;
    let state: T;
    let listeners: Listener<T>[] = []; // maintain a array of listener callbacks

    function setState(partial: setStateArg<T>, replace: boolean = false) {
        // ---snip---

        // when state change, loop through listeners array and trigger each listener
        listeners.forEach((listener) => {
            listener(state, prevState);
        });
    }

    // ---snip---

    // subscribe method takes listern as callback and returns a unsubscriber callback
    function subscribe(listener: Listener<T>) {
        listeners.push(listener);

        return () => {
            listeners = listeners.filter((li) => li !== listener);
        };
    }

    // ---snip---
}