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:
1const user = { 2 name: "Oliver", 3 age: 20, 4 address: { 5 city: "Miami", 6 state: "Florida", 7 country: "USA", 8 }, 9};
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.
1// Create a Context 2const UserContext = React.createContext(); 3 4// Wrap the parent component with the UserContext provider 5const App = () => ( 6 <UserContext.Provider value={user}> 7 {/* Other components that use the user Context */} 8 </UserContext.Provider> 9);
1// In a child component, access the user Context using useContext hook 2const UserProfile = () => { 3 const user = React.useContext(UserContext); 4 5 return ( 6 <div> 7 <p>{user.name}</p> 8 <p>{user.age}</p> 9 </div> 10 ); 11};
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.
1import create from "zustand"; 2 3// Create a user store using zustand 4const useUserStore = create(set => ({ 5 user: { 6 name: "Oliver", 7 age: 20, 8 address: { 9 city: "Miami", 10 state: "Florida", 11 country: "USA", 12 }, 13 } 14 setUser: set, 15}));
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.
1// Access the user via the useUserStore hook 2const UserProfile = () => { 3 const user = useUserStore(store => store.user); 4 5 return ( 6 <div> 7 <p>{user.name}</p> 8 <p>{user.age}</p> 9 </div> 10 ); 11};
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:
1import { shallow } from "zustand/shallow"; 2 3const { name, age } = useUserStore( 4 ({ user }) => ({ name: user.name, age: user.age }), 5 shallow 6);
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).
1const useUserStore = create(() => ({ name: "Oliver", age: 20 })); 2 3// Getting non-reactive fresh state 4const handleUpdate = () => { 5 if (useUserStore.getState().age === 20) { 6 // Our code here 7 } 8};
Working with Zustand
1. Installing zustand
1yarn 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.
1const useUserStore = create( 2 withImmutableActions(set => ({ 3 name: 10, 4 age: 20, 5 address: { 6 city: "Miami", 7 state: "Florida", 8 country: "USA", 9 }, 10 setName: ({ name }) => set({ name }), 11 setGlobalState: set, 12 })) 13);
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.
1setName: ({ 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.
1// throws an error 2setGlobalState({ name: 0, setName: () => {} });
Here is the source code of withImmutableActions:
1import { isEmpty, keys } from "ramda"; 2 3const setWithoutModifyingActions = set => partial => 4 set(previous => { 5 if (typeof partial === "function") partial = partial(previous); 6 7 const overwrittenActions = keys(partial).filter( 8 key => 9 typeof previous?.[key] === "function" && partial[key] !== previous[key] 10 ); 11 if (!isEmpty(overwrittenActions)) { 12 throw new Error( 13 `Actions should not be modified. Touched action(s): ${overwrittenActions.join( 14 ", " 15 )}` 16 ); 17 } 18 19 return partial; 20 }, false); 21 22const withImmutableActions = config => (set, get, api) => 23 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:
1setGlobalState({ value: 0 }, true); 2setGlobalState({ 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:
1const useUserStore = create(set => ({ 2 name: "", 3 subjects: [], 4 address: { 5 city: "", 6 country: "", 7 }, 8 setUser: set, 9}));
If we only need the city value, we can do:
1const 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
1// Not recommended 2const { 3 address: { city, country }, 4 setAddress, 5} = useUserStore();
We can replace the above code with the following approach:
1const { city, country } = useUserStore( 2 store => pick(["city", "country"], store.address), 3 shallow 4); 5const setAddress = useUserStore(prop("setAddress")); 6// `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)
1const handleUpdate = () => { 2 if (useUserStore.getState().role === "admin") { 3 // Our code here 4 } 5};
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.
1// useUserStore.js 2import { create } from "zustand"; 3 4const useUserStore = create(set => ({ 5 name: "", 6 subjects: [], 7 address: { 8 city: "", 9 country: "", 10 }, 11 setUser: set 12})); 13 14export default useUserStore; 15 16// App.jsx 17import React from "react"; 18import Profile from "./Profile"; 19 20const App = () => ( 21 <div> 22 <Profile role="Teacher" /> 23 <Profile role="Student" /> 24 </div> 25); 26 27export default App; 28 29// Profile.jsx 30import React from "react"; 31import { prop } from "ramda" 32 33import useUserStore from "./stores/useUserStore"; 34 35const Profile = ({ role }) => { 36 const name = useUserStore(prop("name")); 37 const setName = useUserStore(prop("setName")); 38 39 return ( 40 <div> 41 <p>{`Enter the ${role}'s name`}</p> 42 <input 43 value={name} 44 onChange={(e) => setName(e.target.value)} 45 /> 46 </div> 47 ); 48}; 49 50export 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.
1// Create a Context 2import { createContext } from "react"; 3 4const UserContext = createContext(null); 5 6export default UserContext; 7 8// Modify useUserStore using createStore 9import { createStore } from "zustand"; 10 11const useUserStore = () => 12 createStore((set) => ({ 13 name: "", 14 subjects: [], 15 address: { 16 city: "", 17 country: "" 18 }, 19 20 setUser: set 21 })); 22 23export default useUserStore; 24 25// Add changes to Profile.jsx 26import React, { useContext, useMemo } from "react"; 27import { pick } from "ramda"; 28import useUserStore from "./stores/useUserStore"; 29import { useStore } from "zustand"; 30import { shallow } from "zustand/shallow"; 31import UserContext from "./contexts/User"; 32 33const Profile = ({ role }) => { 34 const userStore = useContext(UserContext); 35 36 const { name, setName } = useStore( 37 userStore, 38 pick(["name", "setName"]), 39 shallow 40 ); 41 42 return ( 43 <div> 44 <p>{`Enter the ${role}'s name`}</p> 45 <input value={name} onChange={(e) => setName(e.target.value)} /> 46 </div> 47 ); 48}; 49 50const ProfileWithState = (props) => { 51 const stateStore = useMemo(useUserStore, []); 52 53 return ( 54 <UserContext.Provider value={stateStore}> 55 <Profile {...props} /> 56 </UserContext.Provider> 57 ); 58}; 59 60export 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:
1// Before 2import { shallow } from "zustand/shallow"; 3 4const { order, customer } = useGlobalStore( 5 store => ({ 6 order: store[sessionId]?.globals.order, 7 customer: store[sessionId]?.globals.customer, 8 }), 9 shallow 10);
With the babel plugin, the above code can be written as:
1//After 2const { 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.