refactor(core): use connection lookup to check for selected elements (#1737)

* refactor(core): use connection lookup to check for selected elements

Signed-off-by: braks <78412429+bcakmakoglu@users.noreply.github.com>

* chore(changeset): add

Signed-off-by: braks <78412429+bcakmakoglu@users.noreply.github.com>

---------

Signed-off-by: braks <78412429+bcakmakoglu@users.noreply.github.com>

.changeset/flat-impalas-look.md

@@ -0,0 +1,5 @@
+---
+"@vue-flow/core": minor
+---
+
+Use connection lookup to check what elements should be selected by a user selection box.

packages/core/src/container/Pane/Pane.vue

@@ -2,10 +2,10 @@
 import { ref, toRef, watch } from 'vue'
 import UserSelection from '../../components/UserSelection/UserSelection.vue'
 import NodesSelection from '../../components/NodesSelection/NodesSelection.vue'
-import type { NodeChange } from '../../types'
+import type { EdgeChange, NodeChange } from '../../types'
 import { SelectionMode } from '../../types'
 import { useKeyPress, useVueFlow } from '../../composables'
-import { getEventPosition, getNodesInside, getSelectionChanges } from '../../utils'
+import { areSetsEqual, getEventPosition, getNodesInside, getSelectionChanges } from '../../utils'
 import { getMousePosition } from './utils'
 
 const { isSelecting, selectionKeyPressed } = defineProps<{ isSelecting: boolean; selectionKeyPressed: boolean }>()
@@ -30,18 +30,18 @@ const {
   multiSelectionActive,
   edgeLookup,
   nodeLookup,
+  connectionLookup,
+  defaultEdgeOptions,
 } = useVueFlow()
 
 const container = ref<HTMLDivElement | null>(null)
 
-const prevSelectedNodesCount = ref(0)
+const selectedNodeIds = ref<Set<string>>(new Set())
 
-const prevSelectedEdgesCount = ref(0)
+const selectedEdgeIds = ref<Set<string>>(new Set())
 
 const containerBounds = ref<DOMRect>()
 
-const edgeIdLookup = ref<Map<string, Set<string>>>(new Map())
-
 const hasActiveSelection = toRef(() => elementsSelectable.value && (isSelecting || userSelectionActive.value))
 
 // Used to prevent click events when the user lets go of the selectionKey during a selection
@@ -78,14 +78,6 @@ function wrapHandler(handler: Function, containerRef: HTMLDivElement | null) {
   }
 }
 
-function resetUserSelection() {
-  userSelectionActive.value = false
-  userSelectionRect.value = null
-
-  prevSelectedNodesCount.value = 0
-  prevSelectedEdgesCount.value = 0
-}
-
 function onClick(event: MouseEvent) {
   if (selectionInProgress) {
     selectionInProgress = false
@@ -129,12 +121,6 @@ function onPointerDown(event: PointerEvent) {
 
   selectionStarted = true
   selectionInProgress = false
-  edgeIdLookup.value = new Map()
-
-  for (const [id, edge] of edgeLookup.value) {
-    edgeIdLookup.value.set(edge.source, edgeIdLookup.value.get(edge.source)?.add(id) || new Set([id]))
-    edgeIdLookup.value.set(edge.target, edgeIdLookup.value.get(edge.target)?.add(id) || new Set([id]))
-  }
 
   removeSelectedElements()
 
@@ -169,38 +155,39 @@ function onPointerMove(event: PointerEvent) {
     height: Math.abs(mouseY - startY),
   }
 
-  const selectedNodes = getNodesInside(
-    nodes.value,
-    nextUserSelectRect,
-    viewport.value,
-    selectionMode.value === SelectionMode.Partial,
-    true,
+  const prevSelectedNodeIds = selectedNodeIds.value
+  const prevSelectedEdgeIds = selectedEdgeIds.value
+
+  selectedNodeIds.value = new Set(
+    getNodesInside(nodes.value, nextUserSelectRect, viewport.value, selectionMode.value === SelectionMode.Partial, true).map(
+      (node) => node.id,
+    ),
   )
 
-  const selectedEdgeIds = new Set<string>()
-  const selectedNodeIds = new Set<string>()
+  selectedEdgeIds.value = new Set()
+  const edgesSelectable = defaultEdgeOptions.value?.selectable ?? true
 
-  for (const selectedNode of selectedNodes) {
-    selectedNodeIds.add(selectedNode.id)
-
-    const edgeIds = edgeIdLookup.value.get(selectedNode.id)
-
-    if (edgeIds) {
-      for (const edgeId of edgeIds) {
-        selectedEdgeIds.add(edgeId)
+  // We look for all edges connected to the selected nodes
+  for (const nodeId of selectedNodeIds.value) {
+    const connections = connectionLookup.value.get(nodeId)
+    if (!connections) {
+      continue
+    }
+    for (const { edgeId } of connections.values()) {
+      const edge = edgeLookup.value.get(edgeId)
+      if (edge && (edge.selectable ?? edgesSelectable)) {
+        selectedEdgeIds.value.add(edgeId)
       }
     }
   }
 
-  if (prevSelectedNodesCount.value !== selectedNodeIds.size) {
-    prevSelectedNodesCount.value = selectedNodeIds.size
-    const changes = getSelectionChanges(nodeLookup.value, selectedNodeIds, true) as NodeChange[]
+  if (!areSetsEqual(prevSelectedNodeIds, selectedNodeIds.value)) {
+    const changes = getSelectionChanges(nodeLookup.value, selectedNodeIds.value, true) as NodeChange[]
     emits.nodesChange(changes)
   }
 
-  if (prevSelectedEdgesCount.value !== selectedEdgeIds.size) {
-    prevSelectedEdgesCount.value = selectedEdgeIds.size
-    const changes = getSelectionChanges(edgeLookup.value, selectedEdgeIds)
+  if (!areSetsEqual(prevSelectedEdgeIds, selectedEdgeIds.value)) {
+    const changes = getSelectionChanges(edgeLookup.value, selectedEdgeIds.value) as EdgeChange[]
     emits.edgesChange(changes)
   }
 
@@ -222,11 +209,9 @@ function onPointerUp(event: PointerEvent) {
     onClick(event)
   }
 
-  if (prevSelectedNodesCount.value > 0) {
-    nodesSelectionActive.value = true
-  }
-
-  resetUserSelection()
+  userSelectionActive.value = false
+  userSelectionRect.value = null
+  nodesSelectionActive.value = selectedNodeIds.value.size > 0
 
   emits.selectionEnd(event)
 

packages/core/src/utils/store.ts

@@ -233,6 +233,23 @@ export function areConnectionMapsEqual(a?: Map<string, Connection>, b?: Map<stri
   return true
 }
 
+/**
+ * @internal
+ */
+export function areSetsEqual(a: Set<string>, b: Set<string>) {
+  if (a.size !== b.size) {
+    return false
+  }
+
+  for (const item of a) {
+    if (!b.has(item)) {
+      return false
+    }
+  }
+
+  return true
+}
+
 /**
  * @internal
  */