From 0d6bc39e4920a3e54031cca6cc47b82df49d886d Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Tue, 29 Mar 2022 10:52:51 +0200
Subject: [PATCH] fix(nodes,edges): try to resolve component before falling
 back to default

---
 src/components/Edges/EdgeWrapper.vue        | 86 ++++++++++++---------
 src/components/Nodes/NodeWrapper.vue        | 86 ++++++++++++---------
 src/container/EdgeRenderer/EdgeRenderer.vue | 14 +---
 src/container/NodeRenderer/NodeRenderer.vue | 19 +----
 4 files changed, 104 insertions(+), 101 deletions(-)

diff --git a/src/components/Edges/EdgeWrapper.vue b/src/components/Edges/EdgeWrapper.vue
index bfea5782..fbf4ab12 100644
--- a/src/components/Edges/EdgeWrapper.vue
+++ b/src/components/Edges/EdgeWrapper.vue
@@ -1,47 +1,46 @@
 <script lang="ts" setup>
 import { useHandle, useVueFlow } from '../../composables'
-import { ConnectionMode, Position, GraphEdge, EdgeComponent } from '../../types'
+import { ConnectionMode, Position, EdgeComponent } from '../../types'
 import { getEdgePositions, getHandle, getMarkerId } from '../../utils'
 import EdgeAnchor from './EdgeAnchor.vue'
 
 interface EdgeWrapper {
-  edge: GraphEdge
-  component?: EdgeComponent
+  id: string
   selectable?: boolean
   updatable?: boolean
 }
 
 const props = defineProps<EdgeWrapper>()
-const emit = defineEmits(['update:edge'])
-const edge = useVModel(props, 'edge', emit)
 const { store } = useVueFlow()
+const edge = store.getEdge(props.id)
+if (!edge) throw new Error(`Edge with ${props.id} not found!`)
 
 const updating = ref(false)
 const { onMouseDown } = useHandle()
 
 const onEdgeClick = (event: MouseEvent) => {
-  const data = { event, edge: edge.value }
+  const data = { event, edge }
   if (props.selectable) {
     store.nodesSelectionActive = false
-    store.addSelectedEdges([edge.value])
+    store.addSelectedEdges([edge])
   }
   store.hooks.edgeClick.trigger(data)
 }
-const onEdgeContextMenu = (event: MouseEvent) => store.hooks.edgeContextMenu.trigger({ event, edge: edge.value })
-const onDoubleClick = (event: MouseEvent) => store.hooks.edgeDoubleClick.trigger({ event, edge: edge.value })
-const onEdgeMouseEnter = (event: MouseEvent) => store.hooks.edgeMouseEnter.trigger({ event, edge: edge.value })
-const onEdgeMouseMove = (event: MouseEvent) => store.hooks.edgeMouseMove.trigger({ event, edge: edge.value })
-const onEdgeMouseLeave = (event: MouseEvent) => store.hooks.edgeMouseLeave.trigger({ event, edge: edge.value })
+const onEdgeContextMenu = (event: MouseEvent) => store.hooks.edgeContextMenu.trigger({ event, edge })
+const onDoubleClick = (event: MouseEvent) => store.hooks.edgeDoubleClick.trigger({ event, edge })
+const onEdgeMouseEnter = (event: MouseEvent) => store.hooks.edgeMouseEnter.trigger({ event, edge })
+const onEdgeMouseMove = (event: MouseEvent) => store.hooks.edgeMouseMove.trigger({ event, edge })
+const onEdgeMouseLeave = (event: MouseEvent) => store.hooks.edgeMouseLeave.trigger({ event, 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 handleEdgeUpdater = (event: MouseEvent, isSourceHandle: boolean) => {
-  const nodeId = isSourceHandle ? edge.value.target : edge.value.source
-  const handleId = (isSourceHandle ? edge.value.targetHandle : edge.value.sourceHandle) ?? ''
+  const nodeId = isSourceHandle ? edge.target : edge.source
+  const handleId = (isSourceHandle ? edge.targetHandle : edge.sourceHandle) ?? ''
 
-  store.hooks.edgeUpdateStart.trigger({ event, edge: edge.value })
+  store.hooks.edgeUpdateStart.trigger({ event, edge })
   onMouseDown(
     event,
     handleId,
@@ -49,22 +48,22 @@ const handleEdgeUpdater = (event: MouseEvent, isSourceHandle: boolean) => {
     isSourceHandle,
     undefined,
     isSourceHandle ? 'target' : 'source',
-    (connection) => store.hooks.edgeUpdate.trigger({ edge: edge.value, connection }),
-    () => store.hooks.edgeUpdateEnd.trigger({ event, edge: edge.value }),
+    (connection) => store.hooks.edgeUpdate.trigger({ edge, connection }),
+    () => store.hooks.edgeUpdateEnd.trigger({ event, edge }),
   )
 }
 
 // when connection type is loose we can define all handles as sources
 const targetNodeHandles = computed(() =>
   store.connectionMode === ConnectionMode.Strict
-    ? edge.value.targetNode.handleBounds.target
-    : edge.value.targetNode.handleBounds.target ?? edge.value.targetNode.handleBounds.source,
+    ? edge.targetNode.handleBounds.target
+    : edge.targetNode.handleBounds.target ?? edge.targetNode.handleBounds.source,
 )
 const sourceHandle = controlledComputed(
-  () => edge.value.sourceNode.handleBounds,
-  () => getHandle(edge.value.sourceNode.handleBounds.source, edge.value.sourceHandle),
+  () => edge.sourceNode.handleBounds,
+  () => getHandle(edge.sourceNode.handleBounds.source, edge.sourceHandle),
 )
-const targetHandle = computed(() => getHandle(targetNodeHandles.value, edge.value.targetHandle))
+const targetHandle = computed(() => getHandle(targetNodeHandles.value, edge.targetHandle))
 const sourcePosition = computed(() => (sourceHandle.value ? sourceHandle.value.position : Position.Bottom))
 const targetPosition = computed(() => (targetHandle.value ? targetHandle.value.position : Position.Top))
 const edgeUpdaterRadius = computed(() => store.edgeUpdaterRadius)
@@ -74,30 +73,44 @@ 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,
+      () => edge.sourceNode.position,
+      () => edge.targetNode.position,
+      () => edge.sourceNode.computedPosition,
+      () => edge.targetNode.computedPosition,
+      () => edge.sourceNode.dimensions,
+      () => edge.targetNode.dimensions,
     ],
     () => {
       const { sourceX, sourceY, targetY, targetX } = getEdgePositions(
-        edge.value.sourceNode,
+        edge.sourceNode,
         sourceHandle.value,
         sourcePosition.value,
-        edge.value.targetNode,
+        edge.targetNode,
         targetHandle.value,
         targetPosition.value,
       )
-      if (edge.value.sourceX !== sourceX) edge.value.sourceX = sourceX
-      if (edge.value.sourceY !== sourceY) edge.value.sourceY = sourceY
-      if (edge.value.targetX !== targetX) edge.value.targetX = targetX
-      if (edge.value.targetY !== targetY) edge.value.targetY = targetY
+      if (edge.sourceX !== sourceX) edge.sourceX = sourceX
+      if (edge.sourceY !== sourceY) edge.sourceY = sourceY
+      if (edge.targetX !== targetX) edge.targetX = targetX
+      if (edge.targetY !== targetY) edge.targetY = targetY
     },
     { immediate: true, deep: true },
   )
 })
+
+const name = ref(edge.type ?? 'default')
+const type = computed(() => {
+  let edgeType = store.getEdgeTypes[name.value]
+  if (typeof edgeType === 'string') edgeType = resolveComponent(name.value, false) as EdgeComponent
+  if (typeof edgeType !== 'string') return edgeType
+
+  const slot = useSlots()?.[name.value]?.({})
+  if (!slot || !slot[0].key?.toString().includes(name.value)) {
+    console.warn(`Node type "${edge.type}" not found and no slot detected. Using fallback type "default".`)
+    name.value = 'default'
+    return store.getEdgeTypes.default
+  }
+})
 </script>
 <script lang="ts">
 export default {
@@ -109,7 +122,7 @@ export default {
     :key="`edge-${edge.id}`"
     :class="[
       'vue-flow__edge',
-      `vue-flow__edge-${edge.type || 'default'}`,
+      `vue-flow__edge-${name}`,
       store.noPanClassName,
       {
         selected: edge.selected,
@@ -157,7 +170,8 @@ export default {
       }"
     >
       <component
-        :is="props.component ?? edge.type"
+        :is="type"
+        v-if="type"
         v-bind="{
           id: edge.id,
           sourceNode: edge.sourceNode,
diff --git a/src/components/Nodes/NodeWrapper.vue b/src/components/Nodes/NodeWrapper.vue
index 72e5c4d9..765af292 100644
--- a/src/components/Nodes/NodeWrapper.vue
+++ b/src/components/Nodes/NodeWrapper.vue
@@ -1,13 +1,12 @@
 <script lang="ts" setup>
 import { DraggableEventListener, DraggableCore } from '@braks/revue-draggable'
 import { useVueFlow } from '../../composables'
-import { GraphNode, NodeComponent, SnapGrid } from '../../types'
+import { NodeComponent, SnapGrid } from '../../types'
 import { NodeId } from '../../context'
 import { getXYZPos } from '../../utils'
 
 interface NodeWrapperProps {
-  node: GraphNode
-  type: NodeComponent
+  id: string
   draggable: boolean
   selectable: boolean
   connectable: boolean
@@ -15,74 +14,72 @@ interface NodeWrapperProps {
 }
 
 const props = defineProps<NodeWrapperProps>()
-const node = useVModel(props, 'node')
 const { store } = useVueFlow()
-provide(NodeId, props.node.id)
+const node = store.getNode(props.id)
+if (!node) throw new Error(`Node with ${props.id} not found!`)
+provide(NodeId, props.id)
 
 const nodeElement = templateRef<HTMLDivElement>('node-element', null)
 
-const onMouseEnterHandler = () =>
-  node.value.dragging && ((event: MouseEvent) => store.hooks.nodeMouseEnter.trigger({ event, node: node.value }))
-const onMouseMoveHandler = () =>
-  node.value.dragging && ((event: MouseEvent) => store.hooks.nodeMouseMove.trigger({ event, node: node.value }))
-const onMouseLeaveHandler = () =>
-  node.value.dragging && ((event: MouseEvent) => store.hooks.nodeMouseLeave.trigger({ event, node: node.value }))
+const onMouseEnterHandler = () => node.dragging && ((event: MouseEvent) => store.hooks.nodeMouseEnter.trigger({ event, node }))
+const onMouseMoveHandler = () => node.dragging && ((event: MouseEvent) => store.hooks.nodeMouseMove.trigger({ event, node }))
+const onMouseLeaveHandler = () => node.dragging && ((event: MouseEvent) => store.hooks.nodeMouseLeave.trigger({ event, node }))
 const onContextMenuHandler = () => (event: MouseEvent) =>
   store.hooks.nodeContextMenu.trigger({
     event,
-    node: node.value,
+    node,
   })
-const onDoubleClick = () => (event: MouseEvent) => store.hooks.nodeDoubleClick.trigger({ event, node: node.value })
+const onDoubleClick = () => (event: MouseEvent) => store.hooks.nodeDoubleClick.trigger({ event, node })
 const onSelectNodeHandler = (event: MouseEvent) => {
   if (!props.draggable) {
     if (props.selectable) {
       store.nodesSelectionActive = false
-      if (!node.value.selected) store.addSelectedNodes([node.value])
+      if (!node.selected) store.addSelectedNodes([node])
     }
-    store.hooks.nodeClick.trigger({ event, node: node.value })
+    store.hooks.nodeClick.trigger({ event, node })
   }
 }
 const onDragStart: DraggableEventListener = ({ event }) => {
   store.addSelectedNodes([])
-  store.hooks.nodeDragStart.trigger({ event, node: node.value })
+  store.hooks.nodeDragStart.trigger({ event, node })
 
   if (store.selectNodesOnDrag && props.selectable) {
     store.nodesSelectionActive = false
 
-    if (!node.value.selected) store.addSelectedNodes([node.value])
-  } else if (!store.selectNodesOnDrag && !node.value.selected && props.selectable) {
+    if (!node.selected) store.addSelectedNodes([node])
+  } else if (!store.selectNodesOnDrag && !node.selected && props.selectable) {
     store.nodesSelectionActive = false
     store.addSelectedNodes([])
   }
 }
 const onDrag: DraggableEventListener = ({ event, data: { deltaX, deltaY } }) => {
-  nextTick(() => store.updateNodePosition({ id: node.value.id, diff: { x: deltaX, y: deltaY }, dragging: true }))
-  store.hooks.nodeDrag.trigger({ event, node: node.value })
+  nextTick(() => store.updateNodePosition({ id: node.id, diff: { x: deltaX, y: deltaY }, dragging: true }))
+  store.hooks.nodeDrag.trigger({ event, node })
 }
 const onDragStop: DraggableEventListener = ({ event, data: { deltaX, deltaY } }) => {
   // 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 (!node.value.dragging) {
-    if (props.selectable && !store.selectNodesOnDrag && !node.value.selected) {
-      store.addSelectedNodes([node.value])
+  if (!node.dragging) {
+    if (props.selectable && !store.selectNodesOnDrag && !node.selected) {
+      store.addSelectedNodes([node])
     }
-    store.hooks.nodeClick.trigger({ event, node: node.value })
+    store.hooks.nodeClick.trigger({ event, node })
     return
   }
-  store.updateNodePosition({ id: node.value.id, diff: { x: deltaX, y: deltaY }, dragging: false })
-  store.hooks.nodeDragStop.trigger({ event, node: node.value })
+  store.updateNodePosition({ id: node.id, diff: { x: deltaX, y: deltaY }, dragging: false })
+  store.hooks.nodeDragStop.trigger({ event, node })
 }
 
 onMounted(() => {
   useResizeObserver(nodeElement, () =>
-    store.updateNodeDimensions([{ id: node.value.id, nodeElement: nodeElement.value, forceUpdate: true }]),
+    store.updateNodeDimensions([{ id: node.id, nodeElement: nodeElement.value, forceUpdate: true }]),
   )
 
-  watch([() => node.value.type, () => node.value.sourcePosition, () => node.value.targetPosition], () =>
-    nextTick(() => store.updateNodeDimensions([{ id: node.value.id, nodeElement: nodeElement.value }])),
+  watch([() => node.type, () => node.sourcePosition, () => node.targetPosition], () =>
+    nextTick(() => store.updateNodeDimensions([{ id: node.id, nodeElement: nodeElement.value }])),
   )
 
-  store.updateNodeDimensions([{ id: node.value.id, nodeElement: nodeElement.value }])
+  store.updateNodeDimensions([{ id: node.id, nodeElement: nodeElement.value }])
 })
 
 const scale = controlledComputed(
@@ -91,22 +88,36 @@ const scale = controlledComputed(
 )
 
 watch(
-  [() => node.value.position, () => store.getNode(node.value.parentNode!)],
+  [() => node.position, () => store.getNode(node.parentNode!)],
   ([pos, parent]) => {
     const xyzPos = {
       ...pos,
-      z: node.value.dragging || node.value.selected ? 1000 : node.value.computedPosition.z,
+      z: node.dragging || node.selected ? 1000 : node.computedPosition.z,
     }
     if (parent) {
-      node.value.computedPosition = getXYZPos(parent, xyzPos)
+      node.computedPosition = getXYZPos(parent, xyzPos)
     } else {
-      node.value.computedPosition = xyzPos
+      node.computedPosition = xyzPos
     }
   },
   { deep: true },
 )
 
-store.updateNodePosition({ id: node.value.id, diff: { x: 0, y: 0 } })
+store.updateNodePosition({ id: node.id, diff: { x: 0, y: 0 } })
+
+const name = ref(node.type ?? 'default')
+const type = computed(() => {
+  let nodeType = store.getNodeTypes[name.value]
+  if (typeof nodeType === 'string') nodeType = resolveComponent(name.value, false) as NodeComponent
+  if (typeof nodeType !== 'string') return nodeType
+
+  const slot = useSlots()?.[name.value]?.({})
+  if (!slot || !slot[0].key?.toString().includes(name.value)) {
+    console.warn(`Node type "${node.type}" not found and no slot detected. Using fallback type "default".`)
+    name.value = 'default'
+    return store.getNodeTypes.default
+  }
+})
 </script>
 <script lang="ts">
 export default {
@@ -130,7 +141,7 @@ export default {
       :key="`node-${node.id}`"
       :class="[
         'vue-flow__node',
-        `vue-flow__node-${node.type}`,
+        `vue-flow__node-${name}`,
         store.noPanClassName,
         {
           dragging: node.dragging,
@@ -176,7 +187,8 @@ export default {
         }"
       >
         <component
-          :is="props.type ?? node.type"
+          :is="type"
+          v-if="type"
           v-bind="{
             nodeElement,
             id: node.id,
diff --git a/src/container/EdgeRenderer/EdgeRenderer.vue b/src/container/EdgeRenderer/EdgeRenderer.vue
index 366f06d8..e8cbe1a9 100644
--- a/src/container/EdgeRenderer/EdgeRenderer.vue
+++ b/src/container/EdgeRenderer/EdgeRenderer.vue
@@ -6,15 +6,6 @@ import { groupEdgesByZLevel } from '../../utils'
 import MarkerDefinitions from './MarkerDefinitions.vue'
 
 const { store } = useVueFlow()
-const getType = (type?: string) => {
-  const t = type ?? 'default'
-  let edgeType = store.getEdgeTypes[t]
-  if (!edgeType) {
-    edgeType = store.getEdgeTypes.default
-    console.warn(`Edge type "${type}" not found. Using fallback type "default".`)
-  }
-  return edgeType
-}
 
 const sourceNode = controlledComputed(
   () => store.connectionNodeId,
@@ -52,14 +43,13 @@ export default {
     <g>
       <EdgeWrapper
         v-for="edge of group.edges"
+        :id="edge.id"
         :key="edge.id"
-        :edge="edge"
-        :component="getType(edge.type)"
         :selectable="typeof edge.selectable === 'undefined' ? store.elementsSelectable : edge.selectable"
         :updatable="typeof edge.updatable === 'undefined' ? store.edgesUpdatable : edge.updatable"
       >
         <template #default="edgeProps">
-          <slot :name="`edge-${edge.type}`" v-bind="edgeProps"></slot>
+          <slot v-if="edge.type" :name="`edge-${edge.type}`" v-bind="edgeProps"></slot>
         </template>
       </EdgeWrapper>
       <ConnectionLine v-if="connectionLineVisible && sourceNode" :source-node="sourceNode">
diff --git a/src/container/NodeRenderer/NodeRenderer.vue b/src/container/NodeRenderer/NodeRenderer.vue
index 4c0d4e6b..bbe8e380 100644
--- a/src/container/NodeRenderer/NodeRenderer.vue
+++ b/src/container/NodeRenderer/NodeRenderer.vue
@@ -1,20 +1,8 @@
 <script lang="ts" setup>
-import { GraphNode } from '../../types'
 import NodeWrapper from '../../components/Nodes/NodeWrapper.vue'
 import { useVueFlow } from '../../composables'
 
 const { store } = useVueFlow()
-const getType = (type?: string) => {
-  const t = type ?? 'default'
-  let nodeType = store.getNodeTypes[t]
-  if (!nodeType) {
-    nodeType = store.getNodeTypes.default
-    console.warn(`Node type "${type}" not found. Using fallback type "default".`)
-  }
-  return nodeType
-}
-
-const snapGrid = (node: GraphNode) => node.snapGrid ?? (store.snapToGrid ? store.snapGrid : undefined)
 </script>
 <script lang="ts">
 export default {
@@ -25,16 +13,15 @@ export default {
   <div class="vue-flow__nodes vue-flow__container">
     <NodeWrapper
       v-for="node of store.getNodes"
+      :id="node.id"
       :key="node.id"
-      :node="node"
-      :type="getType(node.type)"
       :draggable="typeof node.draggable === 'undefined' ? store.nodesDraggable : !!node.draggable"
       :selectable="typeof node.selectable === 'undefined' ? store.elementsSelectable : !!node.selectable"
       :connectable="typeof node.connectable === 'undefined' ? store.nodesConnectable : !!node.connectable"
-      :snap-grid="snapGrid(node)"
+      :snap-grid="node.snapGrid ?? (store.snapToGrid ? store.snapGrid : undefined)"
     >
       <template #default="nodeProps">
-        <slot :name="`node-${node.type}`" v-bind="nodeProps"></slot>
+        <slot v-if="node.type" :name="`node-${node.type}`" v-bind="nodeProps"></slot>
       </template>
     </NodeWrapper>
   </div>