Writings
frontend
Evergreen

State management isn't magic: building Zustand from scratch

10 min read

Based on a hands-on workshop I ran as part of Lightricks' Advanced Learning Training, where we traced exactly where React Context runs out of road and implemented our own state engine from scratch. Figured it was worth writing up.

React state libraries look like the kind of thing you shouldn't open.

Zustand, Jotai, Redux Toolkit, Recoil — whole teams, years of work, thousands of commits. From the outside it reads as deep complexity: reactivity, invisible subscriptions, some reconciler that knows exactly who needs to re-render. The honest instinct is to use the thing and not look inside.

Here's the surprise: the core is tiny. A Zustand-like store is a closure with three functions. The React bridge is one official hook, about ten lines. Selectivity is a plain comparator. All those years went into ergonomics and edge cases — not into the mechanism.

And the mechanism is the fun part. Watching a black box turn back into familiar pieces — observer, a closure, a subscription — is one of the better feelings in this job. That's what this is, in about ten minutes: we build a working store from scratch, and by the end you can see how every library in that list actually works.

The path:

  • Why Context — the tool React already hands you — doesn't cover this
  • A minimal store in vanilla JS: three primitives
  • The bridge into React, through useSyncExternalStore
  • Atomic subscriptions: selectivity
  • What the real libraries add on top

Why we start with Context

When state has to be shared across the tree — theme, user, cart — React's built-in answer is Context. Drop a value in at the top, read it anywhere below, skip the prop drilling. It's the first tool you reach for, which makes it the first place the problem shows up.

Context doesn't hold state. It injects it. The Provider holds a value, the tree reads it, props get skipped. That's where its job ends: no updates, no subscriptions, no selectivity.

The problem shows up on the first real case. useContext subscribes to the reference of the value, not to any slice of it. Any change to the value marks every consumer dirty — no matter which part each one actually reads.

const AppCtx = createContext<AppState | null>(null);

function ThemeBadge() {
  const ctx = useContext(AppCtx)!;
  return <span>{ctx.user.theme}</span>;
}

function CartTotal() {
  const ctx = useContext(AppCtx)!;
  return <span>{ctx.cart.total}</span>;
}

<AppCtx.Provider value={state}>
  <ThemeBadge />
  <CartTotal />
</AppCtx.Provider>

Change state.cart.total and ThemeBadge re-renders too. It doesn't read cart. React doesn't know that it doesn't.

The two obvious reflexes both miss, for different reasons:

  • memo? memo skips a re-render when a component's props haven't changed. But Context isn't props — a consumer re-renders on every context value change, memo or not. The subscription bypasses it.
  • A selector? useContext has no selector argument. You subscribe to the whole value or to nothing; there's no API to say "wake me only for cart.total."

And the Provider almost always hands down a fresh object ({ ...state, cart: ... }), so even plain reference equality is gone before you start.

Heavier workarounds exist: useMemo on the value, splitting into ten separate contexts, hand-rolled layers of Providers. On a small tree it's tolerable. On a real one it's a sprawling pile of Providers and coupling you can't debug.

Context is an atomic bomb: one change, and the whole blast radius wakes up. State management is a nuclear plant — the same reaction, contained, where only the consumers that need it light up. Getting there takes a different mechanism: subscriptions.

Context broadcasts. State management subscribes. Different models.

The store is a closure with three primitives

A store isn't magic. It's a closed loop: I hold state in a closure, I change it, I call everyone who subscribed.

Three primitives. Nothing else lives in there.

   ┌──────────────── store (closure) ─────────────────┐
   │                                                   │
   │   state ───────────────────► getState() ──► read │
   │     ▲                                             │
   │     └── setState(partial) ──► merge ──► notify ──┐ │
   │                                                 │ │
   │   listeners: Set ◄──── subscribe(fn) ───────────┘ │
   │         └────────────► fn() on every change       │
   │                                                   │
   └───────────────────────────────────────────────────┘

getState — return the current value. setState — write the next one and walk the subscribers. subscribe — add a function to a Set and return an unsubscribe.

A minimal implementation in vanilla TypeScript, with no React for a mile in any direction:

type Listener = () => void;

export function createStore<T extends object>(initial: T) {
  let state = initial;
  const listeners = new Set<Listener>();

  const getState = () => state;

  const setState = (update: Partial<T> | ((prev: T) => Partial<T>)) => {
    const partial = typeof update === 'function' ? update(state) : update;
    state = { ...state, ...partial };
    listeners.forEach((fn) => fn());
  };

  const subscribe = (fn: Listener) => {
    listeners.add(fn);
    return () => { listeners.delete(fn); };
  };

  return { getState, setState, subscribe };
}

That's the whole store. That's all of Zustand before selectors and middleware.

The pattern is ancient — observer. A subject holds a set of observers and notifies them on change. One small choice worth naming: listeners is a Set, not an array, because unsubscribing with delete is cheaper than filtering a listener out by reference.

getState and subscribe are one-liners. The only line doing real work is setState — so here it is, step by step.

It accepts input in two shapes. Either you hand it the next values directly, or a function that computes them from the previous state:

store.setState({ count: 5 });                           // object
store.setState((prev) => ({ count: prev.count + 1 }));  // function of prev

That's the entire meaning of the type Partial<T> | ((prev: T) => Partial<T>): a partial object, or a function returning one. Partial<T>, not T, because you update a slice, never the whole state.

Normalize to one shape:

const partial = typeof update === 'function' ? update(state) : update;

If it's a function, call it with the current state to get the partial. If it's already an object, use it as is. After this line there's only partial to think about — the two call forms are gone.

Merge, don't replace:

state = { ...state, ...partial };

Spread the old state, overlay the changed keys. A fresh object every time, so reference checks downstream work — and partial only carries what changed. setState({ count: 5 }) leaves user and cart untouched. Replacing instead of merging would force every caller to know the entire state.

Notify:

listeners.forEach((fn) => fn());

Walk the subscribers, call each one. Each pulls the new value itself through getState.

No reconciler, no Proxy, no automatic invalidation. Change state, call the listeners; don't, and nobody finds out.

This loop is predictable: the only way to change state is setState, the only way to hear about it is to be in listeners.

It tests without React. createStore is a pure function — create it, call setState, check getState. No act(), no renderHook, no jsdom.

It debugs in a normal debugger. listeners is a plain Set you can inspect; setState is a plain function you can breakpoint. No fiber trees in the stack trace.

What's left is introducing the store to React.

React is a client, not the owner

React 18 shipped an API for exactly this case — useSyncExternalStore:

useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot?)

Look at what it asks for: a subscribe function and a getter. That's the store interface we just built — subscribe and getState. The hook was designed for the exact shape an external store already has, so wiring it up is almost nothing:

import { useSyncExternalStore } from 'react';

type Store<T> = {
  getState: () => T;
  subscribe: (cb: () => void) => () => void;
};

export function useStore<T>(store: Store<T>): T {
  return useSyncExternalStore(store.subscribe, store.getState);
}

What actually happens: we hand React our subscribe, and React registers its own callback through it. Subscribing and unsubscribing to our store becomes React's job, tied to the component lifecycle — mount subscribes, unmount cleans up. We manage none of it.

So why also pass getState? Why isn't the subscription enough? Think of a restaurant pager. subscribe hands React the pager: it'll buzz when something changes. But the pager only buzzes — it carries no food. When it goes off, React still has to walk to the counter, read the current value, and compare it to the last one to decide whether a re-render is even warranted. That walk is getSnapshot — a synchronous pull, "give me the value right now." Two channels: the buzz is push (when), the pickup is pull (what). The third argument, getServerSnapshot, is the same pull for SSR; on the client you can skip it.

Before useSyncExternalStore, Redux and early Zustand leaned on useState + useEffect. It worked until React learned to render asynchronously. In concurrent mode React can pause a render partway through. If the external store mutates during that pause, the components rendered before it see the old value and the ones rendered after see the new one — one screen, two truths. That's tearing. useSyncExternalStore closes it: it detects that the snapshot changed mid-render and forces a synchronous re-render before it paints, so the whole tree shows one consistent value.

The point isn't tearing. The point is which way the arrow points.

In the useState + Context world, React owned the state: useState owned the value, the component owned the useState, the tree owned the component. To share, you lift up, drill down, wrap in Context.

A store with useSyncExternalStore runs the other way. It's a plain JS object living in module memory. React subscribes to it. Unmount and remount — the store doesn't blink. Drop a Vue island or a vanilla script next to it, and they subscribe through the same subscribe and read the same value.

React became a client of your state, not its owner.

Selectivity: atomic subscriptions

What actually happens when setState fires? It calls every listener — and each listener is React's own callback, one per subscribed component. Each fires: React re-runs that component's getSnapshot, compares the result to the last one with Object.is, and where the value is unchanged it bails out — no re-render, no reconciliation, nothing reaches the DOM. Where the value did change, it re-renders the component, diffs the output against the previous fiber tree, and commits the change.

That bail-out is the whole game. The buzz is a broadcast — it reaches everyone holding a pager. Selectivity isn't about who gets notified; it's about who survives the Object.is check. And the current useStore returns the entire state from getSnapshot, so every consumer's snapshot is a new value on every setState. Nobody bails out. We're back to Context's broadcast — only now it's the store doing the waking.

The fix is a selector. Each component states which slice of state it cares about, and survives the check only when that slice changes.

The idea is simple: between the store and React, add a selector: (state) => S. On every render React calls getSnapshot, which applies the selector. If the selected value didn't change, React doesn't re-render the component.

The catch hides in one word — "change." useSyncExternalStore compares getSnapshot results with Object.is. If the selector returns a new object every time (state => ({ name: state.user.name })), Object.is always returns false and React re-renders every time. No selectivity.

You need to keep the previous selected value and compare it to the new one with a custom equality. If they're equal, return the previous reference so Object.is passes.

import { useRef, useSyncExternalStore } from 'react';

export function useStore<T, S>(
  store: Store<T>,
  selector: (state: T) => S,
  isEqual: (a: S, b: S) => boolean = Object.is,
): S {
  const lastRef = useRef<{ state: T; selected: S } | null>(null);

  const getSnapshot = (): S => {
    const state = store.getState();
    const last = lastRef.current;

    if (last && last.state === state) return last.selected;

    const selected = selector(state);

    if (last && isEqual(last.selected, selected)) {
      lastRef.current = { state, selected: last.selected };
      return last.selected;
    }

    lastRef.current = { state, selected };
    return selected;
  };

  return useSyncExternalStore(store.subscribe, getSnapshot);
}

Default is Object.is, which is enough for primitives and stable references. Pass shallowEqual and the selector can return a fresh object — as long as its fields don't change, React stays put.

These are atomic subscriptions. One store, a thousand consumers, each waking on exactly its slice of state.

In real code you rarely hand-roll this. The same logic is packaged as useSyncExternalStoreWithSelector (use-sync-external-store/shim/with-selector), maintained by the React team — and Zustand's equality-fn API (createWithEqualityFn) uses that exact shim. The default create skips it: plain useSyncExternalStore, equality left to you. That's why useShallow exists — it hands a shallow comparator back in where the default does nothing but a reference check. Returning a fresh object from a selector without it is the single most common way real Zustand code quietly leaks re-renders.

Selectors are what turn a bare observer into state management for a UI.

What real libraries add

If the core is this simple, why do Zustand, Jotai, and Redux Toolkit exist?

They handle everything around the core.

  • Selectors with deep equality and hooks like useShallow — so you don't write shallowEqual by hand every time.
  • Middleware: persist for localStorage and async storage, devtools for Redux DevTools, immer for mutable syntax, subscribeWithSelector for subscriptions outside React.
  • API ergonomics: set as the first argument in the factory, the slice pattern for modular stores, combine for partial types.
  • TypeScript inference — the most tedious part. Real libraries carry dense generic types so set infers the shape of your state without you spelling it out.
  • SSR and hydration: the third useSyncExternalStore argument (getServerSnapshot) feeds a value during server render and initial hydration, plus Suspense-aware paths so server and client agree on the first paint.
  • React Compiler-friendly patterns: stable references, pure selectors, so the compiler doesn't fight you.

None of this makes the core more complex.

Knowing the core changes three things. You pick a library by which layers you actually need, not by brand. You debug by stepping into a Set of listeners, not a reconciler. You argue about specific wrappers, not "Redux or Zustand."

Magic is a layer you haven't read yet. This store just stopped being one.