refactor(connectionlookup): use maps instead of arrays

examples/react/src/examples/UseHandleConnectionStatus/SingleHandleNode.tsx

@@ -23,7 +23,7 @@ function CustomHandle({ nodeId, ...handleProps }: HandleComponentProps & { nodeI
   });
 
   useEffect(() => {
-    //  console.log('useEffect, node id:', nodeId, handleProps.type, status);
+    console.log('useEffect, node id:', nodeId, handleProps.type, status);
   }, [status]);
 
   return <Handle {...handleProps} />;

packages/react/src/hooks/useHandleConnectionStatus.ts

@@ -3,7 +3,6 @@ import { Connection, HandleType } from '@xyflow/system';
 
 import { useStore } from './useStore';
 import { useNodeId } from '../contexts/NodeIdContext';
-import { areConnectionsEqual, isSameConnection } from '../utils/general';
 
 type useHandleConnectionStatusParams = {
   handleType: HandleType;
@@ -13,10 +12,69 @@ type useHandleConnectionStatusParams = {
   onDisconnect?: (connections: Connection[]) => void;
 };
 
+function areConnectionMapsEqual(a?: Map<string, Connection>, b?: Map<string, Connection>) {
+  if (!a && !b) {
+    return true;
+  }
+
+  if (!a || !b || a.size !== b.size) {
+    return false;
+  }
+
+  if (!a.size && !b.size) {
+    return true;
+  }
+
+  for (const key of a.keys()) {
+    if (!b.has(key)) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+/**
+ * We call the callback for all connections in a that are not in b
+ * @internal
+ */
+function handleConnectionChange(
+  a: Map<string, Connection>,
+  b: Map<string, Connection>,
+  cb?: (diff: Connection[]) => void
+) {
+  if (!cb) {
+    return;
+  }
+
+  const diff: Connection[] = [];
+
+  a.forEach((connection, key) => {
+    if (!b?.has(key)) {
+      diff.push(connection);
+    }
+  });
+
+  if (diff.length) {
+    cb(diff);
+  }
+}
+
+/**
+ *  Hook to check if a <Handle /> is connected to another <Handle /> and get the connections.
+ *
+ * @public
+ * @param param.handleType - 'source' or 'target'
+ * @param param.handleId - the handle id (this is only needed if the node has multiple handles of the same type)
+ * @param param.nodeId - node id - if not provided, the node id from the NodeIdContext is used
+ * @param param.onConnect - gets called when a connection is established
+ * @param param.onDisconnect - gets called when a connection is removed
+ * @returns a `connected` boolean and a connections array
+ */
 export function useHandleConnectionStatus({
   handleType,
-  nodeId,
   handleId = null,
+  nodeId,
   onConnect,
   onDisconnect,
 }: useHandleConnectionStatusParams): {
@@ -24,41 +82,29 @@ export function useHandleConnectionStatus({
   connections: Connection[] | null;
 } {
   const _nodeId = useNodeId();
-  const prevConnections = useRef<Connection[] | null>(null);
+  const prevConnections = useRef<Map<string, Connection> | null>(null);
   const currentNodeId = nodeId || _nodeId;
 
   const connections = useStore(
-    (state) => state.connectionLookup.get(`${currentNodeId}-${handleType}-${handleId}`) || null,
-    areConnectionsEqual
+    (state) => state.connectionLookup.get(`${currentNodeId}-${handleType}-${handleId}`),
+    areConnectionMapsEqual
   );
 
   useEffect(() => {
     // @todo dicuss if onConnect/onDisconnect should be called when the component mounts/unmounts
     if (prevConnections.current && prevConnections.current !== connections) {
-      const disconnectedConnections = prevConnections.current.filter(
-        (prevConnection) => !connections?.find((connection) => isSameConnection(connection, prevConnection))
-      );
-
-      const newConnections = connections?.filter(
-        (connection) => !prevConnections.current?.find((prevConnection) => isSameConnection(prevConnection, connection))
-      );
-
-      if (disconnectedConnections.length) {
-        onDisconnect?.(disconnectedConnections);
-      }
-
-      if (newConnections?.length) {
-        onConnect?.(newConnections);
-      }
+      const _connections = connections ?? new Map();
+      handleConnectionChange(prevConnections.current, _connections, onDisconnect);
+      handleConnectionChange(_connections, prevConnections.current, onConnect);
     }
 
-    prevConnections.current = connections ?? [];
+    prevConnections.current = connections ?? new Map();
   }, [connections, onConnect, onDisconnect]);
 
   return useMemo(
     () => ({
       connected: !!connections,
-      connections,
+      connections: Array.from(connections?.values() ?? []),
     }),
     [connections]
   );

packages/react/src/store/initialState.ts

@@ -25,7 +25,7 @@ const getInitialState = ({
   fitView?: boolean;
 } = {}): ReactFlowStore => {
   const nodeLookup = new Map<string, Node>();
-  const connectionLookup = updateConnectionLookup(new Map<string, Connection[]>(), edges);
+  const connectionLookup = updateConnectionLookup(new Map<string, Map<string, Connection>>(), edges);
   const nextNodes = updateNodes(nodes, nodeLookup, { nodeOrigin: [0, 0], elevateNodesOnSelect: false });
 
   let transform: Transform = [0, 0, 1];

packages/react/src/store/utils.ts

@@ -44,7 +44,7 @@ export function updateNodesAndEdgesSelections({ changedNodes, changedEdges, get,
   }
 }
 
-export function updateConnectionLookup(lookup: Map<string, Connection[]>, edges: Edge[]) {
+export function updateConnectionLookup(lookup: Map<string, Map<string, Connection>>, edges: Edge[]) {
   lookup.clear();
 
   edges.forEach(({ source, target, sourceHandle = null, targetHandle = null }) => {
@@ -52,13 +52,12 @@ export function updateConnectionLookup(lookup: Map<string, Connection[]>, edges:
       const sourceKey = `${source}-source-${sourceHandle}`;
       const targetKey = `${target}-target-${targetHandle}`;
 
-      const prevSource = lookup.get(sourceKey);
-      const prevTarget = lookup.get(targetKey);
-
+      const prevSource = lookup.get(sourceKey) || new Map();
+      const prevTarget = lookup.get(targetKey) || new Map();
       const connection = { source, target, sourceHandle, targetHandle };
 
-      lookup.set(sourceKey, prevSource ? [...prevSource, connection] : [connection]);
-      lookup.set(targetKey, prevTarget ? [...prevTarget, connection] : [connection]);
+      lookup.set(sourceKey, prevSource.set(`${target}-${targetHandle}`, connection));
+      lookup.set(targetKey, prevTarget.set(`${source}-${sourceHandle}`, connection));
     }
   });
 

packages/react/src/types/store.ts

@@ -50,7 +50,7 @@ export type ReactFlowStore = {
   nodes: Node[];
   nodeLookup: Map<string, Node>;
   edges: Edge[];
-  connectionLookup: Map<string, Connection[]>;
+  connectionLookup: Map<string, Map<string, Connection>>;
 
   onNodesChange: OnNodesChange | null;
   onEdgesChange: OnEdgesChange | null;

packages/react/src/utils/general.ts

@@ -6,7 +6,6 @@ import {
   getIncomersBase,
   updateEdgeBase,
   getConnectedEdgesBase,
-  Connection,
 } from '@xyflow/system';
 
 import type { Edge, Node } from '../types';
@@ -18,32 +17,3 @@ export const getIncomers = getIncomersBase<Node, Edge>;
 export const addEdge = addEdgeBase<Edge>;
 export const updateEdge = updateEdgeBase<Edge>;
 export const getConnectedEdges = getConnectedEdgesBase<Node, Edge>;
-
-export function isSameConnection(a: Connection, b: Connection) {
-  return (
-    a.source === b.source &&
-    a.target === b.target &&
-    a.sourceHandle === b.sourceHandle &&
-    a.targetHandle === b.targetHandle
-  );
-}
-
-export function areConnectionsEqual(a: Connection[] | null | undefined, b: Connection[] | null | undefined) {
-  if (!a && !b) {
-    return true;
-  }
-
-  if (!a || !b) {
-    return false;
-  }
-
-  if (a.length !== b.length) {
-    return false;
-  }
-
-  if (!a.length && !b.length) {
-    return true;
-  }
-
-  return !a.some((connA) => !b.find((connB) => isSameConnection(connA, connB)));
-}