Here you've the source code with final result.

Creating Own State Manager For React

I'm inspired after working with Zustand in recent months. It's the tool that does the same for state management as Tailwind does for application styles. From my perspective, it's close to being perfect. The size of the library is tiny - just 1.3kB gzipped, and when you compare it with long-existing rivals like Redux, it's quickly visible that this library has something special, and life-changing.

This library is a huge step forward. Because it shows just like Tailwind, that simplicity has a lot of positives, especially in simple applications. However, I'm always learning by curiosity or if I need something. So, in this article we'll write our state manager, similar to Zustand, but with some minor changes (to fulfill our curiosity).

Please, take into account that it will not be production-ready - creating such own state manager may take a lot of time and article place - so we'll create a basic form (it's topic for video or course rather than article). The most important to know is how to get started and what features may be implemented.

The Design Phase

Before writing anything generic and something that behaves well with any code, it's nice to put ourselves into the other developer's shoes. So, here is the API signature for our future state manager (we'll call it Morph).

import { store } from 'morph/store';

type User = {
  id: string;
  name: string;
};

type State = {
  loading: boolean;
  error: string;
  users: User[];
};

// Hook to interact.
const useUsersStore = store<State>(
  // Initial state.
  {
    loading: false,
    error: ``,
    users: [],
  }
);

The proposed API may be used in the following way.

// Read outside of components.
const usersState = useUsersStore.get();

// Read inside components.
const UsersList = () => {
  const usersState = useUsersStore();
};

Now it's time for an update logic. It will be pretty similar to what Zustand offers.

import { useUsersStore } from 'store/users';
import { getUsers } from 'services/users';

const loadUsers = async (id: number): Promise<void> => {
  const { get, set } = useUsersStore;

  const state = get();

  if (state.loading) return;

  try {
    set({ loading: true, error: `` });

    const users = await getUsers(id);

    set({ users, loading: false });
  } catch {
    set({ loading: false, error: `Something went wrong` });
  }
};

export { loadUsers };

Nothing fancy right? It has almost the same signature as Zustand but differs in creation aspects. I want to have our state manager a little bit different from Zustand, but still, for understanding purposes, it will be similar.

Here are all the differences that we'll apply when compared with Zustand:

1. Separated read/update - actions cannot be added with the state (we'll use CQS principle).
2. Some functions will have different names - getState = get, and setState = set.
3. Shorter setup - not function, but the object with the initial state.

Implementing Core Features

The design is ready, we've what is required, so let's start implementation step by step. We'll use the React useSyncExternalStore hook as Zustand, and we'll wrap it with our own factory function that will provide additional features.

We need to start with a typical boilerplate - I'll add empty type definitions and some function mocks that we'll attach to our state management mechanism.

// Generic parameter that determines state type.
const store = <TState>(state: TState) => {
  // Some empty functions to illustrate the structure of the "store" factory.
  const get = () => {};

  const subscribe = () => {};

  const emit = () => {};

  const set = () => {};

  const useStore = () => {};
  
  // We're assigning functions to the "useStore" hook prototype.
  useStore.get = get;
  useStore.set = set;
  useStore.subscribe = subscribe;

  return useStore;
};

export { store };

What is interesting here is an assignment of function expressions to other functions (to be precise to its prototype). It's the same as what Zustand does for its setState, getState, and subscribe. It will allow us to use important mechanisms without additional imports.

The first function to implement will be the get. The easiest to deal with. This one always returns the current version of the state. It must be a function, and cannot be just a { state }, because if we do that, we'll always have the state defined at the beginning of the store(initialState) invocation - it will be damn buggy, so that's last what you want to have.

// Yes that's all... It's the whole implementation (～￣▽￣)～.
const get = (): TState => state;

Now something more complex. The useSyncExternalStore requires two functions useSyncExternalStore(subscribe, getSnapshot). As you probably see, the getSnapshot is our get function. So, 50% of the work is done.

The subscribe one is more tricky. It should take a callback that will be called on state change - this is the missing part. We'll use a JavaScript Set to store a list of functions (subscribers), and when there is a set function called, we'll iterate through and notify each subscriber. This is a typical observer implementation, that glues React rendering with an external state. Thus, here is the code:

type Listener<TState> = (state: TState) => void;
type Unsubscribe = () => void;

const store = <TState>(state: TState) => {
  const listeners = new Set<Listener<TState>>();

  const get = (): TState => state;

  const subscribe = (callback: Listener<TState>): Unsubscribe => {
    listeners.add(callback);
    // Function for component or hook to unsubscribe.
    return () => listeners.delete(callback);
  };

  const useStore = (): TState => useSyncExternalStore(subscribe, get);

  return useStore;
};

We've used Set due to its built-in mechanism to identify any data we pass and be able to remove it without providing unnecessary IDs.

We've added some utility types like Listener and Unsubscribe to define the signature of functions. The Listener is a callback passed by React (watch useSyncExternalStore) or us (when using useStore.subscribe in the component).

The anonymous function for removing the listener is returned, and that's all. Usually, it will be called in useEffect return block.

The missing part is a function that will trigger a state change, and walk through the listeners set. In each iteration, each function will be called with the current state. I've added a local emit function for transparency and readability.

import { useSyncExternalStore } from 'react';

type Listener<TState> = (state: TState) => void;
type Unsubscribe = () => void;
// Utility type for "set" function parameter.
type StateSetter<TState> =
  | Partial<TState>
  | ((state: TState) => Partial<TState>);

const store = <TState>(state: TState) => {
  const listeners = new Set<Listener<TState>>();

  const get = (): TState => state;

  const subscribe = (callback: Listener<TState>): Unsubscribe => {
    listeners.add(callback);
    return () => listeners.delete(callback);
  };
  // Contains only iteration logic.
  const emit = (state: TState): void => {
    listeners.forEach((listener) => listener(state));
  };
  // Sets state and triggers emission for subscribers.
  const set = (setter: StateSetter<TState>): void => {
    state = {
      ...state,
      ...(typeof setter === `function` ? setter(get()) : setter),
    };
    emit(state);
  };

  const useStore = (): TState => useSyncExternalStore(subscribe, get);

  useStore.get = get;
  useStore.set = set;
  useStore.subscribe = subscribe;

  return useStore;
};

export { store };

The set function is interesting from the perspective of parameter overload. It takes (as Zustand), a setter that is an object or function that takes the current state and returns a new state. We wrapped both with the Partial utility type. It means we'll have the same behavior as Zustand, the incoming state will be merged with the current one, and set will not require a whole state object.

Using Our State Manager

Let's try to implement typical add user to the list logic, and use our APIs.

Quick State Manager Demo

We need to start from a store invocation. It must be done outside of the components tree.

Under this pull request you've complete code.

import { store } from "morph/store";

type StoreState = {
  username: string;
  users: { id: string; name: string }[];
};

const useUsersStore = store<StoreState>({ users: [], username: `` });

Now actions, something that will change the state according to our application requirements.

// Both are outside of the component body.
const setUsername = (username: string): void => {
  useUsersStore.set({ username });
};

const addUser: FormEventHandler<HTMLFormElement> = (e): void => {
  e.preventDefault();
  useUsersStore.set(({ users, username }) => ({
    users: [...users, { id: new Date().toISOString(), name: username }],
    username: ``,
  }));
};

Last, the time of integration with our component - reading the state, triggering re-renders, and observing state changes with useEffect and subscribe combo.

const UsersPage = () => {
  const { users, username } = useUsersStore();

  React.useEffect(() => {
    // Snapshot reading.
    let currentState = useUsersStore.get();

    const unsubscribe = useUsersStore.subscribe((state) => {
      // Show alert only when the user array length changes.
      if (currentState.users.length !== state.users.length) {
        alert(
          `New user has been added! - current users count: ${state.users.length}`,
        );
        currentState = state;
      }
    });
    // Triggering unsubscribe when the component unmounts.
    return () => unsubscribe();
  }, []);

  return (
    <div className="flex bg-green-100 flex-col justify-center max-w-[350px] p-10 mx-auto">
      <form
        className="mx-auto p-10 flex flex-col justify-center gap-4"
        onSubmit={addUser}
      >
        <input
          className="p-2 rounded-md border border-gray-300"
          type="text"
          name="name"
          value={username}
          onChange={(e) => setUsername(e.target.value)}
        />
        <Button auto s={1} i={1} className="mx-auto" type="submit">
          Add user
        </Button>
      </form>

      {users.length === 0 && <p className="text-center">No users added yet</p>}
      {users.length > 0 && (
        <ul className="p-2 flex flex-col mx-auto">
          {users.map((user) => (
            <li key={user.id}>{user.name}</li>
          ))}
        </ul>
      )}
    </div>
  );
};

Adding State Replacement

Zustand has more features, here they are:

1. Middlewares.
2. State replacement - not always merging is the preferred way.
3. Nice support for SSR/SSG.
4. Selectors.

We don't have a place in this article to implement all of them, but we can implement the simplest one. Maybe others will be good candidates for separate articles. Below, is the code that allows us to replace the state.

type StateReplacer<TState> = TState extends any
  ? TState | ((state: TState) => TState)
  : never;

const store = <TState>(state: TState) => {
  const replace = (replacer: StateReplacer<TState>): void => {
    state = typeof replacer === `function` ? replacer(get()) : replacer;
    emit(state);
  };

  const useStore = (): TState => useSyncExternalStore(subscribe, get);

  useStore.replace = replace;
};

export { store };

The idea is the same as with the set function. However, there is an interesting type definition: TState extends any. I don't want to interrupt this article space to explain it, the best way to understand that is by reading the following thread on GitHub.

This code is a replacement for the set(state, true|false) implementation that exists in Zustand. The flag determines the behavior of the setter. I've split it into separate functions to avoid weird type definitions behavior - in Zustand if you try to replace the state, the state is still typed as Partial, and it provides a place for bugs - you can forget one property and bang! Especially nested ones.

Here is how you may use it:

// All state properties are required.
usersStore.replace({ users: [], activeId: -1 });
usersStore.replace((state) => ({
  users: state.users.filter((_) => _.id !== userId),
  activeId: -1
}));

Summary

As you saw the custom state manager implementation is not as hard as it may look. The beauty of it is guaranteed by useSyncExternalStore. Before React 18 it was much harder to implement stuff like that.

Still, our implementation has some core features missing that Zustand implements, and the code has not even been tested! It's not a candidate for production-ready usage. But, it's a good start, a base for creating something more advanced that matches your needs.

The low-level API that React provides is great and lovely for every place where high performance and independence from third parties are needed. I cannot say anything more - try to build something on your own for learning purposes!