From d2b140ce5e828e00a96b1c7f5e073fd636a3966e Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Tue, 31 Oct 2023 18:31:24 +0100
Subject: [PATCH] fix(svelte): repair selection, add multi selection closes
 #3557

---
 .../svelte/src/routes/customnode/+page.svelte |  40 +------
 .../src/routes/customnode/CustomNode.svelte   |  45 ++++---
 .../react/src/components/Nodes/wrapNode.tsx   |   3 +-
 packages/react/src/hooks/useDrag.ts           |  14 +--
 packages/svelte/src/lib/actions/drag/index.ts |   6 +-
 .../components/EdgeWrapper/EdgeWrapper.svelte |  36 ++++--
 .../components/KeyHandler/KeyHandler.svelte   |  33 +++++-
 .../src/lib/components/KeyHandler/types.ts    |   1 +
 .../components/NodeWrapper/NodeWrapper.svelte |  20 ++--
 packages/svelte/src/lib/store/index.ts        | 110 ++++++++++++------
 packages/svelte/src/lib/store/types.ts        |   3 +-
 packages/system/src/xydrag/XYDrag.ts          |  10 +-
 12 files changed, 204 insertions(+), 117 deletions(-)

diff --git a/examples/svelte/src/routes/customnode/+page.svelte b/examples/svelte/src/routes/customnode/+page.svelte
index e2654938..14130a47 100644
--- a/examples/svelte/src/routes/customnode/+page.svelte
+++ b/examples/svelte/src/routes/customnode/+page.svelte
@@ -23,28 +23,6 @@
 
 	const bgColor = writable('#1A192B');
 
-	const onChange: ChangeEventHandler<HTMLInputElement> = (event) => {
-		nodes.update((nds) =>
-			nds.map((node) => {
-				if (node.type !== 'colorNode') {
-					return node;
-				}
-
-				const color = (event.target as HTMLInputElement)?.value;
-
-				bgColor.set(color);
-
-				return {
-					...node,
-					data: {
-						...node.data,
-						color
-					}
-				};
-			})
-		);
-	};
-
 	const nodes = writable<Node[]>([
 		{
 			id: '1',
@@ -56,17 +34,9 @@
 		{
 			id: '2',
 			type: 'colorNode',
-			data: { onChange: onChange, color: $bgColor },
-			style: 'border: 1px solid #777; padding: 10px',
+			data: { colorStore: bgColor },
 			position: { x: 250, y: 50 }
 		},
-		{
-			id: '2a',
-			type: 'colorNode',
-			data: { onChange: onChange, color: $bgColor },
-			style: 'border: 1px solid #777; padding: 10px',
-			position: { x: 250, y: 520 }
-		},
 		{
 			id: '3',
 			type: 'output',
@@ -115,7 +85,7 @@
 	{nodes}
 	{edges}
 	{nodeTypes}
-	style="--bgcolor: {$bgColor}"
+	style="--background-color: {$bgColor}"
 	fitView
 	on:connect={onConnect}
 >
@@ -123,9 +93,3 @@
 	<Background variant={BackgroundVariant.Dots} />
 	<MiniMap />
 </SvelteFlow>
-
-<style>
-	:global(.svelte-flow) {
-		background: var(--bgcolor);
-	}
-</style>
diff --git a/examples/svelte/src/routes/customnode/CustomNode.svelte b/examples/svelte/src/routes/customnode/CustomNode.svelte
index 7a40a742..25d37d32 100644
--- a/examples/svelte/src/routes/customnode/CustomNode.svelte
+++ b/examples/svelte/src/routes/customnode/CustomNode.svelte
@@ -1,21 +1,40 @@
 <script lang="ts">
+	import type { Writable } from 'svelte/store';
 	import { Handle, Position, type NodeProps } from '@xyflow/svelte';
 
 	type $$Props = NodeProps;
 
-	export let data: { color: string; onChange: () => void } = { color: '#111', onChange: () => {} };
+	export let data: { colorStore: Writable<string> };
+
+	const { colorStore } = data;
 </script>
 
-<Handle type="target" position={Position.Left} on:connect />
-<div>
-	Custom Color Picker Node: <strong>{data.color}</strong>
+<div class="custom">
+	<Handle type="target" position={Position.Left} on:connect />
+	<div>
+		Custom Color Picker Node: <strong>{$colorStore}</strong>
+	</div>
+	<input
+		class="nodrag"
+		type="color"
+		on:input={(evt) => colorStore.set(evt.currentTarget.value)}
+		value={$colorStore}
+	/>
+	<Handle type="source" position={Position.Right} id="a" style="top: 20px;" on:connect />
+	<Handle
+		type="source"
+		position={Position.Right}
+		id="b"
+		style="top: auto; bottom: 10px;"
+		on:connect
+	/>
 </div>
-<input class="nodrag" type="color" on:input={data.onChange} value={data.color} />
-<Handle type="source" position={Position.Right} id="a" style="top: 10px;" on:connect />
-<Handle
-	type="source"
-	position={Position.Right}
-	id="b"
-	style="top: auto; bottom: 10px;"
-	on:connect
-/>
+
+<style>
+	.custom {
+		background-color: white;
+		padding: 10px;
+		border: 1px solid #777;
+		border-radius: 20px;
+	}
+</style>
diff --git a/packages/react/src/components/Nodes/wrapNode.tsx b/packages/react/src/components/Nodes/wrapNode.tsx
index 4ebead71..f96361ce 100644
--- a/packages/react/src/components/Nodes/wrapNode.tsx
+++ b/packages/react/src/components/Nodes/wrapNode.tsx
@@ -72,7 +72,8 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
       const { selectNodesOnDrag, nodeDragThreshold } = store.getState();
 
       if (isSelectable && (!selectNodesOnDrag || !isDraggable || nodeDragThreshold > 0)) {
-        // this handler gets called within the drag start event when selectNodesOnDrag=true
+        // this handler gets called by XYDrag on drag start when selectNodesOnDrag=true
+        // here we only need to call it when selectNodesOnDrag=false
         handleNodeClick({
           id,
           store,
diff --git a/packages/react/src/hooks/useDrag.ts b/packages/react/src/hooks/useDrag.ts
index 1cee64f3..3b332d8a 100644
--- a/packages/react/src/hooks/useDrag.ts
+++ b/packages/react/src/hooks/useDrag.ts
@@ -23,14 +23,12 @@ function useDrag({ nodeRef, disabled = false, noDragClassName, handleSelector, n
       xyDrag.current = XYDrag({
         domNode: nodeRef.current,
         getStoreItems: () => store.getState(),
-        onNodeClick: () => {
-          if (nodeId) {
-            handleNodeClick({
-              id: nodeId,
-              store,
-              nodeRef: nodeRef as RefObject<HTMLDivElement>,
-            });
-          }
+        onNodeMouseDown: (id: string) => {
+          handleNodeClick({
+            id,
+            store,
+            nodeRef: nodeRef as RefObject<HTMLDivElement>,
+          });
         },
         onDragStart: () => {
           setDragging(true);
diff --git a/packages/svelte/src/lib/actions/drag/index.ts b/packages/svelte/src/lib/actions/drag/index.ts
index bf93b041..e936d3a5 100644
--- a/packages/svelte/src/lib/actions/drag/index.ts
+++ b/packages/svelte/src/lib/actions/drag/index.ts
@@ -13,15 +13,17 @@ export type UseDragParams = {
   onDrag?: OnDrag;
   onDragStart?: OnDrag;
   onDragStop?: OnDrag;
+  onNodeMouseDown?: (id: string) => void;
 };
 
 export default function drag(domNode: Element, params: UseDragParams) {
-  const { store, onDrag, onDragStart, onDragStop } = params;
+  const { store, onDrag, onDragStart, onDragStop, onNodeMouseDown } = params;
   const dragInstance = XYDrag({
     domNode,
     onDrag,
     onDragStart,
     onDragStop,
+    onNodeMouseDown,
     getStoreItems: () => {
       const snapGrid = get(store.snapGrid);
       const vp = get(store.viewport);
@@ -33,7 +35,7 @@ export default function drag(domNode: Element, params: UseDragParams) {
         snapGrid: snapGrid ? snapGrid : [0, 0],
         snapToGrid: !!snapGrid,
         nodeOrigin: [0, 0],
-        multiSelectionActive: false,
+        multiSelectionActive: get(store.multiselectionKeyPressed),
         domNode: get(store.domNode),
         transform: [vp.x, vp.y, vp.zoom],
         autoPanOnNodeDrag: get(store.autoPanOnNodeDrag),
diff --git a/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte b/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte
index 59c8134c..7a8df68c 100644
--- a/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte
+++ b/packages/svelte/src/lib/components/EdgeWrapper/EdgeWrapper.svelte
@@ -3,11 +3,12 @@
 <script lang="ts">
   import cc from 'classcat';
   import { createEventDispatcher } from 'svelte';
-  import { getMarkerId } from '@xyflow/system';
+  import { errorMessages, getMarkerId } from '@xyflow/system';
 
   import { useStore } from '$lib/store';
   import BezierEdge from '$lib/components/edges/BezierEdge.svelte';
   import type { EdgeLayouted, Edge } from '$lib/types';
+  import { get } from 'svelte/store';
 
   type $$Props = EdgeLayouted;
 
@@ -41,7 +42,16 @@
   let className: string = '';
   export { className as class };
 
-  const { edges, edgeTypes, flowId, addSelectedEdges } = useStore();
+  const {
+    edges,
+    edgeTypes,
+    flowId,
+    selectionRect,
+    selectionRectMode,
+    multiselectionKeyPressed,
+    addSelectedEdges,
+    unselectNodesAndEdges
+  } = useStore();
   const dispatch = createEventDispatcher<{
     edgeclick: { edge: Edge; event: MouseEvent | TouchEvent };
     edgecontextmenu: { edge: Edge; event: MouseEvent };
@@ -52,15 +62,25 @@
   $: markerEndUrl = markerEnd ? `url(#${getMarkerId(markerEnd, $flowId)})` : undefined;
 
   function onClick(event: MouseEvent | TouchEvent) {
-    if (selectable) {
-      addSelectedEdges([id]);
-    }
-
     const edge = $edges.find((e) => e.id === id);
 
-    if (edge) {
-      dispatch('edgeclick', { event, edge });
+    if (!edge) {
+      console.warn('012', errorMessages['error012'](id));
+      return;
     }
+
+    if (selectable) {
+      selectionRect.set(null);
+      selectionRectMode.set(null);
+
+      if (!edge.selected) {
+        addSelectedEdges([id]);
+      } else if (edge.selected && get(multiselectionKeyPressed)) {
+        unselectNodesAndEdges({ nodes: [], edges: [edge] });
+      }
+    }
+
+    dispatch('edgeclick', { event, edge });
   }
 
   function onContextMenu(event: MouseEvent) {
diff --git a/packages/svelte/src/lib/components/KeyHandler/KeyHandler.svelte b/packages/svelte/src/lib/components/KeyHandler/KeyHandler.svelte
index a7c852ab..77195ea2 100644
--- a/packages/svelte/src/lib/components/KeyHandler/KeyHandler.svelte
+++ b/packages/svelte/src/lib/components/KeyHandler/KeyHandler.svelte
@@ -4,14 +4,21 @@
   import { useStore } from '$lib/store';
   import type { KeyHandlerProps } from './types';
   import type { KeyDefinition, KeyDefinitionObject } from '$lib/types';
+  import { isMacOs } from '@xyflow/system';
 
   type $$Props = KeyHandlerProps;
 
   export let selectionKey: $$Props['selectionKey'] = 'Shift';
+  export let multiSelectionKey: $$Props['multiSelectionKey'] = isMacOs() ? 'Meta' : 'Control';
   export let deleteKey: $$Props['deleteKey'] = 'Backspace';
   export let panActivationKey: $$Props['panActivationKey'] = ' ';
 
-  const { selectionKeyPressed, deleteKeyPressed, panActivationKeyPressed } = useStore();
+  const {
+    selectionKeyPressed,
+    multiselectionKeyPressed,
+    deleteKeyPressed,
+    panActivationKeyPressed
+  } = useStore();
 
   function isKeyObject(key?: KeyDefinition): key is KeyDefinitionObject {
     return typeof key === 'object';
@@ -20,6 +27,10 @@
   $: selectionKeyString = typeof selectionKey === 'string' ? selectionKey : selectionKey!.key;
   $: selectionKeyModifier = isKeyObject(selectionKey) ? selectionKey?.modifier : [];
 
+  $: multiSelectionKeyString =
+    typeof multiSelectionKey === 'string' ? multiSelectionKey : multiSelectionKey!.key;
+  $: multiSelectionKeyModifier = isKeyObject(multiSelectionKey) ? multiSelectionKey?.modifier : [];
+
   $: deleteKeyString = typeof deleteKey === 'string' ? deleteKey : deleteKey!.key;
   $: deleteKeyModifier = isKeyObject(deleteKey) ? deleteKey?.modifier : [];
 
@@ -49,6 +60,26 @@
     ],
     type: 'keyup'
   }}
+  use:shortcut={{
+    trigger: [
+      {
+        key: multiSelectionKeyString,
+        modifier: multiSelectionKeyModifier,
+        callback: () => multiselectionKeyPressed.set(true)
+      }
+    ],
+    type: 'keydown'
+  }}
+  use:shortcut={{
+    trigger: [
+      {
+        key: multiSelectionKeyString,
+        modifier: multiSelectionKeyModifier,
+        callback: () => multiselectionKeyPressed.set(false)
+      }
+    ],
+    type: 'keyup'
+  }}
   use:shortcut={{
     trigger: [
       {
diff --git a/packages/svelte/src/lib/components/KeyHandler/types.ts b/packages/svelte/src/lib/components/KeyHandler/types.ts
index 70cc95a5..de81d1c3 100644
--- a/packages/svelte/src/lib/components/KeyHandler/types.ts
+++ b/packages/svelte/src/lib/components/KeyHandler/types.ts
@@ -2,6 +2,7 @@ import type { KeyDefinition } from '$lib/types';
 
 export type KeyHandlerProps = {
   selectionKey?: KeyDefinition;
+  multiSelectionKey?: KeyDefinition;
   deleteKey?: KeyDefinition;
   panActivationKey?: KeyDefinition;
 };
diff --git a/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte b/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte
index fa1cd154..95dcb933 100644
--- a/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte
+++ b/packages/svelte/src/lib/components/NodeWrapper/NodeWrapper.svelte
@@ -8,8 +8,8 @@
     SvelteComponent,
     type ComponentType
   } from 'svelte';
-  import cc from 'classcat';
   import { get, writable } from 'svelte/store';
+  import cc from 'classcat';
   import { errorMessages, Position, type NodeProps } from '@xyflow/system';
 
   import drag from '$lib/actions/drag';
@@ -44,7 +44,13 @@
   export { className as class };
 
   const store = useStore();
-  const { nodeTypes, nodeDragThreshold, addSelectedNodes, updateNodeDimensions } = store;
+  const {
+    nodeTypes,
+    nodeDragThreshold,
+    selectNodesOnDrag,
+    handleNodeSelection,
+    updateNodeDimensions
+  } = store;
   const nodeType = type || 'default';
 
   let nodeRef: HTMLDivElement;
@@ -56,7 +62,6 @@
 
   const nodeComponent: ComponentType<SvelteComponent<NodeProps>> =
     $nodeTypes[nodeType] || DefaultNode;
-  const selectNodesOnDrag = false;
   const dispatch = createEventDispatcher<{
     nodeclick: { node: Node; event: MouseEvent | TouchEvent };
     nodecontextmenu: { node: Node; event: MouseEvent | TouchEvent };
@@ -113,12 +118,12 @@
   });
 
   function onSelectNodeHandler(event: MouseEvent | TouchEvent) {
-    if (selectable && (!selectNodesOnDrag || !draggable || get(nodeDragThreshold) > 0)) {
-      // this handler gets called within the drag start event when selectNodesOnDrag=true
-      addSelectedNodes([id]);
+    if (selectable && (!get(selectNodesOnDrag) || !draggable || get(nodeDragThreshold) > 0)) {
+      // this handler gets called by XYDrag on drag start when selectNodesOnDrag=true
+      // here we only need to call it when selectNodesOnDrag=false
+      handleNodeSelection(id);
     }
 
-    // @todo: support multiselection
     dispatch('nodeclick', { node, event });
   }
 
@@ -135,6 +140,7 @@
       disabled: false,
       handleSelector: dragHandle,
       noDragClass: 'nodrag',
+      onNodeMouseDown: handleNodeSelection,
       onDrag: (event, _, node, nodes) => {
         dispatch('nodedrag', { event, node, nodes });
       },
diff --git a/packages/svelte/src/lib/store/index.ts b/packages/svelte/src/lib/store/index.ts
index 0c49c337..b78e6fb2 100644
--- a/packages/svelte/src/lib/store/index.ts
+++ b/packages/svelte/src/lib/store/index.ts
@@ -15,7 +15,8 @@ import {
   type CoordinateExtent,
   type UpdateConnection,
   type NodeBase,
-  type NodeDragItem
+  type NodeDragItem,
+  errorMessages
 } from '@xyflow/system';
 
 import { addEdge as addEdgeUtil } from '$lib/utils';
@@ -173,25 +174,29 @@ export function createStore({
     }
   }
 
-  function resetSelectedItem<T extends Node | Edge>(item: T) {
-    if (item.selected) {
-      return {
-        ...item,
-        selected: false
-      };
-    }
+  function resetSelectedItem<T extends Node | Edge>(ids: string[]) {
+    return (item: T) => {
+      if (item.selected && ids.includes(item.id)) {
+        return {
+          ...item,
+          selected: false
+        };
+      }
 
-    return item;
+      return item;
+    };
   }
 
-  function unselectNodesAndEdges() {
-    if (get(store.nodes).some((node) => node.selected)) {
-      store.nodes.update((ns) => ns.map(resetSelectedItem));
-    }
+  function unselectNodesAndEdges(params?: { nodes?: Node[]; edges?: Edge[] }) {
+    const nodeIdsToUnselect = (params?.nodes ? params.nodes : get(store.nodes)).map(
+      (item) => item.id
+    );
+    const edgeIdsToUnselect = (params?.edges ? params.edges : get(store.edges)).map(
+      (item) => item.id
+    );
 
-    if (get(store.edges).some((edge) => edge.selected)) {
-      store.edges.update((es) => es.map(resetSelectedItem));
-    }
+    store.nodes.update((ns) => ns.map(resetSelectedItem(nodeIdsToUnselect)));
+    store.edges.update((es) => es.map(resetSelectedItem(edgeIdsToUnselect)));
   }
 
   store.deleteKeyPressed.subscribe((deleteKeyPressed) => {
@@ -220,36 +225,74 @@ export function createStore({
   });
 
   function addSelectedNodes(ids: string[]) {
-    store.selectionRect.set(null);
-    store.selectionRectMode.set(null);
-
-    if (get(store.multiselectionKeyPressed)) {
-      // @todo handle multiselection key
-    }
+    const isMultiSelection = get(store.multiselectionKeyPressed);
 
     store.nodes.update((ns) =>
       ns.map((node) => {
-        return {
-          ...node,
-          selected: ids.includes(node.id)
-        };
+        const nodeWillBeSelected = ids.includes(node.id);
+        const selected = isMultiSelection
+          ? node.selected || nodeWillBeSelected
+          : nodeWillBeSelected;
+
+        // we need to mutate the node here in order to have the correct selected state in the drag handler
+        node.selected = selected;
+
+        return node;
       })
     );
+
+    if (!isMultiSelection) {
+      store.edges.update((es) =>
+        es.map((edge) => {
+          edge.selected = false;
+          return edge;
+        })
+      );
+    }
   }
 
   function addSelectedEdges(ids: string[]) {
-    if (get(store.multiselectionKeyPressed)) {
-      // @todo handle multiselection key
-    }
+    const isMultiSelection = get(store.multiselectionKeyPressed);
 
     store.edges.update((edges) =>
       edges.map((edge) => {
-        return {
-          ...edge,
-          selected: ids.includes(edge.id)
-        };
+        const edgeWillBeSelected = ids.includes(edge.id);
+        const selected = isMultiSelection
+          ? edge.selected || edgeWillBeSelected
+          : edgeWillBeSelected;
+
+        edge.selected = selected;
+
+        return edge;
       })
     );
+
+    if (!isMultiSelection) {
+      store.nodes.update((ns) =>
+        ns.map((node) => {
+          node.selected = false;
+          return node;
+        })
+      );
+    }
+  }
+
+  function handleNodeSelection(id: string) {
+    const node = get(store.nodes)?.find((n) => n.id === id);
+
+    if (!node) {
+      console.warn('012', errorMessages['error012'](id));
+      return;
+    }
+
+    store.selectionRect.set(null);
+    store.selectionRectMode.set(null);
+
+    if (!node.selected) {
+      addSelectedNodes([id]);
+    } else if (node.selected && get(store.multiselectionKeyPressed)) {
+      unselectNodesAndEdges({ nodes: [node], edges: [] });
+    }
   }
 
   function panBy(delta: XYPosition) {
@@ -326,6 +369,7 @@ export function createStore({
     unselectNodesAndEdges,
     addSelectedNodes,
     addSelectedEdges,
+    handleNodeSelection,
     panBy,
     updateConnection,
     cancelConnection,
diff --git a/packages/svelte/src/lib/store/types.ts b/packages/svelte/src/lib/store/types.ts
index 49213f7a..ae854cfd 100644
--- a/packages/svelte/src/lib/store/types.ts
+++ b/packages/svelte/src/lib/store/types.ts
@@ -28,9 +28,10 @@ export type SvelteFlowStoreActions = {
   fitView: (options?: FitViewOptions) => boolean;
   updateNodePositions: UpdateNodePositions;
   updateNodeDimensions: (updates: NodeDimensionUpdate[]) => void;
-  unselectNodesAndEdges: () => void;
+  unselectNodesAndEdges: (params?: { nodes?: Node[]; edges?: Edge[] }) => void;
   addSelectedNodes: (ids: string[]) => void;
   addSelectedEdges: (ids: string[]) => void;
+  handleNodeSelection: (id: string) => void;
   panBy: (delta: XYPosition) => boolean;
   updateConnection: UpdateConnection;
   cancelConnection: () => void;
diff --git a/packages/system/src/xydrag/XYDrag.ts b/packages/system/src/xydrag/XYDrag.ts
index 5381d0ef..17f4f99d 100644
--- a/packages/system/src/xydrag/XYDrag.ts
+++ b/packages/system/src/xydrag/XYDrag.ts
@@ -46,7 +46,7 @@ type StoreItems = {
   selectNodesOnDrag: boolean;
   nodeDragThreshold: number;
   panBy: PanBy;
-  unselectNodesAndEdges: () => void;
+  unselectNodesAndEdges: (params?: { nodes?: NodeBase[]; edges?: EdgeBase[] }) => void;
   onError?: OnError;
   onNodeDragStart?: OnNodeDrag;
   onNodeDrag?: OnNodeDrag;
@@ -63,7 +63,7 @@ export type XYDragParams = {
   onDragStart?: OnDrag;
   onDrag?: OnDrag;
   onDragStop?: OnDrag;
-  onNodeClick?: () => void;
+  onNodeMouseDown?: (id: string) => void;
 };
 
 export type XYDragInstance = {
@@ -81,7 +81,7 @@ export type DragUpdateParams = {
 
 export function XYDrag({
   domNode,
-  onNodeClick,
+  onNodeMouseDown,
   getStoreItems,
   onDragStart,
   onDrag,
@@ -217,8 +217,8 @@ export function XYDrag({
         }
       }
 
-      if (isSelectable && selectNodesOnDrag) {
-        onNodeClick?.();
+      if (isSelectable && selectNodesOnDrag && nodeId) {
+        onNodeMouseDown?.(nodeId);
       }
 
       const pointerPos = getPointerPosition(event.sourceEvent, { transform, snapGrid, snapToGrid });