refactor(edge-wrapper): cleanup

packages/react/src/components/EdgeWrapper/EdgeUpdateAnchors.tsx

@@ -1,12 +1,11 @@
 // Updatable edges have a anchors around their handles to update the edge.
-import { XYHandle, type Connection, getEdgePosition } from '@xyflow/system';
+import { XYHandle, type Connection, EdgePosition } from '@xyflow/system';
 
 import { EdgeAnchor } from '../Edges/EdgeAnchor';
 import type { EdgeWrapperProps, Edge } from '../../types/edges';
 import { useStoreApi } from '../../hooks/useStore';
 
 type EdgeUpdateAnchorsProps = {
-  edgePosition: NonNullable<ReturnType<typeof getEdgePosition>>;
   edge: Edge;
   isUpdatable: boolean | 'source' | 'target';
   edgeUpdaterRadius: EdgeWrapperProps['edgeUpdaterRadius'];
@@ -17,20 +16,25 @@ type EdgeUpdateAnchorsProps = {
   onEdgeUpdateEnd: EdgeWrapperProps['onEdgeUpdateEnd'];
   setUpdateHover: (hover: boolean) => void;
   setUpdating: (updating: boolean) => void;
-};
+} & EdgePosition;
 
 function EdgeUpdateAnchors({
   isUpdatable,
   edgeUpdaterRadius,
+  edge,
+  targetHandleId,
+  sourceHandleId,
+  sourceX,
+  sourceY,
+  targetX,
+  targetY,
+  sourcePosition,
+  targetPosition,
   onEdgeUpdate,
   onEdgeUpdateStart,
   onEdgeUpdateEnd,
   setUpdating,
   setUpdateHover,
-  edgePosition,
-  edge,
-  targetHandleId,
-  sourceHandleId,
 }: EdgeUpdateAnchorsProps) {
   const store = useStoreApi();
 
@@ -104,9 +108,9 @@ function EdgeUpdateAnchors({
     <>
       {(isUpdatable === 'source' || isUpdatable === true) && (
         <EdgeAnchor
-          position={edgePosition.sourcePosition}
-          centerX={edgePosition.sourceX}
-          centerY={edgePosition.sourceY}
+          position={sourcePosition}
+          centerX={sourceX}
+          centerY={sourceY}
           radius={edgeUpdaterRadius}
           onMouseDown={onEdgeUpdaterSourceMouseDown}
           onMouseEnter={onEdgeUpdaterMouseEnter}
@@ -116,9 +120,9 @@ function EdgeUpdateAnchors({
       )}
       {(isUpdatable === 'target' || isUpdatable === true) && (
         <EdgeAnchor
-          position={edgePosition.targetPosition}
-          centerX={edgePosition.targetX}
-          centerY={edgePosition.targetY}
+          position={targetPosition}
+          centerX={targetX}
+          centerY={targetY}
           radius={edgeUpdaterRadius}
           onMouseDown={onEdgeUpdaterTargetMouseDown}
           onMouseEnter={onEdgeUpdaterMouseEnter}

packages/react/src/components/EdgeWrapper/index.tsx

@@ -1,12 +1,18 @@
 import { memo, useState, useMemo, useRef, type KeyboardEvent, useCallback } from 'react';
 import cc from 'classcat';
 import { shallow } from 'zustand/shallow';
-import { getMarkerId, elementSelectionKeys, getEdgePosition, errorMessages, getEdgeZIndex } from '@xyflow/system';
+import {
+  getMarkerId,
+  elementSelectionKeys,
+  getEdgePosition,
+  errorMessages,
+  getElevatedEdgeZIndex,
+} from '@xyflow/system';
 
 import { useStoreApi, useStore } from '../../hooks/useStore';
 import { ARIA_EDGE_DESC_KEY } from '../A11yDescriptions';
-import type { EdgeWrapperProps, Node } from '../../types';
-import { builtinEdgeTypes } from './utils';
+import type { EdgeWrapperProps } from '../../types';
+import { builtinEdgeTypes, nullPosition } from './utils';
 import EdgeUpdateAnchors from './EdgeUpdateAnchors';
 
 function EdgeWrapper({
@@ -53,45 +59,44 @@ function EdgeWrapper({
   const [updateHover, setUpdateHover] = useState<boolean>(false);
   const [updating, setUpdating] = useState<boolean>(false);
   const store = useStoreApi();
-  const prevSourceNode = useRef<Node | undefined>();
-  const prevTargetNode = useRef<Node | undefined>();
-  const prevZIndex = useRef<number | undefined>(edge.zIndex);
-  const prevEdgePosition = useRef<ReturnType<typeof getEdgePosition> | null>(null);
 
-  const { edgePosition, zIndex } = useStore(
+  const { zIndex, sourceX, sourceY, targetX, targetY, sourcePosition, targetPosition } = useStore(
     useCallback(
-      (state) => {
-        const sourceNode = state.nodeLookup.get(edge.source);
-        const targetNode = state.nodeLookup.get(edge.target);
+      (store) => {
+        const sourceNode = store.nodeLookup.get(edge.source);
+        const targetNode = store.nodeLookup.get(edge.target);
 
         if (!sourceNode || !targetNode) {
-          return { edgePosition: null, zIndex: edge.zIndex };
+          return {
+            zIndex: edge.zIndex,
+            ...nullPosition,
+          };
         }
 
-        const nodesChanged = prevSourceNode.current !== sourceNode || prevTargetNode.current !== targetNode;
+        const edgePosition = getEdgePosition({
+          id,
+          sourceNode,
+          targetNode,
+          sourceHandle: sourceHandleId || null,
+          targetHandle: targetHandleId || null,
+          connectionMode: store.connectionMode,
+          onError,
+        });
 
-        prevSourceNode.current = sourceNode;
-        prevTargetNode.current = targetNode;
-
-        prevEdgePosition.current = nodesChanged
-          ? getEdgePosition({
-              id,
-              sourceNode,
-              targetNode,
-              sourceHandle: sourceHandleId || null,
-              targetHandle: targetHandleId || null,
-              connectionMode: state.connectionMode,
-              onError: state.onError,
-            })
-          : prevEdgePosition.current;
-        prevZIndex.current = getEdgeZIndex(edge.selected, edge.zIndex, sourceNode, targetNode, elevateEdgesOnSelect);
+        const zIndex = getElevatedEdgeZIndex({
+          selected: edge.selected,
+          zIndex: edge.zIndex,
+          sourceNode,
+          targetNode,
+          elevateOnSelect: elevateEdgesOnSelect,
+        });
 
         return {
-          edgePosition: prevEdgePosition.current,
-          zIndex: prevZIndex.current,
+          zIndex,
+          ...(edgePosition || nullPosition),
         };
       },
-      [edge.source, edge.target, edge.selected, edge.zIndex]
+      [edge.source, edge.target, edge.selected, edge.zIndex, elevateEdgesOnSelect]
     ),
     shallow
   );
@@ -100,12 +105,13 @@ function EdgeWrapper({
     () => (edge.markerStart ? `url(#${getMarkerId(edge.markerStart, rfId)})` : undefined),
     [edge.markerStart, rfId]
   );
+
   const markerEndUrl = useMemo(
     () => (edge.markerEnd ? `url(#${getMarkerId(edge.markerEnd, rfId)})` : undefined),
     [edge.markerEnd, rfId]
   );
 
-  if (edge.hidden || !edgePosition) {
+  if (edge.hidden || !sourceX || !sourceY || !targetX || !targetY) {
     return null;
   }
 
@@ -213,14 +219,14 @@ function EdgeWrapper({
             labelBgStyle={edge.labelBgStyle}
             labelBgPadding={edge.labelBgPadding}
             labelBgBorderRadius={edge.labelBgBorderRadius}
+            sourceX={sourceX}
+            sourceY={sourceY}
+            targetX={targetX}
+            targetY={targetY}
+            sourcePosition={sourcePosition}
+            targetPosition={targetPosition}
             data={edge.data}
             style={edge.style}
-            sourceX={edgePosition.sourceX}
-            sourceY={edgePosition.sourceY}
-            targetX={edgePosition.targetX}
-            targetY={edgePosition.targetY}
-            sourcePosition={edgePosition.sourcePosition}
-            targetPosition={edgePosition.targetPosition}
             sourceHandleId={sourceHandleId}
             targetHandleId={targetHandleId}
             markerStart={markerStartUrl}
@@ -237,7 +243,12 @@ function EdgeWrapper({
             onEdgeUpdate={onEdgeUpdate}
             onEdgeUpdateStart={onEdgeUpdateStart}
             onEdgeUpdateEnd={onEdgeUpdateEnd}
-            edgePosition={edgePosition}
+            sourceX={sourceX}
+            sourceY={sourceY}
+            targetX={targetX}
+            targetY={targetY}
+            sourcePosition={sourcePosition}
+            targetPosition={targetPosition}
             setUpdateHover={setUpdateHover}
             setUpdating={setUpdating}
             sourceHandleId={sourceHandleId}

packages/react/src/components/EdgeWrapper/utils.ts

@@ -15,3 +15,12 @@ export const builtinEdgeTypes: EdgeTypes = {
   smoothstep: SmoothStepEdgeInternal as ComponentType<EdgeProps>,
   simplebezier: SimpleBezierEdgeInternal as ComponentType<EdgeProps>,
 };
+
+export const nullPosition = {
+  sourceX: null,
+  sourceY: null,
+  targetX: null,
+  targetY: null,
+  sourcePosition: null,
+  targetPosition: null,
+};