From a0323064081c6051429ef1b1fb5d3986ee387e75 Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Fri, 26 Nov 2021 14:39:22 +0100
Subject: [PATCH] feat(hooks): add useNodesState and useEdgesState

---
 example/src/Basic/index.tsx         | 18 ++++--------------
 example/src/CustomNode/index.tsx    | 22 ++++++----------------
 example/src/FloatingEdges/index.tsx | 22 +++++-----------------
 example/src/Hidden/index.tsx        | 16 +++++++++++++---
 example/src/Layouting/index.tsx     | 18 ++++++------------
 example/src/NestedNodes/index.tsx   | 18 ++++--------------
 src/hooks/useNodesEdgesState.ts     | 27 +++++++++++++++++++++++++++
 src/index.ts                        |  1 +
 8 files changed, 66 insertions(+), 76 deletions(-)
 create mode 100644 src/hooks/useNodesEdgesState.ts

diff --git a/example/src/Basic/index.tsx b/example/src/Basic/index.tsx
index c49da47b..eee70449 100644
--- a/example/src/Basic/index.tsx
+++ b/example/src/Basic/index.tsx
@@ -3,17 +3,15 @@ import { useState, MouseEvent, useCallback } from 'react';
 import ReactFlow, {
   addEdge,
   Background,
-  applyNodeChanges,
-  applyEdgeChanges,
   MiniMap,
   Controls,
   Node,
   Edge,
-  NodeChange,
-  EdgeChange,
   OnLoadParams,
   Connection,
   MarkerType,
+  useNodesState,
+  useEdgesState,
 } from 'react-flow-renderer';
 import DebugNode from './DebugNode';
 
@@ -108,8 +106,8 @@ const nodeTypes = {
 
 const BasicFlow = () => {
   const [rfInstance, setRfInstance] = useState<OnLoadParams | null>(null);
-  const [nodes, setNodes] = useState<Node[]>(initialNodes);
-  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
+  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);
 
   const onConnect = useCallback((connection: Connection) => {
     setEdges((eds) => {
@@ -152,14 +150,6 @@ const BasicFlow = () => {
     });
   };
 
-  const onNodesChange = useCallback((changes: NodeChange[]) => {
-    setNodes((ns) => applyNodeChanges(changes, ns));
-  }, []);
-
-  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
-    setEdges((es) => applyEdgeChanges(changes, es));
-  }, []);
-
   return (
     <ReactFlow
       nodes={nodes}
diff --git a/example/src/CustomNode/index.tsx b/example/src/CustomNode/index.tsx
index 9ac2aa52..c47b0794 100644
--- a/example/src/CustomNode/index.tsx
+++ b/example/src/CustomNode/index.tsx
@@ -1,4 +1,4 @@
-import { useState, useEffect, MouseEvent, useCallback } from 'react';
+import { useState, useEffect, MouseEvent } from 'react';
 import { ChangeEvent } from 'react';
 
 import ReactFlow, {
@@ -10,11 +10,8 @@ import ReactFlow, {
   Position,
   SnapGrid,
   Connection,
-  Edge,
-  NodeChange,
-  applyNodeChanges,
-  applyEdgeChanges,
-  EdgeChange,
+  useNodesState,
+  useEdgesState,
 } from 'react-flow-renderer';
 
 import ColorSelectorNode from './ColorSelectorNode';
@@ -35,8 +32,9 @@ const nodeTypes = {
 };
 
 const CustomNodeFlow = () => {
-  const [nodes, setNodes] = useState<Node[]>([]);
-  const [edges, setEdges] = useState<Edge[]>([]);
+  const [nodes, setNodes, onNodesChange] = useNodesState([]);
+  const [edges, setEdges, onEdgesChange] = useEdgesState([]);
+
   const [bgColor, setBgColor] = useState<string>(initBgColor);
 
   useEffect(() => {
@@ -103,14 +101,6 @@ const CustomNodeFlow = () => {
   const onConnect = (connection: Connection) =>
     setEdges((eds) => addEdge({ ...connection, animated: true, style: { stroke: '#fff' } }, eds));
 
-  const onNodesChange = useCallback((changes: NodeChange[]) => {
-    setNodes((ns) => applyNodeChanges(changes, ns));
-  }, []);
-
-  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
-    setEdges((es) => applyEdgeChanges(changes, es));
-  }, []);
-
   return (
     <ReactFlow
       nodes={nodes}
diff --git a/example/src/FloatingEdges/index.tsx b/example/src/FloatingEdges/index.tsx
index 82168d2c..c489ec0b 100644
--- a/example/src/FloatingEdges/index.tsx
+++ b/example/src/FloatingEdges/index.tsx
@@ -1,17 +1,13 @@
-import { useState, useCallback } from 'react';
+import { useCallback } from 'react';
 
 import ReactFlow, {
   addEdge,
   Background,
   OnLoadParams,
   EdgeTypesType,
-  Node,
   Connection,
-  Edge,
-  applyNodeChanges,
-  applyEdgeChanges,
-  NodeChange,
-  EdgeChange,
+  useNodesState,
+  useEdgesState,
 } from 'react-flow-renderer';
 
 import './style.css';
@@ -29,21 +25,13 @@ const edgeTypes: EdgeTypesType = {
 };
 
 const FloatingEdges = () => {
-  const [nodes, setNodes] = useState<Node[]>(initialNodes);
-  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
+  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);
 
   const onConnect = useCallback((connection: Connection) => {
     setEdges((eds) => addEdge(connection, eds));
   }, []);
 
-  const onNodesChange = useCallback((changes: NodeChange[]) => {
-    setNodes((ns) => applyNodeChanges(changes, ns));
-  }, []);
-
-  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
-    setEdges((es) => applyEdgeChanges(changes, es));
-  }, []);
-
   return (
     <div className="floatingedges">
       <ReactFlow
diff --git a/example/src/Hidden/index.tsx b/example/src/Hidden/index.tsx
index ac3e578c..384bbe4c 100644
--- a/example/src/Hidden/index.tsx
+++ b/example/src/Hidden/index.tsx
@@ -1,7 +1,16 @@
 import { useState, useCallback } from 'react';
 
 import { useEffect } from 'react';
-import ReactFlow, { addEdge, MiniMap, Controls, Connection, Edge, Node } from 'react-flow-renderer';
+import ReactFlow, {
+  addEdge,
+  MiniMap,
+  Controls,
+  Connection,
+  Edge,
+  Node,
+  useNodesState,
+  useEdgesState,
+} from 'react-flow-renderer';
 
 const initialNodes: Node[] = [
   { id: '1', type: 'input', hidden: true, data: { label: 'Node 1' }, position: { x: 250, y: 5 } },
@@ -23,8 +32,9 @@ const setHidden = (hidden: boolean) => (els: any[]) =>
   });
 
 const HiddenFlow = () => {
-  const [nodes, setNodes] = useState<Node[]>(initialNodes);
-  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const [nodes, setNodes] = useNodesState(initialNodes);
+  const [edges, setEdges] = useEdgesState(initialEdges);
+
   const [isHidden, setIsHidden] = useState<boolean>(true);
 
   const onConnect = useCallback((connection: Connection) => {
diff --git a/example/src/Layouting/index.tsx b/example/src/Layouting/index.tsx
index e2616672..a4fc53c7 100644
--- a/example/src/Layouting/index.tsx
+++ b/example/src/Layouting/index.tsx
@@ -1,4 +1,4 @@
-import { useState, useCallback } from 'react';
+import { useCallback } from 'react';
 import ReactFlow, {
   ReactFlowProvider,
   addEdge,
@@ -7,15 +7,15 @@ import ReactFlow, {
   Controls,
   NodeChange,
   EdgeChange,
-  Node,
   Connection,
-  Edge,
   CoordinateExtent,
   Position,
+  useNodesState,
+  useEdgesState,
 } from 'react-flow-renderer';
 import dagre from 'dagre';
 
-import initialNodesAndEdges from './initial-elements';
+import initialItems from './initial-elements';
 
 import './layouting.css';
 
@@ -28,8 +28,8 @@ const nodeExtent: CoordinateExtent = [
 ];
 
 const LayoutFlow = () => {
-  const [nodes, setNodes] = useState<Node[]>(initialNodesAndEdges.nodes);
-  const [edges, setEdges] = useState<Edge[]>(initialNodesAndEdges.edges);
+  const [nodes, setNodes, onNodesChange] = useNodesState(initialItems.nodes);
+  const [edges, setEdges, onEdgesChange] = useEdgesState(initialItems.edges);
 
   const onConnect = useCallback((connection: Connection) => {
     setEdges((eds) => addEdge(connection, eds));
@@ -63,12 +63,6 @@ const LayoutFlow = () => {
     setNodes(layoutedNodes);
   };
 
-  const onNodesChange = useCallback((changes: NodeChange[]) => setNodes((ns) => applyNodeChanges(changes, ns)), []);
-
-  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
-    setEdges((es) => applyEdgeChanges(changes, es));
-  }, []);
-
   return (
     <div className="layoutflow">
       <ReactFlowProvider>
diff --git a/example/src/NestedNodes/index.tsx b/example/src/NestedNodes/index.tsx
index 27b4caae..bf46d158 100644
--- a/example/src/NestedNodes/index.tsx
+++ b/example/src/NestedNodes/index.tsx
@@ -3,14 +3,12 @@ import { useState, MouseEvent, useCallback } from 'react';
 import ReactFlow, {
   addEdge,
   Background,
-  applyNodeChanges,
-  applyEdgeChanges,
+  useNodesState,
+  useEdgesState,
   MiniMap,
   Controls,
   Node,
   Edge,
-  NodeChange,
-  EdgeChange,
   OnLoadParams,
   Connection,
 } from 'react-flow-renderer';
@@ -87,8 +85,8 @@ const initialEdges: Edge[] = [
 
 const NestedFlow = () => {
   const [rfInstance, setRfInstance] = useState<OnLoadParams | null>(null);
-  const [nodes, setNodes] = useState<Node[]>(initialNodes);
-  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
+  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);
 
   const onConnect = useCallback((connection: Connection) => {
     setEdges((eds) => addEdge(connection, eds));
@@ -129,14 +127,6 @@ const NestedFlow = () => {
     });
   };
 
-  const onNodesChange = useCallback((changes: NodeChange[]) => {
-    setNodes((ns) => applyNodeChanges(changes, ns));
-  }, []);
-
-  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
-    setEdges((es) => applyEdgeChanges(changes, es));
-  }, []);
-
   return (
     <ReactFlow
       nodes={nodes}
diff --git a/src/hooks/useNodesEdgesState.ts b/src/hooks/useNodesEdgesState.ts
new file mode 100644
index 00000000..d64abdb7
--- /dev/null
+++ b/src/hooks/useNodesEdgesState.ts
@@ -0,0 +1,27 @@
+import { useState, useCallback, SetStateAction, Dispatch } from 'react';
+
+import { applyNodeChanges, applyEdgeChanges } from '../utils/changes';
+import { Node, NodeChange, Edge, EdgeChange } from '../types';
+
+type ApplyChanges<ItemType, ChangesType> = (changes: ChangesType[], items: ItemType[]) => ItemType[];
+type OnChange<ChangesType> = (changes: ChangesType[]) => void;
+
+// returns a hook that can be used liked this:
+// const [nodes, setNodes, onNodesChange] = useNodesState(intialNodes);
+function createUseItemsState<ItemType, ChangesType>(
+  applyChangesFunction: ApplyChanges<ItemType, ChangesType>
+): (initialItems: ItemType[]) => [ItemType[], Dispatch<SetStateAction<ItemType[]>>, OnChange<ChangesType>] {
+  return (initialItems: ItemType[]) => {
+    const [items, setItems] = useState<ItemType[]>(initialItems);
+
+    const onItemsChange = useCallback(
+      (changes: ChangesType[]) => setItems((items) => applyChangesFunction(changes, items)),
+      []
+    );
+
+    return [items, setItems, onItemsChange];
+  };
+}
+
+export const useNodesState = createUseItemsState<Node, NodeChange>(applyNodeChanges as ApplyChanges<Node, NodeChange>);
+export const useEdgesState = createUseItemsState<Edge, EdgeChange>(applyEdgeChanges as ApplyChanges<Edge, EdgeChange>);
diff --git a/src/index.ts b/src/index.ts
index 6aa465d3..71856c24 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -27,6 +27,7 @@ export { applyNodeChanges, applyEdgeChanges } from './utils/changes';
 
 export { default as useZoomPanHelper } from './hooks/useZoomPanHelper';
 export { default as useUpdateNodeInternals } from './hooks/useUpdateNodeInternals';
+export * from './hooks/useNodesEdgesState';
 
 export * from './additional-components';
 export { useStore, useStoreApi } from './store';