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 me:

Hey, my name is Ramu Narasinga. I study large open-source projects and create content about their codebase architecture and best practices, sharing it through articles, videos.

I am open to work on interesting projects. Send me an email at ramu.narasinga@gmail.com

My Github — https://github.com/ramu-narasinga

My website — https://ramunarasinga.com

My Youtube channel — https://www.youtube.com/@thinkthroo

Learning platform — https://thinkthroo.com

Codebase Architecture — https://app.thinkthroo.com/architecture

Best practices — https://app.thinkthroo.com/best-practices

Production-grade projects — https://app.thinkthroo.com/production-grade-projects

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