feat(react): connection lookup draft

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;
+}