From 21f56db09aae6c32b1b4b00494cca9e7ccba65a4 Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Sun, 22 May 2022 10:41:31 +0200
Subject: [PATCH] refactor(store): add function for updating selection changes

---
 src/store/index.ts | 119 +++++++++++----------------------------------
 src/store/utils.ts |  45 +++++++++++------
 2 files changed, 58 insertions(+), 106 deletions(-)

diff --git a/src/store/index.ts b/src/store/index.ts
index ee7d90ec..419ecbab 100644
--- a/src/store/index.ts
+++ b/src/store/index.ts
@@ -3,7 +3,6 @@ import createContext from 'zustand/context';
 
 import { clampPosition, getDimensions } from '../utils';
 import { applyNodeChanges } from '../utils/changes';
-
 import {
   ReactFlowState,
   Node,
@@ -18,12 +17,7 @@ import {
 } from '../types';
 import { getHandleBounds } from '../components/Nodes/utils';
 import { createSelectionChange, getSelectionChanges } from '../utils/changes';
-import {
-  createNodeInternals,
-  handleControlledEdgeSelectionChange,
-  handleControlledNodeSelectionChange,
-  fitView,
-} from './utils';
+import { createNodeInternals, fitView, updateNodesAndEdgesSelections } from './utils';
 import initialState from './initialState';
 
 const { Provider, useStore, useStoreApi } = createContext<ReactFlowState>();
@@ -120,17 +114,8 @@ const createStore = () =>
         }
       }
     },
-    // @TODO: can we unify addSelectedNodes and addSelectedEdges somehow?
     addSelectedNodes: (selectedNodeIds: string[]) => {
-      const {
-        multiSelectionActive,
-        onNodesChange,
-        nodeInternals,
-        hasDefaultNodes,
-        onEdgesChange,
-        hasDefaultEdges,
-        edges,
-      } = get();
+      const { multiSelectionActive, nodeInternals, edges } = get();
       let changedNodes: NodeSelectionChange[];
       let changedEdges: EdgeSelectionChange[] | null = null;
 
@@ -141,32 +126,15 @@ const createStore = () =>
         changedEdges = getSelectionChanges(edges, []);
       }
 
-      if (changedNodes.length) {
-        if (hasDefaultNodes) {
-          set({ nodeInternals: handleControlledNodeSelectionChange(changedNodes, nodeInternals) });
-        }
-
-        onNodesChange?.(changedNodes);
-      }
-
-      if (changedEdges?.length) {
-        if (hasDefaultEdges) {
-          set({ edges: handleControlledEdgeSelectionChange(changedEdges, edges) });
-        }
-
-        onEdgesChange?.(changedEdges);
-      }
+      updateNodesAndEdgesSelections({
+        changedNodes,
+        changedEdges,
+        get,
+        set,
+      });
     },
     addSelectedEdges: (selectedEdgeIds: string[]) => {
-      const {
-        multiSelectionActive,
-        onEdgesChange,
-        edges,
-        hasDefaultEdges,
-        nodeInternals,
-        hasDefaultNodes,
-        onNodesChange,
-      } = get();
+      const { multiSelectionActive, edges, nodeInternals } = get();
       let changedEdges: EdgeSelectionChange[];
       let changedNodes: NodeSelectionChange[] | null = null;
 
@@ -177,25 +145,15 @@ const createStore = () =>
         changedNodes = getSelectionChanges(Array.from(nodeInternals.values()), []);
       }
 
-      if (changedEdges.length) {
-        if (hasDefaultEdges) {
-          set({
-            edges: handleControlledEdgeSelectionChange(changedEdges, edges),
-          });
-        }
-        onEdgesChange?.(changedEdges);
-      }
-
-      if (changedNodes?.length) {
-        if (hasDefaultNodes) {
-          set({ nodeInternals: handleControlledNodeSelectionChange(changedNodes, nodeInternals) });
-        }
-
-        onNodesChange?.(changedNodes);
-      }
+      updateNodesAndEdgesSelections({
+        changedNodes,
+        changedEdges,
+        get,
+        set,
+      });
     },
     unselectNodesAndEdges: () => {
-      const { nodeInternals, edges, onNodesChange, onEdgesChange, hasDefaultNodes, hasDefaultEdges } = get();
+      const { nodeInternals, edges } = get();
       const nodes = Array.from(nodeInternals.values());
 
       const nodesToUnselect = nodes.map((n) => {
@@ -204,22 +162,13 @@ const createStore = () =>
       }) as NodeSelectionChange[];
       const edgesToUnselect = edges.map((edge) => createSelectionChange(edge.id, false)) as EdgeSelectionChange[];
 
-      if (nodesToUnselect.length) {
-        if (hasDefaultNodes) {
-          set({ nodeInternals: handleControlledNodeSelectionChange(nodesToUnselect, nodeInternals) });
-        }
-        onNodesChange?.(nodesToUnselect);
-      }
-      if (edgesToUnselect.length) {
-        if (hasDefaultEdges) {
-          set({
-            edges: handleControlledEdgeSelectionChange(edgesToUnselect, edges),
-          });
-        }
-        onEdgesChange?.(edgesToUnselect);
-      }
+      updateNodesAndEdgesSelections({
+        changedNodes: nodesToUnselect,
+        changedEdges: edgesToUnselect,
+        get,
+        set,
+      });
     },
-
     setMinZoom: (minZoom: number) => {
       const { d3Zoom, maxZoom } = get();
       d3Zoom?.scaleExtent([minZoom, maxZoom]);
@@ -239,7 +188,7 @@ const createStore = () =>
       set({ translateExtent });
     },
     resetSelectedElements: () => {
-      const { nodeInternals, edges, onNodesChange, onEdgesChange, hasDefaultNodes, hasDefaultEdges } = get();
+      const { nodeInternals, edges } = get();
       const nodes = Array.from(nodeInternals.values());
 
       const nodesToUnselect = nodes
@@ -249,22 +198,12 @@ const createStore = () =>
         .filter((e) => e.selected)
         .map((e) => createSelectionChange(e.id, false)) as EdgeSelectionChange[];
 
-      if (nodesToUnselect.length) {
-        if (hasDefaultNodes) {
-          set({
-            nodeInternals: handleControlledNodeSelectionChange(nodesToUnselect, nodeInternals),
-          });
-        }
-        onNodesChange?.(nodesToUnselect);
-      }
-      if (edgesToUnselect.length) {
-        if (hasDefaultEdges) {
-          set({
-            edges: handleControlledEdgeSelectionChange(edgesToUnselect, edges),
-          });
-        }
-        onEdgesChange?.(edgesToUnselect);
-      }
+      updateNodesAndEdgesSelections({
+        changedNodes: nodesToUnselect,
+        changedEdges: edgesToUnselect,
+        get,
+        set,
+      });
     },
     setNodeExtent: (nodeExtent: CoordinateExtent) => {
       const { nodeInternals } = get();
diff --git a/src/store/utils.ts b/src/store/utils.ts
index 6cd542c6..d7c5921e 100644
--- a/src/store/utils.ts
+++ b/src/store/utils.ts
@@ -1,22 +1,17 @@
-// @ts-nocheck
 import { zoomIdentity } from 'd3-zoom';
-import { GetState } from 'zustand';
+import { GetState, SetState } from 'zustand';
 
-import { clampPosition, isNumeric } from '../utils';
+import { isNumeric } from '../utils';
 import { getD3Transition, getRectOfNodes, getTransformForBounds } from '../utils/graph';
 import {
-  CoordinateExtent,
   Edge,
   EdgeSelectionChange,
   Node,
   NodeInternals,
-  NodePositionChange,
   NodeSelectionChange,
   ReactFlowState,
-  XYPosition,
   XYZPosition,
   FitViewOptions,
-  SnapGrid,
 } from '../types';
 
 type ParentNodes = Record<string, boolean>;
@@ -92,15 +87,6 @@ export function createNodeInternals(nodes: Node[], nodeInternals: NodeInternals)
   return nextNodeInternals;
 }
 
-type CreatePostionChangeParams = {
-  node: Node;
-  nodeExtent: CoordinateExtent;
-  nodeInternals: NodeInternals;
-  diff?: XYPosition;
-  snapToGrid?: boolean;
-  snapGrid?: SnapGrid;
-};
-
 type InternalFitViewOptions = {
   initial?: boolean;
 } & FitViewOptions;
@@ -166,3 +152,30 @@ export function handleControlledEdgeSelectionChange(edgeChanges: EdgeSelectionCh
     return e;
   });
 }
+
+type UpdateNodesAndEdgesParams = {
+  changedNodes: NodeSelectionChange[] | null;
+  changedEdges: EdgeSelectionChange[] | null;
+  get: GetState<ReactFlowState>;
+  set: SetState<ReactFlowState>;
+};
+
+export function updateNodesAndEdgesSelections({ changedNodes, changedEdges, get, set }: UpdateNodesAndEdgesParams) {
+  const { nodeInternals, edges, onNodesChange, onEdgesChange, hasDefaultNodes, hasDefaultEdges } = get();
+
+  if (changedNodes?.length) {
+    if (hasDefaultNodes) {
+      set({ nodeInternals: handleControlledNodeSelectionChange(changedNodes, nodeInternals) });
+    }
+
+    onNodesChange?.(changedNodes);
+  }
+
+  if (changedEdges?.length) {
+    if (hasDefaultEdges) {
+      set({ edges: handleControlledEdgeSelectionChange(changedEdges, edges) });
+    }
+
+    onEdgesChange?.(changedEdges);
+  }
+}