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

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,110 +50,39 @@ 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[]))[];
-  };
+  const setNodes = useCallback<Instance.SetNodes<NodeType>>(
+    (payload) => {
+      batchContext?.nodeQueue.push(payload as NodeType[]);
+    },
+    [batchContext]
+  );
 
-  // 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);
+  const setEdges = useCallback<Instance.SetEdges<EdgeType>>(
+    (payload) => {
+      batchContext?.edgeQueue.push(payload as EdgeType[]);
+    },
+    [batchContext]
+  );
 
-  // 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;
-    }
+  const addNodes = useCallback<Instance.AddNodes<NodeType>>(
+    (payload) => {
+      const newNodes = Array.isArray(payload) ? payload : [payload];
 
-    if (setElementsQueue.current.nodes.length) {
-      const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState();
+      // Queueing a functional update means that we won't worry about other calls
+      // to `setNodes` that might happen elsewhere.
+      batchContext?.nodeQueue.push((nodes) => [...nodes, ...newNodes]);
+    },
+    [batchContext]
+  );
 
-      // 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;
-      }
+  const addEdges = useCallback<Instance.AddEdges<EdgeType>>(
+    (payload) => {
+      const newEdges = Array.isArray(payload) ? payload : [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);
-    setShouldFlushQueue(true);
-  }, []);
-
-  const setEdges = useCallback<Instance.SetEdges<EdgeType>>((payload) => {
-    setElementsQueue.current.edges.push(payload);
-    setShouldFlushQueue(true);
-  }, []);
-
-  const addNodes = useCallback<Instance.AddNodes<NodeType>>((payload) => {
-    const newNodes = Array.isArray(payload) ? payload : [payload];
-
-    // 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];
-
-    setElementsQueue.current.edges.push((edges) => [...edges, ...newEdges]);
-    setShouldFlushQueue(true);
-  }, []);
+      batchContext?.edgeQueue.push((edges) => [...edges, ...newEdges]);
+    },
+    [batchContext]
+  );
 
   const toObject = useCallback<Instance.ToObject<NodeType, EdgeType>>(() => {
     const { nodes = [], edges = [], transform } = store.getState();