refactor: access edge-label-renderer by storing the root element

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

@@ -1,19 +1,26 @@
-import { useRef } from 'react';
 import type { ReactNode } from 'react';
 import { createPortal } from 'react-dom';
 import { useStoreApi } from '../../hooks/useStore';
-import { getEdgeLabelRendererId } from '../../utils/graph';
+import { EDGE_LABEL_RENDERER_MAIN_CLASS } from '../../constants/component';
 
 function EdgeLabelRenderer({ children }: { children: ReactNode }) {
-  const store = useStoreApi()
-  const state = store.getState()
-  const wrapperRef = useRef(document.getElementById(getEdgeLabelRendererId(state.rfId)));
+  const store = useStoreApi();
+  const {
+    rootElementRef: { current: rootElement },
+  } = store.getState();
 
-  if (!wrapperRef.current) {
+  if (!rootElement) {
     return null;
   }
 
-  return createPortal(children, wrapperRef.current);
+  const collection = rootElement.getElementsByClassName(EDGE_LABEL_RENDERER_MAIN_CLASS);
+
+  const edgeLabelRendererElement = collection.item(0);
+  if (!edgeLabelRendererElement) {
+    return null;
+  }
+
+  return createPortal(children, edgeLabelRendererElement);
 }
 
 export default EdgeLabelRenderer;

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

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

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

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

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

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

@@ -1,4 +1,4 @@
-import { forwardRef } from 'react';
+import { forwardRef, useRef } from 'react';
 import type { CSSProperties } from 'react';
 import cc from 'classcat';
 
@@ -164,11 +164,26 @@ 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={ref}
+        ref={setRootRef}
         className={cc(['react-flow', className])}
         data-testid="rf__wrapper"
         id={id}
@@ -275,6 +290,7 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
             noPanClassName={noPanClassName}
             nodeOrigin={nodeOrigin}
             rfId={rfId}
+            rootElementRef={rootElementRef}
           />
           <SelectionListener onSelectionChange={onSelectionChange} />
           {children}

packages/core/src/store/initialState.ts

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

packages/core/src/utils/graph.ts

@@ -224,6 +224,3 @@ export const getD3Transition = (selection: D3Selection<Element, unknown, null, u
   return selection.transition().duration(duration);
 };
 
-export function getEdgeLabelRendererId(id: string) {
-  return `edgelabel-portal-${id}`
-}