From React context to zustand: A seamless transition

Global state refers to data that needs to be accessible and shared across different parts of an application. Unlike local or component-specific state, global state is not confined to a particular component but is available throughout the entire application.

Let's dive into a real-world scenario to understand the need for a global state in a React application.

Imagine we're building a sophisticated e-commerce platform with various components, such as a product catalog, a shopping cart, and a user profile. Each of these components requires access to shared data, like the user's authentication status and the contents of their shopping cart.

In the application, the user logs in on the homepage and starts adding products to their shopping cart. As the user navigates through different sections, such as the product catalog or the user profile, we need to decide how do we seamlessly share and manage the user's authentication status and the contents of the shopping cart across these disparate components? This is where the concept of a global state comes into the picture.

In the early stages of our application development, we might adopt React Context to manage this global state.

In this blog post, we'll discuss the process of upgrading from traditional React Context to Zustand, a state management library that offers simplicity, efficiency, and improved performance.

The Pitfalls of React Context

In our initial setup, we relied on React contexts for managing global states. However, as our application grew, we encountered performance issues and cumbersome boilerplate code. Let's consider a typical scenario where we need a global user state:

const user = {
  name: "Oliver",
  age: 20,
  address: {
    city: "Miami",
    state: "Florida",
    country: "USA",
  },
};

To use this global state, we had to create a Context, wrap the child components within a provider, and use the useContext hook in the child components. This led to unnecessary re-renders and increased boilerplate.

// Create a Context
const UserContext = React.createContext();

// Wrap the parent component with the UserContext provider
const App = () => (
  <UserContext.Provider value={user}>
    {/* Other components that use the user Context */}
  </UserContext.Provider>
);

// In a child component, access the user Context using useContext hook
const UserProfile = () => {
  const user = React.useContext(UserContext);

  return (
    <div>
      <p>{user.name}</p>
      <p>{user.age}</p>
    </div>
  );
};

Components that listen to the Context will trigger a re-render whenever any value within the Context changes, even if those changes are unrelated to the specific component.

For example, in the UserProfile component, if the value of city changes in the Context, the component will re-render, even if the address values aren't actually utilized within UserProfile. This can have a noticeable impact on performance. Furthermore, the usage of Context involves a lot of boilerplate code.

Enter Zustand: A Breath of Fresh Air

Zustand emerged as our solution to these challenges. It offered a more streamlined approach to global state management, addressing the performance concerns.

The useUserStore hook is created using zustand's create function. It initializes a store with initial state values and actions to update the state.

import create from "zustand";

// Create a user store using zustand
const useUserStore = create(set => ({
  user: {
    name: "Oliver",
    age: 20,
    address: {
      city: "Miami",
      state: "Florida",
      country: "USA",
    },
  }
  setUser: set,
}));

The UserProfile component uses the useUserStore hook to access the user state. The store => store.user function is passed as an argument to the hook, which retrieves the user object from the store.

// Access the user via the useUserStore hook
const UserProfile = () => {
  const user = useUserStore(store => store.user);

  return (
    <div>
      <p>{user.name}</p>
      <p>{user.age}</p>
    </div>
  );
};

In this component, useUserStore is used to access the entire user object from the store. Any change in the user object, even if it's a nested property like age, will trigger a re-render of the UserProfile component. This behavior is similar to how React Contexts work.

The first argument to the useUserStore hook is a selector function. Using the selector function, we can specify what to pick from the store. Zustand compares the previous and current values of the selected data and if the current and previous values are different, zustand triggers a re-render.

In the above example, store => store.user is the selector function. Zustand will compare the previous value of user with the current value and will trigger a re-render if the values are different. But inside this component, we need the values of only name and age properties of the user object.

This is where Zustand's ability to selectively pick specific parts of the state for a component comes into play, offering potential performance optimizations.

If we want to construct a single object with multiple state-picks inside, we can use shallow function to prevent unnecessary rerenders.

For example, we can be more specific by picking only name and age values from user store:

import { shallow } from "zustand/shallow";

const { name, age } = useUserStore(
  ({ user }) => ({ name: user.name, age: user.age }),
  shallow
);

Without shallow, the function ({ user }) => ({ name: user.name, age: user.age }) recreates the object { name: user.name, age: user.age } everytime it is called.

shallow is a function of comparison which checks for equality at the top level of the object, without performing a deep comparison of nested properties. Zustand's default behavior is to use Object.is for comparisons of the current and previous values. Even though the current and previous objects can have the same properties with equal values, they are not considered equal when compared using the strict equality operator ( Object.is ), same in the case of arrays. By adding shallow it will dig into the array/object and compare its key values or elements in array and if any one is different it triggers again.

In the above case, shallow ensures that the UserProfile component will re-render only if the name or age properties of the user object change.

Zustand also provides the getState function as a way to directly access the state of a store. This function can be particularly useful when we want to access the state outside of the component rendering cycle.

When using a value within a specific function, the getState() retrieves the latest value at the time of calling. It is useful to avoid having the value loaded using the hook (which will trigger a re-render when this value changes).

const useUserStore = create(() => ({ name: "Oliver", age: 20 }));

// Getting non-reactive fresh state
const handleUpdate = () => {
  if (useUserStore.getState().age === 20) {
    // Our code here
  }
};

Working with Zustand

1. Installing zustand

yarn add zustand

2. Replacing all contexts with zustand

During the initial migration, we replaced all React contexts with Zustand. This involved copying data and replacing Context hooks with Zustand stores. Our focus was on the migration itself, deferring performance enhancements for a later phase.

In the context of Zustand, "actions" refer to functions that are responsible for updating the state. In other words, actions are methods that modify the data within the state container.

const useUserStore = create(
  withImmutableActions(set => ({
    name: 10,
    age: 20,
    address: {
      city: "Miami",
      state: "Florida",
      country: "USA",
    },
    setName: ({ name }) => set({ name }),
    setGlobalState: set,
  }))
);

In the provided code snippet, setName and setGlobalState are examples of actions. Let's break it down:

setName: This action takes an object as an argument, specifically { name }, and updates the name property of the state with the provided value.

setName: ({ name }) => set({ name }),

setGlobalState: Similarly, this action takes an argument, and in this case, it merges the state with the provided argument. It's a more generic action that allows modifying multiple properties of the state at once.

To safeguard against actions being overwritten, we introduced a middleware function called withImmutableActions

This middleware ensures that attempts to overwrite Zustand store actions result in an error, providing a safeguard against unintended behavior.

The withImmutableActions throws an error because we are trying to overwrite the zustand store's actions.

// throws an error
setGlobalState({ name: 0, setName: () => {} });

Here is the source code of withImmutableActions:

import { isEmpty, keys } from "ramda";

const setWithoutModifyingActions = set => partial =>
  set(previous => {
    if (typeof partial === "function") partial = partial(previous);

    const overwrittenActions = keys(partial).filter(
      key =>
        typeof previous?.[key] === "function" && partial[key] !== previous[key]
    );
    if (!isEmpty(overwrittenActions)) {
      throw new Error(
        `Actions should not be modified. Touched action(s): ${overwrittenActions.join(
          ", "
        )}`
      );
    }

    return partial;
  }, false);

const withImmutableActions = config => (set, get, api) =>
  config(setWithoutModifyingActions(set), get, api);

Unlike zustand's default behavior, this middleware disregards the second argument of the set function which is used to overwrite the entire state when set to true. Hence, the following lines of code work identical to each other:

setGlobalState({ value: 0 }, true);
setGlobalState({ value: 0 });

3. Performance optimization strategies

We identified key strategies to optimize performance while using Zustand:

Selective Data Usage

Instead of using the entire state, components can selectively choose the data they need. This ensures that re-renders occur only when relevant data changes.

Consider the following user store:

const useUserStore = create(set => ({
  name: "",
  subjects: [],
  address: {
    city: "",
    country: "",
  },
  setUser: set,
}));

If we only need the city value, we can do:

const city = useUserStore(store => store.address.city);

In this case, the usage of shallow is not necessary because the selected data is a primitive value (city), not a complex object with nested properties. shallow is not needed when the returned object can be compared using Object.is operator.

Avoiding importing values like contexts

// Not recommended
const {
  address: { city, country },
  setAddress,
} = useUserStore();

We can replace the above code with the following approach:

const { city, country } = useUserStore(
  store => pick(["city", "country"], store.address),
  shallow
);
const setAddress = useUserStore(prop("setAddress"));
// `pick` and `prop` are imported from ramda

Avoiding prop drilling

Directly accessing Zustand values within the intended component eliminates the need for prop drilling, improving code clarity and maintainability.

Utilizing getState Method

When used within a function, the getState() retrieves the latest value at the time of calling the function. It is useful to avoid having the value loaded using the hook (which will trigger a re-render when this value changes)

const handleUpdate = () => {
  if (useUserStore.getState().role === "admin") {
    // Our code here
  }
};

Challenges and Solutions

Shared State Instances Across Components

Zustand's design maintains a single instance of state and actions. When using the same store hook across multiple components, values are shared. To address this, we combined Zustand with React Context, achieving a balance between efficient state management and isolation.

When we call the store hook (useUserStore) from different components which need separate states, the values returned by the hook will be the same across those components.

This behavior is a consequence of zustand's design. It maintains a single instance of the state and actions, ensuring that all components using the same hook share the same state and actions.

To illustrate this, consider an example where we have two input components on a form page: one for the Student profile and another for the Teacher profile. Both components are utilizing the same useUserStore to manage both student and teacher details.

// useUserStore.js
import { create } from "zustand";

const useUserStore = create(set => ({
  name: "",
  subjects: [],
  address: {
    city: "",
    country: "",
  },
  setUser: set
}));

export default useUserStore;

// App.jsx
import React from "react";
import Profile from "./Profile";

const App = () => (
  <div>
    <Profile role="Teacher" />
    <Profile role="Student" />
  </div>
);

export default App;

// Profile.jsx
import React from "react";
import { prop } from "ramda"

import useUserStore from "./stores/useUserStore";

const Profile = ({ role }) => {
  const name = useUserStore(prop("name"));
  const setName = useUserStore(prop("setName"));

  return (
    <div>
      <p>{`Enter the ${role}'s name`}</p>
      <input
        value={name}
        onChange={(e) => setName(e.target.value)}
      />
    </div>
  );
};

export default Profile;

In this setup, since both the Student and Teacher profiles are using the same store (useUserStore), the input fields in both components will display the same value.

We combined zustand with React Context to address this challenge of shared state instances across different components on the same page. By doing so, we have achieved a balance between the benefits of zustand's efficient state management and the isolation provided by React Context.

// Create a Context
import { createContext } from "react";

const UserContext = createContext(null);

export default UserContext;

// Modify useUserStore using createStore
import { createStore } from "zustand";

const useUserStore = () =>
  createStore((set) => ({
    name: "",
    subjects: [],
    address: {
      city: "",
      country: ""
    },

   setUser: set
  }));

export default useUserStore;

// Add changes to Profile.jsx
import React, { useContext, useMemo } from "react";
import { pick } from "ramda";
import useUserStore from "./stores/useUserStore";
import { useStore } from "zustand";
import { shallow } from "zustand/shallow";
import UserContext from "./contexts/User";

const Profile = ({ role }) => {
  const userStore = useContext(UserContext);

  const { name, setName } = useStore(
    userStore,
    pick(["name", "setName"]),
    shallow
  );

  return (
    <div>
      <p>{`Enter the ${role}'s name`}</p>
      <input value={name} onChange={(e) => setName(e.target.value)} />
    </div>
  );
};

const ProfileWithState = (props) => {
  const stateStore = useMemo(useUserStore, []);

  return (
    <UserContext.Provider value={stateStore}>
      <Profile {...props} />
    </UserContext.Provider>
  );
};

export default ProfileWithState;

With this implementation, each component gets its own isolated state, avoiding the issue of shared state instances.

Tackling Boilerplate Code

In the codebase, there was a recurring pattern of boilerplate code when trying to pick specific properties from a Zustand store with nested values. This involved using shallow and manually accessing nested properties, resulting in verbose code.

To simplify this process and reduce boilerplate, a custom babel plugin was developed. This plugin provides a cleaner syntax for property picking from Zustand stores.

Without the plugin, to pick specific values from the store, we needed to write:

// Before
import { shallow } from "zustand/shallow";

const { order, customer } = useGlobalStore(
  store => ({
    order: store[sessionId]?.globals.order,
    customer: store[sessionId]?.globals.customer,
  }),
  shallow
);

With the babel plugin, the above code can be written as:

//After
const { order, customer } = useGlobalStore.pick([sessionId, "globals"]);

The babel transformer will transform this code to the one shown above to achieve the same result.

A transformer is a module with a specific goal that is run against our code to transform it. The Babel plugin operates during the code compilation process. By using the Babel plugin, developers can achieve the same functionality with fewer lines of code, reducing the code verbosity.

The useGlobalStore.pick syntax provides a more streamlined and expressive way of picking properties. It abstracts away the need for manual property access and the use of shallow.

Conclusion

Upgrading to Zustand has proven to be a wise decision, addressing performance concerns and streamlining our state management. By combining Zustand with React Context and tackling challenges with innovative solutions, we've achieved a robust and efficient state management system in our React applications.