Creating your own store | Small Recipes for Disaster

Introduction

Creating your own custom store is an exciting and educational journey that can teach you a lot about state management and React. But, just a heads up before you dive in – it’s usually best to start with established solutions to keep things simple, especially in a professional setting. Experimenting and tinkering are fantastic, and I encourage you to do so, but using existing tools can save you time and headaches. Only build your own store if you have specific needs that other options don't cover (and those specific needs probably don't exist).

Alright, now that we’ve got that out of the way, let's dive in and explore the possibilities!

Getting started

First, let's think about the problem we're trying to solve. There are plenty of solutions out there, so let's figure out why our store should exist in the first place. We will go with a minimalistic approach. This is not just for simplicity's sake, but because the more I work on stores, the more I believe they shouldn't dictate how your solution works. Instead, they should simply persist useful client state across components and provide a convenient way to modify that state.

To sum up what's the big idea? We want a store that takes an initial state, which will be an object with properties but no methods or functions. We need an easy way to modify each top-level property individually and the ability to subscribe to specific changes.

A bit of background first

Here's a quick rundown. If you ever want to create a store and make it work in React, you can do it using the useSyncExternalStore React hook. It’s all about connecting your custom store to React using a publish-subscribe pattern. For more details, I highly recommend checking out this awesome video by Jack Herrington or diving into the official React documentation.

Starting implementing

So, while we wait for React to give us a better way to subscribe to specific state update, let's roll with a selector approach like zustand do.

First, we want to specify the types around our store. We know that it will take an object as initial state, and we know that we want a setter for each property. Let's start with the setter part.

import type { Dispatch, SetStateAction } from 'react'

type SetterKey<Key extends string> = `set${Capitalize<Key>}`
type Store<InitialState extends object> = InitialState & {
  [K in keyof InitialState & string as SetterKey<K>]: Dispatch<
    SetStateAction<InitialState[K]>
  >
}

So now, everytime we have a property, we get a setter for it the way useState does. For example, if the initial state has a property foo of type string, we get a method setFoo that takes either a string as a parameter or a function that receives the previous string value and returns a string.

You can see an example in this Typescript playground.

The store itself

We'll create our store inside a hook. This allows us to freely use useRef and useCallback. We will use useRef to store the current store value, so that mutating it won't trigger any refresh unless we subscribe to the specific changes. We'll also use this hook to manage our subscriptions.

function useLocalStore<InitialState extends object>(initialState: InitialState) {
    type MyStore = Store<InitialState>
    const store = useRef<MyStore | null>(null)
    const subscribers = useRef(new Set<() => void>())
    ...
}

We then create our publish function inside the hook by just calling each subscribers:

function publish() {
  subscribers.current.forEach((callback) => callback())
}

Now for the tricky part: we need to create an object that contains both the current store values and the setters for each property. We will create a function injectSetters alongside publish which will do precisely that:

function injectSetters() {
  return Object.fromEntries(
    Object.entries(initialState).flatMap(([key, value]) => [
      [key, value],
      [
        `set${upperFirst(key)}`,
        (v: Dispatch<SetStateAction<InitialState[keyof InitialState]>>) => {
          if (!ref.current) {
            return
          }

          const previousValue = ref.current[key as keyof InitialState]
          const newValue = typeof v === 'function' ? v(previousValue) : v

          if (Object.is(newValue, previousValue)) {
            // Do not update the store if the values are the same
            return
          }

          ref.current = { ...ref.current, [key]: newValue }
          publish()
        },
      ],
    ]),
  ) as MyStore
}

Basically, we iterate over each entry and create a new setter for each one of them. We check the nature of the value passed to see if it's a function. If it is, we call it with the previous value to get the new value. We use Object.is to check that the value has been modified before actually updating the value and publishing the changes. I had to cast some variables here and there, and I'd love to know if there's a better way to avoid this. Feel free to share any tips!

Now that we have our store builder function, let's build our store:

if (ref.current === null) {
  ref.current = injectSetters()
}

Alright, our store is initialized with its setters¹.

Now, let's resolve the subscription part. Instead of returning the store, we'll return a new function called useStore This function allows us to subscribe to specific parts of the store, either a setter (which won't change) or a property.

  type Selector<T, SelectorOutput> = (store: T) => SelectorOutput

  function useStore<SelectorOutput>(selector: Selector<T, SelectorOutput>): SelectorOutput {
    const subscribe = useCallback((callback: () => void) => {
      subscribers.current.add(callback)
      return () => subscribers.current.delete(callback)
    }, [])

    return useSyncExternalStore(
      subscribe,
      () => (selector(ref.current!)),
      () => (selector(injectSetters())),
    )
  }

  return useStore
}

Again, feel free to investigate the references talked about in A bit of background first.

Make it available and easy to use in your components

Alright, our store is created, but it's not yet easy to use across all components. What we can do now is create a utility function which will provide for us a context with our store ready to be consumed. We first need to create a context, we will use createStrictContext utility function which we already talked about on this blog. You can check it out here if you are unsure about it. We'll remove the defautValues because we want to store a function (useStore) in the context instead of an object².

function createLocalStore<InitialState extends object>(name: string) {
  const [Provider, useContext] = createStrictContext<
    <V>(selector: Selector<Store<InitialState>, V>) => V
  >({
    name,
  })

  ...
}

It might look complicated, but it's basically saying that it will store a function that takes a selector as a parameter and returns a specific part of the store.

Then, we "override" the provider so the user doesn't need to be aware of the useLocalStore hook.

function LocalStoreProvider({
  children,
  initialState,
}: PropsWithChildren<{ initialState: InitialState }>) {
  return <Provider value={useLocalStore(initialState)}>{children}</Provider>
}

Finally, we create a utility function that abstracts the context from the developer.

function useStore<V>(selector: Selector<Store<InitialState>, V>): V {
  return useContext()(selector)
}

Now, we return both the custom provider and the utility function.

return [LocalStoreProvider, useStore] as const

And that's it! We're all done! Now you have a full working store, that you can use everywhere in your application once under the Provider with selector ability and automatically generated setters.

type MyStore = {
  foo: string
  bar: string
}

const [MyStoreProvider, useMyStore] = createLocalStore<MyStore>('MyStore')

function Children() {
  const foo = useMyStore((state) => state.foo)
  const setFoo = useMyStore((state) => state.setFoo)

  return <input value={foo} onChange={(event) => setFoo(event.currentTarget.value)} />
}

function App() {
  return (
    <MyStoreProvider initialState={{ foo: 'foo', bar: 'bar' }}>
      <Children />
    </MyStoreProvider>
  )

You can see a working example here.

If you've made it this far, congratulations! I hope you learned something useful and wish you fun tinkering!

Honestly, I'd love an initializer function (like useState) for useRef, but it doesn't seem possible since useRef can also store a function as a value. ↩
Which probably means that this whole defaultValues option isn't well thought out, and you should skip it entirely if you used it as defined in the previous blog post. ↩

Small RecipesFor Disaster