From 3d1b85fce23da9662d33f4f197933b8c65e393d6 Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Tue, 7 Jan 2025 17:57:22 +0100
Subject: [PATCH] refactor(core): allow setting handleId in useNodeConnections
 (#1732)

* refactor(core): allow setting handleId in useNodeConnections

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

* chore(core): deprecate useHandleConnections

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

* chore(examples): cleanup

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

* chore(docs): cleanup

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

* chore(core): cleanup

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

* chore(core): cleanup

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

* chore(docs): cleanup

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

---------

Signed-off-by: braks <78412429+bcakmakoglu@users.noreply.github.com>
---
 docs/examples/custom-node/OutputNode.vue      |   4 +-
 docs/examples/layout/ProcessNode.vue          |   6 +-
 docs/examples/math/ResultNode.vue             |   5 +-
 docs/src/examples/math.md                     |   4 +-
 docs/src/guide/composables.md                 |   7 +-
 examples/vite/src/Layouting/ProcessNode.vue   |   6 +-
 examples/vite/src/Math/ResultNode.vue         |   4 +-
 .../src/composables/useHandleConnections.ts   |   1 +
 .../src/composables/useNodeConnections.ts     | 114 ++++++------------
 packages/core/src/store/actions.ts            |   6 +-
 packages/core/src/types/connection.ts         |  12 +-
 packages/core/src/utils/store.ts              |  67 +++++++---
 12 files changed, 122 insertions(+), 114 deletions(-)

diff --git a/docs/examples/custom-node/OutputNode.vue b/docs/examples/custom-node/OutputNode.vue
index d04a719e..84045efa 100644
--- a/docs/examples/custom-node/OutputNode.vue
+++ b/docs/examples/custom-node/OutputNode.vue
@@ -1,7 +1,7 @@
 <script setup>
-import { Handle, Position, useHandleConnections, useNodesData } from '@vue-flow/core'
+import { Handle, Position, useNodeConnections, useNodesData } from '@vue-flow/core'
 
-const connections = useHandleConnections({
+const connections = useNodeConnections({
   type: 'target',
 })
 
diff --git a/docs/examples/layout/ProcessNode.vue b/docs/examples/layout/ProcessNode.vue
index 7641d874..ec2c0f2c 100644
--- a/docs/examples/layout/ProcessNode.vue
+++ b/docs/examples/layout/ProcessNode.vue
@@ -1,6 +1,6 @@
 <script setup>
 import { computed, toRef } from 'vue'
-import { Handle, useHandleConnections } from '@vue-flow/core'
+import { Handle, useNodeConnections } from '@vue-flow/core'
 
 const props = defineProps({
   data: {
@@ -15,11 +15,11 @@ const props = defineProps({
   },
 })
 
-const sourceConnections = useHandleConnections({
+const sourceConnections = useNodeConnections({
   type: 'target',
 })
 
-const targetConnections = useHandleConnections({
+const targetConnections = useNodeConnections({
   type: 'source',
 })
 
diff --git a/docs/examples/math/ResultNode.vue b/docs/examples/math/ResultNode.vue
index 999f1ab8..c9cc27f5 100644
--- a/docs/examples/math/ResultNode.vue
+++ b/docs/examples/math/ResultNode.vue
@@ -1,6 +1,6 @@
 <script setup>
 import { computed } from 'vue'
-import { Handle, Position, useHandleConnections, useNodeConnections, useNodesData } from '@vue-flow/core'
+import { Handle, Position, useNodeConnections, useNodesData } from '@vue-flow/core'
 
 defineProps(['id'])
 
@@ -12,7 +12,7 @@ const mathFunctions = {
 }
 
 // Get the source connections of the result node. In this example it's only one operator node.
-const sourceConnections = useHandleConnections({
+const sourceConnections = useNodeConnections({
   // type target means all connections where *this* node is the target
   // that means we go backwards in the graph to find the source of the connection(s)
   type: 'target',
@@ -49,7 +49,6 @@ const result = computed(() => {
 </script>
 
 <template>
-  {{ operatorSourceConnections }}
   <div class="calculation">
     <template v-for="(value, i) in valueData" :key="`${value.id}-${value.data}`">
       <span>
diff --git a/docs/src/examples/math.md b/docs/src/examples/math.md
index dd483ad5..d3cdfb25 100644
--- a/docs/src/examples/math.md
+++ b/docs/src/examples/math.md
@@ -1,8 +1,8 @@
 # Math Operation
 
-This example demonstrates how to use the different composables like `useHandleConnections` or `useNodesData` to create a data flow between nodes.
+This example demonstrates how to use the different composables like `useNodeConnections` or `useNodesData` to create a data flow between nodes.
 
-In this example we use nodes to create a simple math operation, where the user can input two numbers and select an operation to perform.
+We use nodes to create a simple math operation, where the user can input two numbers and select an operation to perform.
 
 <div class="mt-6">
   <Repl example="math"></Repl>
diff --git a/docs/src/guide/composables.md b/docs/src/guide/composables.md
index 1714575c..336a9058 100644
--- a/docs/src/guide/composables.md
+++ b/docs/src/guide/composables.md
@@ -102,8 +102,7 @@ const connections = useHandleConnections({
 ## [useNodeConnections](/typedocs/functions/useNodeConnections)
 
 `useNodeConnections` provides you with an array of connections that are connected to a specific node.
-This composable is especially useful when you want to get all connections (of either type `source` or `target`) of a node 
-instead of just the connections of a specific `<Handle>`.
+This composable is especially useful when you want to get all connections (of either type `source` or `target`) of a node.
 
 ```ts
 import { type HandleConnection, useNodeConnections } from '@vue-flow/core'
@@ -119,6 +118,10 @@ const sourceConnections = useNodeConnections({
   type: 'source',
 })
 
+const handleConnections = useNodeConnections({
+  handleId: 'handle-1', // you can explicitly pass a handle id if you want to get connections of a specific handle
+})
+
 const connections = useNodeConnections({
   nodeId: '1', // you can explicitly pass a node id, otherwise it's used from the `NodeId  injection
   type: 'target',
diff --git a/examples/vite/src/Layouting/ProcessNode.vue b/examples/vite/src/Layouting/ProcessNode.vue
index b3110e53..3e590fe5 100644
--- a/examples/vite/src/Layouting/ProcessNode.vue
+++ b/examples/vite/src/Layouting/ProcessNode.vue
@@ -1,15 +1,15 @@
 <script setup lang="ts">
 import type { NodeProps } from '@vue-flow/core'
-import { Handle, useHandleConnections } from '@vue-flow/core'
+import { Handle, useNodeConnections } from '@vue-flow/core'
 import type { ProcessNodeData } from './types'
 
 const props = defineProps<NodeProps<ProcessNodeData>>()
 
-const sourceConnections = useHandleConnections({
+const sourceConnections = useNodeConnections({
   type: 'target',
 })
 
-const targetConnections = useHandleConnections({
+const targetConnections = useNodeConnections({
   type: 'source',
 })
 
diff --git a/examples/vite/src/Math/ResultNode.vue b/examples/vite/src/Math/ResultNode.vue
index b9c9fea9..1e28fa35 100644
--- a/examples/vite/src/Math/ResultNode.vue
+++ b/examples/vite/src/Math/ResultNode.vue
@@ -1,13 +1,13 @@
 <script setup lang="ts">
 import type { GraphNode } from '@vue-flow/core'
-import { Handle, Position, useHandleConnections, useNodeConnections, useNodesData } from '@vue-flow/core'
+import { Handle, Position, useNodeConnections, useNodesData } from '@vue-flow/core'
 import type { OperatorNodeData, ValueNodeData } from './types'
 import { mathFunctions } from './utils'
 
 defineProps<{ id: string }>()
 
 // Get the source connections of the result node. In this example it's only one operator node.
-const sourceConnections = useHandleConnections({
+const sourceConnections = useNodeConnections({
   // type target means all connections where *this* node is the target
   // that means we go backwards in the graph to find the source of the connection(s)
   type: 'target',
diff --git a/packages/core/src/composables/useHandleConnections.ts b/packages/core/src/composables/useHandleConnections.ts
index dd71d5a1..06a91a6d 100644
--- a/packages/core/src/composables/useHandleConnections.ts
+++ b/packages/core/src/composables/useHandleConnections.ts
@@ -16,6 +16,7 @@ export interface UseHandleConnectionsParams {
 /**
  * Composable that returns existing connections of a `<Handle />`.
  *
+ * @deprecated use `useNodeConnections` instead
  * @public
  * @param params
  * @param params.type - handle type `source` or `target`
diff --git a/packages/core/src/composables/useNodeConnections.ts b/packages/core/src/composables/useNodeConnections.ts
index 9ef53ab2..125953f1 100644
--- a/packages/core/src/composables/useNodeConnections.ts
+++ b/packages/core/src/composables/useNodeConnections.ts
@@ -1,120 +1,80 @@
-import type { ComputedRef, MaybeRefOrGetter } from 'vue'
-import { computed, ref, toRef, toValue, watch } from 'vue'
-import type { HandleConnection, HandleElement, HandleType } from '../types'
+import type { MaybeRefOrGetter } from 'vue'
+import { computed, ref, toValue, watch } from 'vue'
+import type { HandleType, NodeConnection } from '../types'
+import { areConnectionMapsEqual, handleConnectionChange } from '../utils'
 import { useNodeId } from './useNodeId'
 import { useVueFlow } from './useVueFlow'
 
 export interface UseNodeConnectionsParams {
-  type: MaybeRefOrGetter<HandleType>
+  type?: MaybeRefOrGetter<HandleType>
+  handleId?: MaybeRefOrGetter<string | null>
   nodeId?: MaybeRefOrGetter<string | null>
-  onConnect?: (connections: HandleConnection[]) => void
-  onDisconnect?: (connections: HandleConnection[]) => void
+  onConnect?: (connections: NodeConnection[]) => void
+  onDisconnect?: (connections: NodeConnection[]) => void
 }
 
 /**
- * Composable that returns existing connections of a node by handle type.
- * This is useful when you want to get all connections of a node by a specific handle type.
+ * Hook to retrieve all edges connected to a node. Can be filtered by handle type and id.
  *
  * @public
  * @param params
  * @param params.type - handle type `source` or `target`
  * @param params.nodeId - node id - if not provided, the node id from the `useNodeId` (meaning, the context-based injection) is used
+ * @param params.handleId - the handle id (this is required if the node has multiple handles of the same type)
  * @param params.onConnect - gets called when a connection is created
  * @param params.onDisconnect - gets called when a connection is removed
  *
  * @returns An array of connections
  */
-export function useNodeConnections(params: UseNodeConnectionsParams): ComputedRef<HandleConnection[]> {
-  const { type, nodeId, onConnect, onDisconnect } = params
+export function useNodeConnections(params: UseNodeConnectionsParams = {}) {
+  const { type, handleId, nodeId, onConnect, onDisconnect } = params
 
-  const { connectionLookup, findNode } = useVueFlow()
+  const { connectionLookup } = useVueFlow()
 
   const _nodeId = useNodeId()
 
-  const currentNodeId = toRef(() => toValue(nodeId) ?? _nodeId)
+  const prevConnections = ref<Map<string, NodeConnection> | null>(null)
 
-  const handleType = toRef(() => toValue(type))
+  const connections = ref<Map<string, NodeConnection>>()
 
-  const node = computed(() => findNode(currentNodeId.value))
+  const lookupKey = computed(() => {
+    const currNodeId = toValue(nodeId) ?? _nodeId
+    const handleType = toValue(type)
+    const currHandleId = toValue(handleId)
 
-  const handleIds = computed(() => {
-    if (!node.value) {
-      return []
-    }
-
-    const handles: HandleElement['id'][] = []
-    for (const handle of node.value?.handleBounds?.[handleType.value] ?? []) {
-      handles.push(handle.id)
-    }
-
-    return handles
-  })
-
-  const prevConnections = ref<Map<string, HandleConnection> | null>(null)
-
-  const connectionsFromLookup = computed(() => {
-    const nodeConnections = [] as Map<string, HandleConnection>[]
-
-    for (const handleId of handleIds.value) {
-      const connectionMap = connectionLookup.value.get(`${currentNodeId.value}-${handleType.value}-${handleId}`)
-      if (connectionMap) {
-        nodeConnections.push(connectionMap)
-      }
-    }
-
-    return nodeConnections
+    return `${currNodeId}${handleType ? (currHandleId ? `-${handleType}-${currHandleId}` : `-${handleType}`) : ''}`
   })
 
   watch(
-    [connectionsFromLookup, () => typeof onConnect !== 'undefined', () => typeof onDisconnect !== 'undefined'],
-    ([currentConnections]) => {
-      if (!currentConnections) {
+    () => connectionLookup.value.get(lookupKey.value),
+    (nextConnections) => {
+      if (areConnectionMapsEqual(connections.value, nextConnections)) {
         return
       }
 
-      const newConnections = new Map<string, HandleConnection>()
+      connections.value = nextConnections
+    },
+    { immediate: true },
+  )
 
-      for (const connectionMap of currentConnections) {
-        for (const [key, connection] of connectionMap) {
-          newConnections.set(key, connection)
-        }
+  watch(
+    [connections, () => typeof onConnect !== 'undefined', () => typeof onDisconnect !== 'undefined'],
+    ([currentConnections = new Map<string, NodeConnection>()]) => {
+      if (prevConnections.value && prevConnections.value !== currentConnections) {
+        handleConnectionChange(prevConnections.value, currentConnections, onDisconnect)
+        handleConnectionChange(currentConnections, prevConnections.value, onConnect)
       }
 
-      if (!prevConnections.value) {
-        prevConnections.value = new Map(newConnections)
-        return
-      }
-
-      const prevConnectionsValue = prevConnections.value
-
-      const addedConnections = Array.from(newConnections.keys()).filter((key) => !prevConnectionsValue.has(key))
-
-      const removedConnections = Array.from(prevConnectionsValue.keys()).filter((key) => !newConnections.has(key))
-
-      if (addedConnections.length && onConnect) {
-        const added = addedConnections.map((key) => newConnections.get(key)!)
-        onConnect(added)
-      }
-
-      if (removedConnections.length && onDisconnect) {
-        const removed = removedConnections.map((key) => prevConnectionsValue.get(key)!)
-        onDisconnect(removed)
-      }
-
-      prevConnections.value = new Map(newConnections)
+      prevConnections.value = currentConnections
     },
     { immediate: true },
   )
 
   return computed(() => {
-    const connections = [] as HandleConnection[]
-
-    for (const connectionMap of connectionsFromLookup.value) {
-      for (const connection of connectionMap.values()) {
-        connections.push(connection)
-      }
+    if (!connections.value) {
+      return []
     }
 
-    return connections
+    return Array.from(connections.value.values())
   })
 }
diff --git a/packages/core/src/store/actions.ts b/packages/core/src/store/actions.ts
index 6c29b135..882aac2e 100644
--- a/packages/core/src/store/actions.ts
+++ b/packages/core/src/store/actions.ts
@@ -349,7 +349,7 @@ export function useActions(state: State, nodeLookup: ComputedRef<NodeLookup>, ed
       state.edges,
     )
 
-    updateConnectionLookup(state.connectionLookup, validEdges)
+    updateConnectionLookup(state.connectionLookup, edgeLookup.value, validEdges)
 
     state.edges = validEdges
   }
@@ -527,7 +527,7 @@ export function useActions(state: State, nodeLookup: ComputedRef<NodeLookup>, ed
 
       state.edges.splice(state.edges.indexOf(prevEdge), 1, validEdge)
 
-      updateConnectionLookup(state.connectionLookup, [validEdge])
+      updateConnectionLookup(state.connectionLookup, edgeLookup.value, [validEdge])
 
       return validEdge
     }
@@ -554,7 +554,7 @@ export function useActions(state: State, nodeLookup: ComputedRef<NodeLookup>, ed
   const applyEdgeChanges: Actions['applyEdgeChanges'] = (changes) => {
     const changedEdges = applyChanges(changes, state.edges)
 
-    updateConnectionLookup(state.connectionLookup, changedEdges)
+    updateConnectionLookup(state.connectionLookup, edgeLookup.value, changedEdges)
 
     return changedEdges
   }
diff --git a/packages/core/src/types/connection.ts b/packages/core/src/types/connection.ts
index 51f1906c..ef623036 100644
--- a/packages/core/src/types/connection.ts
+++ b/packages/core/src/types/connection.ts
@@ -33,11 +33,19 @@ export interface Connection {
   targetHandle?: string | null
 }
 
-/** Connection with edge id */
+/**
+ * Connection with edge id
+ * @deprecated
+ * todo: remove in next major release
+ */
 export interface HandleConnection extends Connection {
   edgeId: string
 }
 
+export type NodeConnection = Connection & {
+  edgeId: string
+}
+
 export type Connector = (
   params: Connection,
 ) => Promise<(Connection & Partial<Edge>) | false> | ((Connection & Partial<Edge>) | false)
@@ -89,4 +97,4 @@ export interface ConnectionLineProps {
   connectionStatus: ConnectionStatus | null
 }
 
-export type ConnectionLookup = Map<string, Map<string, HandleConnection>>
+export type ConnectionLookup = Map<string, Map<string, NodeConnection>>
diff --git a/packages/core/src/utils/store.ts b/packages/core/src/utils/store.ts
index 71f2b3a4..0cbad7f0 100644
--- a/packages/core/src/utils/store.ts
+++ b/packages/core/src/utils/store.ts
@@ -1,14 +1,15 @@
-import { markRaw, unref } from 'vue'
+import { unref } from 'vue'
 import type {
   Actions,
   Connection,
   ConnectionLookup,
   DefaultEdgeOptions,
   Edge,
+  EdgeLookup,
   GraphEdge,
   GraphNode,
-  HandleConnection,
   Node,
+  NodeConnection,
   State,
   ValidConnectionFunc,
   VueFlowStore,
@@ -126,21 +127,57 @@ export function createGraphNodes(nodes: Node[], findNode: Actions['findNode'], t
   return nextNodes
 }
 
-export function updateConnectionLookup(connectionLookup: ConnectionLookup, edges: Edge[]) {
+/**
+ * this function adds the connection to the connectionLookup
+ * at the following keys: nodeId-type-handleId, nodeId-type and nodeId
+ * @param type type of the connection
+ * @param connection connection that should be added to the lookup
+ * @param connectionKey at which key the connection should be added
+ * @param connectionLookup reference to the connection lookup
+ * @param nodeId nodeId of the connection
+ * @param handleId handleId of the conneciton
+ */
+function addConnectionToLookup(
+  type: 'source' | 'target',
+  connection: NodeConnection,
+  connectionKey: string,
+  connectionLookup: ConnectionLookup,
+  nodeId: string,
+  handleId: string | null,
+) {
+  // We add the connection to the connectionLookup at the following keys
+  // 1. nodeId, 2. nodeId-type, 3. nodeId-type-handleId
+  // If the key already exists, we add the connection to the existing map
+  let key = nodeId
+  const nodeMap = connectionLookup.get(key) || new Map()
+  connectionLookup.set(key, nodeMap.set(connectionKey, connection))
+
+  key = `${nodeId}-${type}`
+  const typeMap = connectionLookup.get(key) || new Map()
+  connectionLookup.set(key, typeMap.set(connectionKey, connection))
+
+  if (handleId) {
+    key = `${nodeId}-${type}-${handleId}`
+    const handleMap = connectionLookup.get(key) || new Map()
+    connectionLookup.set(key, handleMap.set(connectionKey, connection))
+  }
+}
+
+export function updateConnectionLookup(connectionLookup: ConnectionLookup, edgeLookup: EdgeLookup, edges: GraphEdge[]) {
   connectionLookup.clear()
+  edgeLookup.clear()
 
   for (const edge of edges) {
-    const { id: edgeId, source, target, sourceHandle = null, targetHandle = null } = edge
+    const { source: sourceNode, target: targetNode, sourceHandle = null, targetHandle = null } = edge
 
-    const sourceKey = `${source}-source-${sourceHandle}`
-    const targetKey = `${target}-target-${targetHandle}`
+    const connection = { edgeId: edge.id, source: sourceNode, target: targetNode, sourceHandle, targetHandle }
+    const sourceKey = `${sourceNode}-${sourceHandle}`
+    const targetKey = `${targetNode}-${targetHandle}`
 
-    const prevSource = connectionLookup.get(sourceKey) || new Map()
-    const prevTarget = connectionLookup.get(targetKey) || new Map()
-    const connection = markRaw({ edgeId, source, target, sourceHandle, targetHandle })
+    addConnectionToLookup('source', connection, targetKey, connectionLookup, sourceNode, sourceHandle)
+    addConnectionToLookup('target', connection, sourceKey, connectionLookup, targetNode, targetHandle)
 
-    connectionLookup.set(sourceKey, prevSource.set(`${target}-${targetHandle}`, connection))
-    connectionLookup.set(targetKey, prevTarget.set(`${source}-${sourceHandle}`, connection))
+    edgeLookup.set(edge.id, edge)
   }
 }
 
@@ -150,15 +187,15 @@ export function updateConnectionLookup(connectionLookup: ConnectionLookup, edges
  * @internal
  */
 export function handleConnectionChange(
-  a: Map<string, HandleConnection>,
-  b: Map<string, HandleConnection>,
-  cb?: (diff: HandleConnection[]) => void,
+  a: Map<string, NodeConnection>,
+  b: Map<string, NodeConnection>,
+  cb?: (diff: NodeConnection[]) => void,
 ) {
   if (!cb) {
     return
   }
 
-  const diff: HandleConnection[] = []
+  const diff: NodeConnection[] = []
 
   for (const key of a.keys()) {
     if (!b.has(key)) {