Merge pull request #4209 from xyflow/fix/batching-provider

fix(react): handle setNodes batching in provider

examples/react/src/examples/UseNodesData/index.tsx

@@ -45,6 +45,12 @@ const initNodes: MyNode[] = [
     data: { text: '' },
     position: { x: 100, y: 0 },
   },
+  {
+    id: '1b',
+    type: 'uppercase',
+    data: { text: '' },
+    position: { x: 100, y: -100 },
+  },
   {
     id: '2',
     type: 'text',
@@ -53,9 +59,14 @@ const initNodes: MyNode[] = [
     },
     position: { x: 0, y: 100 },
   },
-
   {
-    id: '3',
+    id: '3a',
+    type: 'result',
+    data: {},
+    position: { x: 300, y: -75 },
+  },
+  {
+    id: '3b',
     type: 'result',
     data: {},
     position: { x: 300, y: 50 },
@@ -69,14 +80,24 @@ const initEdges: Edge[] = [
     target: '1a',
   },
   {
-    id: 'e1a-3',
+    id: 'e1a-3a',
-    source: '1a',
+    source: '1b',
-    target: '3',
+    target: '3a',
   },
   {
-    id: 'e2-3',
+    id: 'e1-1b',
+    source: '1',
+    target: '1b',
+  },
+  {
+    id: 'e1a-3b',
+    source: '1a',
+    target: '3b',
+  },
+  {
+    id: 'e2-3b',
     source: '2',
-    target: '3',
+    target: '3b',
   },
 ];
 

packages/react/CHANGELOG.md

@@ -1,5 +1,11 @@
 # @xyflow/react
 
+## 12.0.0-next.16
+
+## Patch changes
+
+- fix batching for `setNodes`, `updateNode`, `updateNodeData` etc.
+
 ## 12.0.0-next.15
 
 ## Patch changes

packages/react/src/components/BatchProvider/index.tsx

@@ -0,0 +1,88 @@
+import { createContext, ReactNode, useCallback, useContext, useMemo } from 'react';
+import { EdgeChange, NodeChange } from '@xyflow/system';
+
+import { useStoreApi } from '../../hooks/useStore';
+import { getElementsDiffChanges } from '../../utils';
+import { Queue, QueueItem } from './types';
+import type { Edge, Node } from '../../types';
+import { useQueue } from './useQueue';
+
+const BatchContext = createContext<{
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  nodeQueue: Queue<any>;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  edgeQueue: Queue<any>;
+} | null>(null);
+
+/**
+ * This is a context provider that holds and processes the node and edge update queues
+ * that are needed to handle setNodes, addNodes, setEdges and addEdges.
+ *
+ * @internal
+ */
+export function BatchProvider<NodeType extends Node = Node, EdgeType extends Edge = Edge>({
+  children,
+}: {
+  children: ReactNode;
+}) {
+  const store = useStoreApi<NodeType, EdgeType>();
+
+  const nodeQueueHandler = useCallback((queueItems: QueueItem<NodeType>[]) => {
+    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 queueItems) {
+      next = typeof payload === 'function' ? payload(next) : payload;
+    }
+
+    if (hasDefaultNodes) {
+      setNodes(next);
+    } else if (onNodesChange) {
+      onNodesChange(
+        getElementsDiffChanges({
+          items: next,
+          lookup: nodeLookup,
+        }) as NodeChange<NodeType>[]
+      );
+    }
+  }, []);
+  const nodeQueue = useQueue<NodeType>(nodeQueueHandler);
+
+  const edgeQueueHandler = useCallback((queueItems: QueueItem<EdgeType>[]) => {
+    const { edges = [], setEdges, hasDefaultEdges, onEdgesChange, edgeLookup } = store.getState();
+
+    let next = edges as EdgeType[];
+    for (const payload of queueItems) {
+      next = typeof payload === 'function' ? payload(next) : payload;
+    }
+
+    if (hasDefaultEdges) {
+      setEdges(next);
+    } else if (onEdgesChange) {
+      onEdgesChange(
+        getElementsDiffChanges({
+          items: next,
+          lookup: edgeLookup,
+        }) as EdgeChange<EdgeType>[]
+      );
+    }
+  }, []);
+  const edgeQueue = useQueue<EdgeType>(edgeQueueHandler);
+
+  const value = useMemo(() => ({ nodeQueue, edgeQueue }), []);
+
+  return <BatchContext.Provider value={value}>{children}</BatchContext.Provider>;
+}
+
+export function useBatchContext() {
+  const batchContext = useContext(BatchContext);
+
+  if (!batchContext) {
+    throw new Error('useBatchContext must be used within a BatchProvider');
+  }
+
+  return batchContext;
+}

packages/react/src/components/BatchProvider/types.ts

@@ -0,0 +1,7 @@
+export type QueueItem<T> = T[] | ((items: T[]) => T[]);
+
+export type Queue<T> = {
+  get: () => QueueItem<T>[];
+  reset: () => void;
+  push: (item: QueueItem<T>) => void;
+};

packages/react/src/components/BatchProvider/useQueue.ts

@@ -0,0 +1,67 @@
+import { useState } from 'react';
+
+import { useIsomorphicLayoutEffect } from '../../hooks/useIsomorphicLayoutEffect';
+import { Queue, QueueItem } from './types';
+
+/**
+ * This hook returns a queue that can be used to batch updates.
+ *
+ * @param runQueue - a function that gets called when the queue is flushed
+ * @internal
+ *
+ * @returns a Queue object
+ */
+export function useQueue<T>(runQueue: (items: QueueItem<T>[]) => void) {
+  // 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 [shouldFlush, setShouldFlush] = useState(false);
+
+  // A reference of all the batched updates to process before the next render. We
+  // want a reference here so multiple synchronous calls to `setNodes` etc can be
+  // batched together.
+  const [queue] = useState(() => createQueue<T>(() => setShouldFlush(true)));
+
+  // 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`).
+  useIsomorphicLayoutEffect(() => {
+    // 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 (!shouldFlush) {
+      queue.reset();
+      return;
+    }
+
+    const queueItems = queue.get();
+
+    if (queueItems.length) {
+      runQueue(queueItems);
+
+      queue.reset();
+    }
+
+    // Beacuse we're using reactive state to trigger this effect, we need to flip
+    // it back to false.
+    setShouldFlush(false);
+  }, [shouldFlush]);
+
+  return queue;
+}
+
+function createQueue<T>(cb: () => void): Queue<T> {
+  let queue: QueueItem<T>[] = [];
+
+  return {
+    get: () => queue,
+    reset: () => {
+      queue = [];
+    },
+    push: (item) => {
+      queue.push(item);
+      cb();
+    },
+  };
+}

packages/react/src/components/ReactFlowProvider/index.tsx

@@ -2,6 +2,7 @@ import { useState, type ReactNode } from 'react';
 
 import { Provider } from '../../contexts/StoreContext';
 import { createStore } from '../../store';
+import { BatchProvider } from '../BatchProvider';
 import type { Node, Edge } from '../../types';
 
 export type ReactFlowProviderProps = {
@@ -37,5 +38,9 @@ export function ReactFlowProvider({
     })
   );
 
-  return <Provider value={store}>{children}</Provider>;
+  return (
+    <Provider value={store}>
+      <BatchProvider>{children}</BatchProvider>
+    </Provider>
+  );
 }

packages/react/src/hooks/useReactFlow.ts

@@ -1,4 +1,4 @@
-import { useCallback, useMemo, useRef, useState } from 'react';
+import { useCallback, useMemo } from 'react';
 import {
   evaluateAbsolutePosition,
   getElementsToRemove,
@@ -10,9 +10,9 @@ import {
 
 import useViewportHelper from './useViewportHelper';
 import { useStoreApi } from './useStore';
+import { useBatchContext } from '../components/BatchProvider';
+import { isNode } from '../utils';
 import type { ReactFlowInstance, Instance, Node, Edge, InternalNode } from '../types';
-import { getElementsDiffChanges, isNode } from '../utils';
-import { useIsomorphicLayoutEffect } from './useIsomorphicLayoutEffect';
 
 /**
  * Hook for accessing the ReactFlow instance.
@@ -26,6 +26,7 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
 > {
   const viewportHelper = useViewportHelper();
   const store = useStoreApi();
+  const batchContext = useBatchContext();
 
   const getNodes = useCallback<Instance.GetNodes<NodeType>>(
     () => store.getState().nodes.map((n) => ({ ...n })) as NodeType[],
@@ -49,109 +50,22 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
 
   const getEdge = useCallback<Instance.GetEdge<EdgeType>>((id) => store.getState().edgeLookup.get(id) as EdgeType, []);
 
-  type SetElementsQueue = {
-    nodes: (NodeType[] | ((nodes: NodeType[]) => NodeType[]))[];
-    edges: (EdgeType[] | ((edges: EdgeType[]) => EdgeType[]))[];
-  };
-
-  // 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);
-
-  // 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`).
-  useIsomorphicLayoutEffect(() => {
-    // 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 (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;
-      }
-
-      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);
+    batchContext.nodeQueue.push(payload as NodeType[]);
-    setShouldFlushQueue(true);
   }, []);
 
   const setEdges = useCallback<Instance.SetEdges<EdgeType>>((payload) => {
-    setElementsQueue.current.edges.push(payload);
+    batchContext.edgeQueue.push(payload as EdgeType[]);
-    setShouldFlushQueue(true);
   }, []);
 
   const addNodes = useCallback<Instance.AddNodes<NodeType>>((payload) => {
     const newNodes = Array.isArray(payload) ? payload : [payload];
-
+    batchContext.nodeQueue.push((nodes) => [...nodes, ...newNodes]);
-    // 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 newEdges = Array.isArray(payload) ? payload : [payload];
-
+    batchContext.edgeQueue.push((edges) => [...edges, ...newEdges]);
-    setElementsQueue.current.edges.push((edges) => [...edges, ...newEdges]);
-    setShouldFlushQueue(true);
   }, []);
 
   const toObject = useCallback<Instance.ToObject<NodeType, EdgeType>>(() => {

packages/react/src/utils/changes.ts

@@ -236,7 +236,8 @@ export function getElementsDiffChanges({
   const itemsLookup = new Map<string, any>(items.map((item) => [item.id, item]));
 
   for (const item of items) {
-    const storeItem = lookup.get(item.id);
+    const lookupItem = lookup.get(item.id);
+    const storeItem = lookupItem?.internals?.userNode ?? lookupItem;
 
     if (storeItem !== undefined && storeItem !== item) {
       changes.push({ id: item.id, item: item, type: 'replace' });