From b15d3012e0de60bdb9046f8db3c59931a5f2b153 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christopher=20M=C3=B6ller?= <chrtze@gmail.com>
Date: Thu, 21 Oct 2021 14:08:39 +0200
Subject: [PATCH] fix(nested-nodes): fix getSourceTargetNode for childNodes

---
 example/src/Basic/index.tsx         | 40 ++++++++++++++++++++++++++---
 src/container/EdgeRenderer/utils.ts | 36 ++++++++++++++++----------
 2 files changed, 58 insertions(+), 18 deletions(-)

diff --git a/example/src/Basic/index.tsx b/example/src/Basic/index.tsx
index a9802e6f..513c7621 100644
--- a/example/src/Basic/index.tsx
+++ b/example/src/Basic/index.tsx
@@ -31,7 +31,13 @@ const initialNodes: Node[] = [
     className: 'light',
     style: { backgroundColor: 'rgba(255, 0, 0, .2)' },
     childNodes: [
-      { id: '4a', data: { label: 'Node 4a' }, position: { x: 400, y: 400 }, className: 'light' },
+      {
+        id: '4a',
+        draggable: false,
+        data: { label: 'Node 4a', isNested: true },
+        position: { x: 400, y: 400 },
+        className: 'light',
+      },
       {
         id: '4b',
         data: { label: 'Node 4b' },
@@ -39,8 +45,20 @@ const initialNodes: Node[] = [
         className: 'light',
         style: { backgroundColor: 'rgba(255, 0, 0, .2)' },
         childNodes: [
-          { id: '4b1', data: { label: 'Node 4b1' }, position: { x: 450, y: 450 }, className: 'light' },
-          { id: '4b2', data: { label: 'Node 4b2' }, position: { x: 550, y: 550 }, className: 'light' },
+          {
+            id: '4b1',
+            draggable: false,
+            data: { label: 'Node 4b1', isNested: true },
+            position: { x: 450, y: 450 },
+            className: 'light',
+          },
+          {
+            id: '4b2',
+            draggable: false,
+            data: { label: 'Node 4b2', isNested: true },
+            position: { x: 550, y: 550 },
+            className: 'light',
+          },
         ],
       },
     ],
@@ -58,7 +76,9 @@ const BasicFlow = () => {
   const [edges, setEdges] = useState<Edge[]>(initialEdges);
 
   const onConnect = useCallback((params: Edge | Connection) => {
-    setEdges((eds) => addEdge(params, eds));
+    setEdges((eds) => {
+      return addEdge(params, eds);
+    });
   }, []);
   const onLoad = useCallback((reactFlowInstance: OnLoadParams) => setRfInstance(reactFlowInstance), []);
 
@@ -87,6 +107,15 @@ const BasicFlow = () => {
     });
   };
 
+  const toggleChildNodes = () => {
+    setNodes((nds) => {
+      return nodeHelper(nds).map((n) => {
+        n.isHidden = n.data.isNested && !n.isHidden;
+        return n;
+      });
+    });
+  };
+
   const onNodesChange = useCallback((changes: NodeChange[]) => {
     setNodes((ns) => applyNodeChanges(changes, ns));
   }, []);
@@ -125,6 +154,9 @@ const BasicFlow = () => {
         <button onClick={toggleClassnames} style={{ marginRight: 5 }}>
           toggle classnames
         </button>
+        <button style={{ marginRight: 5 }} onClick={toggleChildNodes}>
+          toggleChildNodes
+        </button>
         <button onClick={logToObject}>toObject</button>
       </div>
     </ReactFlow>
diff --git a/src/container/EdgeRenderer/utils.ts b/src/container/EdgeRenderer/utils.ts
index f6ad4360..c9c3ddc8 100644
--- a/src/container/EdgeRenderer/utils.ts
+++ b/src/container/EdgeRenderer/utils.ts
@@ -172,18 +172,26 @@ type SourceTargetNode = {
 };
 
 export const getSourceTargetNodes = (edge: Edge, nodes: Node[]): 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
-    );
+  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
+  //   );
 };