State management in React Flow: How to manage it wisely?

Discover how to manage state in React Flow apps. Learn when
to use global and uncontrolled state, and optimize performance
with controlled state.

Introduction
When to use global state
Uncontrolled state
Controlled state
Summary

Introduction

State management in React Flow applications is a critical factor that directly impacts user interface performance and responsiveness. Poorly structured state can cause unnecessary re-renders, slowing down the app and frustrating users – especially in complex systems with large data sets and lots of interactions.
This article highlights best practices for state management, along with practical use cases for different scenarios.

When to use global state

When you need communication between components that aren’t directly nested in the component tree. Global state lets them share data and stay in sync even if they’re not structurally related.
When you want to enable interactions between nodes that affect the entire diagram. For instance, setting a node color that applies globally and updates across the app.
When you want to simplify components. With global state, components can just dispatch actions to update data, without passing event handlers down the component tree. This helps prevent tangled data flows.

Uncontrolled state

The simplest way to manage state in React Flow is to rely on the library’s built-in mechanisms – known as uncontrolled state. In this model, all interactions with nodes and edges are handled internally by React Flow. This makes implementation easier, as you don’t need to pass down action handlers or install any extra libraries.

useReactFlow() is a hook provided by the library. It’s a handy tool for accessing and controlling the state of your diagram – including nodes, edges, and viewport settings. It gives you access to the internal React Flow instance, which comes with a variety of useful methods and properties.

Adding nodes

Here’s a basic example of a component that adds a new node to the diagram:

1// AddNewNodeButton.tsx
2import { useReactFlow } from '@xyflow/react';
3const defaultNode = {
4  position: { x: 0, y: 0 },
5  type: 'default',
6  data: { label: 'New Node' },
7};
8export const AddNewNodeButton = () => {
9  const { addNodes } = useReactFlow();
10  const handleAddNode = () => {
11    addNodes({ ...defaultNode, id: crypto.randomUUID() });
12  };
13  return <button onClick={handleAddNode}>Add new node</button>;
14};

Getting nodes and edges

The useReactFlow() hook includes two helpful methods for accessing node data:

getNodes() – Returns all nodes currently in the diagram,
getNode(id: string) – Returns the specific node with the given ID.

The getNode() offers constant O(1) access time to a node, regardless of how many nodes exist. That makes it far more efficient than manually filtering arrays.

Similarly, you can use getEdges() and getEdge(id: string) to access edges.

Removing nodes and edges

To remove elements from the diagram, you can use the built-in deleteElements method returned by useReactFlow hook. This method accepts an object like this:

1{
2  nodes?: (Partial<Node> & { id: Node['id'] })[],
3  edges?: (Partial<Edge> & { id: Edge['id'] })[]
4}

It allows you to remove either nodes or edges. At a minimum, you need to pass an array with the id of the elements to delete.

When removing nodes, the method automatically deletes their connected edges and any child elements in grouped structures.

1// RemoveNodeWithEdgesButton.tsx
2import { useReactFlow } from '@xyflow/react';
3export const RemoveNodeWithEdgesButton = ({ nodeId } : { nodeId: string }) => {
4  const { deleteElements } = useReactFlow();
5  const handleRemoveNodeWithEdges = () => {
6    deleteElements({ nodes: [{ id: nodeId }] }); 
7  };
8  return <button onClick={handleRemoveNodeWithEdges}>Remove node</button>;
9};

Updating nodes and edges

To update nodes and edges, you can use the built-in hooks useNodesState() and useEdgeState() from the React Flow library. They’re similar to React’s native useState() hook, but extended with an additional helper:

1const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
2const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);

Both hooks return the current array of elements (like nodes), a setter function (like setNodes), and a callback that handles changes (like onNodesChange). When initializing the hook, you provide the initial array of elements – for example, initialNodes.

Here’s a sample implementation using these hooks to update nodes and edges:

1// App.tsx
2
3import { ReactFlow, useNodesState, useEdgesState, Panel, Node, Edge } from '@xyflow/react';
4import { useMemo } from 'react';
5
6type MyNode = Node & {
7  data: {
8    label: string;
9  };
10};
11
12const initialNodes = [
13  { id: '1', data: { label: 'New Node' }, position: { x: 100, y: 100 } },
14  { id: '2', data: { label: 'New Node 2' }, position: { x: 100, y: 200 } },
15] as MyNode[];
16const initialEdges = [{ id: 'e1-2', source: '1', target: '2' }] as Edge[];
17
18export function App() {
19  const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
20  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);
21
22  const selectedNode = useMemo(() => {
23    const selected = nodes.filter((node) => node.selected);
24    return selected?.[0] ?? null;
25  }, [nodes]);
26
27  const updateNodeLabel = (newLabel: string) => {
28    if (!selectedNode) return;
29    setNodes((nds) =>
30      nds.map((node) => (node.id === selectedNode.id ? { ...node, data: { ...node.data, label: newLabel } } : node)),
31    );
32  };
33
34  const updateNodeHiddenState = (newHidden: boolean) => {
35    if (!selectedNode) return;
36    setNodes((nds) => nds.map((node) => (node.id === selectedNode.id ? { ...node, hidden: newHidden } : node)));
37    setEdges((eds) =>
38      eds.map((edge) =>
39        edge.source === selectedNode.id || edge.target === selectedNode.id ? { ...edge, hidden: newHidden } : edge,
40      ),
41    );
42  };
43
44  return (
45    <ReactFlow nodes={nodes} edges={edges} onNodesChange={onNodesChange} onEdgesChange={onEdgesChange}>
46      {selectedNode && (
47        <Panel position="top-left" style={{ width: 200 }}>
48          <label>Label:</label>
49          <input value={selectedNode.data.label} onChange={(evt) => updateNodeLabel(evt.target.value)} />
50          <label>Hidden:</label>
51          <input
52            type="checkbox"
53            checked={selectedNode.hidden ?? false}
54            onChange={(evt) => updateNodeHiddenState(evt.target.checked)}
55          />
56        </Panel>
57      )}
58    </ReactFlow>
59  );
60}
61

In the example above, we update the label and hidden properties of a node, and also toggle the visibility of an edge based on the node’s visibility.

Controlled state

In the Uncontrolled state section, we covered how React Flow can work with local component state to manage nodes and edges in the diagram. However, as your application grows, managing that state can become more complex. To make things easier in larger apps, you can use React Context or integrate a state management library.

When opting for an external library for a React app, there are several options to choose from. For React Flow, Zustand and Redux are among the most popular choices. If your state is less complex, a lightweight library like Jotai may also be a good fit.

In the following examples, we’ll use Zustand – which is also used internally by React Flow itself.

Creating the store

First, define the type that describes your state:

1// use-diagram-store.ts
2
3import { Node, Edge, NodeChange, EdgeChange, Connection } from '@xyflow/react';
4
5export type DiagramState = {
6  nodes: Node[];
7  edges: Edge[];
8  reactFlowInstance: ReactFlowInstance | null;
9  selectedNodes: Node[];
10  setSelectedNodes: (nodes: Node[]) => void;
11  onNodesChange: (changes: NodeChange[]) => void;
12  onEdgesChange: (changes: EdgeChange[]) => void;
13  onConnect: (connection: Connection) => void;
14  onInit: (instance: ReactFlowInstance) => void;
15};

Then, use the create method to initialize your store with default values:

1// use-diagram-store.ts
2
3import { create } from 'zustand';
4import { applyNodeChanges, applyEdgeChanges, addEdge } from @xyflow/react;
5
6export const useDiagramStore = create<DiagramState>((set) => ({
7  nodes: [],
8  edges: [],
9  reactFlowInstance: null;
10  selectedNodes: [],
11  setSelectedNodes: (nodes) => set({ selectedNodes: nodes }),
12  onNodesChange: (changes) => set((state) => ({ nodes: applyNodeChanges(changes, state.nodes) })),
13  onEdgesChange: (changes) => set((state) => ({ edges: applyEdgeChanges(changes, state.edges) })),
14  onConnect: (connection) => set((state) => ({ edges: addEdge(connection, state.edges) })),
15  onInit: (instance) => {
16    (window as any).diagram = instance; // optionally, for testing purposes
17    set({ reactFlowInstance: instance });
18  },
19}));

Using the store

The create method returns a custom hook you can use across your app. Here’s how to access individual methods or values from the store:

1const selectedNodes = useDiagramStore((state) => state.selectedNodes);

To access multiple values from the store at once, use this approach:

1const [selectedNodes, someOtherProperty] = useDiagramStore((state) => [
2  state.selectedNodes,
3  state.someOtherProperty,
4]);

Just make sure your store is properly configured so it doesn't hurt your app’s performance. For more on this topic, see The ultimate guide to optimize React Flow project performance and the Optimization section.

Integrating the store with the ReactFlow component

To make our Zustand store respond to diagram changes – and vice versa – we need to pass its methods and values into the ReactFlow component:

1// Workspace.tsx
2
3import { useDiagramStore } from '@/store';
4import { ReactFlow, Background } from '@xyflow/react'; 
5  
6export const Workspace = () => {
7  // ... rest of the component
8  
9  const [
10    nodes, edges, onInit, onNodesChange, onEdgesChange, onConnect
11  ] = useDiagramStore((state) => [
12    state.nodes, state.edges, state.onInit, state.onNodesChange, state.onEdgesChange, state.onConnect
13  ]);
14  
15  return (
16    <ReactFlow
17      onInit={onInit}
18      nodes={nodes}
19      edges={edges}
20      onNodesChange={onNodesChange}
21      onEdgesChange={onEdgesChange}
22      onConnect={onConnect}
23      // some other props
24    >
25      <Background />
26    </ReactFlow>
27  );
28};

The ReactFlow component comes with many more methods and properties. You can find the full list here.

Adding nodes

To add nodes through the store, define a method that modifies the nodes array:

1// use-diagram-store.ts 
2
3import { getNodeAddChange } from '@/utils/get-node-add-change'; 
4 
5export type DiagramState = {
6  // ... rest of the type
7  addNodes: (nodes: Node[]) => void;
8};
9
10export const useDiagramStore = create<DiagramState>((set, get) => ({
11  // ... rest of the store
12  addNodes: (nodes) => set((state) => ({ nodes: [...state.nodes, ...nodes] })) 
13}));

Alternatively, instead of using a callback function, you can use the get method:

1addNodes: (nodes) => set({ nodes: get().nodes.concat(nodes) })

If your app fires off many rapid, concurrent actions to add nodes, the callback approach offers better protection against missing updates. For simpler, less frequent calls, both approaches should behave the same way.

Here’s an example component that adds a new node to the diagram:

1// AddNewNodeButton.tsx
2
3const defaultNode = {
4  position: { x: 0, y: 0 },
5  type: 'default',
6  data: { label: 'New Node' },
7};
8  
9export const AddNewNodeButton = () => {
10  const addNodes = useDiagramStore((store) => store.addNodes);
11
12  const handleAddNode = () => {
13    addNodes([{ ...defaultNode, id: crypto.randomUUID() }]);
14  };
15
16  return <button onClick={handleAddNode}>Add New Node</button>;
17};

Removing nodes and edges

When setting up the ReactFlow component, we’ve already connected the onNodesChange and onEdgesChange handlers. That means we don’t need to define extra methods for removing elements. We can simply reuse the deleteElements method – just like in the Uncontrolled state. One key advantage here is that React Flow automatically removes any orphaned edges and child elements from deleted groups. This setup also enables support for other deletion methods, like pressing the Backspace key.

The example below works the same way as in the uncontrolled state section – the removal is now fully integrated with our state.

1// AddNewNodeButton.tsx
2
3const defaultNode = {
4  position: { x: 0, y: 0 },
5  type: 'default',
6  data: { label: 'New Node' },
7};
8  
9export const AddNewNodeButton = () => {
10  const addNodes = useDiagramStore((store) => store.addNodes);
11
12  const handleAddNode = () => {
13    addNodes([{ ...defaultNode, id: crypto.randomUUID() }]);
14  };
15
16  return <button onClick={handleAddNode}>Add New Node</button>;
17};

Slicing the store into parts

Let’s look at the current definition of our state:

1// use-diagram-store.ts
2
3export type DiagramState = {
4  nodes: Node[];
5  edges: Edge[];
6  reactFlowInstance: ReactFlowInstance | null;
7  selectedNodes: Node[];
8  setSelectedNodes: (nodes: Node[]) => void;
9  onNodesChange: (changes: NodeChange[]) => void;
10  onEdgesChange: (changes: EdgeChange[]) => void;
11  onConnect: (connection: Connection) => void;
12  onInit: (instance: ReactFlowInstance) => void;
13  addNodes: (nodes: Node[]) => void;
14};

We’ll now add a few more values – like the currently dragged element from the palette, whether the palette is expanded, and the selection state of nodes and edges.

1// use-diagram-store.ts
2
3export type DiagramState = {
4  nodes: Node[];
5  edges: Edge[];
6  reactFlowInstance: ReactFlowInstance | null;
7  selectedNodes: Node[];
8  setSelectedNodes: (nodes: Node[]) => void;
9  onNodesChange: (changes: NodeChange[]) => void;
10  onEdgesChange: (changes: EdgeChange[]) => void;
11  onConnect: (connection: Connection) => void;
12  onInit: (instance: ReactFlowInstance) => void;
13  addNodes: (nodes: Node[]) => void;
14  // ------
15  isPaletteExpanded: boolean;
16  setIsPaletteExpanded: (isExpanded: boolean) => void;
17  draggedItemType: DraggedItem | null;
18  setDraggedItemType: (item: DraggedItem | null) => void;
19  selectedElements: (Node | Edge)[];
20  setSelectedElements: (elements: (Node | Edge)[]) => void;
21  addToSelection: (elements: (Node | Edge)[]) => void;
22  removeFromSelection: (elements: (Node | Edge)[] | string[]) => void;
23  hoveredElementId: string | null;
24};

As the state grows, every new method or value makes it more complex. To avoid creating large, hard-to-maintain files, it’s worth breaking the state into smaller slices.

In Zustand, slices help you organize large and complex state into manageable sections. Each slice represents a distinct part of the application’s state, typically tied to a specific feature or domain.

Start by defining types for the set and get methods:

1// use-diagram-store.ts
2
3export type SetDiagramState = (
4  partial:
5    | DiagramState
6    | Partial<DiagramState>
7    | ((
8        state: DiagramState,
9      ) => DiagramState | Partial<DiagramState>),
10  replace?: false | undefined,
11) => void;
12
13export type GetDiagramState = () => DiagramState;

A sample slice structure might look like this:

1// create-workspace-slice.ts
2
3export type WorkspaceState = {
4  nodes: Node[];
5  edges: Edge[];
6  reactFlowInstance: ReactFlowInstance | null;
7  onNodesChange: (changes: NodeChange[]) => void;
8  onEdgesChange: (changes: EdgeChange[]) => void;
9  onConnect: (connection: Connection) => void;
10  onInit: (instance: ReactFlowInstance) => void;
11  addNodes: (nodes: Node[]) => void;
12};
13
14export const createWorkspaceSlice = (set: SetDiagramState, get: GetDiagramState): WorkspaceState => ({
15  nodes: [],
16  edges: [],
17  // ... rest of the slice
18});

1// create-palette-slice.ts
2
3export type PaletteState = {
4  isPaletteExpanded: boolean;
5  setIsPaletteExpanded: (isExpanded: boolean) => void;
6  draggedItemType: DraggedItem | null;
7  setDraggedItemType: (item: DraggedItem | null) => void;
8};
9
10export const createPaletteSlice = (set: SetDiagramState, get: GetDiagramState): PaletteState => ({
11  isPaletteExpanded: true,
12  setIsPaletteExpanded: (isPaletteExpanded) => set({ isPaletteExpanded }),
13  draggedItemType: null,
14  setDraggedItemType: (draggedItemType) => set({ draggedItemType }),
15});

1// create-selection-slice.ts
2
3export type SelectionState = {
4  selectedElements: (Node | Edge)[];
5  setSelectedElements: (elements: (Node | Edge)[]) => void;
6  addToSelection: (elements: (Node | Edge)[]) => void;
7  removeFromSelection: (elements: (Node | Edge)[] | string[]) => void;
8  hoveredElementId: string | null;
9};
10
11export const createSelectionSlice = (set: SetDiagramState, get: GetDiagramState): SelectionState => ({
12  selectedElements: [],
13  setSelectedElements: (selectedElements) => set({ selectedElements }),
14  // ... rest of the store
15});

Back into the use-diagram-store.ts file:

1// use-diagram-store.ts
2export type DiagramState = WorkspaceState & PaletteState & SelectionState;
3
4export const useDiagramStore = create<DiagramState>((set, get) => {
5  ...createWorkspaceSlice(set, get),
6  ...createPaletteSlice(set, get),
7  ...createSelectionSlice(set, get),
8});

Using a sliced store doesn’t change how it behaves – it simply improves code clarity and organization. You can still access all values and methods from the same hook:

1const [selectedElements, isPaletteExpanded, reactFlowInstance] = useDiagramStore((state) => [
2  state.selectedElements,
3  state.isPaletteExpanded,
4  state.reactFlowInstance,
5]);

For more on using slices in Zustand, check out Slices Pattern.

Optimization

As mentioned earlier, when creating the store, our current approach can lead to performance issues. That’s because the array returned by useStore is recreated every time the state changes. React treats this as a new reference, which triggers component re-renders. To reduce unnecessary re-renders when working with our store, we can use Zustand’s useShallow hook:

1const [selectedNodes, someOtherProperty] = useDiagramStore(useShallow((state) => [
2  state.selectedNodes,
3  state.someOtherProperty,
4]));

Or we can adjust how we create the store using createWithEqualityFn, and pass the shallow comparison function as a parameter. This way, all selectors will automatically use shallow comparison, which helps minimize re-renders of components.

1import { shallow } from 'zustand/shallow';
2import { createWithEqualityFn } from 'zustand/traditional';
3
4// ...
5
6export const useDiagramStore = createWithEqualityFn<DiagramState>((set) => ({
7  // ... no changes
8}), shallow);

Initializing state with default values

If the default values, you want to use in the state are static – like a fixed diagram template – you can define them as a variable and pass them directly when creating the store.

1// use-diagram-store.ts
2
3const defaultNodes: Node[] = [
4  { id: 'default-1', data: { label: 'Default node' }, position: { x: 0, y: 0 }},
5  { id: 'default-2', data: { label: 'Second default node' }, { x: 0, y: 100 }}
6];
7
8const defaultEdges = Edge[] = [
9  { id: 'default-1=>default-2', source: 'default-1', target: 'default-2' },
10];
11
12const useDiagramStore = create<DiagramState>((set) => ({
13  nodes: defaultNodes,
14  edges: defaultEdges,
15  // ...
16}));

If you’re using slices, the default values should be defined inside their respective slices.

Saving and loading state

When working with diagrams in React Flow, you’ll often need to save the state and load it later from an external source like a database or API. To keep things simple, the example here uses localStorage.

Let’s create a hook for saving and loading the diagram:

1import { ReactFlowJsonObject, useReactFlow } from '@xyflow/react';
2import { localStorageSaveKey } from '@/config';
3import { useCallback } from 'react';
4
5export function useDiagramDataPersistence() {
6  const { setNodes, setEdges, toObject, setViewport } = useReactFlow();
7
8  const saveDiagram = useCallback(() => {
9    try {
10      const diagramData = toObject();
11      const json = JSON.stringify(diagramData);
12      localStorage.setItem(localStorageSaveKey, json);
13    } catch (e) {
14      console.error('Failed to save diagram:', e);
15    }
16  }, [toObject]);
17
18  const loadDiagram = useCallback(() => {
19    try {
20      const data = localStorage.getItem(localStorageSaveKey);
21      if (!data) return;
22
23      const diagram = JSON.parse(data) as ReactFlowJsonObject;
24      const { nodes, edges, viewport } = diagram;
25
26      setNodes(nodes);
27      setEdges(edges);
28      setViewport(viewport);
29    } catch (e) {
30      console.error('Failed to load diagram:', e);
31    }
32  }, [setNodes, setEdges, setViewport]);
33
34  return { saveDiagram, loadDiagram };
35}
36

Here’s a sample use case inside a component:

1// Workspace.tsx
2
3import { useDiagramStore } from '@/store';
4import { ReactFlow, Background } from '@xyflow/react'; 
5import { useDiagramDataPersistence } from '@/hooks'
6  
7export const Workspace = () => {
8  // ... rest of the component
9  
10  const { saveDiagram, loadDiagram } = useDiagramDataPersistence()
11
12  return (
13    <ReactFlow
14      // ... props
15    >
16      <button onClick={saveDiagram}>Save diagram</button>
17      <button onClick={loadDiagram}>Load diagram</button>
18      <Background />
19    </ReactFlow>
20  );
21};

You can extend the loading function to accept input data in a specific format and load it directly – for example, when fetching from an external source.

Defining actions outside the state

Let’s revisit our state definition before we introduce slices:

1// use-diagram-store.ts
2
3export type DiagramState = {
4  nodes: Node[];
5  edges: Edge[];
6  reactFlowInstance: ReactFlowInstance | null;
7  selectedNodes: Node[];
8  setSelectedNodes: (nodes: Node[]) => void;
9  addNodes: (nodes: Node[]) => void;
10  // ...
11};
12
13export const useDiagramStore = create<DiagramState>((set) => ({
14  nodes: [],
15  edges: [],
16  reactFlowInstance: null,
17  selectedNodes: [],
18  setSelectedNodes: (nodes) => set({ selectedNodes: nodes }),
19  addNodes: (nodes) => set((state) => ({ nodes: [...state.nodes, ...nodes] })),
20  // ... 
21}));

An alternative to storing everything inside the store is to move methods and actions outside.

1// use-diagram-store.ts
2
3export type DiagramState = {
4  nodes: Node[];
5  edges: Edge[];
6  reactFlowInstance: ReactFlowInstance | null;
7  selectedNodes: Node[];
8  // ...
9};
10
11export const useDiagramStore = create<DiagramState>((set) => ({
12  nodes: [],
13  edges: [],
14  reactFlowInstance: null,
15  selectedNodes: [],
16  // ... 
17}));
18
19export const addNodes = (nodes: Node[]) => 
20  useDiagramStore.setState((state) => ({ nodes: [...state.nodes, ...nodes] }));
21  
22export const setSeletedNodes = (nodes: Node[]) => 
23  useDiagramStore.setState((state) => ({ selectedNodes: nodes }));

This approach eliminates the need to call actions through the hook and can simplify your code. For example, the addNodes method must be added to the useCallback dependency array:

1import { useDiagramStore } from '@/store';
2
3export const SomeComponent = () => {
4  const addNodes = useDiagramStore(store => store.addNodes);
5
6  const handleSomething = useCallback(() => {
7    // ... 
8    addNodes(nodes);  
9  }, [addNodes]); // <-- required in dependency array  
10};

But if addNodes is defined outside the useDiagramStore hook, there’s no need to add it to the dependencies:

1import { addNodes } from '@/store';
2
3export const SomeComponent = () => {
4  const handleSomething = useCallback(() => { 
5    addNodes(nodes);  
6  }, []); // <-- not required
7};

In short – if you want full encapsulation of state and logic, follow Zustand’s recommended pattern and keep everything inside the store.

Alternatively, separating state from actions offers some benefits:

Separation of state and logic: the store holds values only, while functions live in standalone modules,
No need to use the hook when calling an action: you can trigger addNodes directly without useDiagramStore,
Code splitting: it allows easier loading of chunks of functionality and doesn’t require using useCallback – making it easier to take care of performance.

Selectors

To retrieve data stored in multiple fields of the store, or to process that data at the moment of access (e.g., by filtering it), we can use selectors – a familiar concept from working with React stores.

Example selectors:

1// @/store/selectors/diagram-selectors.ts
2
3export const diagramStateSelector = ({
4  nodes,
5  edges,
6  selectedNodes,
7}: DiagramState) => ({
8  nodes,
9  edges,
10})

1// @/store/selectors/palette-selector.ts
2 
3export const paletteSelector = ({
4  isPaletteExpanded,
5  setIsPaletteExpanded
6}: PaletteState) => ({
7  isPaletteExpanded,
8  setIsPaletteExpanded
9})

Sample usage:

1// Diagram.tsx
2
3import { diagramStateSelector } from '@/store/selectors/diagram-selectors'
4
5export const Diagram = () => {
6  const { nodes, edges } = useDiagramStore(diagramStateSelector);
7  
8  return (
9    <ReactFlow 
10      nodes={nodes}
11      edges={edges}
12    />
13  );
14}

1// Palette.tsx
2
3import { paletteSelector } from '@/store/selectors/palette-selectors'
4
5export const Palette = () => {
6  const {  isPaletteExpanded, setIsPaletteExpanded } = useDiagramStore(paletteSelector);
7  
8  return (
9    <button onClick={() => setIsPaletteExpanded(!isPaletteExpanded)}>Toggle Palette</button>
10    {isPaletteExpanded && 
11      <div>
12        ...
13      </div>
14    }
15  );
16}

Zustand also supports automatic selector generation when creating a store using the createSelectors function. You can find more details about this feature in the documentation.

Memoized selectors

Be careful when using selectors like the ones shown above. They often return new references, which bypass Zustand’s default caching mechanism. They can also perform heavy computations, and frequent calls may affect performance. To handle both issues, you can use memoized selectors. These cache the results and return the same reference if the input hasn’t changed.

A popular library for memoized selectors is Reselect. Its core function is createSelector, which takes an array of input selectors and a transform function. The transform function receives the output from those input selectors and calculates a final value.

1import { createSelector } from 'reselect';
2
3// Simple selectors
4const selectAllNodes = (state: DiagramState) => state.nodes;
5const selectSelectedNodeIds = (state: DiagramState) => state.selectedNodeIds;
6
7// Standard selector
8export const getSelectedNodes = (state: DiagramState) => {
9  console.log('Getting selected nodes...');
10     
11  const allNodes = selectAllnodes(state);
12  return allNodes.filter((node) => state.selectedNodeIds.includes(node.id)); 
13};
14
15getSelectedNodes(state); // 'Getting selected nodes';
16getSelectedNodes(state); // 'Getting selected nodes';
17getSelectedNodes(state); // 'Getting selected nodes';
18 
19// Memoized reselect selector
20export const getMemoizedSelectedNodes = createSelector(
21  [selectAllNodes, selectSelectedNodeIds],
22  (allNodes, selectedIds) => {           
23    console.log('Getting memoized selected nodes');
24    
25    return allNodes.filter((node) => selectedIds.includes(node.id));
26  }
27);
28
29getMemoizedSelectedNodes(state); // 'Getting memoized selected nodes'
30getMemoizedSelectedNodes(state);
31getMemoizedSelectedNodes(state);
32
33console.log(getSelectedNodes(state) === getSelectedNodes(state)); // false
34console.log(getMemoizedSelectedNodes(state) === getMemoizedSelectedNodes(state)); // true

As shown in the example, getMemoizedSelectedNodes is called only once, and its value is a reference to the same object – unlike the standard getSelectedNodes.

Summary

This article presented different ways to manage state in a diagram-based application built with React Flow, with a special focus on optimization techniques. We covered two main approaches.

The first – uncontrolled state – relies entirely on React Flow’s built-in hooks and methods. It combines simplicity with solid performance and is ideal for smaller projects or proof-of-concept (PoC) work.

The second approach uses a controlled state, where an external store plays the key role. While the app still uses some built-in hooks and features from React Flow, they support the logic handled by the store. This method is more complex and requires careful implementation to avoid performance bottlenecks.

To address those concerns, it’s worth using the techniques mentioned earlier – like the useShallow hook or the createWithEqualityFn function with the shallow comparator while defining store. Another optimization option is to use memoized selectors, for example with the Reselect library.

In conclusion, with proper attention to optimization, the controlled state can work extremely well in large-scale, long-term projects. It offers better scalability, improved data flow between unrelated components, and full control over every part of the application.

State management in React Flow: How to manage it wisely?

Table of contents

Introduction

When to use global state

Uncontrolled state

Adding nodes

Getting nodes and edges

Removing nodes and edges

Updating nodes and edges

Controlled state

Creating the store

Using the store

Integrating the store with the ReactFlow component

Adding nodes

Removing nodes and edges

Slicing the store into parts

Optimization

Initializing state with default values

Saving and loading state

Defining actions outside the state

Selectors

Memoized selectors

Summary