From 2b065b4e564eab2bcc06c9cd2ad6d4233e6ce66c Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Mon, 22 Nov 2021 13:35:47 +0100
Subject: [PATCH] update: pass store state as props to edge-renderer and edge

Signed-off-by: Braks <78412429+bcakmakoglu@users.noreply.github.com>
---
 examples/Edges/CustomEdge2.vue              |   3 +-
 examples/Edges/CustomLabel.vue              |   4 +-
 examples/Edges/EdgesExample.vue             |   2 +-
 src/components/Edges/EdgeWrapper.vue        | 121 +++++++-------------
 src/components/Nodes/NodeWrapper.vue        |  28 ++---
 src/container/EdgeRenderer/EdgeRenderer.vue |  57 +++++++--
 src/container/NodeRenderer/NodeRenderer.vue |   2 +-
 src/container/VueFlow/VueFlow.vue           |   9 ++
 src/container/ZoomPane/ZoomPane.vue         |   8 +-
 src/store/stateStore.ts                     |   2 +-
 src/types/components.ts                     |   6 +-
 src/types/edge.ts                           |   6 +-
 src/types/node.ts                           |   5 +
 src/utils/edge.ts                           |  42 ++++---
 14 files changed, 152 insertions(+), 143 deletions(-)

diff --git a/examples/Edges/CustomEdge2.vue b/examples/Edges/CustomEdge2.vue
index c6c30af7..8513d7a9 100644
--- a/examples/Edges/CustomEdge2.vue
+++ b/examples/Edges/CustomEdge2.vue
@@ -42,6 +42,7 @@ const center = computed(() =>
     targetY: props.targetY,
   }),
 )
+const onClick = () => console.log(props.data)
 </script>
 <script lang="ts">
 export default {
@@ -59,6 +60,6 @@ export default {
     :label-bg-style="{ fill: 'red' }"
     :label-bg-padding="[2, 4]"
     :label-bg-border-radius="2"
-    @click="() => console.log(props.data)"
+    @click="onClick"
   />
 </template>
diff --git a/examples/Edges/CustomLabel.vue b/examples/Edges/CustomLabel.vue
index 67073d7d..3f9c0c05 100644
--- a/examples/Edges/CustomLabel.vue
+++ b/examples/Edges/CustomLabel.vue
@@ -1,10 +1,10 @@
 <script lang="ts" setup>
 interface CustomLabelProps {
-  text: string
+  label: string
 }
 
 const props = defineProps<CustomLabelProps>()
 </script>
 <template>
-  <tspan dy="10" x="0">{{ props.text }}</tspan>
+  <tspan dy="10" x="0">{{ props.label }}</tspan>
 </template>
diff --git a/examples/Edges/EdgesExample.vue b/examples/Edges/EdgesExample.vue
index e24106d7..2c0b5094 100644
--- a/examples/Edges/EdgesExample.vue
+++ b/examples/Edges/EdgesExample.vue
@@ -42,7 +42,7 @@ const initialElements: Elements = [
     label: {
       component: markRaw(CustomLabel),
       props: {
-        text: 'custom label text',
+        label: 'custom label text',
       },
     },
     labelStyle: { fill: 'red', fontWeight: 700 },
diff --git a/src/components/Edges/EdgeWrapper.vue b/src/components/Edges/EdgeWrapper.vue
index cb7c46b5..e7263531 100644
--- a/src/components/Edges/EdgeWrapper.vue
+++ b/src/components/Edges/EdgeWrapper.vue
@@ -1,14 +1,16 @@
 <script lang="ts" setup>
-import { CSSProperties, DefineComponent } from 'vue'
+import { CSSProperties } from 'vue'
 import { useHandle, useStore } from '../../composables'
-import { ArrowHeadType, ConnectionMode, Edge, EdgePositions, EdgeTextProps, ElementId, Position } from '../../types'
-import { isEdge, getEdgePositions, getHandle, getSourceTargetNodes, isEdgeVisible } from '../../utils'
+import { ArrowHeadType, ConnectionMode, Edge, EdgeTextProps, ElementId, SourceTargetNode, Position } from '../../types'
+import { isEdge, getEdgePositions, getHandle } from '../../utils'
 import EdgeAnchor from './EdgeAnchor.vue'
 
 interface EdgeWrapper {
   id: ElementId
   edge: Edge
+  sourceTargetNodes: SourceTargetNode
   type?: string
+  component?: any
   source: ElementId
   target: ElementId
   sourceHandle?: ElementId | null
@@ -19,7 +21,7 @@ interface EdgeWrapper {
   label?:
     | string
     | {
-        component: DefineComponent<EdgeTextProps>
+        component: any
         props?: EdgeTextProps
       }
   labelStyle?: any
@@ -35,23 +37,15 @@ interface EdgeWrapper {
   class?: string
   isHidden?: boolean
   edgeUpdaterRadius?: number
+  selectable?: boolean
 }
 
 const props = withDefaults(defineProps<EdgeWrapper>(), {})
 
 const store = useStore()
 
-const type = computed(() => {
-  const t = props.type ?? 'default'
-  let edge = store.getEdgeTypes[t]
-  if (!edge) {
-    edge = store.getEdgeTypes.default
-    console.warn(`Edge type "${t}" not found. Using fallback type "default".`)
-  }
-  return edge
-})
-
 const updating = ref<boolean>(false)
+const handler = useHandle()
 
 const onEdgeClick = (event: MouseEvent) => {
   if (store.elementsSelectable) {
@@ -61,24 +55,24 @@ const onEdgeClick = (event: MouseEvent) => {
   store.hooks.elementClick.trigger({ event, element: props.edge })
   store.hooks.edgeClick.trigger({ event, edge: props.edge })
 }
-
-const onDoubleClick = (event: MouseEvent) => {
-  store.hooks.edgeDoubleClick.trigger({ event, edge: props.edge })
-}
-
 const onEdgeContextMenu = (event: MouseEvent) =>
   store.hooks.edgeContextMenu.trigger({
     event,
     edge: props.edge,
   })
-
+const onDoubleClick = (event: MouseEvent) => store.hooks.edgeDoubleClick.trigger({ event, edge: props.edge })
 const onEdgeMouseEnter = (event: MouseEvent) => store.hooks.edgeMouseEnter.trigger({ event, edge: props.edge })
-
 const onEdgeMouseMove = (event: MouseEvent) => store.hooks.edgeMouseMove.trigger({ event, edge: props.edge })
-
 const onEdgeMouseLeave = (event: MouseEvent) => store.hooks.edgeMouseLeave.trigger({ event, edge: props.edge })
+const onEdgeUpdaterMouseEnter = () => (updating.value = true)
+const onEdgeUpdaterMouseOut = () => (updating.value = false)
+const onEdgeUpdaterSourceMouseDown = (event: MouseEvent) => {
+  handleEdgeUpdater(event, true)
+}
+const onEdgeUpdaterTargetMouseDown = (event: MouseEvent) => {
+  handleEdgeUpdater(event, false)
+}
 
-const handler = useHandle()
 const handleEdgeUpdater = (event: MouseEvent, isSourceHandle: boolean) => {
   const nodeId = isSourceHandle ? props.target : props.source
   const handleId = (isSourceHandle ? props.targetHandle : props.sourceHandle) ?? ''
@@ -91,75 +85,45 @@ const handleEdgeUpdater = (event: MouseEvent, isSourceHandle: boolean) => {
   )
 }
 
-const onEdgeUpdaterSourceMouseDown = (event: MouseEvent) => {
-  handleEdgeUpdater(event, true)
-}
-
-const onEdgeUpdaterTargetMouseDown = (event: MouseEvent) => {
-  handleEdgeUpdater(event, false)
-}
-
-const onEdgeUpdaterMouseEnter = () => (updating.value = true)
-
-const onEdgeUpdaterMouseOut = () => (updating.value = false)
-
-const isVisible = ({ sourceX, sourceY, targetX, targetY }: EdgePositions) => {
-  return store.onlyRenderVisibleElements
-    ? isEdgeVisible({
-        sourcePos: { x: sourceX, y: sourceY },
-        targetPos: { x: targetX, y: targetY },
-        width: store.dimensions.width,
-        height: store.dimensions.height,
-        transform: store.transform,
-      })
-    : true
-}
-
-const nodes = computed(() => {
-  const n = getSourceTargetNodes(props.edge, store.nodes)
-
-  if (!n.sourceNode) {
-    console.warn(`couldn't create edge for source id: ${props.source}; edge id: ${props.id}`)
-  }
-
-  if (!n.targetNode) {
-    console.warn(`couldn't create edge for target id: ${props.target}; edge id: ${props.id}`)
-  }
-
-  return n
-})
-
 // when connection type is loose we can define all handles as sources
-const targetNodeHandles = computed(() =>
-  store.connectionMode === ConnectionMode.Strict
-    ? nodes.value.targetNode?.__vf?.handleBounds?.target
-    : nodes.value.targetNode?.__vf?.handleBounds?.target ?? nodes.value.targetNode?.__vf?.handleBounds?.source,
+const targetNodeHandles = controlledComputed(
+  () => props.sourceTargetNodes,
+  () =>
+    store.connectionMode === ConnectionMode.Strict
+      ? props.sourceTargetNodes.targetNode.__vf.handleBounds.target
+      : props.sourceTargetNodes.targetNode.__vf.handleBounds.target ??
+        props.sourceTargetNodes.targetNode.__vf.handleBounds.source,
 )
 
-const sourceHandle = computed(() => {
-  if (nodes.value.sourceNode && nodes.value.sourceNode.__vf?.handleBounds?.source)
-    return getHandle(nodes.value.sourceNode.__vf.handleBounds.source, props.sourceHandle ?? null)
-  else return null
-})
+const sourceHandle = controlledComputed(
+  () => props.sourceTargetNodes,
+  () => {
+    if (props.sourceTargetNodes.sourceNode && props.sourceTargetNodes.sourceNode.__vf.handleBounds.source)
+      return getHandle(props.sourceTargetNodes.sourceNode.__vf.handleBounds.source, props.sourceHandle ?? null)
+    else return null
+  },
+)
 const targetHandle = computed(() => {
   if (targetNodeHandles.value) return getHandle(targetNodeHandles.value, props.targetHandle ?? null)
   else return null
 })
-const sourcePosition = computed(() => (sourceHandle.value ? sourceHandle.value.position : Position.Bottom))
-const targetPosition = computed(() => (targetHandle.value ? targetHandle.value.position : Position.Top))
+const sourcePosition = eagerComputed(() => (sourceHandle.value ? sourceHandle.value.position : Position.Bottom))
+const targetPosition = eagerComputed(() => (targetHandle.value ? targetHandle.value.position : Position.Top))
 
-const isSelected = computed(() => store.selectedElements?.some((elm) => isEdge(elm) && elm.id === props.id) ?? false)
+const isSelected = controlledComputed(
+  () => store.selectedElements,
+  () => (props.selectable && store.selectedElements?.some((elm) => isEdge(elm) && elm.id === props.id)) ?? false,
+)
 const edgePos = computed(() =>
   getEdgePositions(
-    nodes.value.sourceNode,
+    props.sourceTargetNodes.sourceNode,
     sourceHandle.value,
     sourcePosition.value,
-    nodes.value.targetNode,
+    props.sourceTargetNodes.targetNode,
     targetHandle.value,
     targetPosition.value,
   ),
 )
-const elementsSelectable = computed(() => store.elementsSelectable)
 </script>
 <script lang="ts">
 export default {
@@ -168,14 +132,13 @@ export default {
 </script>
 <template>
   <g
-    v-show="!props.isHidden && isVisible(edgePos)"
     :class="[
       'vue-flow__edge',
       `vue-flow__edge-${props.type || 'default'}`,
       {
         selected: isSelected,
         animated: props.animated,
-        inactive: !elementsSelectable,
+        inactive: !props.selectable,
         updating,
       },
       props.class,
@@ -216,7 +179,7 @@ export default {
       }"
     >
       <component
-        :is="type"
+        :is="props.component ?? props.type"
         v-if="typeof type !== 'boolean'"
         v-bind="{
           id: props.id,
diff --git a/src/components/Nodes/NodeWrapper.vue b/src/components/Nodes/NodeWrapper.vue
index 8de1afbf..2aacce22 100644
--- a/src/components/Nodes/NodeWrapper.vue
+++ b/src/components/Nodes/NodeWrapper.vue
@@ -1,6 +1,6 @@
 <script lang="ts" setup>
-import { DraggableEventListener, DraggableCore } from '@braks/revue-draggable'
 import { CSSProperties } from 'vue'
+import { DraggableEventListener, DraggableCore } from '@braks/revue-draggable'
 import { useStore } from '../../composables'
 import { ElementId, FlowEvents, GraphNode, Position, SnapGrid, VFInternals, XYPosition } from '../../types'
 import { NodeId } from '../../context'
@@ -47,11 +47,7 @@ provide(NodeId, props.id)
 
 const nodeElement = templateRef<HTMLDivElement>('node-element', null)
 
-const selectable = computed(() => (typeof props.selectable === 'undefined' ? store.elementsSelectable : props.selectable))
-const draggable = computed(() => (typeof props.draggable === 'undefined' ? store.nodesDraggable : props.draggable))
-const connectable = computed(() => (typeof props.connectable === 'undefined' ? store.nodesConnectable : props.connectable))
-const scale = computed(() => props.scale)
-const selected = computed(() => selectable.value && store.selectedElements?.some(({ id }) => id === props.id))
+const selected = computed(() => props.selectable && store.selectedElements?.some(({ id }) => id === props.id))
 
 const onMouseEnterHandler = () =>
   props.vf.isDragging &&
@@ -91,8 +87,8 @@ const onDoubleClick = () => (event: MouseEvent) => {
 
 const onSelectNodeHandler = (event: MouseEvent) => {
   const n = props.node
-  if (!draggable.value) {
-    if (selectable.value) {
+  if (!props.draggable) {
+    if (props.selectable) {
       store.unsetNodesSelection()
       if (!selected.value) store.addSelectedElements([n])
     }
@@ -106,11 +102,11 @@ const onDragStart: DraggableEventListener = ({ event }) => {
   emit('dragStart', { event, node: n })
   store.hooks.nodeDragStart.trigger({ event, node: n })
 
-  if (props.selectNodesOnDrag && selectable) {
+  if (props.selectNodesOnDrag && props.selectable) {
     store.unsetNodesSelection()
 
     if (!selected.value) store.addSelectedElements([n])
-  } else if (!props.selectNodesOnDrag && !selected.value && selectable.value) {
+  } else if (!props.selectNodesOnDrag && !selected.value && props.selectable) {
     store.unsetNodesSelection()
     store.addSelectedElements([])
   }
@@ -138,7 +134,7 @@ const onDragStop: DraggableEventListener = ({ event }) => {
   // onDragStop also gets called when user just clicks on a node.
   // Because of that we set dragging to true inside the onDrag handler and handle the click here
   if (!props.vf.isDragging) {
-    if (selectable.value && !props.selectNodesOnDrag && !selected.value) {
+    if (props.selectable && !props.selectNodesOnDrag && !selected.value) {
       store.addSelectedElements([n])
     }
     store.hooks.nodeClick.trigger({ event, node: n })
@@ -193,7 +189,7 @@ export default {
     cancel=".nodrag"
     :handle="props.dragHandle"
     :disabled="!draggable"
-    :scale="scale"
+    :scale="props.scale"
     :grid="props.snapGrid"
     :enable-user-select-hack="false"
     @start="onDragStart"
@@ -207,14 +203,14 @@ export default {
         `vue-flow__node-${props.type}`,
         {
           selected,
-          selectable: selectable,
+          selectable: props.selectable,
         },
         props.class,
       ]"
       :style="{
         zIndex: selected ? 10 : 3,
         transform: `translate(${props.vf.position.x}px,${props.vf.position.y}px)`,
-        pointerEvents: selectable || draggable ? 'all' : 'none',
+        pointerEvents: props.selectable || props.draggable ? 'all' : 'none',
         opacity: props.vf.width !== null && props.vf.height !== null ? 1 : 0,
         ...props.style,
       }"
@@ -234,7 +230,7 @@ export default {
           xPos: props.vf.position.x,
           yPos: props.vf.position.y,
           selected,
-          connectable,
+          connectable: props.connectable,
           sourcePosition: props.sourcePosition,
           targetPosition: props.targetPosition,
           dragging: props.vf.isDragging,
@@ -249,7 +245,7 @@ export default {
             xPos: props.vf.position.x,
             yPos: props.vf.position.y,
             selected,
-            connectable,
+            connectable: props.connectable,
             sourcePosition: props.sourcePosition,
             targetPosition: props.targetPosition,
             dragging: props.vf.isDragging,
diff --git a/src/container/EdgeRenderer/EdgeRenderer.vue b/src/container/EdgeRenderer/EdgeRenderer.vue
index ca73d672..d3d6c03e 100644
--- a/src/container/EdgeRenderer/EdgeRenderer.vue
+++ b/src/container/EdgeRenderer/EdgeRenderer.vue
@@ -1,17 +1,26 @@
 <script lang="ts" setup>
 import { CSSProperties } from 'vue'
-import { ConnectionLineType } from '../../types'
-import { useStore } from '../../composables'
-import Edge from '../../components/Edges/EdgeWrapper.vue'
+import { ConnectionLineType, Dimensions, Edge, EdgeComponent, GraphNode, HandleType, Transform } from '../../types'
+import { getSourceTargetNodes } from '../../utils'
+import EdgeWrapper from '../../components/Edges/EdgeWrapper.vue'
 import ConnectionLine from '../../components/ConnectionLine/ConnectionLine.vue'
 import MarkerDefinitions from './MarkerDefinitions.vue'
 
 interface EdgeRendererProps {
+  edges: Edge[]
+  nodes: GraphNode[]
+  edgeTypes: Record<string, EdgeComponent>
+  transform: Transform
+  dimensions: Dimensions
   connectionLineType?: ConnectionLineType
   connectionLineStyle?: CSSProperties
   arrowHeadColor?: string
   markerEndId?: string
+  elementsSelectable?: boolean
   edgeUpdaterRadius?: number
+  connectionNodeId?: string
+  connectionHandleType?: HandleType
+  nodesConnectable?: boolean
 }
 
 const props = withDefaults(defineProps<EdgeRendererProps>(), {
@@ -20,14 +29,35 @@ const props = withDefaults(defineProps<EdgeRendererProps>(), {
   edgeUpdaterRadius: 10,
 })
 
-const store = useStore()
+const getType = (type: string) => {
+  const t = type ?? 'default'
+  let edgeType = props.edgeTypes[t]
+  if (!edgeType) {
+    edgeType = props.edgeTypes.default
+    console.warn(`Edge type "${type}" not found. Using fallback type "default".`)
+  }
+  return edgeType
+}
 
-const sourceNode = computed(() => store.nodes.find((n) => n.id === store.connectionNodeId))
+const sourceNode = computed(() => props.nodes.find((n) => n.id === props.connectionNodeId))
 const connectionLineVisible = computed(
-  () => !!(store.nodesConnectable && sourceNode.value && store.connectionNodeId && store.connectionHandleType),
+  () =>
+    !!(
+      sourceNode.value &&
+      (typeof sourceNode.value.selectable === 'undefined' ? props.nodesConnectable : sourceNode.value.selectable) &&
+      props.connectionNodeId &&
+      props.connectionHandleType
+    ),
 )
-const dimensions = computed(() => store.dimensions)
-const transform = computed(() => `translate(${store.transform[0]},${store.transform[1]}) scale(${store.transform[2]})`)
+
+const transform = computed(() => `translate(${props.transform[0]},${props.transform[1]}) scale(${props.transform[2]})`)
+
+const sourceTargetNodes = (edge: Edge) => {
+  const { sourceNode, targetNode } = getSourceTargetNodes(edge, props.nodes)
+  if (!sourceNode) console.warn(`couldn't create edge for source id: ${edge.source}; edge id: ${edge.id}`)
+  if (!targetNode) console.warn(`couldn't create edge for target id: ${edge.target}; edge id: ${edge.id}`)
+  return { sourceNode, targetNode }
+}
 </script>
 <script lang="ts">
 export default {
@@ -35,20 +65,23 @@ export default {
 }
 </script>
 <template>
-  <svg :width="dimensions.width" :height="dimensions.height" class="vue-flow__edges">
+  <svg :width="props.dimensions.width" :height="props.dimensions.height" class="vue-flow__edges">
     <MarkerDefinitions :color="props.arrowHeadColor" />
     <g :transform="transform">
-      <Edge
-        v-for="(edge, i) of store.getEdges"
+      <EdgeWrapper
+        v-for="(edge, i) of props.edges"
         :id="edge.id"
         :key="`${edge.id}-${i}`"
         :edge="edge"
+        :source-target-nodes="sourceTargetNodes(edge)"
         :type="edge.type"
+        :component="getType(edge.type)"
         :source="edge.source"
         :target="edge.target"
         :source-handle="edge.sourceHandle"
         :target-handle="edge.targetHandle"
         :selected="edge.selected"
+        :selectable="props.elementsSelectable"
         :source-position="edge.sourcePosition"
         :target-position="edge.targetPosition"
         :label="edge.label"
@@ -69,7 +102,7 @@ export default {
         <template #default="edgeProps">
           <slot :name="`edge-${edge.type}`" v-bind="edgeProps"></slot>
         </template>
-      </Edge>
+      </EdgeWrapper>
       <ConnectionLine
         v-if="connectionLineVisible && sourceNode"
         :source-node="sourceNode"
diff --git a/src/container/NodeRenderer/NodeRenderer.vue b/src/container/NodeRenderer/NodeRenderer.vue
index 7b960605..9e687b76 100644
--- a/src/container/NodeRenderer/NodeRenderer.vue
+++ b/src/container/NodeRenderer/NodeRenderer.vue
@@ -38,7 +38,7 @@ const snapGrid = computed(() => (props.snapToGrid ? props.snapGrid : undefined))
 </script>
 <script lang="ts">
 export default {
-  name: 'Node',
+  name: 'Nodes',
 }
 </script>
 <template>
diff --git a/src/container/VueFlow/VueFlow.vue b/src/container/VueFlow/VueFlow.vue
index 91a06e1e..e6991556 100644
--- a/src/container/VueFlow/VueFlow.vue
+++ b/src/container/VueFlow/VueFlow.vue
@@ -231,6 +231,15 @@ const transitionName = computed(() => {
                 <slot name="edge-renderer" v-bind="{ edges: store.getEdges, edgeTypes: store.getEdgeTypes }">
                   <EdgeRenderer
                     :key="`edge-renderer-${store.$id}`"
+                    :transform="store.transform"
+                    :dimensions="store.dimensions"
+                    :edges="store.getEdges"
+                    :edge-types="store.getEdgeTypes"
+                    :nodes="store.getNodes"
+                    :elements-selectable="store.elementsSelectable"
+                    :connection-node-id="store.connectionNodeId"
+                    :connection-handle-type="store.connectionHandleType"
+                    :nodes-connectable="store.nodesConnectable"
                     :connection-line-type="store.connectionLineType"
                     :connection-line-style="store.connectionLineStyle"
                     :arrow-head-color="store.arrowHeadColor"
diff --git a/src/container/ZoomPane/ZoomPane.vue b/src/container/ZoomPane/ZoomPane.vue
index 57a31b07..1f1eb480 100644
--- a/src/container/ZoomPane/ZoomPane.vue
+++ b/src/container/ZoomPane/ZoomPane.vue
@@ -176,10 +176,6 @@ invoke(async () => {
 })
 
 const { width, height } = useElementBounding(zoomPaneEl)
-store.dimensions = {
-  width: width.value,
-  height: height.value,
-}
 
 // skip waiting for ssr
 const window = useWindow()
@@ -191,7 +187,7 @@ watch(
     store.dimensions.width = newWidth
     store.dimensions.height = newHeight
   },
-  { flush: 'pre' },
+  { flush: 'sync', immediate: true },
 )
 watch(
   transform,
@@ -201,7 +197,7 @@ watch(
       clamp(val.y, store.translateExtent[0][0], store.translateExtent[1][0]),
       clamp(val.zoom, store.translateExtent[0][0], store.translateExtent[1][0]),
     ]),
-  { flush: 'pre' },
+  { flush: 'sync', immediate: true },
 )
 </script>
 <script lang="ts">
diff --git a/src/store/stateStore.ts b/src/store/stateStore.ts
index 0a3e6323..b6c757b9 100644
--- a/src/store/stateStore.ts
+++ b/src/store/stateStore.ts
@@ -1,5 +1,5 @@
-import { setActivePinia, createPinia, defineStore, StoreDefinition, acceptHMRUpdate } from 'pinia'
 import microDiff from 'microdiff'
+import { setActivePinia, createPinia, defineStore, StoreDefinition, acceptHMRUpdate } from 'pinia'
 import { FlowState, Node, FlowActions, Elements, FlowGetters, Edge, GraphNode, NextElements } from '~/types'
 import {
   clampPosition,
diff --git a/src/types/components.ts b/src/types/components.ts
index d5b97074..dfb434fb 100644
--- a/src/types/components.ts
+++ b/src/types/components.ts
@@ -58,12 +58,12 @@ export interface MiniMapNodeProps {
 }
 
 export interface EdgeTextProps {
-  x: number
-  y: number
+  x?: number
+  y?: number
   label?:
     | string
     | {
-        component: Component | DefineComponent
+        component: any
         props?: Record<string, any>
       }
   labelStyle?: CSSProperties
diff --git a/src/types/edge.ts b/src/types/edge.ts
index 35e76e74..eb6cbf0e 100644
--- a/src/types/edge.ts
+++ b/src/types/edge.ts
@@ -1,6 +1,6 @@
 import { Component, CSSProperties, DefineComponent } from 'vue'
 import { ArrowHeadType, ElementId, Position } from './flow'
-import { EdgeTextProps } from '~/types/components'
+import { EdgeTextProps } from './components'
 
 export interface Edge<T = any> {
   id: ElementId
@@ -15,8 +15,8 @@ export interface Edge<T = any> {
   label?:
     | string
     | {
-        component: DefineComponent<EdgeTextProps>
-        props?: EdgeTextProps
+        component: DefineComponent<EdgeTextProps> | any
+        props?: EdgeTextProps | any
       }
   labelStyle?: any
   labelShowBg?: boolean
diff --git a/src/types/node.ts b/src/types/node.ts
index 645edf53..66966235 100644
--- a/src/types/node.ts
+++ b/src/types/node.ts
@@ -33,6 +33,11 @@ export interface GraphNode<T = any> extends Node<T> {
   __vf: VFInternals
 }
 
+export type SourceTargetNode = {
+  sourceNode: GraphNode
+  targetNode: GraphNode
+}
+
 export type NodePosUpdate = {
   id: ElementId
   pos: XYPosition
diff --git a/src/utils/edge.ts b/src/utils/edge.ts
index 0a07e262..9b8e194c 100644
--- a/src/utils/edge.ts
+++ b/src/utils/edge.ts
@@ -1,5 +1,15 @@
 import { rectToBox } from './graph'
-import { Edge, EdgePositions, ElementId, GraphNode, HandleElement, Position, Transform, XYPosition } from '~/types'
+import {
+  Edge,
+  EdgePositions,
+  ElementId,
+  GraphNode,
+  HandleElement,
+  Position,
+  SourceTargetNode,
+  Transform,
+  XYPosition,
+} from '~/types'
 
 export function getHandlePosition(position: Position, node: GraphNode, handle: any | null = null): XYPosition {
   const x = (handle?.x || 0) + node.__vf.position.x
@@ -103,22 +113,18 @@ export function isEdgeVisible({ sourcePos, targetPos, width, height, transform }
   return overlappingArea > 0
 }
 
-type SourceTargetNode = {
-  sourceNode: GraphNode
-  targetNode: GraphNode
-}
-
 export const getSourceTargetNodes = (edge: Edge, nodes: GraphNode[]): SourceTargetNode => {
-  return nodes.reduce(
-    (res, node) => {
-      if (node.id === edge.source) {
-        res.sourceNode = node
-      }
-      if (node.id === edge.target) {
-        res.targetNode = node
-      }
-      return res
-    },
-    { sourceNode: null, targetNode: null } as any,
-  )
+  let { sourceNode, targetNode }: any = {
+    sourceNode: null,
+    targetNode: null,
+  }
+  for (const node of nodes) {
+    if (!sourceNode || !targetNode) {
+      if (node.id === edge.source) sourceNode = node
+      if (node.id === edge.target) targetNode = node
+    } else {
+      break
+    }
+  }
+  return { sourceNode, targetNode }
 }