From b2ac2199121f7c9a98da819619785f99ac9abbdf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christopher=20M=C3=B6ller?= <chrtze@gmail.com>
Date: Thu, 25 Nov 2021 16:04:08 +0100
Subject: [PATCH] refactor(edge): pass onEdgeUpdate instead of creating
 onConnect function

---
 src/components/Edges/wrapEdge.tsx    | 15 ++++++++++++---
 src/container/EdgeRenderer/index.tsx | 15 ++++-----------
 src/types/edges.ts                   |  4 ++--
 3 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/src/components/Edges/wrapEdge.tsx b/src/components/Edges/wrapEdge.tsx
index a89eb28e..e0af3934 100644
--- a/src/components/Edges/wrapEdge.tsx
+++ b/src/components/Edges/wrapEdge.tsx
@@ -3,7 +3,7 @@ import cc from 'classcat';
 import shallow from 'zustand/shallow';
 
 import { useStore, useStoreApi } from '../../store';
-import { Edge, EdgeProps, WrapEdgeProps, ReactFlowState } from '../../types';
+import { Edge, EdgeProps, WrapEdgeProps, ReactFlowState, Connection } from '../../types';
 import { onMouseDown } from '../../components/Handle/handler';
 import { EdgeAnchor } from './EdgeAnchor';
 import { getMarkerId } from '../../utils/graph';
@@ -46,12 +46,12 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
     sourceHandleId,
     targetHandleId,
     handleEdgeUpdate,
-    onConnectEdge,
     onContextMenu,
     onMouseEnter,
     onMouseMove,
     onMouseLeave,
     edgeUpdaterRadius,
+    onEdgeUpdate,
     onEdgeUpdateStart,
     onEdgeUpdateEnd,
     markerEnd,
@@ -156,6 +156,15 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
           ? (evt: MouseEvent): void => onEdgeUpdateEnd(evt, edgeElement)
           : undefined;
 
+        const onConnectEdge = (connection: Connection) => {
+          const { edges } = store.getState();
+          const edge = edges.find((e) => e.id === id);
+
+          if (edge && onEdgeUpdate) {
+            onEdgeUpdate(edge, connection);
+          }
+        };
+
         onMouseDown(
           event,
           handleId,
@@ -181,7 +190,7 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
         setConnectionNodeId,
         setPosition,
         edgeElement,
-        onConnectEdge,
+        onEdgeUpdate,
       ]
     );
 
diff --git a/src/container/EdgeRenderer/index.tsx b/src/container/EdgeRenderer/index.tsx
index 8a3d45d2..7ee45e5f 100644
--- a/src/container/EdgeRenderer/index.tsx
+++ b/src/container/EdgeRenderer/index.tsx
@@ -9,7 +9,6 @@ import { getEdgePositions, getHandle } from './utils';
 import {
   Position,
   Edge,
-  Connection,
   ConnectionLineType,
   ConnectionLineComponent,
   ConnectionMode,
@@ -92,12 +91,6 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
               const sourceHandleId = edge.sourceHandle || null;
               const targetHandleId = edge.targetHandle || null;
 
-              const onConnectEdge = (connection: Connection) => {
-                props.onEdgeUpdate?.(edge, connection);
-              };
-
-              const sourceNodeWidth = sourceNode?.width;
-              const sourceNodeHeight = sourceNode?.height;
               const sourceNodeX = sourceNode?.positionAbsolute?.x;
               const sourceNodeY = sourceNode?.positionAbsolute?.y;
               const sourceNodeHandleBounds = sourceNode?.handleBounds;
@@ -113,8 +106,8 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
               }
 
               if (
-                !sourceNodeWidth ||
-                !sourceNodeHeight ||
+                !sourceNode?.width ||
+                !sourceNode?.height ||
                 typeof sourceNodeX === 'undefined' ||
                 typeof sourceNodeY === 'undefined'
               ) {
@@ -155,7 +148,7 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
               }
 
               const { sourceX, sourceY, targetX, targetY } = getEdgePositions(
-                { x: sourceNodeX, y: sourceNodeY, width: sourceNodeWidth, height: sourceNodeHeight },
+                { x: sourceNodeX, y: sourceNodeY, width: sourceNode?.width, height: sourceNode?.height },
                 sourceHandle,
                 sourcePosition,
                 { x: targetNodeX, y: targetNodeY, width: targetNodeWidth, height: targetNodeHeight },
@@ -193,8 +186,8 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
                   sourcePosition={sourcePosition}
                   targetPosition={targetPosition}
                   elementsSelectable={elementsSelectable}
-                  onConnectEdge={onConnectEdge}
                   handleEdgeUpdate={typeof props.onEdgeUpdate !== 'undefined'}
+                  onEdgeUpdate={props.onEdgeUpdate}
                   onContextMenu={props.onEdgeContextMenu}
                   onMouseEnter={props.onEdgeMouseEnter}
                   onMouseMove={props.onEdgeMouseMove}
diff --git a/src/types/edges.ts b/src/types/edges.ts
index c2b98ec3..6fbd9d07 100644
--- a/src/types/edges.ts
+++ b/src/types/edges.ts
@@ -1,7 +1,7 @@
 import { CSSProperties, ReactNode, HTMLAttributes } from 'react';
 
 import { Position } from './utils';
-import { OnConnect, Connection } from './general';
+import { Connection } from './general';
 import { HandleElement } from './handles';
 import { Node } from './nodes';
 
@@ -90,7 +90,7 @@ export interface WrapEdgeProps<T = any> {
   elementsSelectable?: boolean;
   hidden?: boolean;
   handleEdgeUpdate: boolean;
-  onConnectEdge: OnConnect;
+  onEdgeUpdate: OnEdgeUpdateFunc;
   onContextMenu?: EdgeMouseHandler;
   onMouseEnter?: EdgeMouseHandler;
   onMouseMove?: EdgeMouseHandler;