Merge pull request #4725 from xyflow/feat/use-handle-connection

add `useNodeConnections`, deprecate `useHandleConnections`

packages/system/src/constants.ts

@@ -24,6 +24,8 @@ export const errorMessages = {
     `Node with id "${id}" does not exist, it may have been removed. This can happen when a node is deleted before the "onNodeClick" handler is called.`,
   error013: (lib: string = 'react') =>
     `It seems that you haven't loaded the styles. Please import '@xyflow/${lib}/dist/style.css' or base.css to make sure everything is working properly.`,
+  error014: () =>
+    'useNodeConnections: No node ID found. Call useNodeConnections inside a custom Node or provide a node ID.',
 };
 
 export const infiniteExtent: CoordinateExtent = [

packages/system/src/types/general.ts

@@ -33,10 +33,15 @@ export type Connection = {
   targetHandle: string | null;
 };
 
+// TODO: remove in next version
 export type HandleConnection = Connection & {
   edgeId: string;
 };
 
+export type NodeConnection = Connection & {
+  edgeId: string;
+};
+
 export enum ConnectionMode {
   Strict = 'strict',
   Loose = 'loose',

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,56 @@ export async function panBy({
   return Promise.resolve(transformChanged);
 }
 
+/**
+ * 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: HandleConnection,
+  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: 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));
   }
 }