From 996b1b51f5aab400b4b47765c8450bad6c51960c Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Thu, 4 Nov 2021 18:13:21 +0100
Subject: [PATCH] refactor(store): rename nodeLookup to nodeInternals

---
 example/src/Basic/index.tsx                 |  2 +-
 src/additional-components/MiniMap/index.tsx | 12 ++-
 src/components/ConnectionLine/index.tsx     |  8 +-
 src/components/Nodes/wrapNode.tsx           |  4 +-
 src/container/EdgeRenderer/index.tsx        | 10 +--
 src/container/NodeRenderer/index.tsx        | 14 ++--
 src/hooks/useNodeInternalsRef.ts            | 14 ++++
 src/hooks/useNodeLookupRef.ts               | 14 ----
 src/hooks/useVisibleEdges.ts                | 18 ++---
 src/store/index.ts                          | 90 +++------------------
 src/store/utils.ts                          | 77 ++++++++++++++++++
 src/style.css                               |  1 +
 src/types/index.ts                          |  6 +-
 13 files changed, 142 insertions(+), 128 deletions(-)
 create mode 100644 src/hooks/useNodeInternalsRef.ts
 delete mode 100644 src/hooks/useNodeLookupRef.ts
 create mode 100644 src/store/utils.ts

diff --git a/example/src/Basic/index.tsx b/example/src/Basic/index.tsx
index 40ed203c..0e53b216 100644
--- a/example/src/Basic/index.tsx
+++ b/example/src/Basic/index.tsx
@@ -1,4 +1,4 @@
-import { useState, useMemo, MouseEvent, useCallback } from 'react';
+import { useState, MouseEvent, useCallback } from 'react';
 
 import ReactFlow, {
   addEdge,
diff --git a/src/additional-components/MiniMap/index.tsx b/src/additional-components/MiniMap/index.tsx
index 7cc82df9..9cc5a0c1 100644
--- a/src/additional-components/MiniMap/index.tsx
+++ b/src/additional-components/MiniMap/index.tsx
@@ -29,7 +29,7 @@ const selector = (s: ReactFlowState) => ({
   height: s.height,
   transform: s.transform,
   nodes: s.nodes,
-  nodeLookup: s.nodeLookup,
+  nodeInternals: s.nodeInternals,
 });
 
 const MiniMap = ({
@@ -42,7 +42,13 @@ const MiniMap = ({
   nodeStrokeWidth = 2,
   maskColor = 'rgb(240, 242, 243, 0.7)',
 }: MiniMapProps) => {
-  const { width: containerWidth, height: containerHeight, transform, nodes, nodeLookup } = useStore(selector, shallow);
+  const {
+    width: containerWidth,
+    height: containerHeight,
+    transform,
+    nodes,
+    nodeInternals,
+  } = useStore(selector, shallow);
   const [tX, tY, tScale] = transform;
 
   const mapClasses = cc(['react-flow__minimap', className]);
@@ -85,7 +91,7 @@ const MiniMap = ({
       {nodes
         .filter((node) => !node.isHidden && node.width && node.height)
         .map((node) => {
-          const positionAbsolute = nodeLookup.get(node.id)?.positionAbsolute;
+          const positionAbsolute = nodeInternals.get(node.id)?.positionAbsolute;
 
           return (
             <MiniMapNode
diff --git a/src/components/ConnectionLine/index.tsx b/src/components/ConnectionLine/index.tsx
index 55143f72..7a40dd71 100644
--- a/src/components/ConnectionLine/index.tsx
+++ b/src/components/ConnectionLine/index.tsx
@@ -6,7 +6,7 @@ import { getBezierPath } from '../Edges/BezierEdge';
 import { getSmoothStepPath } from '../Edges/SmoothStepEdge';
 import {
   ElementId,
-  NodeLookupItem,
+  NodeInternalsItem,
   HandleElement,
   Position,
   ConnectionLineType,
@@ -28,7 +28,7 @@ interface ConnectionLineProps {
   CustomConnectionLineComponent?: ConnectionLineComponent;
 }
 
-const selector = (s: ReactFlowState) => ({ nodeLookup: s.nodeLookup, nodes: s.nodes, transform: s.transform });
+const selector = (s: ReactFlowState) => ({ nodeInternals: s.nodeInternals, nodes: s.nodes, transform: s.transform });
 
 export default ({
   connectionNodeId,
@@ -44,8 +44,8 @@ export default ({
   const nodeId = connectionNodeId;
   const handleId = connectionHandleId;
 
-  const { nodeLookup, nodes, transform } = useStore(selector, shallow);
-  const sourceNodeInternals = useRef<NodeLookupItem | undefined>(nodeLookup.get(nodeId));
+  const { nodeInternals, nodes, transform } = useStore(selector, shallow);
+  const sourceNodeInternals = useRef<NodeInternalsItem | undefined>(nodeInternals.get(nodeId));
   const sourceNode = useRef<Node | undefined>(nodes.find((n) => n.id === nodeId));
 
   if (
diff --git a/src/components/Nodes/wrapNode.tsx b/src/components/Nodes/wrapNode.tsx
index dfc9cc71..b1cd90fe 100644
--- a/src/components/Nodes/wrapNode.tsx
+++ b/src/components/Nodes/wrapNode.tsx
@@ -66,12 +66,11 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
     const nodeStyle: CSSProperties = useMemo(
       () => ({
         zIndex,
-        // zIndex: isSelected ? zIndex + 1 : zIndex,
         transform: `translate(${xPos}px,${yPos}px)`,
         pointerEvents:
           isSelectable || isDraggable || onClick || onMouseEnter || onMouseMove || onMouseLeave ? 'all' : 'none',
         // prevents jumping of nodes on start
-        // opacity: isInitialized ? 1 : 0,
+        opacity: isInitialized ? 1 : 0,
         ...style,
       }),
       [
@@ -87,6 +86,7 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
         onMouseMove,
         onMouseLeave,
         isParentNode,
+        zIndex,
       ]
     );
 
diff --git a/src/container/EdgeRenderer/index.tsx b/src/container/EdgeRenderer/index.tsx
index eea579d2..738aa8e9 100644
--- a/src/container/EdgeRenderer/index.tsx
+++ b/src/container/EdgeRenderer/index.tsx
@@ -216,7 +216,7 @@ const selector = (s: ReactFlowState) => ({
   width: s.width,
   height: s.height,
   connectionMode: s.connectionMode,
-  nodeLookup: s.nodeLookup,
+  nodeInternals: s.nodeInternals,
 });
 
 const EdgeRenderer = (props: EdgeRendererProps) => {
@@ -230,9 +230,9 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
     width,
     height,
     connectionMode,
-    nodeLookup,
+    nodeInternals,
   } = useStore(selector, shallow);
-  const edgeTree = useVisibleEdges(props.onlyRenderVisibleElements, nodeLookup);
+  const edgeTree = useVisibleEdges(props.onlyRenderVisibleElements, nodeInternals);
 
   if (!width) {
     return null;
@@ -254,8 +254,8 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
           {isMaxLevel && <MarkerDefinitions defaultColor={defaultMarkerColor} />}
           <g>
             {edges.map((edge: Edge) => {
-              const sourceNode = nodeLookup.get(edge.source);
-              const targetNode = nodeLookup.get(edge.target);
+              const sourceNode = nodeInternals.get(edge.source);
+              const targetNode = nodeInternals.get(edge.target);
 
               return (
                 <Edge
diff --git a/src/container/NodeRenderer/index.tsx b/src/container/NodeRenderer/index.tsx
index 30d045c4..a34fdc46 100644
--- a/src/container/NodeRenderer/index.tsx
+++ b/src/container/NodeRenderer/index.tsx
@@ -4,7 +4,7 @@ import shallow from 'zustand/shallow';
 import { useStore } from '../../store';
 import { Node, NodeTypesType, ReactFlowState, WrapNodeProps } from '../../types';
 import useVisibleNodes from '../../hooks/useVisibleNodes';
-import useNodeLookupRef from '../../hooks/useNodeLookupRef';
+import useNodeInternalsRef from '../../hooks/useNodeInternalsRef';
 
 interface NodeRendererProps {
   nodeTypes: NodeTypesType;
@@ -34,7 +34,7 @@ const selector = (s: ReactFlowState) => ({
 const NodeRenderer = (props: NodeRendererProps) => {
   const { scale, nodesDraggable, nodesConnectable, elementsSelectable, updateNodeDimensions, snapGrid, snapToGrid } =
     useStore(selector, shallow);
-  const nodeLookup = useNodeLookupRef();
+  const nodeInternals = useNodeInternalsRef();
   const nodes = useVisibleNodes(props.onlyRenderVisibleElements);
 
   const resizeObserver = useMemo(() => {
@@ -56,7 +56,7 @@ const NodeRenderer = (props: NodeRendererProps) => {
     <div className="react-flow__nodes react-flow__container">
       {nodes.map((node) => {
         const nodeType = node.type || 'default';
-        const lookupNode = nodeLookup.current.get(node.id);
+        const internals = nodeInternals.current.get(node.id);
 
         if (!props.nodeTypes[nodeType]) {
           console.warn(`Node type "${nodeType}" not found. Using fallback type "default".`);
@@ -83,8 +83,8 @@ const NodeRenderer = (props: NodeRendererProps) => {
             sourcePosition={node.sourcePosition}
             targetPosition={node.targetPosition}
             isHidden={node.isHidden}
-            xPos={lookupNode?.positionAbsolute?.x || 0}
-            yPos={lookupNode?.positionAbsolute?.y || 0}
+            xPos={internals?.positionAbsolute?.x || 0}
+            yPos={internals?.positionAbsolute?.y || 0}
             isDragging={node.isDragging}
             isInitialized={isInitialized}
             snapGrid={snapGrid}
@@ -106,8 +106,8 @@ const NodeRenderer = (props: NodeRendererProps) => {
             isConnectable={isConnectable}
             resizeObserver={resizeObserver}
             dragHandle={node.dragHandle}
-            zIndex={lookupNode?.treeLevel || 0}
-            isParentNode={!!lookupNode?.isParentNode}
+            zIndex={internals?.treeLevel || 0}
+            isParentNode={!!internals?.isParentNode}
           />
         );
       })}
diff --git a/src/hooks/useNodeInternalsRef.ts b/src/hooks/useNodeInternalsRef.ts
new file mode 100644
index 00000000..d6b9b2d9
--- /dev/null
+++ b/src/hooks/useNodeInternalsRef.ts
@@ -0,0 +1,14 @@
+import { useRef, useEffect } from 'react';
+
+import { useStoreApi } from '../store';
+
+function useNodeInternalsRef() {
+  const store = useStoreApi();
+  const nodeInternals = useRef(store.getState().nodeInternals);
+
+  useEffect(() => store.subscribe((state) => (nodeInternals.current = state.nodeInternals)), []);
+
+  return nodeInternals;
+}
+
+export default useNodeInternalsRef;
diff --git a/src/hooks/useNodeLookupRef.ts b/src/hooks/useNodeLookupRef.ts
deleted file mode 100644
index 707751f3..00000000
--- a/src/hooks/useNodeLookupRef.ts
+++ /dev/null
@@ -1,14 +0,0 @@
-import { useRef, useEffect } from 'react';
-
-import { useStoreApi } from '../store';
-
-function useNodeLookupRef() {
-  const store = useStoreApi();
-  const nodeLookup = useRef(store.getState().nodeLookup);
-
-  useEffect(() => store.subscribe((state) => (nodeLookup.current = state.nodeLookup)), []);
-
-  return nodeLookup;
-}
-
-export default useNodeLookupRef;
diff --git a/src/hooks/useVisibleEdges.ts b/src/hooks/useVisibleEdges.ts
index 14cd22e9..33befa1e 100644
--- a/src/hooks/useVisibleEdges.ts
+++ b/src/hooks/useVisibleEdges.ts
@@ -2,15 +2,15 @@ import { useCallback } from 'react';
 
 import { useStore } from '../store';
 import { isEdgeVisible } from '../container/EdgeRenderer/utils';
-import { ReactFlowState, NodeLookup, Edge } from '../types';
+import { ReactFlowState, NodeInternals, Edge } from '../types';
 
-function groupEdgesByTreeLevel(edges: Edge[], nodeLookup: NodeLookup) {
+function groupEdgesByTreeLevel(edges: Edge[], nodeInternals: NodeInternals) {
   let maxLevel = -1;
 
   const levelLookup = edges.reduce<Record<string, Edge[]>>((tree, edge) => {
     const treeLevel = Math.max(
-      nodeLookup.get(edge.source)?.treeLevel || 0,
-      nodeLookup.get(edge.target)?.treeLevel || 0
+      nodeInternals.get(edge.source)?.treeLevel || 0,
+      nodeInternals.get(edge.target)?.treeLevel || 0
     );
     if (tree[treeLevel]) {
       tree[treeLevel].push(edge);
@@ -34,7 +34,7 @@ function groupEdgesByTreeLevel(edges: Edge[], nodeLookup: NodeLookup) {
   });
 }
 
-function useVisibleEdges(onlyRenderVisible: boolean, nodeLookup: NodeLookup) {
+function useVisibleEdges(onlyRenderVisible: boolean, nodeInternals: NodeInternals) {
   const edges = useStore(
     useCallback(
       (s: ReactFlowState) => {
@@ -43,8 +43,8 @@ function useVisibleEdges(onlyRenderVisible: boolean, nodeLookup: NodeLookup) {
         }
 
         return s.edges.filter((e) => {
-          const sourceNode = nodeLookup.get(e.source);
-          const targetNode = nodeLookup.get(e.target);
+          const sourceNode = nodeInternals.get(e.source);
+          const targetNode = nodeInternals.get(e.target);
 
           return (
             sourceNode?.width &&
@@ -65,11 +65,11 @@ function useVisibleEdges(onlyRenderVisible: boolean, nodeLookup: NodeLookup) {
           );
         });
       },
-      [onlyRenderVisible, nodeLookup]
+      [onlyRenderVisible, nodeInternals]
     )
   );
 
-  return groupEdgesByTreeLevel(edges, nodeLookup);
+  return groupEdgesByTreeLevel(edges, nodeInternals);
 }
 
 export default useVisibleEdges;
diff --git a/src/store/index.ts b/src/store/index.ts
index d139d124..b6c3cba0 100644
--- a/src/store/index.ts
+++ b/src/store/index.ts
@@ -26,12 +26,10 @@ import {
   OnEdgesChange,
   EdgeChange,
   NodeDimensionChange,
-  NodeLookup,
-  NodeLookupItem,
-  ElementId,
 } from '../types';
 import { isNode, isEdge, getRectOfNodes, getNodesInside, getConnectedEdges } from '../utils/graph';
 import { getHandleBounds } from '../components/Nodes/utils';
+import { createNodeInternals } from './utils';
 
 const { Provider, useStore, useStoreApi } = createContext<ReactFlowState>();
 
@@ -41,34 +39,6 @@ const createNodeOrEdgeSelectionChange = (isSelected: boolean) => (item: Node | E
   isSelected,
 });
 
-type XYPosAndTreeLevel = XYPosition & { treeLevel: number };
-
-function addPositions(a: XYPosAndTreeLevel, b: XYPosition): XYPosAndTreeLevel {
-  return {
-    x: (a.x ?? 0) + (b.x ?? 0),
-    y: (a.y ?? 0) + (b.y ?? 0),
-    treeLevel: a.treeLevel + 1,
-  };
-}
-
-function getAbsolutePositionAndTreeLevel(
-  node: NodeLookupItem,
-  nodeLookup: NodeLookup,
-  result: XYPosAndTreeLevel
-): XYPosAndTreeLevel {
-  const parentNode = node.parentNode ? nodeLookup.get(node.parentNode) : false;
-
-  if (!parentNode) {
-    return result;
-  }
-
-  return getAbsolutePositionAndTreeLevel(
-    parentNode,
-    nodeLookup,
-    addPositions(result, parentNode.position || { x: 0, y: 0 })
-  );
-}
-
 const createStore = () =>
   create<ReactFlowState>((set, get) => ({
     width: 0,
@@ -125,58 +95,18 @@ const createStore = () =>
 
     reactFlowVersion: typeof __REACT_FLOW_VERSION__ !== 'undefined' ? __REACT_FLOW_VERSION__ : '-',
 
-    nodeLookup: new Map(),
+    nodeInternals: new Map(),
 
     setNodes: (nodes: Node[]) => {
-      const { nodeLookup } = get();
-      const nextNodeLookup = new Map<ElementId, NodeLookupItem>();
+      const nodeInternals = createNodeInternals(nodes, get().nodeInternals);
 
-      nodes.forEach((node) => {
-        const lookupNode: NodeLookupItem = {
-          ...nodeLookup.get(node.id),
-          width: node.width || null,
-          height: node.height || null,
-          position: node.position,
-          positionAbsolute: node.position,
-          treeLevel: node.zIndex || 0,
-        };
-        if (node.parentNode) {
-          lookupNode.parentNode = node.parentNode;
-        }
-        nextNodeLookup.set(node.id, lookupNode);
-      });
-
-      nodes
-        .filter((node) => node.parentNode)
-        .forEach((node) => {
-          const positionAbsoluteAndTreeLevel = getAbsolutePositionAndTreeLevel(node, nextNodeLookup, {
-            ...node.position,
-            treeLevel: node.zIndex || 0,
-          });
-
-          nextNodeLookup.set(node.parentNode!, { ...nextNodeLookup.get(node.parentNode!), isParentNode: true });
-
-          if (positionAbsoluteAndTreeLevel) {
-            const { treeLevel, x, y } = positionAbsoluteAndTreeLevel;
-
-            nextNodeLookup.set(node.id, {
-              ...nextNodeLookup.get(node.id),
-              positionAbsolute: {
-                x,
-                y,
-              },
-              treeLevel,
-            });
-          }
-        });
-
-      set({ nodes, nodeLookup: nextNodeLookup });
+      set({ nodes, nodeInternals });
     },
     setEdges: (edges: Edge[]) => {
       set({ edges });
     },
     updateNodeDimensions: (updates: NodeDimensionUpdate[]) => {
-      const { onNodesChange, nodes, transform, nodeLookup } = get();
+      const { onNodesChange, nodes, transform, nodeInternals } = get();
 
       const nodesToChange: NodeChange[] = updates.reduce<NodeChange[]>((res, update) => {
         const node = nodes.find((n) => n.id === update.id);
@@ -190,8 +120,8 @@ const createStore = () =>
 
           if (doUpdate) {
             const handleBounds = getHandleBounds(update.nodeElement, transform[2]);
-            nodeLookup.set(node.id, {
-              ...nodeLookup.get(node.id),
+            nodeInternals.set(node.id, {
+              ...nodeInternals.get(node.id),
               handleBounds,
               ...dimensions,
             });
@@ -208,7 +138,7 @@ const createStore = () =>
         return res;
       }, []);
 
-      set({ nodeLookup: new Map(nodeLookup) });
+      set({ nodeInternals: new Map(nodeInternals) });
 
       onNodesChange?.(nodesToChange);
     },
@@ -216,11 +146,11 @@ const createStore = () =>
       const { onNodesChange, nodes, nodeExtent } = get();
 
       if (onNodesChange) {
-        const matchingNodes = nodes.filter((n) => !!n.isSelected || n.id === id);
+        const matchingNodes = nodes.filter((n) => !!(n.isSelected || n.id === id));
 
         if (matchingNodes?.length) {
           onNodesChange(
-            matchingNodes.map((node) => {
+            matchingNodes?.map((node) => {
               const change: NodeDimensionChange = {
                 id: node.id,
                 type: 'dimensions',
diff --git a/src/store/utils.ts b/src/store/utils.ts
new file mode 100644
index 00000000..928adf45
--- /dev/null
+++ b/src/store/utils.ts
@@ -0,0 +1,77 @@
+import { ElementId, Node, NodeInternals, NodeInternalsItem, XYPosition } from '../types';
+
+type XYPosAndTreeLevel = XYPosition & { treeLevel: number };
+
+function addPositions(a: XYPosAndTreeLevel, b: XYPosition): XYPosAndTreeLevel {
+  return {
+    x: (a.x ?? 0) + (b.x ?? 0),
+    y: (a.y ?? 0) + (b.y ?? 0),
+    treeLevel: a.treeLevel + 1,
+  };
+}
+
+function getAbsolutePosAndTreeLevel(
+  node: NodeInternalsItem,
+  nodeInternals: NodeInternals,
+  result: XYPosAndTreeLevel
+): XYPosAndTreeLevel {
+  const parentNode = node.parentNode ? nodeInternals.get(node.parentNode) : false;
+
+  if (!parentNode) {
+    return result;
+  }
+
+  return getAbsolutePosAndTreeLevel(
+    parentNode,
+    nodeInternals,
+    addPositions(result, parentNode.position || { x: 0, y: 0 })
+  );
+}
+export function createNodeInternals(nodes: Node[], nodeInternals: NodeInternals): NodeInternals {
+  const nextNodeInternals = new Map<ElementId, NodeInternalsItem>();
+  const parentNodes: Record<ElementId, boolean> = {};
+
+  nodes.forEach((node) => {
+    const internals: NodeInternalsItem = {
+      ...nodeInternals.get(node.id),
+      width: node.width || null,
+      height: node.height || null,
+      position: node.position,
+      positionAbsolute: node.position,
+      treeLevel: node.zIndex || 0,
+    };
+    if (node.parentNode) {
+      internals.parentNode = node.parentNode;
+      parentNodes[node.parentNode] = true;
+    }
+    nextNodeInternals.set(node.id, internals);
+  });
+
+  nodes.forEach((node) => {
+    const updatedInternals: NodeInternalsItem = nextNodeInternals.get(node.id)!;
+
+    if (node.parentNode) {
+      const positionAbsoluteAndTreeLevel = getAbsolutePosAndTreeLevel(node, nextNodeInternals, {
+        ...node.position,
+        treeLevel: node.zIndex || 0,
+      });
+
+      const { treeLevel, x, y } = positionAbsoluteAndTreeLevel;
+
+      nextNodeInternals.set(node.parentNode!, { ...nextNodeInternals.get(node.parentNode!), isParentNode: true });
+      updatedInternals.positionAbsolute = {
+        x,
+        y,
+      };
+      updatedInternals.treeLevel = treeLevel;
+    }
+
+    if ((node.isDragging || node.isSelected) && !parentNodes[node.id]) {
+      nextNodeInternals.set(node.id, { ...updatedInternals, treeLevel: 1000 });
+    } else {
+      nextNodeInternals.set(node.id, { ...updatedInternals, treeLevel: updatedInternals?.treeLevel || 0 });
+    }
+  });
+
+  return nextNodeInternals;
+}
diff --git a/src/style.css b/src/style.css
index fe74d2c0..30290001 100644
--- a/src/style.css
+++ b/src/style.css
@@ -20,6 +20,7 @@
 .react-flow__viewport {
   transform-origin: 0 0;
   z-index: 2;
+  pointer-events: none;
 }
 
 .react-flow__renderer {
diff --git a/src/types/index.ts b/src/types/index.ts
index 378a4666..02cf3d43 100644
--- a/src/types/index.ts
+++ b/src/types/index.ts
@@ -457,7 +457,7 @@ export type InitD3ZoomPayload = {
 export type OnNodesChange = (nodes: NodeChange[]) => void;
 export type OnEdgesChange = (nodes: EdgeChange[]) => void;
 
-export type NodeLookupItem = {
+export type NodeInternalsItem = {
   width?: number | null;
   height?: number | null;
   parentNode?: ElementId;
@@ -468,14 +468,14 @@ export type NodeLookupItem = {
   isParentNode?: boolean;
 };
 
-export type NodeLookup = Map<ElementId, NodeLookupItem>;
+export type NodeInternals = Map<ElementId, NodeInternalsItem>;
 
 export interface ReactFlowState {
   width: number;
   height: number;
   transform: Transform;
   nodes: Node[];
-  nodeLookup: NodeLookup;
+  nodeInternals: NodeInternals;
   edges: Edge[];
   selectedNodesBbox: Rect;
   onNodesChange: OnNodesChange | null;