From 007bdabb897628c68a8ea0fa3fd739b605db0837 Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Tue, 22 Feb 2022 17:38:53 +0100
Subject: [PATCH] fix(selection): deselect nodes when select edge and vice
 versa closes #1899

---
 src/store/index.ts | 49 ++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/store/index.ts b/src/store/index.ts
index 49a202a6..14caeeec 100644
--- a/src/store/index.ts
+++ b/src/store/index.ts
@@ -123,16 +123,25 @@ const createStore = () =>
         }
       }
     },
+    // @TODO: can we unify addSelectedNodes and addSelectedEdges somehow?
     addSelectedNodes: (selectedNodeIds: string[]) => {
-      const { multiSelectionActive, onNodesChange, nodeInternals, hasDefaultNodes } = get();
-      // @TODO: work with nodeInternals instead of converting it to an array
-      const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+      const {
+        multiSelectionActive,
+        onNodesChange,
+        nodeInternals,
+        hasDefaultNodes,
+        onEdgesChange,
+        hasDefaultEdges,
+        edges,
+      } = get();
       let changedNodes: NodeSelectionChange[];
+      let changedEdges: EdgeSelectionChange[] | null = null;
 
       if (multiSelectionActive) {
         changedNodes = selectedNodeIds.map((nodeId) => createSelectionChange(nodeId, true)) as NodeSelectionChange[];
       } else {
-        changedNodes = getSelectionChanges(nodes, selectedNodeIds);
+        changedNodes = getSelectionChanges(Array.from(nodeInternals.values()), selectedNodeIds);
+        changedEdges = getSelectionChanges(edges, []);
       }
 
       if (changedNodes.length) {
@@ -142,16 +151,33 @@ const createStore = () =>
 
         onNodesChange?.(changedNodes);
       }
+
+      if (changedEdges?.length) {
+        if (hasDefaultEdges) {
+          set({ edges: handleControlledEdgeSelectionChange(changedEdges, edges) });
+        }
+
+        onEdgesChange?.(changedEdges);
+      }
     },
     addSelectedEdges: (selectedEdgeIds: string[]) => {
-      const { multiSelectionActive, onEdgesChange, edges, hasDefaultEdges } = get();
-
+      const {
+        multiSelectionActive,
+        onEdgesChange,
+        edges,
+        hasDefaultEdges,
+        nodeInternals,
+        hasDefaultNodes,
+        onNodesChange,
+      } = get();
       let changedEdges: EdgeSelectionChange[];
+      let changedNodes: NodeSelectionChange[] | null = null;
 
       if (multiSelectionActive) {
         changedEdges = selectedEdgeIds.map((edgeId) => createSelectionChange(edgeId, true)) as EdgeSelectionChange[];
       } else {
         changedEdges = getSelectionChanges(edges, selectedEdgeIds);
+        changedNodes = getSelectionChanges(Array.from(nodeInternals.values()), []);
       }
 
       if (changedEdges.length) {
@@ -162,11 +188,18 @@ const createStore = () =>
         }
         onEdgesChange?.(changedEdges);
       }
+
+      if (changedNodes?.length) {
+        if (hasDefaultNodes) {
+          set({ nodeInternals: handleControlledNodeSelectionChange(changedNodes, nodeInternals) });
+        }
+
+        onNodesChange?.(changedNodes);
+      }
     },
     unselectNodesAndEdges: () => {
       const { nodeInternals, edges, onNodesChange, onEdgesChange, hasDefaultNodes, hasDefaultEdges } = get();
-      // @TODO: work with nodeInternals instead of converting it to an array
-      const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+      const nodes = Array.from(nodeInternals.values());
 
       const nodesToUnselect = nodes.map((n) => {
         n.selected = false;