From 9350884836a47b6ab22e66f0564ff2f46a04cc3f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christopher=20M=C3=B6ller?= <chrtze@gmail.com>
Date: Wed, 20 Oct 2021 18:06:41 +0200
Subject: [PATCH] fix(perf): improve flattenNodes

---
 src/store/index.ts | 86 +++++++++++++++++++++++++++++++++++++++++-----
 src/utils/graph.ts | 32 ++++++++++++-----
 2 files changed, 101 insertions(+), 17 deletions(-)

diff --git a/src/store/index.ts b/src/store/index.ts
index eb5555ee..10693808 100644
--- a/src/store/index.ts
+++ b/src/store/index.ts
@@ -38,6 +38,50 @@ const createNodeOrEdgeSelectionChange = (isSelected: boolean) => (item: Node | E
   isSelected,
 });
 
+const findNodeById = (id: string, nodes: Node[]): Node | null => {
+  let res = null;
+
+  for (let i = 0; i < nodes.length; i++) {
+    const n = nodes[i];
+
+    if (n.id === id) {
+      return n;
+    }
+
+    if (n.childNodes) {
+      res = findNodeById(id, n.childNodes);
+
+      if (res) {
+        return res;
+      }
+    }
+  }
+
+  return res;
+};
+
+const findNodes = (condition: (node: Node) => boolean, nodes: Node[]): Node[] => {
+  let res = [];
+
+  for (let i = 0; i < nodes.length; i++) {
+    const n = nodes[i];
+
+    if (condition(n)) {
+      res.push(n);
+    }
+
+    if (n.childNodes) {
+      const matches = findNodes(condition, n.childNodes);
+
+      for (let j = 0; j < matches.length; j++) {
+        res.push(matches[j]);
+      }
+    }
+  }
+
+  return res;
+};
+
 const createStore = () =>
   create<ReactFlowState>((set, get) => ({
     width: 0,
@@ -96,6 +140,7 @@ const createStore = () =>
 
     setNodes: (propNodes: Node[]) => {
       const { nodes } = get();
+
       const nextNodes = propNodes.map((propNode: Node) => {
         const storeNode = nodes.find((node) => node.id === propNode.id);
 
@@ -124,10 +169,10 @@ const createStore = () =>
     updateNodeDimensions: (updates: NodeDimensionUpdate[]) => {
       const { onNodesChange, nodes, transform } = get();
 
-      const initialChanges: NodeChange[] = [];
-      const nodesToChange: NodeChange[] = flattenNodes(nodes).reduce((res, node) => {
-        const update = updates.find((u) => u.id === node.id);
-        if (update) {
+      const nodesToChange: NodeChange[] = updates.reduce<NodeChange[]>((res, update) => {
+        const node = findNodeById(update.id, nodes);
+
+        if (node) {
           const dimensions = getDimensions(update.nodeElement);
           const doUpdate =
             dimensions.width &&
@@ -147,7 +192,31 @@ const createStore = () =>
         }
 
         return res;
-      }, initialChanges);
+      }, []);
+
+      // const nodesToChange: NodeChange[] = flattenNodes(nodes).reduce((res, node) => {
+      //   const update = updates.find((u) => u.id === node.id);
+      //   if (update) {
+      //     const dimensions = getDimensions(update.nodeElement);
+      //     const doUpdate =
+      //       dimensions.width &&
+      //       dimensions.height &&
+      //       (node.width !== dimensions.width || node.height !== dimensions.height || update.forceUpdate);
+
+      //     if (doUpdate) {
+      //       const handleBounds = getHandleBounds(update.nodeElement, transform[2]);
+      //       const change = {
+      //         id: node.id,
+      //         type: 'dimensions',
+      //         dimensions,
+      //         handleBounds,
+      //       } as NodeChange;
+      //       res.push(change);
+      //     }
+      //   }
+
+      //   return res;
+      // }, initialChanges);
 
       onNodesChange?.(nodesToChange);
     },
@@ -155,12 +224,11 @@ const createStore = () =>
       const { onNodesChange, nodes, nodeExtent } = get();
 
       if (onNodesChange) {
-        const matchingNodes = flattenNodes(nodes).filter((n) => n.id === id || n.isSelected);
-        const matchingChildNodes = flattenNodes(matchingNodes); //.filter(n => !!n.childNodes).reduce<Node[]>((result, node) => result.concat(node.childNodes!), []));
+        const matchingNodes = flattenNodes(findNodes((n) => n.id === id || !!n.isSelected, nodes));
 
-        if (matchingChildNodes?.length) {
+        if (matchingNodes?.length) {
           onNodesChange(
-            matchingChildNodes.map((n) => {
+            matchingNodes.map((n) => {
               const change: NodePositionChange = {
                 id: n.id,
                 type: 'position',
diff --git a/src/utils/graph.ts b/src/utils/graph.ts
index 4895966b..7e88a608 100644
--- a/src/utils/graph.ts
+++ b/src/utils/graph.ts
@@ -343,12 +343,28 @@ export function applyEdgeChanges(changes: EdgeChange[], edges: Edge[]): Edge[] {
   return applyChanges(changes, edges) as Edge[];
 }
 
-export function flattenNodes(nodes: Node[] | undefined): Node[] {
-  if (!nodes) {
-    return [];
-  }
-
-  return nodes.reduce<Node[]>((result, node) => {
-    return result.concat([node, ...flattenNodes(node.childNodes)]);
-  }, []);
+function flat(arr: Node[], target: Node[]) {
+  arr.forEach(function (el) {
+    if (el.childNodes) {
+      flat(el.childNodes, target);
+    } else {
+      target.push(el);
+    }
+  });
+}
+
+export function flattenNodes(nodes: Node[]): Node[] {
+  const flattened: Node[] = [];
+  flat(nodes, flattened);
+  return flattened;
+
+  // return nodes.reduce<Node[]>((result, node) => {
+  //   result.push(node);
+
+  //   if (node.childNodes) {
+  //     result.push(...flattenNodes(node.childNodes));
+  //   }
+
+  //   return result;
+  // }, []);
 }