setState in Zustand source code.

In this article, I will provide a review on how setState in Zustand’s source code is written/works. This concept leverages closures in JavaScript and arrow functions.

StoreApi type is straight forward.

export interface StoreApi<T> {
  setState: SetStateInternal<T>
  getState: () => T
  getInitialState: () => T
  subscribe: (listener: (state: T, prevState: T) => void) => () => void
}

setState accepts two parameters

1. partial 

2. replace

Let’s perform an experiment using the example demo app provided in theZustand repo.

I added some console statements in the dist to see what’s in partial and replace.

And this is what the values are when you update the count in the demo example.

Since partial is a function here, 

const nextState = typeof partial === "function" ? partial(state) : partial;

If you look closely, state is initialised when you createStore and is outside the setState function. Does that ring a bell? Refer to Closures in Javascript.

partial is an arrow function

(state)=>({
    count: state.count + 1
})

The beauty is that you can call these functions with a parameter since it returns a function, that is why we have partial(state) and state is outside the setState. setState has access to this state variable, thanks to closures in JavaScript.

You can run the below code snippet in a browser console and it logs what you sent as a parameter.

(a => console.log(a))("test")
// Output: test

I wrote detailed articles about Object.is and Object.assign usage. Since replace is null,

if (!Object.is(nextState, state)) {
  const previousState = state
  state =
    (replace ?? (typeof nextState !== 'object' || nextState === null))
      ? (nextState as TState)
      : Object.assign({}, state, nextState)
  listeners.forEach((listener) => listener(state, previousState))
}

State is updated using Object.assign. We will look at an advanced use case where replace is not null and understand how setState behaves in the future articles.

About us:

At Think Throo, we are on a mission to teach the best practices inspired by open-source projects.

10x your coding skills by practising advanced architectural concepts in Next.js/React, learn the best practices and build production-grade projects.

We are open source — https://github.com/thinkthroo/thinkthroo (Do give us a star!)

Looking to build bespoke web systems for your business? Contact us at hello@thinkthroo.com

References:

1. https://github.com/pmndrs/zustand/blob/main/src/vanilla.ts#L64

2. https://github.com/pmndrs/zustand/blob/main/src/vanilla.ts#L9

3. https://developer.mozilla.org/en-US/docs/Web/JavaScript/Closures

4. https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Functions/Arrow_functions