refactor(nodelookup): use as normale store item

src/components/ConnectionLine/index.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect, useState, CSSProperties } from 'react';
 
+import { useStore } from '../../store';
 import { getBezierPath } from '../Edges/BezierEdge';
 import { getSmoothStepPath } from '../Edges/SmoothStepEdge';
 import {
@@ -12,8 +13,9 @@ import {
   ConnectionLineComponent,
   HandleType,
   Node,
+  ReactFlowState,
 } from '../../types';
-import useNodeLookup from '../../hooks/useNodeLookup';
+
 interface ConnectionLineProps {
   connectionNodeId: ElementId;
   connectionHandleId: ElementId | null;
@@ -27,6 +29,8 @@ interface ConnectionLineProps {
   CustomConnectionLineComponent?: ConnectionLineComponent;
 }
 
+const selector = (s: ReactFlowState) => s.nodeLookup;
+
 export default ({
   connectionNodeId,
   connectionHandleId,
@@ -39,7 +43,7 @@ export default ({
   isConnectable,
   CustomConnectionLineComponent,
 }: ConnectionLineProps) => {
-  const nodeLookup = useNodeLookup();
+  const nodeLookup = useStore(selector);
   const [sourceNode, setSourceNode] = useState<NodeLookupItem | null>(null);
   const nodeId = connectionNodeId;
   const handleId = connectionHandleId;

src/container/EdgeRenderer/index.tsx

@@ -17,7 +17,6 @@ import {
   NodeHandleBounds,
 } from '../../types';
 import useVisibleEdges from '../../hooks/useVisibleEdges';
-import useNodeLookup from '../../hooks/useNodeLookup';
 
 interface EdgeRendererProps {
   edgeTypes: any;
@@ -218,7 +217,7 @@ const selector = (s: ReactFlowState) => ({
   width: s.width,
   height: s.height,
   connectionMode: s.connectionMode,
-  nodes: s.nodes,
+  nodeLookup: s.nodeLookup,
 });
 
 const EdgeRenderer = (props: EdgeRendererProps) => {
@@ -233,8 +232,8 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
     width,
     height,
     connectionMode,
+    nodeLookup,
   } = useStore(selector, shallow);
-  const nodeLookup = useNodeLookup();
   const edgeTree = useVisibleEdges(props.onlyRenderVisibleElements, nodeLookup);
 
   if (!width) {

src/container/NodeRenderer/index.tsx

@@ -4,7 +4,7 @@ import shallow from 'zustand/shallow';
 import { useStore } from '../../store';
 import { Node, NodeTypesType, ReactFlowState, WrapNodeProps, SnapGrid, NodeRendererNode } from '../../types';
 import useVisibleNodes from '../../hooks/useVisibleNodes';
-import useNodeLookup from '../../hooks/useNodeLookup';
+
 interface NodeRendererProps {
   nodeTypes: NodeTypesType;
   selectNodesOnDrag: boolean;
@@ -131,9 +131,9 @@ const NodeRenderer = (props: NodeRendererProps) => {
     updateNodeDimensions,
     snapGrid,
     snapToGrid,
+    nodeLookup,
   } = useStore(selector, shallow);
 
-  const nodeLookup = useNodeLookup();
   const nodes = useVisibleNodes(props.onlyRenderVisibleElements);
 
   const resizeObserver = useMemo(() => {

src/hooks/useNodeLookup.ts

@@ -1,12 +0,0 @@
-import { useRef } from 'react';
-
-import { useStoreApi } from '../store';
-
-function useNodeLookup() {
-  const store = useStoreApi();
-  const nodeLookup = useRef(store.getState().nodeLookup);
-
-  return nodeLookup.current;
-}
-
-export default useNodeLookup;

src/store/index.ts

@@ -28,6 +28,7 @@ import {
   NodeDimensionChange,
   NodeLookup,
   NodeLookupItem,
+  ElementId,
 } from '../types';
 import { isNode, isEdge, getRectOfNodes, getNodesInside, getConnectedEdges } from '../utils/graph';
 import { getHandleBounds } from '../components/Nodes/utils';
@@ -128,9 +129,10 @@ const createStore = () =>
 
     setNodes: (nodes: Node[]) => {
       const { nodeLookup } = get();
+      const nextNodeLookup = new Map<ElementId, NodeLookupItem>();
 
       nodes.forEach((node) => {
-        const lookupNode = {
+        const lookupNode: NodeLookupItem = {
           ...nodeLookup.get(node.id),
           width: node.width || null,
           height: node.height || null,
@@ -141,24 +143,24 @@ const createStore = () =>
         if (node.parentNode) {
           lookupNode.parentNode = node.parentNode;
         }
-        nodeLookup.set(node.id, lookupNode);
+        nextNodeLookup.set(node.id, lookupNode);
       });
 
       nodes
         .filter((node) => node.parentNode)
         .forEach((node) => {
-          const positionAbsoluteAndTreeLevel = getAbsolutePositionAndTreeLevel(node, nodeLookup, {
+          const positionAbsoluteAndTreeLevel = getAbsolutePositionAndTreeLevel(node, nextNodeLookup, {
             ...node.position,
             treeLevel: node.zIndex || 0,
           });
 
-          nodeLookup.set(node.parentNode!, { ...nodeLookup.get(node.parentNode!), isParentNode: true });
+          nextNodeLookup.set(node.parentNode!, { ...nextNodeLookup.get(node.parentNode!), isParentNode: true });
 
           if (positionAbsoluteAndTreeLevel) {
             const { treeLevel, x, y } = positionAbsoluteAndTreeLevel;
 
-            nodeLookup.set(node.id, {
+            nextNodeLookup.set(node.id, {
-              ...nodeLookup.get(node.id),
+              ...nextNodeLookup.get(node.id),
               positionAbsolute: {
                 x,
                 y,
@@ -168,7 +170,7 @@ const createStore = () =>
           }
         });
 
-      set({ nodes });
+      set({ nodes, nodeLookup: nextNodeLookup });
     },
     setEdges: (edges: Edge[]) => {
       set({ edges });
@@ -188,13 +190,16 @@ const createStore = () =>
 
           if (doUpdate) {
             const handleBounds = getHandleBounds(update.nodeElement, transform[2]);
-            nodeLookup.set(node.id, { ...nodeLookup.get(node.id), handleBounds });
+            nodeLookup.set(node.id, {
+              ...nodeLookup.get(node.id),
+              handleBounds,
+              ...dimensions,
+            });
 
             const change = {
               id: node.id,
               type: 'dimensions',
               dimensions,
-              handleBounds,
             } as NodeChange;
             res.push(change);
           }
@@ -203,6 +208,8 @@ const createStore = () =>
         return res;
       }, []);
 
+      set({ nodeLookup: new Map(nodeLookup) });
+
       onNodesChange?.(nodesToChange);
     },
     updateNodePosition: ({ id, diff, isDragging }: NodeDiffUpdate) => {