From 266f8744ad3e7bbaca599f99294b08d2da313c9b Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Sat, 29 Jun 2024 12:25:24 +0200
Subject: [PATCH] chore(core): cleanup pane (#1507)

* chore(core): cleanup pane

* chore(core): cleanup actions
---
 packages/core/src/container/Pane/Pane.vue | 126 +++++++++++++---------
 packages/core/src/store/actions.ts        | 121 ++++++---------------
 packages/core/src/types/hooks.ts          |   6 +-
 packages/core/src/types/store.ts          |  12 ++-
 packages/core/src/utils/changes.ts        |  41 +++----
 packages/core/src/utils/storage.ts        |   2 +
 6 files changed, 144 insertions(+), 164 deletions(-)

diff --git a/packages/core/src/container/Pane/Pane.vue b/packages/core/src/container/Pane/Pane.vue
index 40b98111..b9417950 100644
--- a/packages/core/src/container/Pane/Pane.vue
+++ b/packages/core/src/container/Pane/Pane.vue
@@ -2,9 +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 { SelectionMode } from '../../types'
 import { useKeyPress, useVueFlow } from '../../composables'
-import { getConnectedEdges, getNodesInside } from '../../utils'
+import { getEventPosition, getNodesInside, getSelectionChanges } from '../../utils'
 import { getMousePosition } from './utils'
 
 const { isSelecting, selectionKeyPressed } = defineProps<{ isSelecting: boolean; selectionKeyPressed: boolean }>()
@@ -12,7 +13,6 @@ const { isSelecting, selectionKeyPressed } = defineProps<{ isSelecting: boolean;
 const {
   vueFlowRef,
   getNodes,
-  getEdges,
   viewport,
   emits,
   userSelectionActive,
@@ -21,7 +21,6 @@ const {
   userSelectionRect,
   elementsSelectable,
   nodesSelectionActive,
-  addSelectedElements,
   getSelectedEdges,
   getSelectedNodes,
   removeNodes,
@@ -30,6 +29,8 @@ const {
   deleteKeyCode,
   multiSelectionKeyCode,
   multiSelectionActive,
+  edgeLookup,
+  nodeLookup,
 } = useVueFlow()
 
 const container = ref<HTMLDivElement | null>(null)
@@ -40,6 +41,8 @@ const prevSelectedEdgesCount = ref(0)
 
 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
@@ -65,6 +68,16 @@ watch(multiSelectKeyPressed, (isKeyPressed) => {
   multiSelectionActive.value = isKeyPressed
 })
 
+function wrapHandler(handler: Function, containerRef: HTMLDivElement | null) {
+  return (event: MouseEvent) => {
+    if (event.target !== containerRef) {
+      return
+    }
+
+    handler?.(event)
+  }
+}
+
 function resetUserSelection() {
   userSelectionActive.value = false
   userSelectionRect.value = null
@@ -74,7 +87,7 @@ function resetUserSelection() {
 }
 
 function onClick(event: MouseEvent) {
-  if (selectionInProgress.value || event.target !== container.value || hasActiveSelection.value) {
+  if (selectionInProgress.value) {
     selectionInProgress.value = false
     return
   }
@@ -87,10 +100,6 @@ function onClick(event: MouseEvent) {
 }
 
 function onContextMenu(event: MouseEvent) {
-  if (event.target !== container.value) {
-    return
-  }
-
   if (Array.isArray(panOnDrag.value) && panOnDrag.value?.includes(2)) {
     event.preventDefault()
     return
@@ -100,27 +109,32 @@ function onContextMenu(event: MouseEvent) {
 }
 
 function onWheel(event: WheelEvent) {
-  if (event.target !== container.value) {
-    return
-  }
-
   emits.paneScroll(event)
 }
 
 function onPointerDown(event: PointerEvent) {
-  if (!hasActiveSelection.value) {
-    return emits.paneMouseMove(event)
-  }
-
   containerBounds.value = vueFlowRef.value?.getBoundingClientRect()
   container.value?.setPointerCapture(event.pointerId)
 
-  if (!elementsSelectable || !isSelecting || event.button !== 0 || event.target !== container.value || !containerBounds.value) {
+  if (
+    !elementsSelectable.value ||
+    !isSelecting ||
+    event.button !== 0 ||
+    event.target !== container.value ||
+    !containerBounds.value
+  ) {
     return
   }
 
   const { x, y } = getMousePosition(event, containerBounds.value)
 
+  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()
 
   userSelectionRect.value = {
@@ -138,43 +152,62 @@ function onPointerDown(event: PointerEvent) {
 }
 
 function onPointerMove(event: PointerEvent) {
-  if (!hasActiveSelection.value) {
-    return emits.paneMouseMove(event)
-  }
-
   if (!containerBounds.value || !userSelectionRect.value) {
     return
   }
 
   selectionInProgress.value = true
 
-  const mousePos = getMousePosition(event, containerBounds.value)
+  const { x: mouseX, y: mouseY } = getEventPosition(event, containerBounds.value)
   const { startX = 0, startY = 0 } = userSelectionRect.value
 
   const nextUserSelectRect = {
-    ...userSelectionRect.value,
-    x: mousePos.x < startX ? mousePos.x : startX,
-    y: mousePos.y < startY ? mousePos.y : startY,
-    width: Math.abs(mousePos.x - startX),
-    height: Math.abs(mousePos.y - startY),
+    startX,
+    startY,
+    x: mouseX < startX ? mouseX : startX,
+    y: mouseY < startY ? mouseY : startY,
+    width: Math.abs(mouseX - startX),
+    height: Math.abs(mouseY - startY),
   }
 
   const selectedNodes = getNodesInside(
     getNodes.value,
-    userSelectionRect.value,
+    nextUserSelectRect,
     viewport.value,
     selectionMode.value === SelectionMode.Partial,
     true,
   )
 
-  const selectedEdges = getConnectedEdges(selectedNodes, getEdges.value)
+  const selectedEdgeIds = new Set<string>()
+  const selectedNodeIds = new Set<string>()
 
-  prevSelectedNodesCount.value = selectedNodes.length
-  prevSelectedEdgesCount.value = selectedEdges.length
+  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)
+      }
+    }
+  }
+
+  if (prevSelectedNodesCount.value !== selectedNodeIds.size) {
+    prevSelectedNodesCount.value = selectedNodeIds.size
+    const changes = getSelectionChanges(nodeLookup.value, selectedNodeIds, true) as NodeChange[]
+    emits.nodesChange(changes)
+  }
+
+  if (prevSelectedEdgesCount.value !== selectedEdgeIds.size) {
+    prevSelectedEdgesCount.value = selectedEdgeIds.size
+    const changes = getSelectionChanges(edgeLookup.value, selectedEdgeIds)
+    emits.edgesChange(changes)
+  }
 
   userSelectionRect.value = nextUserSelectRect
-
-  addSelectedElements([...selectedNodes, ...selectedEdges])
+  userSelectionActive.value = true
+  nodesSelectionActive.value = false
 }
 
 function onPointerUp(event: PointerEvent) {
@@ -184,10 +217,6 @@ function onPointerUp(event: PointerEvent) {
 
   container.value?.releasePointerCapture(event.pointerId)
 
-  if (!hasActiveSelection.value) {
-    return
-  }
-
   // We only want to trigger click functions when in selection mode if
   // the user did not move the mouse.
   if (!userSelectionActive.value && userSelectionRect.value && event.target === container.value) {
@@ -206,14 +235,6 @@ function onPointerUp(event: PointerEvent) {
     selectionInProgress.value = false
   }
 }
-
-function onPointerEnter(event: PointerEvent) {
-  if (hasActiveSelection.value) {
-    return
-  }
-
-  emits.paneMouseEnter(event)
-}
 </script>
 
 <script lang="ts">
@@ -228,13 +249,14 @@ export default {
     ref="container"
     class="vue-flow__pane vue-flow__container"
     :class="{ selection: isSelecting }"
-    @click="onClick"
-    @contextmenu="onContextMenu"
-    @wheel.passive="onWheel"
-    @pointerenter="onPointerEnter"
-    @pointerdown="onPointerDown"
-    @pointermove="onPointerMove"
-    @pointerup="onPointerUp"
+    @click="(event) => (hasActiveSelection ? undefined : wrapHandler(onClick, container)(event))"
+    @contextmenu="wrapHandler(onContextMenu, container)($event)"
+    @wheel.passive="wrapHandler(onWheel, container)($event)"
+    @pointerenter="(event) => (hasActiveSelection ? undefined : emits.paneMouseEnter(event))"
+    @pointerdown="(event) => (hasActiveSelection ? onPointerDown(event) : emits.paneMouseMove(event))"
+    @pointermove="(event) => (hasActiveSelection ? onPointerMove(event) : emits.paneMouseMove(event))"
+    @pointerup="(event) => (hasActiveSelection ? onPointerUp(event) : undefined)"
+    @pointerleave="emits.paneMouseLeave($event)"
   >
     <slot />
     <UserSelection v-if="userSelectionActive && userSelectionRect" :user-selection-rect="userSelectionRect" />
diff --git a/packages/core/src/store/actions.ts b/packages/core/src/store/actions.ts
index 9a58c743..f0428a6e 100644
--- a/packages/core/src/store/actions.ts
+++ b/packages/core/src/store/actions.ts
@@ -6,7 +6,6 @@ import type {
   CoordinateExtent,
   Edge,
   EdgeAddChange,
-  EdgeChange,
   EdgeLookup,
   EdgeRemoveChange,
   EdgeSelectionChange,
@@ -16,7 +15,6 @@ import type {
   GraphNode,
   Node,
   NodeAddChange,
-  NodeChange,
   NodeDimensionChange,
   NodeLookup,
   NodePositionChange,
@@ -178,76 +176,20 @@ export function useActions(state: State, nodeLookup: ComputedRef<NodeLookup>, ed
     }
   }
 
-  const nodeSelectionHandler = (nodes: GraphNode[], selected: boolean) => {
-    const nodeIds: string[] = []
-    for (const node of nodes) {
-      nodeIds.push(node.id)
-    }
-
-    let changedNodes: NodeChange[] = []
-    let changedEdges: EdgeChange[] = []
-
-    if (state.multiSelectionActive) {
-      for (const nodeId of nodeIds) {
-        changedNodes.push(createSelectionChange(nodeId, selected))
-      }
-    } else {
-      const selectionChanges = getSelectionChanges([...state.nodes, ...state.edges], nodeIds)
-      changedNodes = selectionChanges.changedNodes
-      changedEdges = selectionChanges.changedEdges
-    }
-
-    if (changedNodes.length) {
-      state.hooks.nodesChange.trigger(changedNodes)
-    }
-
-    if (changedEdges.length) {
-      state.hooks.edgesChange.trigger(changedEdges)
-    }
-  }
-
-  const edgeSelectionHandler = (edges: GraphEdge[], selected: boolean) => {
-    const edgeIds: string[] = []
-
-    for (const edge of edges) {
-      edgeIds.push(edge.id)
-    }
-
-    let changedNodes: NodeChange[] = []
-    let changedEdges: EdgeChange[] = []
-
-    if (state.multiSelectionActive) {
-      for (const edgeId of edgeIds) {
-        changedEdges.push(createSelectionChange(edgeId, selected))
-      }
-    } else {
-      const selectionChanges = getSelectionChanges([...state.nodes, ...state.edges], edgeIds)
-      changedNodes = selectionChanges.changedNodes
-      changedEdges = selectionChanges.changedEdges
-    }
-
-    if (changedNodes.length) {
-      state.hooks.nodesChange.trigger(changedNodes)
-    }
-
-    if (changedEdges.length) {
-      state.hooks.edgesChange.trigger(changedEdges)
-    }
-  }
-
   const elementSelectionHandler = (elements: Elements, selected: boolean) => {
-    const nodeIds: string[] = []
-    const edgeIds: string[] = []
+    const nodeIds = new Set<string>()
+    const edgeIds = new Set<string>()
 
     for (const element of elements) {
       if (isNode(element)) {
-        nodeIds.push(element.id)
+        nodeIds.add(element.id)
       } else if (isEdge(element)) {
-        edgeIds.push(element.id)
+        edgeIds.add(element.id)
       }
     }
 
-    const { changedNodes, changedEdges } = getSelectionChanges([...state.nodes, ...state.edges], [...nodeIds, ...edgeIds])
+    const changedNodes = getSelectionChanges(nodeLookup.value, nodeIds, true)
+    const changedEdges = getSelectionChanges(edgeLookup.value, edgeIds)
 
     if (state.multiSelectionActive) {
       for (const nodeId of nodeIds) {
@@ -269,11 +211,25 @@ export function useActions(state: State, nodeLookup: ComputedRef<NodeLookup>, ed
   }
 
   const addSelectedNodes: Actions['addSelectedNodes'] = (nodes) => {
-    nodeSelectionHandler(nodes, true)
+    if (state.multiSelectionActive) {
+      const nodeChanges = nodes.map((node) => createSelectionChange(node.id, true))
+      state.hooks.nodesChange.trigger(nodeChanges)
+      return
+    }
+
+    state.hooks.nodesChange.trigger(getSelectionChanges(nodeLookup.value, new Set(nodes.map((n) => n.id)), true))
+    state.hooks.edgesChange.trigger(getSelectionChanges(edgeLookup.value))
   }
 
   const addSelectedEdges: Actions['addSelectedEdges'] = (edges) => {
-    edgeSelectionHandler(edges, true)
+    if (state.multiSelectionActive) {
+      const changedEdges = edges.map((edge) => createSelectionChange(edge.id, true))
+      state.hooks.edgesChange.trigger(changedEdges as EdgeSelectionChange[])
+      return
+    }
+
+    state.hooks.edgesChange.trigger(getSelectionChanges(edgeLookup.value, new Set(edges.map((e) => e.id))))
+    state.hooks.nodesChange.trigger(getSelectionChanges(nodeLookup.value, new Set(), true))
   }
 
   const addSelectedElements: Actions['addSelectedElements'] = (elements) => {
@@ -281,34 +237,25 @@ export function useActions(state: State, nodeLookup: ComputedRef<NodeLookup>, ed
   }
 
   const removeSelectedNodes: Actions['removeSelectedNodes'] = (nodes) => {
-    if (!nodes.length) {
-      return nodeSelectionHandler(nodes, false)
-    }
+    const nodesToUnselect = nodes || state.nodes
 
-    const changedNodes: NodeSelectionChange[] = []
+    const nodeChanges = nodesToUnselect.map((n) => {
+      n.selected = false
+      return createSelectionChange(n.id, false)
+    })
 
-    for (const node of nodes) {
-      changedNodes.push(createSelectionChange(node.id, false))
-    }
-
-    if (changedNodes.length) {
-      state.hooks.nodesChange.trigger(changedNodes)
-    }
+    state.hooks.nodesChange.trigger(nodeChanges)
   }
 
   const removeSelectedEdges: Actions['removeSelectedEdges'] = (edges) => {
-    if (!edges.length) {
-      return edgeSelectionHandler(edges, false)
-    }
+    const edgesToUnselect = edges || state.edges
 
-    const changedEdges: EdgeSelectionChange[] = []
-    for (const edge of edges) {
-      changedEdges.push(createSelectionChange(edge.id, false))
-    }
+    const edgeChanges = edgesToUnselect.map((e) => {
+      e.selected = false
+      return createSelectionChange(e.id, false)
+    })
 
-    if (changedEdges.length) {
-      state.hooks.edgesChange.trigger(changedEdges)
-    }
+    state.hooks.edgesChange.trigger(edgeChanges)
   }
 
   const removeSelectedElements: Actions['removeSelectedElements'] = (elements) => {
diff --git a/packages/core/src/types/hooks.ts b/packages/core/src/types/hooks.ts
index ba5091db..03e8f1d1 100644
--- a/packages/core/src/types/hooks.ts
+++ b/packages/core/src/types/hooks.ts
@@ -78,9 +78,9 @@ export interface FlowEvents {
   paneScroll: WheelEvent | undefined
   paneClick: MouseEvent
   paneContextMenu: MouseEvent
-  paneMouseEnter: MouseEvent
-  paneMouseMove: MouseEvent
-  paneMouseLeave: MouseEvent
+  paneMouseEnter: PointerEvent
+  paneMouseMove: PointerEvent
+  paneMouseLeave: PointerEvent
   edgeContextMenu: EdgeMouseEvent
   edgeMouseEnter: EdgeMouseEvent
   edgeMouseMove: EdgeMouseEvent
diff --git a/packages/core/src/types/store.ts b/packages/core/src/types/store.ts
index 358d598a..e1385c06 100644
--- a/packages/core/src/types/store.ts
+++ b/packages/core/src/types/store.ts
@@ -258,7 +258,10 @@ export interface Actions extends Omit<ViewportHelper, 'viewportInitialized'> {
   applyEdgeChanges: (changes: EdgeChange[]) => GraphEdge[]
   /** applies default node change handler */
   applyNodeChanges: (changes: NodeChange[]) => GraphNode[]
-  /** manually select elements and add to state */
+  /**
+   * @deprecated will be removed in the next major, use `addSelectedNodes`/`addSelectedEdges` instead
+   * manually select elements and add to state
+   */
   addSelectedElements: (elements: FlowElements) => void
   /** manually select edges and add to state */
   addSelectedEdges: (edges: GraphEdge[]) => void
@@ -268,7 +271,10 @@ export interface Actions extends Omit<ViewportHelper, 'viewportInitialized'> {
   removeSelectedEdges: (edges: GraphEdge[]) => void
   /** manually unselect nodes and remove from state */
   removeSelectedNodes: (nodes: GraphNode[]) => void
-  /** unselect selected elements (if none are passed, all elements are unselected) */
+  /**
+   * @deprecated will be replaced in the next major
+   * unselect selected elements (if none are passed, all elements are unselected)
+   */
   removeSelectedElements: (elements?: Elements) => void
   /** apply min zoom value to d3 */
   setMinZoom: (zoom: number) => void
@@ -371,6 +377,8 @@ export type ComputedGetters = {
 export type VueFlowStore = {
   readonly id: string
   readonly emits: FlowHooksEmit
+  readonly nodeLookup: ComputedRef<NodeLookup>
+  readonly edgeLookup: ComputedRef<EdgeLookup>
   /** current vue flow version you're using */
   readonly vueFlowVersion: string
 } & FlowHooksOn &
diff --git a/packages/core/src/utils/changes.ts b/packages/core/src/utils/changes.ts
index 39f3435a..f2c01b7f 100644
--- a/packages/core/src/utils/changes.ts
+++ b/packages/core/src/utils/changes.ts
@@ -6,7 +6,6 @@ import type {
   EdgeSelectionChange,
   ElementChange,
   FlowElement,
-  FlowElements,
   GraphEdge,
   GraphNode,
   NodeAddChange,
@@ -16,7 +15,7 @@ import type {
   StyleFunc,
   Styles,
 } from '../types'
-import { isDef, isGraphNode } from '.'
+import { isGraphNode } from '.'
 
 function handleParentExpand(updateItem: GraphNode, parent: GraphNode) {
   if (parent) {
@@ -255,25 +254,27 @@ export function createEdgeRemoveChange(
   }
 }
 
-export function getSelectionChanges(elements: FlowElements, selectedIds: string[]) {
-  return elements.reduce(
-    (res, item) => {
-      let willBeSelected = selectedIds.includes(item.id)
+export function getSelectionChanges(
+  items: Map<string, any>,
+  selectedIds: Set<string> = new Set(),
+  mutateItem = false,
+): NodeSelectionChange[] | EdgeSelectionChange[] {
+  const changes: NodeSelectionChange[] | EdgeSelectionChange[] = []
 
-      if (isDef(item.selectable) && !item.selectable) {
-        willBeSelected = false
+  for (const [id, item] of items) {
+    const willBeSelected = selectedIds.has(id)
+
+    // we don't want to set all items to selected=false on the first selection
+    if (!(item.selected === undefined && !willBeSelected) && item.selected !== willBeSelected) {
+      if (mutateItem) {
+        // this hack is needed for nodes. When the user dragged a node, it's selected.
+        // When another node gets dragged, we need to deselect the previous one,
+        // in order to have only one selected node at a time - the onNodesChange callback comes too late here :/
+        item.selected = willBeSelected
       }
+      changes.push(createSelectionChange(item.id, willBeSelected))
+    }
+  }
 
-      const key = isGraphNode(item) ? 'changedNodes' : 'changedEdges'
-
-      if (!item.selected && willBeSelected) {
-        res[key].push(createSelectionChange(item.id, true))
-      } else if (item.selected && !willBeSelected) {
-        res[key].push(createSelectionChange(item.id, false))
-      }
-
-      return res
-    },
-    { changedNodes: [], changedEdges: [] } as { changedNodes: NodeSelectionChange[]; changedEdges: EdgeSelectionChange[] },
-  )
+  return changes
 }
diff --git a/packages/core/src/utils/storage.ts b/packages/core/src/utils/storage.ts
index 4034d5de..0da844ab 100644
--- a/packages/core/src/utils/storage.ts
+++ b/packages/core/src/utils/storage.ts
@@ -87,6 +87,8 @@ export class Storage {
       ...getters,
       ...actions,
       ...toRefs(reactiveState),
+      nodeLookup,
+      edgeLookup,
       emits,
       id,
       vueFlowVersion: typeof __VUE_FLOW_VERSION__ !== 'undefined' ? __VUE_FLOW_VERSION__ : 'UNKNOWN',