Merge pull request #3907 from xyflow/refactor/set-nodes

Refactor(setNodes): improve current batching solution

examples/react/src/App/routes.ts

@@ -28,6 +28,7 @@ import NodeTypesObjectChange from '../examples/NodeTypesObjectChange';
 import Overview from '../examples/Overview';
 import Provider from '../examples/Provider';
 import SaveRestore from '../examples/SaveRestore';
+import SetNodesBatching from '../examples/SetNodesBatching';
 import Stress from '../examples/Stress';
 import Subflow from '../examples/Subflow';
 import SwitchFlow from '../examples/Switch';
@@ -226,6 +227,11 @@ const routes: IRoute[] = [
     path: 'save-restore',
     component: SaveRestore,
   },
+  {
+    name: 'SetNodes Batching',
+    path: 'setnodes-batching',
+    component: SetNodesBatching,
+  },
   {
     name: 'Stress',
     path: 'stress',

examples/react/src/examples/SetNodesBatching/index.tsx

@@ -0,0 +1,70 @@
+import { useCallback } from 'react';
+import {
+  ReactFlow,
+  MiniMap,
+  Background,
+  BackgroundVariant,
+  Controls,
+  ReactFlowProvider,
+  Node,
+  Edge,
+  useReactFlow,
+  Panel,
+} from '@xyflow/react';
+
+const a = { id: 'a', data: { label: 'A' }, position: { x: 250, y: 5 } };
+const b = { id: 'b', data: { label: 'B' }, position: { x: 100, y: 100 } };
+const c = { id: 'c', data: { label: 'C' }, position: { x: 400, y: 100 } };
+
+const SetNotesBatchingFlow = () => {
+  const { setNodes, updateNode } = useReactFlow();
+
+  const triggerMultipleSetNodes = useCallback(() => {
+    setNodes([a]);
+    setNodes((nodes) => [...nodes, b]);
+    setNodes((nodes) => [...nodes, c]);
+    setNodes((nodes) =>
+      nodes.map((node) =>
+        node.id === 'a' ? { ...node, position: { x: node.position.x + 20, y: node.position.y + 20 } } : node
+      )
+    );
+  }, []);
+
+  const triggerMultipleUpdateNodes = useCallback(() => {
+    triggerMultipleSetNodes();
+    updateNode('a', (a) => ({ position: { x: a.position.x + 20, y: a.position.y + 20 } }));
+    updateNode('b', (b) => ({ position: { x: b.position.x + 20, y: b.position.y + 20 } }));
+    updateNode('c', (c) => ({ position: { x: c.position.x + 20, y: c.position.y + 20 } }));
+    updateNode('a', (a) => ({ data: { ...a.data, label: `A ${Date.now()}` } }));
+    updateNode('b', (b) => ({ data: { ...b.data, label: `B ${Date.now()}` } }));
+    updateNode('c', (c) => ({ data: { ...c.data, label: `C ${Date.now()}` } }));
+  }, []);
+
+  return (
+    <ReactFlow
+      defaultNodes={[]}
+      defaultEdges={[]}
+      className="react-flow-basic-example"
+      minZoom={0.2}
+      maxZoom={4}
+      fitView
+    >
+      <Background variant={BackgroundVariant.Dots} />
+      <MiniMap />
+      <Controls />
+
+      <Panel position="top-right">
+        <button onClick={triggerMultipleSetNodes}>queue multiple setNodes calls</button>
+        <button onClick={triggerMultipleUpdateNodes}>queue multiple updateNode calls</button>
+      </Panel>
+    </ReactFlow>
+  );
+};
+
+export default function App() {
+  return (
+    <ReactFlowProvider>
+      <SetNotesBatchingFlow />
+    </ReactFlowProvider>
+  );
+}

examples/react/src/examples/UseNodesData/TextNode.tsx

@@ -14,8 +14,16 @@ function TextNode({ id, data }: NodeProps) {
   return (
     <div style={{ background: '#eee', color: '#222', padding: 10, fontSize: 12, borderRadius: 10 }}>
       <div>node {id}</div>
-      <div>
-        <input onChange={(evt) => updateText(evt.target.value)} value={text} />
+      <div style={{ marginTop: 5 }}>
+        <label style={{ fontSize: 12 }}>useState + updateNodeData</label>
+        <input onChange={(evt) => updateText(evt.target.value)} value={text} style={{ display: 'block' }} />
+
+        <label style={{ fontSize: 12 }}>updateNodeData</label>
+        <input
+          onChange={(evt) => updateNodeData(id, { text: evt.target.value })}
+          value={data.text}
+          style={{ display: 'block' }}
+        />
       </div>
       <Handle type="source" position={Position.Right} />
     </div>

packages/react/src/hooks/useReactFlow.ts

@@ -1,18 +1,9 @@
-import { useCallback, useMemo, useRef } from 'react';
+import { useCallback, useLayoutEffect, useMemo, useRef, useState } from 'react';
 import { getElementsToRemove, getOverlappingArea, isRectObject, nodeToRect, type Rect } from '@xyflow/system';
 
 import useViewportHelper from './useViewportHelper';
 import { useStoreApi } from './useStore';
-import type {
-  ReactFlowInstance,
-  Instance,
-  NodeAddChange,
-  EdgeAddChange,
-  Node,
-  Edge,
-  NodeChange,
-  EdgeChange,
-} from '../types';
+import type { ReactFlowInstance, Instance, Node, Edge } from '../types';
 import { getElementsDiffChanges, isNode } from '../utils';
 
 /**
@@ -46,82 +37,109 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
     return edges.find((e) => e.id === id) as EdgeType;
   }, []);
 
-  // this is used to handle multiple syncronous setNodes calls
-  const setNodesData = useRef<Node[]>();
-  const setNodesTimeout = useRef<ReturnType<typeof setTimeout>>();
-  const setNodes = useCallback<Instance.SetNodes<NodeType>>((payload) => {
-    const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState();
-    const nextNodes = typeof payload === 'function' ? payload((setNodesData.current as NodeType[]) || nodes) : payload;
+  type SetElementsQueue = {
+    nodes: (NodeType[] | ((nodes: NodeType[]) => NodeType[]))[];
+    edges: (EdgeType[] | ((edges: EdgeType[]) => EdgeType[]))[];
+  };
 
-    setNodesData.current = nextNodes;
+  // A reference of all the batched updates to process before the next render. We
+  // want a mutable reference here so multiple synchronous calls to `setNodes` etc
+  // can be batched together.
+  const setElementsQueue = useRef<SetElementsQueue>({ nodes: [], edges: [] });
+  // Because we're using a ref above, we need some way to let React know when to
+  // actually process the queue. We flip this bit of state to `true` any time we
+  // mutate the queue and then flip it back to `false` after flushing the queue.
+  const [shouldFlushQueue, setShouldFlushQueue] = useState(false);
 
-    if (setNodesTimeout.current) {
-      clearTimeout(setNodesTimeout.current);
+  // Layout effects are guaranteed to run before the next render which means we
+  // shouldn't run into any issues with stale state or weird issues that come from
+  // rendering things one frame later than expected (we used to use `setTimeout`).
+  useLayoutEffect(() => {
+    // Because we need to flip the state back to false after flushing, this should
+    // trigger the hook again (!). If the hook is being run again we know that any
+    // updates should have been processed by now and we can safely clear the queue
+    // and bail early.
+    if (!shouldFlushQueue) {
+      setElementsQueue.current = { nodes: [], edges: [] };
+      return;
     }
 
-    // if there are multiple synchronous setNodes calls, we only want to call onNodesChange once
-    // for this, we use a timeout to wait for the last call and store updated nodes in setNodesData
-    // this is not perfect, but should work in most cases
-    setNodesTimeout.current = setTimeout(() => {
-      if (hasDefaultNodes) {
-        setNodes(nextNodes);
-      } else if (onNodesChange) {
-        const changes: NodeChange[] = getElementsDiffChanges({ items: setNodesData.current, lookup: nodeLookup });
-        onNodesChange(changes);
+    if (setElementsQueue.current.nodes.length) {
+      const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState();
+
+      // This is essentially an `Array.reduce` in imperative clothing. Processing
+      // this queue is a relatively hot path so we'd like to avoid the overhead of
+      // array methods where we can.
+      let next = nodes as NodeType[];
+      for (const payload of setElementsQueue.current.nodes) {
+        next = typeof payload === 'function' ? payload(next) : payload;
       }
 
-      setNodesData.current = undefined;
-    }, 0);
+      if (hasDefaultNodes) {
+        setNodes(next);
+      } else if (onNodesChange) {
+        onNodesChange(
+          getElementsDiffChanges({
+            items: next,
+            lookup: nodeLookup,
+          })
+        );
+      }
+
+      setElementsQueue.current.nodes = [];
+    }
+
+    if (setElementsQueue.current.edges.length) {
+      const { edges = [], setEdges, hasDefaultEdges, onEdgesChange, edgeLookup } = store.getState();
+
+      let next = edges as EdgeType[];
+      for (const payload of setElementsQueue.current.edges) {
+        next = typeof payload === 'function' ? payload(next) : payload;
+      }
+
+      if (hasDefaultEdges) {
+        setEdges(next);
+      } else if (onEdgesChange) {
+        onEdgesChange(
+          getElementsDiffChanges({
+            items: next,
+            lookup: edgeLookup,
+          })
+        );
+      }
+
+      setElementsQueue.current.edges = [];
+    }
+
+    // Beacuse we're using reactive state to trigger this effect, we need to flip
+    // it back to false.
+    setShouldFlushQueue(false);
+  }, [shouldFlushQueue]);
+
+  const setNodes = useCallback<Instance.SetNodes<NodeType>>((payload) => {
+    setElementsQueue.current.nodes.push(payload);
+    setShouldFlushQueue(true);
   }, []);
 
-  // this is used to handle multiple syncronous setEdges calls
-  const setEdgesData = useRef<Edge[]>();
-  const setEdgesTimeout = useRef<ReturnType<typeof setTimeout>>();
   const setEdges = useCallback<Instance.SetEdges<EdgeType>>((payload) => {
-    const { edges = [], setEdges, hasDefaultEdges, onEdgesChange, edgeLookup } = store.getState();
-    const nextEdges = typeof payload === 'function' ? payload((setEdgesData.current as EdgeType[]) || edges) : payload;
-
-    setEdgesData.current = nextEdges;
-
-    if (setEdgesTimeout.current) {
-      clearTimeout(setEdgesTimeout.current);
-    }
-
-    setEdgesTimeout.current = setTimeout(() => {
-      if (hasDefaultEdges) {
-        setEdges(nextEdges);
-      } else if (onEdgesChange) {
-        const changes: EdgeChange[] = getElementsDiffChanges({ items: nextEdges, lookup: edgeLookup });
-        onEdgesChange(changes);
-      }
-
-      setEdgesData.current = undefined;
-    }, 0);
+    setElementsQueue.current.edges.push(payload);
+    setShouldFlushQueue(true);
   }, []);
 
   const addNodes = useCallback<Instance.AddNodes<NodeType>>((payload) => {
-    const nodes = Array.isArray(payload) ? payload : [payload];
-    const { nodes: currentNodes, hasDefaultNodes, onNodesChange, setNodes } = store.getState();
+    const newNodes = Array.isArray(payload) ? payload : [payload];
 
-    if (hasDefaultNodes) {
-      const nextNodes = [...currentNodes, ...nodes];
-      setNodes(nextNodes);
-    } else if (onNodesChange) {
-      const changes = nodes.map((node) => ({ item: node, type: 'add' } as NodeAddChange<NodeType>));
-      onNodesChange(changes);
-    }
+    // Queueing a functional update means that we won't worry about other calls
+    // to `setNodes` that might happen elsewhere.
+    setElementsQueue.current.nodes.push((nodes) => [...nodes, ...newNodes]);
+    setShouldFlushQueue(true);
   }, []);
 
   const addEdges = useCallback<Instance.AddEdges<EdgeType>>((payload) => {
-    const nextEdges = Array.isArray(payload) ? payload : [payload];
-    const { edges = [], setEdges, hasDefaultEdges, onEdgesChange } = store.getState();
+    const newEdges = Array.isArray(payload) ? payload : [payload];
 
-    if (hasDefaultEdges) {
-      setEdges([...edges, ...nextEdges]);
-    } else if (onEdgesChange) {
-      const changes = nextEdges.map((edge) => ({ item: edge, type: 'add' } as EdgeAddChange<EdgeType>));
-      onEdgesChange(changes);
-    }
+    setElementsQueue.current.edges.push((edges) => [...edges, ...newEdges]);
+    setShouldFlushQueue(true);
   }, []);
 
   const toObject = useCallback<Instance.ToObject<NodeType, EdgeType>>(() => {