refactor: use existing d3Selection to get the element ref

examples/vite-app/cypress/components/reactflow/basic-props.cy.tsx

@@ -2,7 +2,6 @@ import ReactFlow, { EdgeProps } from 'reactflow';
 
 import ControlledFlow from '../../support/ControlledFlow';
 import * as simpleflow from '../../fixtures/simpleflow';
-import { RefReactFlow } from '../../support/RefReactFlow';
 
 describe('<ReactFlow />: Basic Props', () => {
   describe('uses defaultNodes and defaultEdges', () => {
@@ -195,32 +194,4 @@ describe('<ReactFlow />: Basic Props', () => {
     cy.mount(<ControlledFlow className="custom" />);
     cy.get('.react-flow').should('have.class', 'custom');
   });
-
-  it('uses function ref', () => {
-    let pResolve: ((element: HTMLDivElement | PromiseLike<HTMLDivElement> | null) => void) | null = null;
-    const promise: Promise<HTMLDivElement | null> = new Promise((resolve) => {
-      pResolve = resolve;
-    });
-
-    function ref(root: HTMLDivElement | null) {
-      pResolve && pResolve(root);
-    }
-
-    cy.mount(<ReactFlow ref={ref} />);
-    cy.wrap(promise).should('be.instanceOf', HTMLDivElement);
-  });
-
-  it('use mutable ref', () => {
-    let pResolve: ((element: HTMLDivElement | PromiseLike<HTMLDivElement> | null) => void) | null = null;
-    const promise: Promise<HTMLDivElement | null> = new Promise((resolve) => {
-      pResolve = resolve;
-    });
-
-    function onGetRef(root: HTMLDivElement | null) {
-      pResolve && pResolve(root);
-    }
-
-    cy.mount(<RefReactFlow onGetRef={onGetRef} />);
-    cy.wrap(promise).should('be.instanceOf', HTMLDivElement);
-  });
 });

examples/vite-app/cypress/support/RefReactFlow.tsx

@@ -1,12 +0,0 @@
-import { useEffect, useRef } from 'react';
-import ReactFlow from 'reactflow';
-
-export function RefReactFlow({ onGetRef }: { onGetRef: (element: HTMLDivElement | null) => void }) {
-  const ref = useRef<HTMLDivElement | null>(null);
-
-  useEffect(() => {
-    onGetRef(ref.current);
-  }, []);
-
-  return <ReactFlow ref={ref} />;
-}

packages/core/src/components/EdgeLabelRenderer/index.tsx

@@ -1,22 +1,18 @@
 import type { ReactNode } from 'react';
 import { createPortal } from 'react-dom';
-import { useStoreApi } from '../../hooks/useStore';
-import { EDGE_LABEL_RENDERER_MAIN_CLASS } from '../../constants/component';
+import { useStore } from '../../hooks/useStore';
+import { ReactFlowState } from '../../types';
+
+const selector = (s: ReactFlowState) => ({
+  d3Selection: s.d3Selection,
+});
 
 function EdgeLabelRenderer({ children }: { children: ReactNode }) {
-  const store = useStoreApi();
-  const {
-    rootElementRef: { current: rootElement },
-  } = store.getState();
+  const { d3Selection } = useStore(selector);
 
-  if (!rootElement) {
-    return null;
-  }
+  const edgeLabelRendererElement = d3Selection?.select('.react-flow__edgelabel-renderer').node();
 
-  const collection = rootElement.getElementsByClassName(EDGE_LABEL_RENDERER_MAIN_CLASS);
-
-  const edgeLabelRendererElement = collection.item(0);
-  if (!edgeLabelRendererElement) {
+  if (!edgeLabelRendererElement || !(edgeLabelRendererElement instanceof HTMLDivElement)) {
     return null;
   }
 

packages/core/src/components/StoreUpdater/index.tsx

@@ -3,15 +3,7 @@ import { StoreApi } from 'zustand';
 import shallow from 'zustand/shallow';
 
 import { useStore, useStoreApi } from '../../hooks/useStore';
-import type {
-  Node,
-  Edge,
-  ReactFlowState,
-  CoordinateExtent,
-  ReactFlowProps,
-  ReactFlowStore,
-  RootElementRef,
-} from '../../types';
+import type { Node, Edge, ReactFlowState, CoordinateExtent, ReactFlowProps, ReactFlowStore } from '../../types';
 
 type StoreUpdaterProps = Pick<
   ReactFlowProps,
@@ -52,7 +44,7 @@ type StoreUpdaterProps = Pick<
   | 'onSelectionDragStop'
   | 'noPanClassName'
   | 'nodeOrigin'
-> & { rfId: string; rootElementRef: RootElementRef };
+> & { rfId: string };
 
 const selector = (s: ReactFlowState) => ({
   setNodes: s.setNodes,
@@ -125,7 +117,6 @@ const StoreUpdater = ({
   noPanClassName,
   nodeOrigin,
   rfId,
-  rootElementRef,
 }: StoreUpdaterProps) => {
   const {
     setNodes,
@@ -178,7 +169,6 @@ const StoreUpdater = ({
   useDirectStoreUpdater('noPanClassName', noPanClassName, store.setState);
   useDirectStoreUpdater('nodeOrigin', nodeOrigin, store.setState);
   useDirectStoreUpdater('rfId', rfId, store.setState);
-  useDirectStoreUpdater('rootElementRef', rootElementRef, store.setState);
 
   useStoreUpdater<Node[]>(nodes, setNodes);
   useStoreUpdater<Edge[]>(edges, setEdges);

packages/core/src/constants/component.ts

@@ -1 +0,0 @@
-export const EDGE_LABEL_RENDERER_MAIN_CLASS = 'react-flow__edgelabel-renderer'

packages/core/src/container/GraphView/index.tsx

@@ -6,7 +6,6 @@ import EdgeRenderer from '../EdgeRenderer';
 import ViewportWrapper from '../Viewport';
 import useOnInitHandler from '../../hooks/useOnInitHandler';
 import type { EdgeTypesWrapped, NodeTypesWrapped, ReactFlowProps } from '../../types';
-import { EDGE_LABEL_RENDERER_MAIN_CLASS } from '../../constants/component';
 
 export type GraphViewProps = Omit<ReactFlowProps, 'onSelectionChange' | 'nodes' | 'edges' | 'nodeTypes' | 'edgeTypes'> &
   Required<
@@ -159,7 +158,7 @@ const GraphView = ({
           disableKeyboardA11y={disableKeyboardA11y}
           rfId={rfId}
         />
-        <div className={EDGE_LABEL_RENDERER_MAIN_CLASS} />
+        <div className="react-flow__edgelabel-renderer" />
 
         <NodeRenderer
           nodeTypes={nodeTypes}

packages/core/src/container/ReactFlow/index.tsx

@@ -1,4 +1,4 @@
-import { forwardRef, useRef } from 'react';
+import { forwardRef } from 'react';
 import type { CSSProperties } from 'react';
 import cc from 'classcat';
 
@@ -164,26 +164,11 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
     const edgeTypesWrapped = useNodeOrEdgeTypes(edgeTypes, createEdgeTypes) as EdgeTypesWrapped;
     const rfId = id || '1';
 
-    const rootElementRef = useRef<HTMLDivElement | null>(null);
-
-    function setRootRef(element: HTMLDivElement | null) {
-      rootElementRef.current = element;
-
-      if (typeof ref === 'function') {
-        ref(element);
-        return;
-      }
-
-      if (ref) {
-        ref.current = element;
-      }
-    }
-
     return (
       <div
         {...rest}
         style={{ ...style, ...wrapperStyle }}
-        ref={setRootRef}
+        ref={ref}
         className={cc(['react-flow', className])}
         data-testid="rf__wrapper"
         id={id}
@@ -290,7 +275,6 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
             noPanClassName={noPanClassName}
             nodeOrigin={nodeOrigin}
             rfId={rfId}
-            rootElementRef={rootElementRef}
           />
           <SelectionListener onSelectionChange={onSelectionChange} />
           {children}

packages/core/src/store/initialState.ts

@@ -54,7 +54,6 @@ const initialState: ReactFlowStore = {
   connectOnClick: true,
 
   ariaLiveMessage: '',
-  rootElementRef: { current: null }
 };
 
 export default initialState;

packages/core/src/types/general.ts

@@ -1,5 +1,5 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import type { MouseEvent as ReactMouseEvent, ComponentType, MemoExoticComponent, MutableRefObject } from 'react';
+import type { MouseEvent as ReactMouseEvent, ComponentType, MemoExoticComponent } from 'react';
 import type { Selection as D3Selection, ZoomBehavior } from 'd3';
 
 import type { XYPosition, Rect, Transform, CoordinateExtent } from './utils';
@@ -129,8 +129,6 @@ export type ViewportHelperFunctions = {
   viewportInitialized: boolean;
 };
 
-export type RootElementRef = MutableRefObject<HTMLDivElement | null>
-
 export type ReactFlowStore = {
   rfId: string;
   width: number;
@@ -210,7 +208,6 @@ export type ReactFlowStore = {
   onSelectionChange?: OnSelectionChangeFunc;
 
   ariaLiveMessage: string;
-  rootElementRef: RootElementRef
 };
 
 export type ReactFlowActions = {