feat(react): connection lookup draft

packages/react/src/hooks/useHandleConnectionStatus.ts

@@ -0,0 +1,85 @@
+import { useEffect, useMemo, useRef } from 'react';
+import { Connection, HandleType } from '@xyflow/system';
+
+import { useStore } from './useStore';
+import { useNodeId } from '../contexts/NodeIdContext';
+
+type useHandleConnectionStatusParams = {
+  handleType: HandleType;
+  nodeId?: string;
+  handleId?: string | null;
+  onConnect?: (connections: Connection[]) => void;
+  onDisconnect?: (connections: Connection[]) => void;
+};
+
+function connectionsEqual(a: Connection[] | null, b: Connection[] | null) {
+  if (!a && !b) {
+    return true;
+  }
+
+  if (!a || !b) {
+    return false;
+  }
+
+  if (a.length !== b.length) {
+    return false;
+  }
+
+  return a.every((connA) =>
+    b.find(
+      (connB) =>
+        connA.source === connB.source &&
+        connA.target === connB.target &&
+        connA.sourceHandle === connB.sourceHandle &&
+        connA.targetHandle === connB.targetHandle
+    )
+  );
+}
+
+export function useHandleConnectionStatus({
+  handleType,
+  nodeId,
+  handleId = null,
+  onConnect,
+  onDisconnect,
+}: useHandleConnectionStatusParams): {
+  connected: boolean;
+  connections: Connection[] | null;
+} {
+  const _nodeId = useNodeId();
+  const prevConnections = useRef<Connection[] | null>(null);
+  const currentNodeId = nodeId || _nodeId;
+
+  const connections = useStore(
+    (state) => state.connectionLookup.get(`${currentNodeId}-${handleType}-${handleId}`) || null,
+    connectionsEqual
+  );
+
+  useEffect(() => {
+    // we don't want to trigger the handlers for the initial render
+    if (prevConnections.current && prevConnections.current !== connections) {
+      if (prevConnections.current?.length > (connections?.length ?? 0)) {
+        const disconnect = prevConnections.current.filter(
+          (prevConnection) => !connections?.find((connection) => connection.source === prevConnection.source)
+        );
+        onDisconnect?.(disconnect);
+      } else if (connections?.length) {
+        const connect = connections.filter(
+          (connection) =>
+            !prevConnections.current?.find((prevConnection) => prevConnection.source === connection.source)
+        );
+        onConnect?.(connect);
+      }
+    }
+
+    prevConnections.current = connections ?? [];
+  }, [connections, onConnect, onDisconnect]);
+
+  return useMemo(
+    () => ({
+      connected: !!connections,
+      connections,
+    }),
+    [connections]
+  );
+}

packages/react/src/index.ts

@@ -22,6 +22,7 @@ export { useStore, useStoreApi } from './hooks/useStore';
 export { default as useOnViewportChange, type UseOnViewportChangeOptions } from './hooks/useOnViewportChange';
 export { default as useOnSelectionChange, type UseOnSelectionChangeOptions } from './hooks/useOnSelectionChange';
 export { default as useNodesInitialized, type UseNodesInitializedOptions } from './hooks/useNodesInitialized';
+export { useHandleConnectionStatus } from './hooks/useHandleConnectionStatus';
 export { useNodeId } from './contexts/NodeIdContext';
 
 export { applyNodeChanges, applyEdgeChanges, handleParentExpand } from './utils/changes';

packages/react/src/store/index.ts

@@ -10,7 +10,7 @@ import {
 } from '@xyflow/system';
 
 import { applyNodeChanges, createSelectionChange, getSelectionChanges } from '../utils/changes';
-import { updateNodesAndEdgesSelections } from './utils';
+import { updateConnectionLookup, updateNodesAndEdgesSelections } from './utils';
 import getInitialState from './initialState';
 import type {
   ReactFlowState,
@@ -49,8 +49,12 @@ const createRFStore = ({
         set({ nodes: nextNodes });
       },
       setEdges: (edges: Edge[]) => {
-        const { defaultEdgeOptions = {} } = get();
-        set({ edges: edges.map((e) => ({ ...defaultEdgeOptions, ...e })) });
+        const { defaultEdgeOptions = {}, connectionLookup } = get();
+        const nextEdges = edges.map((e) => ({ ...defaultEdgeOptions, ...e }));
+
+        updateConnectionLookup(connectionLookup, nextEdges);
+
+        set({ edges: nextEdges });
       },
       // when the user works with an uncontrolled flow,
       // we set a flag `hasDefaultNodes` / `hasDefaultEdges`

packages/react/src/store/initialState.ts

@@ -5,9 +5,11 @@ import {
   getNodesBounds,
   getViewportForBounds,
   Transform,
+  Connection,
 } from '@xyflow/system';
 
 import type { Edge, Node, ReactFlowStore } from '../types';
+import { updateConnectionLookup } from './utils';
 
 const getInitialState = ({
   nodes = [],
@@ -23,6 +25,7 @@ const getInitialState = ({
   fitView?: boolean;
 } = {}): ReactFlowStore => {
   const nodeLookup = new Map<string, Node>();
+  const connectionLookup = updateConnectionLookup(new Map<string, Connection[]>(), edges);
   const nextNodes = updateNodes(nodes, nodeLookup, { nodeOrigin: [0, 0], elevateNodesOnSelect: false });
 
   let transform: Transform = [0, 0, 1];
@@ -42,6 +45,7 @@ const getInitialState = ({
     nodes: nextNodes,
     nodeLookup,
     edges: edges,
+    connectionLookup,
     onNodesChange: null,
     onEdgesChange: null,
     hasDefaultNodes: false,

packages/react/src/store/utils.ts

@@ -1,5 +1,6 @@
 import type { StoreApi } from 'zustand';
 import type { Edge, EdgeSelectionChange, Node, NodeSelectionChange, ReactFlowState } from '../types';
+import { Connection } from '@xyflow/system';
 
 export function handleControlledSelectionChange<NodeOrEdge extends Node | Edge>(
   changes: NodeSelectionChange[] | EdgeSelectionChange[],
@@ -42,3 +43,26 @@ export function updateNodesAndEdgesSelections({ changedNodes, changedEdges, get,
     onEdgesChange?.(changedEdges);
   }
 }
+
+export function updateConnectionLookup(lookup: Map<string, Connection[]>, edges: Edge[]) {
+  lookup.clear();
+
+  edges.forEach((edge) => {
+    const { source, target, sourceHandle = null, targetHandle = null } = edge;
+
+    if (source && target) {
+      const sourceKey = `${source}-source-${sourceHandle}`;
+      const targetKey = `${target}-target-${targetHandle}`;
+
+      const prevSource = lookup.get(sourceKey);
+      const prevTarget = lookup.get(targetKey);
+
+      const connection = { source, target, sourceHandle, targetHandle };
+
+      lookup.set(sourceKey, prevSource ? [...prevSource, connection] : [connection]);
+      lookup.set(targetKey, prevTarget ? [...prevTarget, connection] : [connection]);
+    }
+  });
+
+  return lookup;
+}

packages/react/src/types/store.ts

@@ -24,6 +24,7 @@ import {
   type OnMoveEnd,
   type IsValidConnection,
   type UpdateConnection,
+  Connection,
 } from '@xyflow/system';
 
 import type {
@@ -49,6 +50,7 @@ export type ReactFlowStore = {
   nodes: Node[];
   nodeLookup: Map<string, Node>;
   edges: Edge[];
+  connectionLookup: Map<string, Connection[]>;
   onNodesChange: OnNodesChange | null;
   onEdgesChange: OnEdgesChange | null;
   hasDefaultNodes: boolean;