refactor: use existing d3Selection to get the element 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);