refactor(node-dimensions): use store domNode for calculation #2316

src/container/ReactFlow/index.tsx

@@ -1,5 +1,6 @@
-import cc from 'classcat';
 import React, { forwardRef } from 'react';
+import cc from 'classcat';
+
 import Attribution from '../../components/Attribution';
 import { BezierEdge, SmoothStepEdge, StepEdge, StraightEdge, SimpleBezierEdge } from '../../components/Edges';
 import DefaultNode from '../../components/Nodes/DefaultNode';
@@ -148,10 +149,9 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
   ) => {
     const nodeTypesWrapped = useNodeOrEdgeTypes(nodeTypes, createNodeTypes) as NodeTypesWrapped;
     const edgeTypesWrapped = useNodeOrEdgeTypes(edgeTypes, createEdgeTypes) as EdgeTypesWrapped;
-    const reactFlowClasses = cc(['react-flow', className]);
 
     return (
-      <div {...rest} ref={ref} className={reactFlowClasses}>
+      <div {...rest} ref={ref} className={cc(['react-flow', className])}>
         <Wrapper>
           <GraphView
             onInit={onInit}

src/container/ZoomPane/index.tsx

@@ -81,6 +81,7 @@ const ZoomPane = ({
         d3ZoomHandler: selection.on('wheel.zoom'),
         // we need to pass transform because zoom handler is not registered when we set the initial transform
         transform: [clampedX, clampedY, clampedZoom],
+        domNode: selection.node()?.closest('.react-flow') as HTMLElement,
       });
     }
   }, []);

src/hooks/useUpdateNodeInternals.ts

@@ -1,15 +1,21 @@
 import { useCallback } from 'react';
 
-import { useStore } from '../store';
+import { useStore, useStoreApi } from '../store';
 import { UpdateNodeInternals, ReactFlowState } from '../types';
 
 const selector = (state: ReactFlowState) => state.updateNodeDimensions;
 
 function useUpdateNodeInternals(): UpdateNodeInternals {
+  const store = useStoreApi();
   const updateNodeDimensions = useStore(selector);
 
   return useCallback<UpdateNodeInternals>((id: string) => {
-    const nodeElement = document.querySelector(`.react-flow__node[data-id="${id}"]`) as HTMLDivElement;
+    const { domNode } = store.getState();
+    if (!domNode) {
+      return;
+    }
+
+    const nodeElement = domNode.querySelector(`.react-flow__node[data-id="${id}"]`) as HTMLDivElement;
 
     if (nodeElement) {
       updateNodeDimensions([{ id, nodeElement, forceUpdate: true }]);

src/store/index.ts

@@ -43,9 +43,14 @@ const createStore = () =>
       set({ nodeInternals, edges: nextEdges, hasDefaultNodes, hasDefaultEdges });
     },
     updateNodeDimensions: (updates: NodeDimensionUpdate[]) => {
-      const { onNodesChange, nodeInternals, fitViewOnInit, fitViewOnInitDone, fitViewOnInitOptions } = get();
+      const { onNodesChange, nodeInternals, fitViewOnInit, fitViewOnInitDone, fitViewOnInitOptions, domNode } = get();
+      const viewportNode = domNode?.querySelector('.react-flow__viewport');
 
-      const style = window.getComputedStyle(document.querySelector('.react-flow__viewport')!);
+      if (!viewportNode) {
+        return;
+      }
+
+      const style = window.getComputedStyle(viewportNode);
       const { m22: zoom } = new window.DOMMatrixReadOnly(style.transform);
 
       const changes: NodeDimensionChange[] = updates.reduce<NodeDimensionChange[]>((res, update) => {

src/store/initialState.ts

@@ -29,6 +29,7 @@ const initialState: ReactFlowStore = {
   connectionHandleType: 'source',
   connectionPosition: { x: 0, y: 0 },
   connectionMode: ConnectionMode.Strict,
+  domNode: null,
 
   snapGrid: [15, 15],
   snapToGrid: false,

src/types/general.ts

@@ -143,6 +143,7 @@ export type ReactFlowStore = {
   onEdgesChange: OnEdgesChange | null;
   hasDefaultNodes: boolean;
   hasDefaultEdges: boolean;
+  domNode: HTMLElement | null;
 
   d3Zoom: ZoomBehavior<Element, unknown> | null;
   d3Selection: D3Selection<Element, unknown, null, undefined> | null;