Build your own React state management library with TypeScript

Today we’ll build our own state management library with TypeScript involves creating a centralized store that components can subscribe to, dispatching actions to update the state, and using context to provide the store to components.

Understanding useSyncExternalStore

React 18 introduced a new hook called useSyncExternalStore which allowes react to sync to any external store.

useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot?)

Here’s a breakdown of its parameters:

• subscribe receives a callback as a parameter, and subscribes the callback to the external store so that it is called when the store state changes, it needs to return an unsubscribe function.
• getSnapshot gets the current snapshot of the store, this snapshot must be a cached value as react compares this value on each render using Object.is(getSnapshot(), oldSnapshot), providing a new value each time will cause an infinite loop.
• getServerSnapshot (optional) allows us to return a snapshot when rendering on the server, which can be helpful in certain cases where the external store or subscription source can’t run on the server or needs specific handling to run on the server.

Leveraging useSyncExternalStore, we can build a minimalist store tailored to our requirements.

Why not just use React Context?

React Context is a feature in react that allows a component to pass down props to the entire component tree below it, which means that it can be used as a store and is a viable option.

React context requires some boilerplate:

const context = createContext();

const CountProvider = ({ children }) => {
  const [count, setCount] = useState(0);
  return <context.Provider value={{ count, setCount }}>{children}</context.Provider>;
};

export function App() {
  return (
    <CountProvider>
      <Outer />
      <Other />
    </CountProvider>
  );
}

Using Context a lot can lead to “Context Hell”, where a lot of context providers are nested in the App component:

export function App() {
  return (
    <CountProvider>
      <AuthProvider>
        <ThemeProvider>
          <CacheProvider>
            <IntlProvider>
              <TooltipProvider>
                <UserSettingsProvider>
                  <NotificationProvider>
                    <AnalyticsProvider>
                      <Content />
                    </UserSettingsProvider>
                  </NotificationProvider>
                </AnalyticsProvider>
              </TooltipProvider>
            </IntlProvider>
          </CacheProvider>
        </ThemeProvider>
      </AuthProvider>
    </CountProvider>
  );
}

Moreover, using context can inadvertently trigger re-renders across the entire component tree, as demonstrated below:

export function App() {
  const [count, setCount] = useState(0);

  return (
    <context.Provider value={{ count, setCount }}>
      <Outer />
      <Other />
    </context.Provider>
  );
}

Using setCount from a consumer of the context will cause a re-render of the entire app (both Outer and Other will be re-rendered) because the state is set on the App component, and when it re-renders, all of it’s child components are also re-rendered.

Also, using an external store can allow us to more easily sync react with external systems like http requests, with context you’ll resort to using useEffect, while with an external store you can just update the store directly and the change will take effect in the subscribing components.

Building our store

Let’s delve into the implementation of our store. We’ll start with a basic structure and gradually enhance it based on our requirements.

import { useSyncExternalStore } from 'react';

export type Listener = () => void;

function createStore<T>({ initialState }: { initialState: T }) {
  let subscribers: Listener[] = [];
  let state = initialState;

  const notifyStateChanged = () => {
    subscribers.forEach((fn) => fn());
  };

  return {
    subscribe(fn: Listener) {
      subscribers.push(fn);
      return () => {
        subscribers = subscribers.filter((listener) => listener !== fn);
      };
    },
    getSnapshot() {
      return state;
    },
    setState(newState: T) {
      state = newState;
      notifyStateChanged();
    },
  };
}

Subscribers is an array of listeners that our store will notify on every change of the store’s state.
State is the store’s state which we’ll update when setState is called and then notify all of the store’s subscribers of the update.

To use the store in react we’ll create createUseStore which is a helper that wraps createStore and useSyncExternalStore in a convenient way:

export function createUseStore<T>(initialState: T) {
  const store = createStore({ initialState });
  return () => [useSyncExternalStore(store.subscribe, store.getSnapshot), store.setState] as const;
}

Using the store

With our store in place, let’s start by building a Counter component:

import React, { useState } from "react";

export function Counter() {
  const [count, setCount] = useState(0);

  const increment = () => {
    setCount(count + 1);
  };

  const decrement = () => {
    setCount(count - 1);
  };

  return (
    <div>
      <p>Count: {count}</p>
      <button onClick={increment}>Increment</button>
      <button onClick={decrement}>Decrement</button>
    </div>
  );
}

and render it three times in our app:

import React from 'react';
import ReactDOM from 'react-dom/client';
import { Counter } from './Counter.tsx';

ReactDOM.createRoot(document.getElementById('root')!).render(
  <React.StrictMode>
    <Counter />
    <Counter />
    <Counter />
  </React.StrictMode>,
);

We now see three counters in our page, clicking on Increment only increments one of the counters:

Let’s make these 3 counters use the same state using our store, first we’ll create useCountStore using the createUseStore helper that we’ve created before:

export const useCountStore = createUseStore(0);

now let’s use our useCountStore hook in our counter:

import { useCountStore } from "./countStore";

function Counter() {
  const [count, setCount] = useCountStore();

  const increment = () => {
    setCount(count + 1);
  };

  const decrement = () => {
    setCount(count - 1);
  };

  return (
    <div>
      <p>Count: {count}</p>
      <button onClick={increment}>Increment</button>
      <button onClick={decrement}>Decrement</button>
    </div>
  );
}

Now our 3 counters are synced up and all of them increment together:

Thanks to the use of generics, typescript knows that count is a number and setCount is a callback that accepts a number:

Next steps

Next, create a reducer function to handle state changes:

dispatch(action) {
  state = reducer(action);
  notifyStateChanged();
},

// without immer
setState({
    ...state,
    nested: {
        ...state.nested,
        sub: {
            ...state.nested.sub,
            new: true,
        }
    }
});

// with immer
import { produce } from "immer";

const nextState = produce(state, s => {
    s.nested.sub.new = true;
});
setState(nextState);

setState((state) => {
    state.nested.sub.new = true;
    return state;
});