From 268ed0e84ee3d77aab6862ad3003da6e8356c879 Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Mon, 6 Jan 2025 12:30:47 +0100
Subject: [PATCH] feat(core): add `useNodeConnections` composable (#1728)

* feat(core): add `useNodeConnections` composable

* chore(core): export `useNodeConnections`

* chore(docs): update math example

* docs: add guide for `useNodeConnections`

* chore(changeset): add
---
 .changeset/gentle-owls-grin.md                |   5 +
 docs/examples/math/ResultNode.vue             |  12 +-
 docs/src/guide/composables.md                 |  41 +++++-
 examples/vite/src/Math/ResultNode.vue         |  11 +-
 .../src/composables/useNodeConnections.ts     | 120 ++++++++++++++++++
 packages/core/src/index.ts                    |   1 +
 6 files changed, 174 insertions(+), 16 deletions(-)
 create mode 100644 .changeset/gentle-owls-grin.md
 create mode 100644 packages/core/src/composables/useNodeConnections.ts

diff --git a/.changeset/gentle-owls-grin.md b/.changeset/gentle-owls-grin.md
new file mode 100644
index 00000000..d4a7d5cc
--- /dev/null
+++ b/.changeset/gentle-owls-grin.md
@@ -0,0 +1,5 @@
+---
+"@vue-flow/core": minor
+---
+
+Add `useNodeConnections` composable
diff --git a/docs/examples/math/ResultNode.vue b/docs/examples/math/ResultNode.vue
index 2e15a807..999f1ab8 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, useNodesData, useVueFlow } from '@vue-flow/core'
+import { Handle, Position, useHandleConnections, useNodeConnections, useNodesData } from '@vue-flow/core'
 
 defineProps(['id'])
 
@@ -11,8 +11,6 @@ const mathFunctions = {
   '/': (a, b) => a / b,
 }
 
-const { getConnectedEdges } = useVueFlow()
-
 // Get the source connections of the result node. In this example it's only one operator node.
 const sourceConnections = useHandleConnections({
   // type target means all connections where *this* node is the target
@@ -21,9 +19,10 @@ const sourceConnections = useHandleConnections({
 })
 
 // Get the source connections of the operator node
-const operatorSourceConnections = computed(() =>
-  getConnectedEdges(sourceConnections.value[0].source).filter((e) => e.source !== sourceConnections.value[0].source),
-)
+const operatorSourceConnections = useNodeConnections({
+  type: 'target',
+  nodeId: () => sourceConnections.value[0]?.source,
+})
 
 const operatorData = useNodesData(() => sourceConnections.value.map((connection) => connection.source))
 
@@ -50,6 +49,7 @@ 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/guide/composables.md b/docs/src/guide/composables.md
index 63d02df2..1714575c 100644
--- a/docs/src/guide/composables.md
+++ b/docs/src/guide/composables.md
@@ -70,13 +70,14 @@ onInit((instance) => {
 
 ## [useHandleConnections](/typedocs/functions/useHandleConnections)
 
-`useHandleConnections` provides you with an array of connections that are connected to the node you pass to it.
+`useHandleConnections` provides you with an array of connections that are connected to specific `<Handle>`.
 
 ```ts
-import { useHandleConnections } from '@vue-flow/core'
+import { type HandleConnection, useHandleConnections } from '@vue-flow/core'
 
 // get all connections where this node is the target (incoming connections)
 const targetConnections = useHandleConnections({
+  // type is required
   type: 'target',
 })
 
@@ -89,10 +90,42 @@ const connections = useHandleConnections({
   id: 'handle-1', // you can explicitly pass a handle id if there are multiple handles of the same type
   nodeId: '1', // you can explicitly pass a node id, otherwise it's used from the `NodeId  injection
   type: 'target',
-  onConnect: (connections: Connection[]) => {
+  onConnect: (connections: HandleConnection[]) => {
     // do something with the connections
   },
-  onDisconnect: (connections: Connection[]) => {
+  onDisconnect: (connections: HandleConnection[]) => {
+    // do something with the connections
+  },
+})
+```
+
+## [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>`.
+
+```ts
+import { type HandleConnection, useNodeConnections } from '@vue-flow/core'
+
+// get all connections where this node is the target (incoming connections)
+const targetConnections = useNodeConnections({
+  // type is required
+  type: 'target',
+})
+
+// get all connections where this node is the source (outgoing connections)
+const sourceConnections = useNodeConnections({
+  type: 'source',
+})
+
+const connections = useNodeConnections({
+  nodeId: '1', // you can explicitly pass a node id, otherwise it's used from the `NodeId  injection
+  type: 'target',
+  onConnect: (connections: HandleConnection[]) => {
+    // do something with the connections
+  },
+  onDisconnect: (connections: HandleConnection[]) => {
     // do something with the connections
   },
 })
diff --git a/examples/vite/src/Math/ResultNode.vue b/examples/vite/src/Math/ResultNode.vue
index 495e44fe..b9c9fea9 100644
--- a/examples/vite/src/Math/ResultNode.vue
+++ b/examples/vite/src/Math/ResultNode.vue
@@ -1,13 +1,11 @@
 <script setup lang="ts">
 import type { GraphNode } from '@vue-flow/core'
-import { Handle, Position, useHandleConnections, useNodesData, useVueFlow } from '@vue-flow/core'
+import { Handle, Position, useHandleConnections, useNodeConnections, useNodesData } from '@vue-flow/core'
 import type { OperatorNodeData, ValueNodeData } from './types'
 import { mathFunctions } from './utils'
 
 defineProps<{ id: string }>()
 
-const { getConnectedEdges } = useVueFlow()
-
 // Get the source connections of the result node. In this example it's only one operator node.
 const sourceConnections = useHandleConnections({
   // type target means all connections where *this* node is the target
@@ -16,9 +14,10 @@ const sourceConnections = useHandleConnections({
 })
 
 // Get the source connections of the operator node
-const operatorSourceConnections = computed(() =>
-  getConnectedEdges(sourceConnections.value[0].source).filter((e) => e.source !== sourceConnections.value[0].source),
-)
+const operatorSourceConnections = useNodeConnections({
+  type: 'target',
+  nodeId: () => sourceConnections.value[0]?.source,
+})
 
 const operatorData = useNodesData<GraphNode<OperatorNodeData>>(() =>
   sourceConnections.value.map((connection) => connection.source),
diff --git a/packages/core/src/composables/useNodeConnections.ts b/packages/core/src/composables/useNodeConnections.ts
new file mode 100644
index 00000000..9ef53ab2
--- /dev/null
+++ b/packages/core/src/composables/useNodeConnections.ts
@@ -0,0 +1,120 @@
+import type { ComputedRef, MaybeRefOrGetter } from 'vue'
+import { computed, ref, toRef, toValue, watch } from 'vue'
+import type { HandleConnection, HandleElement, HandleType } from '../types'
+import { useNodeId } from './useNodeId'
+import { useVueFlow } from './useVueFlow'
+
+export interface UseNodeConnectionsParams {
+  type: MaybeRefOrGetter<HandleType>
+  nodeId?: MaybeRefOrGetter<string | null>
+  onConnect?: (connections: HandleConnection[]) => void
+  onDisconnect?: (connections: HandleConnection[]) => 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.
+ *
+ * @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.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
+
+  const { connectionLookup, findNode } = useVueFlow()
+
+  const _nodeId = useNodeId()
+
+  const currentNodeId = toRef(() => toValue(nodeId) ?? _nodeId)
+
+  const handleType = toRef(() => toValue(type))
+
+  const node = computed(() => findNode(currentNodeId.value))
+
+  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
+  })
+
+  watch(
+    [connectionsFromLookup, () => typeof onConnect !== 'undefined', () => typeof onDisconnect !== 'undefined'],
+    ([currentConnections]) => {
+      if (!currentConnections) {
+        return
+      }
+
+      const newConnections = new Map<string, HandleConnection>()
+
+      for (const connectionMap of currentConnections) {
+        for (const [key, connection] of connectionMap) {
+          newConnections.set(key, connection)
+        }
+      }
+
+      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)
+    },
+    { immediate: true },
+  )
+
+  return computed(() => {
+    const connections = [] as HandleConnection[]
+
+    for (const connectionMap of connectionsFromLookup.value) {
+      for (const connection of connectionMap.values()) {
+        connections.push(connection)
+      }
+    }
+
+    return connections
+  })
+}
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 371c4dfa..782f2ed2 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -70,6 +70,7 @@ export { useGetPointerPosition } from './composables/useGetPointerPosition'
 export { useNodeId } from './composables/useNodeId'
 export { useConnection } from './composables/useConnection'
 export { useHandleConnections } from './composables/useHandleConnections'
+export { useNodeConnections } from './composables/useNodeConnections'
 export { useNodesData } from './composables/useNodesData'
 export { useEdgesData } from './composables/useEdgesData'
 export { useNodesInitialized } from './composables/useNodesInitialized'