From ee38f4af9be69220c72d72d2d2e5078dff6ea7e6 Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Fri, 15 Apr 2022 20:54:08 +0200
Subject: [PATCH] fix(edges): Remove orphaned edges

# What's changed?

* Remove source/targetNode properties from GraphEdge type, the source/target will be fetched with getters if necessary
* Remove edges when they have no existing source or target nodes that can be found in the state
* Add nodes to custom edge test
---
 .../integration/1-store/setState.spec.ts      | 10 +++
 package/src/components/Edges/EdgeWrapper.vue  | 36 ++++-----
 package/src/components/Nodes/NodeWrapper.vue  |  2 +-
 .../container/EdgeRenderer/EdgeRenderer.vue   |  8 +-
 .../container/NodeRenderer/NodeRenderer.vue   |  2 +-
 package/src/store/actions.ts                  | 13 ++--
 package/src/store/getters.ts                  | 75 ++++++++++++-------
 package/src/types/edge.ts                     |  4 +-
 package/src/utils/changes.ts                  | 17 +++--
 package/src/utils/edge.ts                     | 17 ++---
 10 files changed, 111 insertions(+), 73 deletions(-)

diff --git a/e2e/cypress/integration/1-store/setState.spec.ts b/e2e/cypress/integration/1-store/setState.spec.ts
index 5077333b..8450e5e0 100644
--- a/e2e/cypress/integration/1-store/setState.spec.ts
+++ b/e2e/cypress/integration/1-store/setState.spec.ts
@@ -46,6 +46,16 @@ describe('test store state', () => {
 
   it('gets custom edge types', () => {
     store.setState({
+      nodes: [
+        {
+          id: '1',
+          position: { x: 0, y: 0 },
+        },
+        {
+          id: '2',
+          position: { x: 50, y: 50 },
+        },
+      ],
       edges: [
         {
           id: '1',
diff --git a/package/src/components/Edges/EdgeWrapper.vue b/package/src/components/Edges/EdgeWrapper.vue
index ed8c55e2..559478f2 100644
--- a/package/src/components/Edges/EdgeWrapper.vue
+++ b/package/src/components/Edges/EdgeWrapper.vue
@@ -1,7 +1,7 @@
 <script lang="ts" setup>
 import { CSSProperties } from 'vue'
 import { useHandle, useVueFlow } from '../../composables'
-import { ConnectionMode, EdgeComponent, GraphEdge, Position } from '../../types'
+import { ConnectionMode, EdgeComponent, GraphEdge, GraphNode, Position } from '../../types'
 import { getEdgePositions, getHandle, getMarkerId } from '../../utils'
 import { Slots } from '../../context'
 import EdgeAnchor from './EdgeAnchor.vue'
@@ -9,6 +9,8 @@ import EdgeAnchor from './EdgeAnchor.vue'
 interface EdgeWrapper {
   id: string
   edge: GraphEdge
+  sourceNode: GraphNode
+  targetNode: GraphNode
   selectable?: boolean
   updatable?: boolean
 }
@@ -59,21 +61,21 @@ const handleEdgeUpdater = (event: MouseEvent, isSourceHandle: boolean) => {
 // when connection type is loose we can define all handles as sources
 const targetNodeHandles = computed(() => {
   if (store.connectionMode === ConnectionMode.Strict) {
-    return edge.value.targetNode.handleBounds.target
+    return props.targetNode.handleBounds.target
   }
 
-  const targetBounds = edge.value.targetNode.handleBounds.target
-  const sourceBounds = edge.value.targetNode.handleBounds.source
+  const targetBounds = props.targetNode.handleBounds.target
+  const sourceBounds = props.targetNode.handleBounds.source
   return targetBounds ?? sourceBounds
 })
 
 const sourceNodeHandles = computed(() => {
   if (store.connectionMode === ConnectionMode.Strict) {
-    return edge.value.sourceNode.handleBounds.source
+    return props.sourceNode.handleBounds.source
   }
 
-  const targetBounds = edge.value.sourceNode.handleBounds.target
-  const sourceBounds = edge.value.sourceNode.handleBounds.source
+  const targetBounds = props.sourceNode.handleBounds.target
+  const sourceBounds = props.sourceNode.handleBounds.source
   return sourceBounds ?? targetBounds
 })
 
@@ -90,19 +92,19 @@ onMounted(() => {
     [
       sourcePosition,
       targetPosition,
-      () => edge.value.sourceNode.position,
-      () => edge.value.targetNode.position,
-      () => edge.value.sourceNode.computedPosition,
-      () => edge.value.targetNode.computedPosition,
-      () => edge.value.sourceNode.dimensions,
-      () => edge.value.targetNode.dimensions,
+      () => props.sourceNode.position,
+      () => props.targetNode.position,
+      () => props.sourceNode.computedPosition,
+      () => props.targetNode.computedPosition,
+      () => props.sourceNode.dimensions,
+      () => props.targetNode.dimensions,
     ],
     () => {
       const { sourceX, sourceY, targetY, targetX } = getEdgePositions(
-        edge.value.sourceNode,
+        props.sourceNode,
         sourceHandle.value,
         sourcePosition.value,
-        edge.value.targetNode,
+        props.targetNode,
         targetHandle.value,
         targetPosition.value,
       )
@@ -185,8 +187,8 @@ export default {
     <component
       :is="type"
       :id="edge.id"
-      :source-node="edge.sourceNode"
-      :target-node="edge.targetNode"
+      :source-node="props.sourceNode"
+      :target-node="props.targetNode"
       :source="edge.source"
       :target="edge.target"
       :updatable="props.updatable"
diff --git a/package/src/components/Nodes/NodeWrapper.vue b/package/src/components/Nodes/NodeWrapper.vue
index ca804cd0..893af955 100644
--- a/package/src/components/Nodes/NodeWrapper.vue
+++ b/package/src/components/Nodes/NodeWrapper.vue
@@ -1,6 +1,6 @@
 <script lang="ts" setup>
 import { useDraggableCore } from '@braks/revue-draggable'
-import { CSSProperties, isVNode } from 'vue'
+import { CSSProperties } from 'vue'
 import { useVueFlow } from '../../composables'
 import { GraphNode, NodeComponent, SnapGrid } from '../../types'
 import { NodeId, Slots } from '../../context'
diff --git a/package/src/container/EdgeRenderer/EdgeRenderer.vue b/package/src/container/EdgeRenderer/EdgeRenderer.vue
index 7dbeca3a..071d13c1 100644
--- a/package/src/container/EdgeRenderer/EdgeRenderer.vue
+++ b/package/src/container/EdgeRenderer/EdgeRenderer.vue
@@ -23,7 +23,11 @@ const connectionLineVisible = controlledComputed(
       store.connectionHandleType
     ),
 )
-const groups = computed(() => groupEdgesByZLevel(store.getEdges, store.getNodes))
+
+const getNode = computed(() => (node: string) => store.getNode(node)!)
+
+const memoizedGroups = useMemoize(groupEdgesByZLevel)
+const groups = computed(() => memoizedGroups(store.getEdges, getNode.value))
 </script>
 <script lang="ts">
 export default {
@@ -39,6 +43,8 @@ export default {
         :id="edge.id"
         :key="edge.id"
         :edge="edge"
+        :source-node="getNode(edge.source)"
+        :target-node="getNode(edge.target)"
         :selectable="typeof edge.selectable === 'undefined' ? store.elementsSelectable : edge.selectable"
         :updatable="typeof edge.updatable === 'undefined' ? store.edgesUpdatable : edge.updatable"
       />
diff --git a/package/src/container/NodeRenderer/NodeRenderer.vue b/package/src/container/NodeRenderer/NodeRenderer.vue
index 8f8765fb..f169a38d 100644
--- a/package/src/container/NodeRenderer/NodeRenderer.vue
+++ b/package/src/container/NodeRenderer/NodeRenderer.vue
@@ -14,7 +14,7 @@ export default {
     <NodeWrapper
       v-for="node of store.getNodes"
       :id="node.id"
-      :key="`vue-flow__node-${node.id}`"
+      :key="node.id"
       :node="node"
       :draggable="typeof node.draggable === 'undefined' ? store.nodesDraggable : !!node.draggable"
       :selectable="typeof node.selectable === 'undefined' ? store.elementsSelectable : !!node.selectable"
diff --git a/package/src/store/actions.ts b/package/src/store/actions.ts
index 35de54ed..09514f8a 100644
--- a/package/src/store/actions.ts
+++ b/package/src/store/actions.ts
@@ -16,8 +16,7 @@ import {
   ComputedGetters,
 } from '~/types'
 import {
-  applyNodeChanges as applyNodes,
-  applyEdgeChanges as applyEdges,
+  applyChanges,
   connectionExists,
   createPositionChange,
   createSelectionChange,
@@ -241,6 +240,7 @@ export default (state: State, getters: ComputedGetters): Actions => {
       const missingTarget = !targetNode || typeof targetNode === 'undefined'
       if (missingSource) console.warn(`[vueflow]: Couldn't create edge for source id: ${edge.source}; edge id: ${edge.id}`)
       if (missingTarget) console.warn(`[vueflow]: Couldn't create edge for target id: ${edge.target}; edge id: ${edge.id}`)
+      if (missingSource || missingTarget) return res
 
       const storedEdge = getters.getEdge.value(edge.id)
 
@@ -249,8 +249,6 @@ export default (state: State, getters: ComputedGetters): Actions => {
           ...state.defaultEdgeOptions,
           ...storedEdge,
         }),
-        sourceNode,
-        targetNode,
       })
 
       return res
@@ -284,12 +282,11 @@ export default (state: State, getters: ComputedGetters): Actions => {
         const missingTarget = !targetNode || typeof targetNode === 'undefined'
         if (missingSource) console.warn(`[vueflow]: Couldn't create edge for source id: ${edge.source}; edge id: ${edge.id}`)
         if (missingTarget) console.warn(`[vueflow]: Couldn't create edge for target id: ${edge.target}; edge id: ${edge.id}`)
+        if (missingTarget || missingSource) return
 
         state.edges.push({
           ...state.defaultEdgeOptions,
           ...edge,
-          sourceNode,
-          targetNode,
         })
       }
     })
@@ -298,9 +295,9 @@ export default (state: State, getters: ComputedGetters): Actions => {
   const updateEdge: Actions['updateEdge'] = (oldEdge, newConnection) =>
     updateEdgeAction(oldEdge, newConnection, state.edges, addEdges)
 
-  const applyNodeChanges: Actions['applyNodeChanges'] = (changes) => applyNodes(changes, state.nodes)
+  const applyNodeChanges: Actions['applyNodeChanges'] = (changes) => applyChanges(changes, state.nodes, addNodes)
 
-  const applyEdgeChanges: Actions['applyEdgeChanges'] = (changes) => applyEdges(changes, state.edges)
+  const applyEdgeChanges: Actions['applyEdgeChanges'] = (changes) => applyChanges(changes, state.edges, addEdges)
 
   const setState: Actions['setState'] = (options) => {
     const skip = ['modelValue', 'nodes', 'edges', 'maxZoom', 'minZoom', 'translateExtent']
diff --git a/package/src/store/getters.ts b/package/src/store/getters.ts
index 20e4178c..ee5d2ad0 100644
--- a/package/src/store/getters.ts
+++ b/package/src/store/getters.ts
@@ -3,6 +3,11 @@ import { State, GraphEdge, GraphNode, ComputedGetters } from '~/types'
 import { getNodesInside, isEdgeVisible } from '~/utils'
 
 export default (state: State): ComputedGetters => {
+  const nodeIds = computed(() => state.nodes.map((n) => n.id))
+  const edgeIds = computed(() => state.edges.map((e) => e.id))
+  const getNode: ComputedGetters['getNode'] = computed(() => (id: string) => state.nodes[nodeIds.value.indexOf(id)])
+  const getEdge: ComputedGetters['getEdge'] = computed(() => (id: string) => state.edges[edgeIds.value.indexOf(id)])
+
   const getEdgeTypes = computed(() => {
     const edgeTypes: Record<string, any> = {
       ...defaultEdgeTypes,
@@ -41,31 +46,50 @@ export default (state: State): ComputedGetters => {
       : nodes ?? []
   })
 
+  const edgeHidden = (e: GraphEdge) => {
+    const source = getNode.value(e.source)
+    const target = getNode.value(e.target)
+
+    if (!source || !target) {
+      console.warn(`[vue-flow]: Orphaned edge ${e.id} will be removed.`)
+      state.edges.splice(state.edges.indexOf(e), 1)
+      return true
+    }
+
+    return (
+      !e.hidden &&
+      target &&
+      !target.hidden &&
+      source &&
+      !source.hidden &&
+      source.dimensions.width &&
+      source.dimensions.height &&
+      target.dimensions.width &&
+      target.dimensions.height
+    )
+  }
   const getEdges = computed<GraphEdge[]>(() => {
-    if (!state.onlyRenderVisibleElements)
-      return state.edges.filter((e) => !e.hidden && e.targetNode && !e.targetNode.hidden && e.sourceNode && !e.sourceNode.hidden)
-    else
-      return state.edges.filter(
-        (e) =>
-          !e.hidden &&
-          e.sourceNode.dimensions.width &&
-          e.sourceNode.dimensions.height &&
-          e.targetNode.dimensions.width &&
-          e.targetNode.dimensions.height &&
-          !e.targetNode.hidden &&
-          !e.sourceNode.hidden &&
-          isEdgeVisible({
-            sourcePos: e.sourceNode.computedPosition || { x: 0, y: 0 },
-            targetPos: e.targetNode.computedPosition || { x: 0, y: 0 },
-            sourceWidth: e.sourceNode.dimensions.width,
-            sourceHeight: e.sourceNode.dimensions.height,
-            targetWidth: e.targetNode.dimensions.width,
-            targetHeight: e.targetNode.dimensions.height,
-            width: state.dimensions.width,
-            height: state.dimensions.height,
-            viewport: state.viewport,
-          }),
+    if (!state.onlyRenderVisibleElements) return state.edges.filter(edgeHidden)
+
+    return state.edges.filter((e) => {
+      const source = getNode.value(e.source)!
+      const target = getNode.value(e.target)!
+
+      return (
+        edgeHidden(e) &&
+        isEdgeVisible({
+          sourcePos: source.computedPosition || { x: 0, y: 0 },
+          targetPos: target.computedPosition || { x: 0, y: 0 },
+          sourceWidth: source.dimensions.width,
+          sourceHeight: source.dimensions.height,
+          targetWidth: target.dimensions.width,
+          targetHeight: target.dimensions.height,
+          width: state.dimensions.width,
+          height: state.dimensions.height,
+          viewport: state.viewport,
+        })
       )
+    })
   })
 
   const getSelectedNodes: ComputedGetters['getSelectedNodes'] = computed(() => state.nodes.filter((n) => n.selected))
@@ -75,11 +99,6 @@ export default (state: State): ComputedGetters => {
     ...(getSelectedEdges.value ?? []),
   ])
 
-  const nodeIds = computed(() => state.nodes.map((n) => n.id))
-  const edgeIds = computed(() => state.edges.map((e) => e.id))
-  const getNode: ComputedGetters['getNode'] = computed(() => (id: string) => state.nodes[nodeIds.value.indexOf(id)])
-  const getEdge: ComputedGetters['getEdge'] = computed(() => (id: string) => state.edges[edgeIds.value.indexOf(id)])
-
   return {
     getNode,
     getEdge,
diff --git a/package/src/types/edge.ts b/package/src/types/edge.ts
index 02be03cb..07ed768b 100644
--- a/package/src/types/edge.ts
+++ b/package/src/types/edge.ts
@@ -94,9 +94,7 @@ export interface EdgePositions {
 }
 
 /** Internal edge type */
-export type GraphEdge<Data = ElementData, SourceNodeData = any, TargetNodeData = SourceNodeData> = Edge<Data> & {
-  sourceNode: GraphNode<SourceNodeData>
-  targetNode: GraphNode<TargetNodeData>
+export type GraphEdge<Data = ElementData> = Edge<Data> & {
   selected?: boolean
   z?: number
 } & EdgePositions
diff --git a/package/src/utils/changes.ts b/package/src/utils/changes.ts
index dd39eb1f..83b2d00c 100644
--- a/package/src/utils/changes.ts
+++ b/package/src/utils/changes.ts
@@ -1,9 +1,8 @@
-import { clampPosition, isGraphNode } from './graph'
+import { clampPosition, isGraphEdge, isGraphNode } from './graph'
 import {
   EdgeChange,
   EdgeSelectionChange,
   ElementChange,
-  FlowElement,
   FlowElements,
   GraphNode,
   NodeChange,
@@ -13,6 +12,9 @@ import {
   CoordinateExtent,
   XYPosition,
   GraphEdge,
+  Node,
+  FlowElement,
+  Edge,
 } from '~/types'
 
 type CreatePositionChangeParams = {
@@ -90,20 +92,25 @@ function handleParentExpand(updateItem: GraphNode, curr: GraphNode[]) {
 }
 
 export const applyChanges = <
-  T extends FlowElement = GraphNode,
+  T extends Node | Edge | FlowElement = Node,
   C extends ElementChange = T extends GraphNode ? NodeChange : EdgeChange,
 >(
   changes: C[],
   elements: T[],
+  addElement?: (els: T[]) => void,
 ): T[] => {
   let elementIds = elements.map((el) => el.id)
   changes.forEach((change) => {
-    if (change.type === 'add') return elements.push(change.item as any)
+    if (change.type === 'add') {
+      if (addElement) addElement([change.item as any])
+      else elements.push(change.item as any)
+    }
+
     const i = elementIds.indexOf((<any>change).id)
     const el = elements[i]
     switch (change.type) {
       case 'select':
-        el.selected = change.selected
+        if (isGraphNode(el) || isGraphEdge(el)) el.selected = change.selected
         break
       case 'position':
         if (isGraphNode(el)) {
diff --git a/package/src/utils/edge.ts b/package/src/utils/edge.ts
index 231d53da..684f4dc4 100644
--- a/package/src/utils/edge.ts
+++ b/package/src/utils/edge.ts
@@ -1,5 +1,5 @@
 import { rectToBox } from './graph'
-import { EdgePositions, GraphEdge, GraphNode, HandleElement, Position, Rect, Viewport, XYPosition } from '~/types'
+import { EdgePositions, Getters, GraphEdge, GraphNode, HandleElement, Position, Rect, Viewport, XYPosition } from '~/types'
 
 export const getHandlePosition = (position: Position, rect: Rect, handle?: HandleElement): XYPosition => {
   const x = (handle?.x ?? 0) + rect.x
@@ -126,17 +126,16 @@ export function isEdgeVisible({
   return overlappingArea > 0
 }
 
-export const groupEdgesByZLevel = (edges: GraphEdge[], nodes: GraphNode[]) => {
+export const groupEdgesByZLevel = (edges: GraphEdge[], getNode: Getters['getNode']) => {
   let maxLevel = -1
-  const nodeIds = nodes.map((n) => n.id)
 
   const levelLookup = edges.reduce<Record<string, GraphEdge[]>>((tree, edge) => {
-    const z = edge.z
-      ? edge.z
-      : Math.max(
-          nodes[nodeIds.indexOf(edge.source)]?.computedPosition.z || 0,
-          nodes[nodeIds.indexOf(edge.target)]?.computedPosition.z || 0,
-        )
+    const source = getNode(edge.source)
+    const target = getNode(edge.target)
+
+    if (!source || !target) return tree
+
+    const z = edge.z ? edge.z : Math.max(source.computedPosition.z || 0, target.computedPosition.z || 0)
     if (tree[z]) {
       tree[z].push(edge)
     } else {