refactor(react/handle-connection-state): cleanup

packages/react/src/hooks/useHandleConnectionStatus.ts

@@ -3,6 +3,7 @@ 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;
@@ -12,30 +13,6 @@ type useHandleConnectionStatusParams = {
   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,
@@ -52,23 +29,26 @@ export function useHandleConnectionStatus({
 
   const connections = useStore(
     (state) => state.connectionLookup.get(`${currentNodeId}-${handleType}-${handleId}`) || null,
-    connectionsEqual
+    areConnectionsEqual
   );
 
   useEffect(() => {
-    // we don't want to trigger the handlers for the initial render
+    // @todo dicuss if onConnect/onDisconnect should be called when the component mounts/unmounts
     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);
+      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);
       }
     }
 

packages/react/src/store/index.ts

@@ -330,6 +330,7 @@ const createRFStore = ({
 
         set(currentConnection);
       },
+
       reset: () => {
         // @todo: what should we do about this? Do we still need it?
         // if you are on a SPA with multiple flows, we want to make sure that the store gets resetted

packages/react/src/store/utils.ts

@@ -47,9 +47,7 @@ export function updateNodesAndEdgesSelections({ changedNodes, changedEdges, get,
 export function updateConnectionLookup(lookup: Map<string, Connection[]>, edges: Edge[]) {
   lookup.clear();
 
-  edges.forEach((edge) => {
-    const { source, target, sourceHandle = null, targetHandle = null } = edge;
-
+  edges.forEach(({ source, target, sourceHandle = null, targetHandle = null }) => {
     if (source && target) {
       const sourceKey = `${source}-source-${sourceHandle}`;
       const targetKey = `${target}-target-${targetHandle}`;

packages/react/src/types/store.ts

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

packages/react/src/utils/general.ts

@@ -6,6 +6,7 @@ import {
   getIncomersBase,
   updateEdgeBase,
   getConnectedEdgesBase,
+  Connection,
 } from '@xyflow/system';
 
 import type { Edge, Node } from '../types';
@@ -17,3 +18,32 @@ 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)));
+}