Small RecipesFor Disaster

Introduction

This article is simply about me having fun with zustand, a library that excels at being simple, minimalistic, yet powerful. And I just wanted to explore by myself if I could in a non-obtrusive way modify the way it's being used.

Basically, one of the main benefit of zustand is the selector:

function App() {
  const bears = useBearStore((state) => state.bears)
}

By using (state) => state.bears, you subscribe to the bear store and only react to changes specifically affecting the bears property, not any other properties.

Now, I wanted to explore whether it's possible to achieve a signature similar to useState without writing too much boilerplate. I wanted to be able to write:

function App() {
  const [bears, setBears] = useBearStore((state) => state.bears)
}

Obviously this isn't exactly identical. Unlike useState, a store is persisted outside of the lifecycle of a component, so we can't setup the initial value here, it doesn't make sense. Also we have to provide the selector.

Before going deeper, you have to know that in the destructured array return, the first parameter would be the value, the second parameter would be the dispatch, or set function. What's interesting in this is that the dispatcher allows you to either send a value or to do a functional update. So both approach below work:

function App() {
  const [count, setCount] = useState(initialCount);
  return (
    <>
      Count: {count}
      <button onClick={() => setCount(count + 1)}>+</button> // you set the value directly
      <button onClick={() => setCount(prevCount => prevCount + 1)}>+</button> // you do the functional update
    </>
  )
}

This is particularly interesting because it means that with the functional update, you can use the previous value to determine the new value without needing to explicitly retrieve it — it is automatically provided. But how do you do that in Zustand?

Handling the functional update first

In order to handle the functional update, I first created some type which aren't bound to Zustand but follow the same principles (eg use of getState and setState):

type DispatchValue<T> = T | ((previousValue: T) => T)
type StoreProperty<T> = {
  getState: () => T
  setState: (value: T) => void
}

export function dispatchFactory<T>(store: StoreProperty<T>) {
  return (data: DispatchValue<T>) =>
    store.setState(data instanceof Function ? data(store.getState()) : data)
}

export type Dispatch<T> = (value: T | ((prevState: T) => T)) => void

Now, I could create a store which would return a dispatch function able to either use directly the value or a functional update.

import type { ObjectProperties } from '~/types/object'
import { dispatchFactory, type Dispatch } from '~/libs/zustand/helpers'

type SessionStore = {
  userId: string
  setUserId: Dispatch<string>
}

const INITIAL_STATE: ObjectProperties<SessionStore> = {
  userId: ''
}

export const useSession = create<SessionStore>((set, get) => ({
  ...INITIAL_STATE,
  setUserId: dispatchFactory({
    getState: () => get().userId,
    setState: userId => set({ userId })
  })
}))

ObjectProperties is just a utility type which pick only the properties from an object, not the methods.

This solution works fine. And we could create something à la useState but it would involve a lot of boilerplate. Also, it does not cover the destructuring array part.

Define the property and the dispatch in one go

So now, the idea is to avoid repeating ourselves. We should not have to specify three times the name of the property. So the idea is to have a factory which take an initial state and build the store for us. First, let's define the type of the thing we are trying to build:

type FrankensteinStore<T extends Object> = {
  [key in keyof T]: [T[key], Dispatch<T[key]>]
}

Weirdly, this was pretty simple. We just define a tuple for each key with the first value being the value itself, the second value the dispatcher. Nice!

Now for the factory, sadly it's a bit more complex and it requires some cast as Typescript/Javascript doesn't help much in regard to reduce.

import { StateCreator } from 'zustand'
import R from 'remeda'

function setFrankensteinStore<T extends Object>(initialState: T) {
  return ((set, get) => {
    return R.entries.strict(initialState).reduce<FrankensteinStore<T>>(
      (acc, [key, value]) =>
        Object.assign(acc, {
          [key]: [
            value,
            dispatchFactory({
              getState: () => get()[key][0],
              setState: (newValue) =>
                set({ [key]: [newValue, get()[key][1]] } as Partial<
                  FrankensteinStore<T>
                >),
            }),
          ],
        }),
      {} as never,
    )
  }) satisfies StateCreator<FrankensteinStore<T>>
}

I've been using Remeda strict to just not have to deal with reduce losing the types. Otherwise I think it's pretty much straightforward. We need to return a function which will be called by create of Zustand in order to setup the store, set and get are parameters provided by create. I iterate over each entry of the initial state (which also mean that it won't handle anything else than top-level properties), and for each one of them, I create a new tuple containing the value and then the dispatch function.

Wrapping things up

So how do you use it now? Just like so:

const useBearStore = create(setFrankensteinStore({ bears: 4 }))

function App() {
  const [bears, setBearCount] = useBearStore((state) => state.bears)

  return (
    <>
      Count: {bears}
      <button onClick={() => setBearCount(bears + 1)}>+</button>
      <button onClick={() => setBearCount((v) => v + 1)}>+</button>
    </>
  )
}

You can check the full example here.

Do you understand why we need to cast as Partial<FrankensteinStore<T>> where we do? I don't for now. But I will think about it. And what do you think of this? Would you use it? Would it be a nice addition as a recipe in the official repository?

Also, yes, I know that we could just have write:

const useBearStore = create<{bears: number}>(() => ({ bears: 4 }))
const setBearCount = (bears: number) => useBearStore.setState({ bears })

function App() {
  const bears = useBearStore((state) => state.bears)

  return (
    <>
      Count: {bears}
      <button onClick={() => setBearCount(bears + 1)}>+</button>
      <button onClick={() => setBearCount(useBearStore.getState().bears + 1)}>+</button>
    </>
  )
}

But that wasn't the point of this experiment.

Thanks to @dai_shi and @0xca0a for this library and Poimandres as a whole. I've had so much fun with all the libraries they have built.