chore(BatchProvider): cleanup

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',
-    source: '1a',
-    target: '3',
+    id: 'e1a-3a',
+    source: '1b',
+    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/src/components/BatchProvider/index.tsx

@@ -8,10 +8,18 @@ import type { Edge, Node } from '../../types';
 import { useQueue } from './useQueue';
 
 const BatchContext = createContext<{
-  nodeQueue: Queue<Node>;
-  edgeQueue: Queue<Edge>;
+  // 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,
 }: {

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

@@ -1,8 +1,16 @@
 import { useState } from 'react';
-import { createQueue } from './utils';
-import { useIsomorphicLayoutEffect } from '../../hooks/useIsomorphicLayoutEffect';
-import { QueueItem } from './types';
 
+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
@@ -10,8 +18,8 @@ export function useQueue<T>(runQueue: (items: QueueItem<T>[]) => void) {
   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.
+  // 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
@@ -30,7 +38,7 @@ export function useQueue<T>(runQueue: (items: QueueItem<T>[]) => void) {
     const queueItems = queue.get();
 
     if (queueItems.length) {
-      runQueue?.(queueItems);
+      runQueue(queueItems);
 
       queue.reset();
     }
@@ -42,3 +50,18 @@ export function useQueue<T>(runQueue: (items: QueueItem<T>[]) => void) {
 
   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/hooks/useReactFlow.ts

@@ -50,39 +50,23 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
 
   const getEdge = useCallback<Instance.GetEdge<EdgeType>>((id) => store.getState().edgeLookup.get(id) as EdgeType, []);
 
-  const setNodes = useCallback<Instance.SetNodes<NodeType>>(
-    (payload) => {
-      batchContext?.nodeQueue.push(payload as NodeType[]);
-    },
-    [batchContext]
-  );
+  const setNodes = useCallback<Instance.SetNodes<NodeType>>((payload) => {
+    batchContext.nodeQueue.push(payload as NodeType[]);
+  }, []);
 
-  const setEdges = useCallback<Instance.SetEdges<EdgeType>>(
-    (payload) => {
-      batchContext?.edgeQueue.push(payload as EdgeType[]);
-    },
-    [batchContext]
-  );
+  const setEdges = useCallback<Instance.SetEdges<EdgeType>>((payload) => {
+    batchContext.edgeQueue.push(payload as EdgeType[]);
+  }, []);
 
-  const addNodes = useCallback<Instance.AddNodes<NodeType>>(
-    (payload) => {
-      const newNodes = Array.isArray(payload) ? payload : [payload];
+  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.
-      batchContext?.nodeQueue.push((nodes) => [...nodes, ...newNodes]);
-    },
-    [batchContext]
-  );
-
-  const addEdges = useCallback<Instance.AddEdges<EdgeType>>(
-    (payload) => {
-      const newEdges = Array.isArray(payload) ? payload : [payload];
-
-      batchContext?.edgeQueue.push((edges) => [...edges, ...newEdges]);
-    },
-    [batchContext]
-  );
+  const addEdges = useCallback<Instance.AddEdges<EdgeType>>((payload) => {
+    const newEdges = Array.isArray(payload) ? payload : [payload];
+    batchContext.edgeQueue.push((edges) => [...edges, ...newEdges]);
+  }, []);
 
   const toObject = useCallback<Instance.ToObject<NodeType, EdgeType>>(() => {
     const { nodes = [], edges = [], transform } = store.getState();