refactor(nodelookup): add treelevel

src/container/EdgeRenderer/index.tsx

@@ -251,8 +251,8 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
       <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 = nodeLookup.current.get(edge.source);
-          const targetNode = nodeLookup.current.get(edge.target);
+          const sourceNode = nodeLookup.get(edge.source);
+          const targetNode = nodeLookup.get(edge.target);
 
           return (
             <Edge

src/container/NodeRenderer/index.tsx

@@ -236,11 +236,13 @@ const NodeRenderer = (props: NodeRendererProps) => {
     });
   }, []);
 
+  console.log(nodeLookup);
+
   return (
     <div className="react-flow__nodes" style={transformStyle}>
       <MemoizedNodes
         nodes={nodes}
-        nodeLookup={nodeLookup.current}
+        nodeLookup={nodeLookup}
         snapToGrid={snapToGrid}
         snapGrid={snapGrid}
         nodesDraggable={nodesDraggable}

src/hooks/useNodeLookup.ts

@@ -6,7 +6,7 @@ function useNodeLookup() {
   const store = useStoreApi();
   const nodeLookup = useRef(store.getState().nodeLookup);
 
-  return nodeLookup;
+  return nodeLookup.current;
 }
 
 export default useNodeLookup;

src/store/index.ts

@@ -73,21 +73,32 @@ const createNodeOrEdgeSelectionChange = (isSelected: boolean) => (item: Node | E
 //   return result;
 // }
 
-function addPositions(posA: XYPosition, posB: XYPosition): XYPosition {
+type XYPosAndTreeLevel = XYPosition & { treeLevel: number };
+
+function addPositions(a: XYPosAndTreeLevel, b: XYPosition): XYPosAndTreeLevel {
   return {
-    x: (posA.x ?? 0) + (posB.x ?? 0),
-    y: (posA.y ?? 0) + (posB.y ?? 0),
+    x: (a.x ?? 0) + (b.x ?? 0),
+    y: (a.y ?? 0) + (b.y ?? 0),
+    treeLevel: a.treeLevel + 1,
   };
 }
 
-function getAbsolutePosition(node: NodeLookupItem, nodeLookup: NodeLookup, result: XYPosition): XYPosition {
+function getAbsolutePositionAndTreeLevel(
+  node: NodeLookupItem,
+  nodeLookup: NodeLookup,
+  result: XYPosAndTreeLevel
+): XYPosAndTreeLevel {
   const parentNode = node.parentNode ? nodeLookup.get(node.parentNode) : false;
 
   if (!parentNode) {
     return result;
   }
 
-  return getAbsolutePosition(parentNode, nodeLookup, addPositions(result, parentNode.position || { x: 0, y: 0 }));
+  return getAbsolutePositionAndTreeLevel(
+    parentNode,
+    nodeLookup,
+    addPositions(result, parentNode.position || { x: 0, y: 0 })
+  );
 }
 
 const createStore = () =>
@@ -158,6 +169,7 @@ const createStore = () =>
           height: node.height || null,
           position: node.position,
           positionAbsolute: node.position,
+          treeLevel: 0,
         };
         if (node.parentNode) {
           lookupNode.parentNode = node.parentNode;
@@ -168,12 +180,19 @@ const createStore = () =>
       nodes
         .filter((node) => node.parentNode)
         .forEach((node) => {
-          const positionAbsolute = getAbsolutePosition(node, nodeLookup, node.position);
+          const positionAbsoluteAndTreeLevel = getAbsolutePositionAndTreeLevel(node, nodeLookup, {
+            ...node.position,
+            treeLevel: 0,
+          });
 
-          if (positionAbsolute) {
+          if (positionAbsoluteAndTreeLevel) {
             nodeLookup.set(node.id, {
               ...nodeLookup.get(node.id),
-              positionAbsolute,
+              positionAbsolute: {
+                x: positionAbsoluteAndTreeLevel.x,
+                y: positionAbsoluteAndTreeLevel.y,
+              },
+              treeLevel: positionAbsoluteAndTreeLevel.treeLevel,
             });
           }
         });

src/types/index.ts

@@ -465,6 +465,7 @@ export type NodeLookupItem = {
   position?: XYPosition;
   positionAbsolute?: XYPosition;
   handleBounds?: NodeHandleBounds;
+  treeLevel?: number;
 };
 
 export type NodeLookup = Map<ElementId, NodeLookupItem>;