diff --git a/src/container/EdgeRenderer/index.tsx b/src/container/EdgeRenderer/index.tsx
index 0b83ce99..dc8e960b 100644
--- a/src/container/EdgeRenderer/index.tsx
+++ b/src/container/EdgeRenderer/index.tsx
@@ -249,6 +249,7 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
       <MarkerDefinitions defaultColor={defaultMarkerColor} />
       <g transform={`translate(${transform[0]},${transform[1]}) scale(${transform[2]})`}>
         {edges.map((edge: Edge) => {
+          // @todo: getSourceTargetNodes is called many times during dragging/creating edges
           const { sourceNode, targetNode } = getSourceTargetNodes(edge, nodes);
           return (
             <Edge
diff --git a/src/container/EdgeRenderer/utils.ts b/src/container/EdgeRenderer/utils.ts
index c9c3ddc8..f6ad4360 100644
--- a/src/container/EdgeRenderer/utils.ts
+++ b/src/container/EdgeRenderer/utils.ts
@@ -172,26 +172,18 @@ type SourceTargetNode = {
 };
 
 export const getSourceTargetNodes = (edge: Edge, nodes: Node[]): SourceTargetNode => {
-  const sourceNode = nodeHelper(nodes).find((n) => n.id === edge.source) || null;
-  const targetNode = nodeHelper(nodes).find((n) => n.id === edge.target) || null;
-
-  return {
-    sourceNode,
-    targetNode,
-  };
-
-  // return nodeHelper(nodes)
-  //   .flatten()
-  //   .reduce(
-  //     (res, node) => {
-  //       if (node.id === edge.source) {
-  //         res.sourceNode = node;
-  //       }
-  //       if (node.id === edge.target) {
-  //         res.targetNode = node;
-  //       }
-  //       return res;
-  //     },
-  //     { sourceNode: null, targetNode: null } as SourceTargetNode
-  //   );
+  return nodeHelper(nodes)
+    .flatten()
+    .reduce(
+      (res, node) => {
+        if (node.id === edge.source) {
+          res.sourceNode = node;
+        }
+        if (node.id === edge.target) {
+          res.targetNode = node;
+        }
+        return res;
+      },
+      { sourceNode: null, targetNode: null } as SourceTargetNode
+    );
 };
diff --git a/src/utils/nodes.ts b/src/utils/nodes.ts
index 47a33898..c5520b1b 100644
--- a/src/utils/nodes.ts
+++ b/src/utils/nodes.ts
@@ -2,10 +2,9 @@ import { Node } from '../types';
 
 function flat(arr: Node[], target: Node[]) {
   arr.forEach(function (el) {
+    target.push(el);
     if (el.childNodes) {
       flat(el.childNodes, target);
-    } else {
-      target.push(el);
     }
   });
 }