Zustand EntityAdapter - An EntityAdapter example for Zustand

Hi there! Recently, I came across this amazing library to manage states in React, called Zustand. This library allows you to create simple stores that can be consumed inside the components like a hook. It's simple to learn and use but also very powerful.

What's Zustand?

Zustand is a state library with a Flux-like API. It gave me vibes of what a Service is for Angular's state management, a simple solution to share state between components. It's like creating a Redux store, with zero boilerplate code, that can be accessed anywhere and it doesn't require to configure a central state manager.

The example Zustand has in its introduction page does nothing but show what simple, yet useful it is:

// useBearStore.ts
import { create } from 'zustand'

export const useBearStore = create((set) => ({
  bears: 0,
  increasePopulation: () => set((state) => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 }),
  updateBears: (newBears) => set({ bears: newBears }),
}))

// BearCounter.tsx
import { useBearStore } from '..../useBearStore' // This can be anywhere;

export function BearCounter() {
  const bears = useBearStore((state) => state.bears)
  return <h1>{bears} around here...</h1>
}

// Controls.tsx
import { useBearStore } from '..../useBearStore' // This can be anywhere;

export function Controls() {
  const increasePopulation = useBearStore((state) => state.increasePopulation)

  return <button onClick={increasePopulation}>one up</button>
}

As you can see, we have defined a store somewhere, and then we can consume that store elsewhere the app. This alone can be very helpful to avoid prop drilling and improve your application-wide state management.

You might notice that when I said it gave me vibes of what a Service is for Angular's state management, I didn't say it lightly, this is as simple as creating a Service when you are using Angular. Even simpler because here we don't need classes or DI configurations, just a hook and that's all.

What is an EntityAdapter?

An EntityAdapter is a function that generates the same set of prebuilt actions and selectors, allowing it to interact with a list of objects sharing the same structure and perform efficient CRUD operations. This concept, originally from the ngrx team, has already been ported to RTK, and I thank both teams for their effort.

If you check them @ngrx/entity and @reduxjs/toolkit both their API is very similar. They both have a createEntityAdapter function that will return an object with a state creator, a bunch of actions, and some selectors.

Why should we create our own EntityAdapter for Zustand?

But if this is already been developed for ngrx and RTK, why can we just use those implementations instead? They were developed to be used with each library, and even if they do the same, they do so differently. That's why we can't use either directly with Zustand. But it allows us to learn more from both, the EntityAdapter functionality, and Zustand.

Let's start!

I want us to focus on what we want from this function, a state creator, a bunch of actions, and some selectors.

The initial state

The state that we want to create will be consistent, and we are going to be using the same shape both libraries are using:

type EntityId = string | number;

interface EntityState<Entity, Id extends string | number = EntityId> {
  ids: Id[];
  entities: Record<Id, Entity>;
}

The simpler function we can create to return this shape is:

function stateFactory<Entity>(): EntityState<Entity> {
  return { ids: [], entities: {} };
}

But they don't do that. Instead, both ngrx and RTK allow for more properties to be added, and at first, I wanted to give this ability to the users, but, I think doing so misses the point of using something as simple as Zustand.

Zustand allows us to create encapsulated state managers, and because of this, instead of having our factory to handle this merging, we can do that merging ourselves when we are creating our store. We can even create another store and have one store to handle the entities and other store to handle any additional state.

So, we finish with the stateFactory. Let's move on to the actions!

Random actions go!

Same as we did with the state, it would be better to identify what would be the same that we would need to express better the actions we want. This is where the libraries start to diverge, while both libraries provide methods to add, set, update, upsert, and delete one or many entities, and other methods to delete and replace all entities, each also provides additional functionality.

• ngrx additionally provides a way to map directly through the entities list to update one or many entities.
• RTK also provides a signature overload for each method enabling them to be used directly as a reducer case.

After that review, I think the bare minimum methods that we need are to add, set, update, upsert, and delete one or many entities, one to remove them all, and another one to replace them all. We can follow the same logic to handle the difference between add and set, if we try to add one entity, we verify that the entity doesn't exist, and if we set that value, we introduce or replace that entity. A similar difference can be made with update and upsert methods, but consider that update requires an object with an id property and an update property containing the updates, but the upsert needs the whole entity to be passed.

So, we will need something like this:

interface EntityActions<Entity> {
  addOne(entity: Entity): void;
  addMany(entities: Entity[]): void;

  setOne(entity: Entity): void;
  setMany(entities: Entity[]): void;

  updateOne(entity: Entity): void;
  updateMany(entities: Entity[]): void;

  upsertOne(entity: Entity): void;
  upsertMany(entities: Entity[]): void;

  removeMany(entities: Entity[]): void;
  removeOne(entity: Entity): void;

  removeAll(): void;
  setAll(entities: Entity[]): void;
}

Since we are following the Flux pattern, we don't want our actions to return anything, instead, we need to query for what we need from the main store.

To implement those methods in Zustand, we can leverage Zustand's feature that allows us to return just a slice of the state, so we can focus on what we want to interact with, and we can also consider that to update the state in Zustand, our update method needs to call a setState function. Because of this, we will create our methods to accept two parameters, the current state, and the arguments for what we want to do, an entity most of the time. Those methods will return the updated state.

To add an entity, we can do something like this:

const addOne = (state: S, entity: Entity): S => {
  const id = idSelector(entity);

  if (state.ids.includes(id)) {
    return state;
  }


  const entities = {
    ...state.entities,
    [id]: entity,
  };

  const ids = [...state.ids, id];

  return {
    entities,
    ids,
  };
};

We will talk later more about the idSelector and the State type, but for now, the first one is just a function that gets the value of that property's id, and the State type is a type alias for the store state.

So, to add an entity to the collection, we shallow copy the current state merging it with the new entity and its id. Does the user have additional properties in that state? We don't know, and better than that, we don't care, because as I mentioned, Zustand will use this and update the state merging it with any other properties that the user could have.

But we need to call this using the setState function. To use this we can have a method like:

  addOne(entity: Entity) {
    setState(state => addOne(state, entity))
  }

That method has the same signature as the addOne we want to use as our EntityActions. We can also reuse that method to create the "many" one as well, just by reducing the list of the entities.

  addMany(entities) {
    setState((state) => entities.reduce(addOne, state));
  }

To implement the other methods, we can do something like that, we can create one function for each type of operation we want, and use it to implement both alternatives, the single one and the "many" one.

After we do that, we'll have all basic C*UD operators, but to complete the implementation we want, we would be missing two methods, one to replace all entities, and one to remove them all.

Again, we are going to leverage the simplicity of Zustand for this, and when we want to remove all values from the collection, all we have to do is return an empty state.

  removeAll() {
    setState({ ids: [], entities: {} });
  }

Likewise, for our replace all function, we can still use reduce on the entities list we want to add, but instead of passing the state as the initial value as we are doing with the "many" alternatives, we will pass the empty state.

  setAll(entities: Entity[]){ 
    setState(() => entities.reduce(setOne, { ids: [], entities: {} });
  }

Now we have all the actions we need to interact with our application! Awesome! We can create our actionFactory now. This is when we will talk about the idSelector and setState functions, and the State type.

• The idSelector is a function we will provide when we create each adapter. The idSelector should have this signature:

type IdSelector<Entity> = (model: Entity) => string | number;

• The setState is the Zustand function to set the state. The signature for this one is complex, as the simpler version in Zustand has been deprecated, so we will have to trust the process.

• The State type, is a type alias for a EntityState:
  

type State = EntityState<Entity>;

Considering this, for our actionsFactory we will need the Entity generic, and the setState and the idSelector functions as props.

interface ActionsFactoryProps<Entity extends object> {
  setState: SetState<Entity>;
  idSelector?: IdSelector<Entity>;
}

We'll mark idSelector as optional because we will use a default implementation:

const defaultIdSelector = (entity: any): string | number => entity.id;

With all that, we can create our actionsFactory.

export function actionsFactory<Entity extends object>({
  setState,
  idSelector = defaultIdSelector,
}: ActionFactoryProps<Entity>): EntityActions<Entity> {
  type State = EntityState<Entity>;

  const addOne = (state: State, entity: Entity): State => {
    ...
  }
  ...
  return { 
    addOne(entity: Entity) {
    ...
    }
    ...
 };
}

Unlike the state, I won't be adding the whole factory function for actions, because it's large, and I consider it will create noise in the post. But you can still check the implementation in the linked StackBlitz. Sorry!

Gimme, gimme, gimme... an entity

So far, we have the state and the actions. But, at the end of the day, we need to display that data somewhere, and the best way to handle this is by using selectors. Zustand itself is very unopinionated about the usage of complex selectors, if you want to use memoized selectors, that's on you. But when you use the store, you still need to pass a simple selector function.

Both libraries provide pretty much the same set of selectors, but RTK also provides a selectById selector, that we are going to implement. You need to consider that, because those selectors are going to be used within a Redux-like store, they can memoized them for you.

In our case, what we need is a simple selector function, and we can again leverage Zustand's features to ensure we won't have extra renders.

interface EntitySelectors<
  Entity,
  State extends EntityState<Entity> = EntityState<Entity>
> {
  selectIds(state: State): EntityId[];
  selectEntities(state: State): Dictionary<Entity>;
  selectAll(state: State): Entity[];
  selectTotal(state: State): number;
  selectById(id: EntityId): (state: State) => Entity | undefined;
}

Those are the methods we want to define. Maybe the new thing here is the State extends EntityState<Entity>. We already saw what the EntityState interface is, and what we want to do with this is to allow the user to use another State type, as long as it has the required properties in EntityState. We use a default value, to allow the user to only set the Entity generic type, and we can infer the other generic type.

And readers, it doesn't get any easier than this:

  const selectIds = (state: State) => state.ids;
  const selectEntities = (state: State) => state.entities;
  const selectAll = ({ entities, ids }: State) => ids.map((id) => entities[id]);
  const selectTotal = (state: State) => state.ids.length;
  const selectById = (id: EntityId) => ({ entities }: State) => entities[id];

Five inline functions allow us to get the slice of state that we want.

Be an adapter, my friend

With all the pieces we have, we can already create the adapter function:

type IdSelector<Entity> = (model: Entity) => EntityId;

interface EntityCreatorsProps<Entity extends object> {
  idSelector?: IdSelector<Entity>;
}

export function createEntityAdapter<Entity extends object>({
  idSelector,
}: EntityCreatorsProps<Entity>) {
  return {
    getState() {
      return stateFactory<Entity>();
    },
    getActions(setState: SetState<Entity>) {
      return actionsFactory({ setState, idSelector });
    },
    getSelectors() {
      return selectorsFactory<Entity>();
    },
  };
}

Here we see again our idSelector function, but it's clearer here what it is doing. We have a function, that will be called with an entity, and should return an ID for that entity. We have a default implementation for this, selecting the property id, so we don't have to implement this each time, as long as the default is a valid selector.

And that's it. We have created our version of the EntityAdapter. We are generating our state, our actions, and our selectors. And we have done so, thinking about Zustand first. Funny enough, we are not using any Zustand directly just yet, those are plain JS functions that we are creating!

That's all folks!

We had covered how to implement the main methods and functions we need to recreate our version of the EntityAdaptor for Zustand. In another post, I'm showing you how to use it right here!

Image generated with Microsoft Designer AI using as a prompt "A bear sits on a sofa reading a book in a living room with a firewood, with a 16 bits palette"