feat(nodes): add helper functions to iterate over child nodes

example/src/Basic/index.tsx

@@ -13,6 +13,7 @@ import ReactFlow, {
   EdgeChange,
   OnLoadParams,
   Connection,
+  nodeHelper,
 } from 'react-flow-renderer';
 
 const onNodeDragStop = (_: MouseEvent, node: Node) => console.log('drag stop', node);
@@ -79,9 +80,8 @@ const BasicFlow = () => {
 
   const toggleClassnames = () => {
     setNodes((nds) => {
-      return nds.map((n) => {
+      return nodeHelper(nds).map((n) => {
         n.className = n.className === 'light' ? 'dark' : 'light';
-
         return n;
       });
     });

src/components/ConnectionLine/index.tsx

@@ -14,7 +14,7 @@ import {
   HandleType,
   ReactFlowState,
 } from '../../types';
-import { flattenNodes } from '../../utils/graph';
+import { nodeHelper } from '../../utils/nodes';
 
 interface ConnectionLineProps {
   connectionNodeId: ElementId;
@@ -29,7 +29,7 @@ interface ConnectionLineProps {
   CustomConnectionLineComponent?: ConnectionLineComponent;
 }
 
-const nodesSelector = (s: ReactFlowState) => flattenNodes(s.nodes);
+const nodesSelector = (s: ReactFlowState) => nodeHelper(s.nodes).flatten();
 
 export default ({
   connectionNodeId,

src/container/EdgeRenderer/utils.ts

@@ -2,7 +2,8 @@ import { ComponentType } from 'react';
 
 import { BezierEdge, StepEdge, SmoothStepEdge, StraightEdge } from '../../components/Edges';
 import wrapEdge from '../../components/Edges/wrapEdge';
-import { rectToBox, flattenNodes } from '../../utils/graph';
+import { rectToBox } from '../../utils/graph';
+import { nodeHelper } from '../../utils/nodes';
 
 import {
   EdgeTypesType,
@@ -171,16 +172,18 @@ type SourceTargetNode = {
 };
 
 export const getSourceTargetNodes = (edge: Edge, nodes: Node[]): SourceTargetNode => {
-  return flattenNodes(nodes).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
+    );
 };

src/index.ts

@@ -25,6 +25,9 @@ export {
   applyNodeChanges,
   applyEdgeChanges,
 } from './utils/graph';
+
+export { nodeHelper } from './utils/nodes';
+
 export { default as useZoomPanHelper } from './hooks/useZoomPanHelper';
 export { default as useUpdateNodeInternals } from './hooks/useUpdateNodeInternals';
 

src/store/index.ts

@@ -27,7 +27,8 @@ import {
   EdgeChange,
   NodePositionChange,
 } from '../types';
-import { isNode, isEdge, getRectOfNodes, getNodesInside, getConnectedEdges, flattenNodes } from '../utils/graph';
+import { isNode, isEdge, getRectOfNodes, getNodesInside, getConnectedEdges } from '../utils/graph';
+import { nodeHelper } from '../utils/nodes';
 import { getHandleBounds } from '../components/Nodes/utils';
 
 const { Provider, useStore, useStoreApi } = createContext<ReactFlowState>();
@@ -38,50 +39,6 @@ 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,
@@ -170,7 +127,7 @@ const createStore = () =>
       const { onNodesChange, nodes, transform } = get();
 
       const nodesToChange: NodeChange[] = updates.reduce<NodeChange[]>((res, update) => {
-        const node = findNodeById(update.id, nodes);
+        const node = nodeHelper(nodes).find((n) => n.id === update.id);
 
         if (node) {
           const dimensions = getDimensions(update.nodeElement);
@@ -224,11 +181,12 @@ const createStore = () =>
       const { onNodesChange, nodes, nodeExtent } = get();
 
       if (onNodesChange) {
-        const matchingNodes = flattenNodes(findNodes((n) => n.id === id || !!n.isSelected, nodes));
+        const matchingNodes = nodeHelper(nodes).filter((n) => n.id === id || !!n.isSelected);
+        const changingNodes = nodeHelper(matchingNodes).flatten();
 
-        if (matchingNodes?.length) {
+        if (changingNodes?.length) {
           onNodesChange(
-            matchingNodes.map((n) => {
+            changingNodes.map((n) => {
               const change: NodePositionChange = {
                 id: n.id,
                 type: 'position',
@@ -346,10 +304,12 @@ const createStore = () =>
     unselectNodesAndEdges: () => {
       const { nodes, edges, onNodesChange, onEdgesChange } = get();
 
-      const nodesToUnselect = flattenNodes(nodes).map((n) => {
-        n.isSelected = false;
-        return createNodeOrEdgeSelectionChange(false)(n);
-      }) as NodeChange[];
+      const nodesToUnselect = nodeHelper(nodes)
+        .flatten()
+        .map((n) => {
+          n.isSelected = false;
+          return createNodeOrEdgeSelectionChange(false)(n);
+        }) as NodeChange[];
       const edgesToUnselect = edges.map(createNodeOrEdgeSelectionChange(false)) as EdgeChange[];
 
       if (nodesToUnselect.length) {

src/utils/graph.ts

@@ -342,29 +342,3 @@ export function applyNodeChanges(changes: NodeChange[], nodes: Node[]): Node[] {
 export function applyEdgeChanges(changes: EdgeChange[], edges: Edge[]): Edge[] {
   return applyChanges(changes, edges) as Edge[];
 }
-
-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;
-  // }, []);
-}

src/utils/nodes.ts

@@ -0,0 +1,121 @@
+import { Node } from '../types';
+
+function flat(arr: Node[], target: Node[]) {
+  arr.forEach(function (el) {
+    if (el.childNodes) {
+      flat(el.childNodes, target);
+    } else {
+      target.push(el);
+    }
+  });
+}
+
+const filterNodes = (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 = filterNodes(condition, n.childNodes);
+
+      for (let j = 0; j < matches.length; j++) {
+        res.push(matches[j]);
+      }
+    }
+  }
+
+  return res;
+};
+
+const mapNodes = (accessor: (node: Node) => any, nodes: Node[]): Node[] => {
+  let res = [];
+
+  for (let i = 0; i < nodes.length; i++) {
+    const n = nodes[i];
+    res.push(accessor(n));
+
+    if (n.childNodes) {
+      n.childNodes = mapNodes(accessor, n.childNodes);
+    }
+  }
+
+  return res;
+};
+
+const forEachNode = (accessor: (node: Node) => any, nodes: Node[]): void => {
+  for (let i = 0; i < nodes.length; i++) {
+    const n = nodes[i];
+    accessor(n);
+
+    if (n.childNodes) {
+      forEachNode(accessor, n.childNodes);
+    }
+  }
+};
+
+function findNode(accessor: (node: Node) => boolean, nodes: Node[]): Node | undefined {
+  let res = undefined;
+
+  for (let i = 0; i < nodes.length; i++) {
+    const n = nodes[i];
+
+    if (accessor(n)) {
+      return n;
+    }
+
+    if (n.childNodes) {
+      res = findNode(accessor, n.childNodes);
+
+      if (res) {
+        return res;
+      }
+    }
+  }
+
+  return res;
+}
+
+export interface NodeHelper {
+  filter: (accessor: (node: Node) => boolean) => Node[];
+  map: (accessor: (node: Node) => any) => Node[];
+  find: (accessor: (node: Node) => boolean) => Node | undefined;
+  forEach: (accessor: (node: Node) => void) => void;
+  flatten: () => Node[];
+}
+
+export function nodeHelper(nodes: Node[]): NodeHelper {
+  const flatten: NodeHelper['flatten'] = () => {
+    const flattened: Node[] = [];
+    flat(nodes, flattened);
+    return flattened;
+  };
+
+  const filter: NodeHelper['filter'] = (accessor) => {
+    return filterNodes(accessor, nodes);
+  };
+
+  const forEach: NodeHelper['forEach'] = (accessor) => {
+    return forEachNode(accessor, nodes);
+  };
+
+  const find: NodeHelper['find'] = (accessor) => {
+    return findNode(accessor, nodes);
+  };
+
+  const map: NodeHelper['map'] = (accessor) => {
+    return mapNodes(accessor, nodes);
+  };
+
+  return {
+    filter,
+    forEach,
+    flatten,
+    find,
+    map,
+  };
+}