From 0ae06ce5cbe658501c6c8332036c27f7c824d67b Mon Sep 17 00:00:00 2001
From: peterkogo <peter@xyflow.com>
Date: Thu, 10 Oct 2024 17:20:21 +0200
Subject: [PATCH] implemented more broader tracking of handleConnections

---
 .../react/src/hooks/useHandleConnections.ts   |  6 ++--
 .../src/lib/hooks/useHandleConnections.ts     |  8 +++--
 packages/system/src/utils/store.ts            | 36 +++++++++++++------
 3 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/packages/react/src/hooks/useHandleConnections.ts b/packages/react/src/hooks/useHandleConnections.ts
index 20b65a2d..ff20c909 100644
--- a/packages/react/src/hooks/useHandleConnections.ts
+++ b/packages/react/src/hooks/useHandleConnections.ts
@@ -12,7 +12,7 @@ import { useNodeId } from '../contexts/NodeIdContext';
 
 type useHandleConnectionsParams = {
   type: HandleType;
-  id?: string | null;
+  id?: string;
   nodeId?: string;
   onConnect?: (connections: Connection[]) => void;
   onDisconnect?: (connections: Connection[]) => void;
@@ -31,7 +31,7 @@ type useHandleConnectionsParams = {
  */
 export function useHandleConnections({
   type,
-  id = null,
+  id,
   nodeId,
   onConnect,
   onDisconnect,
@@ -42,7 +42,7 @@ export function useHandleConnections({
   const prevConnections = useRef<Map<string, HandleConnection> | null>(null);
 
   const connections = useStore(
-    (state) => state.connectionLookup.get(`${currentNodeId}-${type}-${id}`),
+    (state) => state.connectionLookup.get(`${currentNodeId}${type ? (id ? `-${type}-${id}` : `-${type}`) : ''}`),
     areConnectionMapsEqual
   );
 
diff --git a/packages/svelte/src/lib/hooks/useHandleConnections.ts b/packages/svelte/src/lib/hooks/useHandleConnections.ts
index 7aec3858..909ecee6 100644
--- a/packages/svelte/src/lib/hooks/useHandleConnections.ts
+++ b/packages/svelte/src/lib/hooks/useHandleConnections.ts
@@ -7,7 +7,7 @@ import { getContext } from 'svelte';
 export type useHandleConnectionsParams = {
   type: HandleType;
   nodeId?: string;
-  id?: string | null;
+  id?: string;
 };
 
 const initialConnections: HandleConnection[] = [];
@@ -21,7 +21,7 @@ const initialConnections: HandleConnection[] = [];
  * @param param.id - the handle id (this is only needed if the node has multiple handles of the same type)
  * @returns an array with connections
  */
-export function useHandleConnections({ type, nodeId, id = null }: useHandleConnectionsParams) {
+export function useHandleConnections({ type, nodeId, id }: useHandleConnectionsParams) {
   const { edges, connectionLookup } = useStore();
 
   const _nodeId = getContext<string>('svelteflow__node_id');
@@ -32,7 +32,9 @@ export function useHandleConnections({ type, nodeId, id = null }: useHandleConne
   return derived(
     [edges, connectionLookup],
     ([, connectionLookup], set) => {
-      const nextConnections = connectionLookup.get(`${currentNodeId}-${type}-${id || null}`);
+      const nextConnections = connectionLookup.get(
+        `${currentNodeId}${type ? (id ? `-${type}-${id}` : `-${type}`) : ''}`
+      );
 
       if (!areConnectionMapsEqual(nextConnections, prevConnections)) {
         prevConnections = nextConnections;
diff --git a/packages/system/src/utils/store.ts b/packages/system/src/utils/store.ts
index ec55759a..d7282e96 100644
--- a/packages/system/src/utils/store.ts
+++ b/packages/system/src/utils/store.ts
@@ -1,4 +1,4 @@
-import { infiniteExtent } from '..';
+import { HandleConnection, infiniteExtent } from '..';
 import {
   NodeBase,
   CoordinateExtent,
@@ -42,7 +42,7 @@ const adoptUserNodesDefaultOptions = {
   checkEquality: true,
 };
 
-function mergeObjects<T extends Record<string, any>>(base: T, incoming?: Partial<T>): T {
+function mergeObjects<T extends Record<string, unknown>>(base: T, incoming?: Partial<T>): T {
   const result = { ...base };
   for (const key in incoming) {
     if (incoming[key] !== undefined) {
@@ -439,22 +439,38 @@ export async function panBy({
   return Promise.resolve(transformChanged);
 }
 
+function addConnectionToLookup(
+  type: 'source' | 'target',
+  connection: HandleConnection,
+  connectionKey: string,
+  connectionLookup: ConnectionLookup,
+  nodeId: string,
+  handleId: string | null
+) {
+  const keyFragments = [nodeId, type, handleId];
+  let key = '';
+  for (const keyFragment of keyFragments) {
+    key += keyFragment;
+    const prevMap = connectionLookup.get(key) || new Map();
+    connectionLookup.set(key, prevMap.set(connectionKey, connection));
+    key += '-';
+  }
+}
+
 export function updateConnectionLookup(connectionLookup: ConnectionLookup, edgeLookup: EdgeLookup, edges: EdgeBase[]) {
   connectionLookup.clear();
   edgeLookup.clear();
 
   for (const edge of edges) {
-    const { 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 = { edgeId: edge.id, source, target, sourceHandle, targetHandle };
+    addConnectionToLookup('source', connection, targetKey, connectionLookup, sourceNode, sourceHandle);
+    addConnectionToLookup('target', connection, sourceKey, connectionLookup, targetNode, targetHandle);
 
     edgeLookup.set(edge.id, edge);
-    connectionLookup.set(sourceKey, prevSource.set(`${target}-${targetHandle}`, connection));
-    connectionLookup.set(targetKey, prevTarget.set(`${source}-${sourceHandle}`, connection));
   }
 }