fix(react): handle setNodes batching in provider #4147

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

@@ -0,0 +1,80 @@
+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<{
+  nodeQueue: Queue<Node>;
+  edgeQueue: Queue<Edge>;
+} | null>(null);
+
+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,44 @@
+import { useState } from 'react';
+import { createQueue } from './utils';
+import { useIsomorphicLayoutEffect } from '../../hooks/useIsomorphicLayoutEffect';
+import { QueueItem } from './types';
+
+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 mutable 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;
+}

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

@@ -0,0 +1,16 @@
+import type { Queue, QueueItem } from './types';
+
+export 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>
+  );
 }