refactor: use existing d3Selection to get the element ref

2022-11-14 08:30:09 +08:00
parent 1993734a99
commit 23aa8bb270
9 changed files with 15 additions and 92 deletions
@@ -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;
  }