From 3f15983f3efab984c0c2819d5aee765804255329 Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Wed, 30 Mar 2022 10:25:55 +0200
Subject: [PATCH] feat(edges): Use slot injection for custom edges

---
 package/src/components/Edges/EdgeWrapper.vue  | 77 ++++++++-----------
 package/src/components/Nodes/NodeWrapper.vue  |  3 +-
 .../container/EdgeRenderer/EdgeRenderer.vue   |  6 +-
 .../container/NodeRenderer/NodeRenderer.vue   |  6 +-
 package/src/index.ts                          |  5 +-
 5 files changed, 40 insertions(+), 57 deletions(-)

diff --git a/package/src/components/Edges/EdgeWrapper.vue b/package/src/components/Edges/EdgeWrapper.vue
index 43e407e4..1f11e04e 100644
--- a/package/src/components/Edges/EdgeWrapper.vue
+++ b/package/src/components/Edges/EdgeWrapper.vue
@@ -2,6 +2,7 @@
 import { useHandle, useVueFlow } from '../../composables'
 import { ConnectionMode, Position, EdgeComponent } from '../../types'
 import { getEdgePositions, getHandle, getMarkerId } from '../../utils'
+import { Slots } from '../../context'
 import EdgeAnchor from './EdgeAnchor.vue'
 
 interface EdgeWrapper {
@@ -98,18 +99,38 @@ onMounted(() => {
   )
 })
 
+const slots = inject(Slots)
+
 const name = ref(edge.value.type ?? 'default')
+watch(
+  () => edge.value.type,
+  (v) => v && (name.value = v),
+)
+
 const type = computed(() => {
+  const instance = getCurrentInstance()
   let edgeType = store.getEdgeTypes[name.value]
-  if (typeof edgeType === 'string') edgeType = resolveComponent(name.value, false) as EdgeComponent
+  if (typeof edgeType === 'string') {
+    if (instance) {
+      const components = Object.keys(instance.appContext.components)
+      if (components && components.includes(name.value)) {
+        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.value.type}" not found and no slot detected. Using fallback type "default".`)
+  const slot = slots?.[`edge-${name.value}`]?.({})
+  if (!slot) {
+    console.warn(`[vueflow]: Edge type "${edge.value.type}" not found and no edge-slot detected. Using fallback type "default".`)
     name.value = 'default'
     return store.getEdgeTypes.default
   }
+  if (slot.length > 1) {
+    console.warn('[vue-flow]: More than one element in edge-slot detected. Using fallback to first slot item.')
+  }
+
+  return slot[0]
 })
 </script>
 <script lang="ts">
@@ -139,16 +160,17 @@ export default {
     @mousemove="onEdgeMouseMove"
     @mouseleave="onEdgeMouseLeave"
   >
-    <slot
+    <component
+      :is="type"
       v-bind="{
         id: edge.id,
         sourceNode: edge.sourceNode,
         targetNode: edge.targetNode,
         source: edge.source,
         target: edge.target,
+        updatable: props.updatable,
         selected: edge.selected,
         animated: edge.animated,
-        updatable: props.updatable,
         label: edge.label,
         labelStyle: edge.labelStyle,
         labelShowBg: edge.labelShowBg,
@@ -157,51 +179,18 @@ export default {
         labelBgBorderRadius: edge.labelBgBorderRadius,
         data: edge.data,
         style: edge.style,
-        sourceX: edge.sourceX,
-        sourceY: edge.sourceY,
-        targetX: edge.targetX,
-        targetY: edge.targetY,
         markerStart: `url(#${getMarkerId(edge.markerStart)})`,
         markerEnd: `url(#${getMarkerId(edge.markerEnd)})`,
         sourcePosition,
         targetPosition,
+        sourceX: edge.sourceX,
+        sourceY: edge.sourceY,
+        targetX: edge.targetX,
+        targetY: edge.targetY,
         sourceHandleId: edge.sourceHandle,
         targetHandleId: edge.targetHandle,
       }"
-    >
-      <component
-        :is="type"
-        v-if="type"
-        v-bind="{
-          id: edge.id,
-          sourceNode: edge.sourceNode,
-          targetNode: edge.targetNode,
-          source: edge.source,
-          target: edge.target,
-          updatable: props.updatable,
-          selected: edge.selected,
-          animated: edge.animated,
-          label: edge.label,
-          labelStyle: edge.labelStyle,
-          labelShowBg: edge.labelShowBg,
-          labelBgStyle: edge.labelBgStyle,
-          labelBgPadding: edge.labelBgPadding,
-          labelBgBorderRadius: edge.labelBgBorderRadius,
-          data: edge.data,
-          style: edge.style,
-          markerStart: `url(#${getMarkerId(edge.markerStart)})`,
-          markerEnd: `url(#${getMarkerId(edge.markerEnd)})`,
-          sourcePosition,
-          targetPosition,
-          sourceX: edge.sourceX,
-          sourceY: edge.sourceY,
-          targetX: edge.targetX,
-          targetY: edge.targetY,
-          sourceHandleId: edge.sourceHandle,
-          targetHandleId: edge.targetHandle,
-        }"
-      />
-    </slot>
+    />
     <g
       v-if="props.updatable"
       @mousedown="onEdgeUpdaterSourceMouseDown"
diff --git a/package/src/components/Nodes/NodeWrapper.vue b/package/src/components/Nodes/NodeWrapper.vue
index 75750e40..009b710f 100644
--- a/package/src/components/Nodes/NodeWrapper.vue
+++ b/package/src/components/Nodes/NodeWrapper.vue
@@ -136,7 +136,7 @@ const type = computed(() => {
     return store.getNodeTypes.default
   }
   if (slot.length > 1) {
-    console.warn('[vue-flow]: More than one element in node-slots detected. Using fallback to first slot item.')
+    console.warn('[vue-flow]: More than one element in node-slot detected. Using fallback to first slot item.')
   }
 
   return slot[0]
@@ -189,7 +189,6 @@ export default {
     >
       <component
         :is="type"
-        v-if="type"
         v-bind="{
           nodeElement,
           id: node.id,
diff --git a/package/src/container/EdgeRenderer/EdgeRenderer.vue b/package/src/container/EdgeRenderer/EdgeRenderer.vue
index 321d08cf..17a42ee6 100644
--- a/package/src/container/EdgeRenderer/EdgeRenderer.vue
+++ b/package/src/container/EdgeRenderer/EdgeRenderer.vue
@@ -47,11 +47,7 @@ export default {
         :key="edge.id"
         :selectable="typeof edge.selectable === 'undefined' ? store.elementsSelectable : edge.selectable"
         :updatable="typeof edge.updatable === 'undefined' ? store.edgesUpdatable : edge.updatable"
-      >
-        <template #default="edgeProps">
-          <slot v-if="edge.type" :name="`edge-${edge.type}`" v-bind="edgeProps"></slot>
-        </template>
-      </EdgeWrapper>
+      />
       <ConnectionLine v-if="connectionLineVisible && sourceNode" :source-node="sourceNode" />
     </g>
   </svg>
diff --git a/package/src/container/NodeRenderer/NodeRenderer.vue b/package/src/container/NodeRenderer/NodeRenderer.vue
index 42304720..6c6d9234 100644
--- a/package/src/container/NodeRenderer/NodeRenderer.vue
+++ b/package/src/container/NodeRenderer/NodeRenderer.vue
@@ -19,10 +19,6 @@ export default {
       :selectable="typeof node.selectable === 'undefined' ? store.elementsSelectable : !!node.selectable"
       :connectable="typeof node.connectable === 'undefined' ? store.nodesConnectable : !!node.connectable"
       :snap-grid="node.snapGrid ?? (store.snapToGrid ? store.snapGrid : undefined)"
-    >
-      <template #default="nodeProps">
-        <slot v-if="node.type" :name="`node-${node.type}`" v-bind="nodeProps"></slot>
-      </template>
-    </NodeWrapper>
+    />
   </div>
 </template>
diff --git a/package/src/index.ts b/package/src/index.ts
index e0639952..5c891486 100644
--- a/package/src/index.ts
+++ b/package/src/index.ts
@@ -6,6 +6,7 @@ export { default as StepEdge } from './components/Edges/StepEdge.vue'
 export { default as BezierEdge } from './components/Edges/BezierEdge.vue'
 export { default as SimpleBezierEdge } from './components/Edges/SimpleBezierEdge.vue'
 export { default as SmoothStepEdge } from './components/Edges/SmoothStepEdge.vue'
+
 export {
   getBezierPath,
   getBezierCenter,
@@ -29,11 +30,13 @@ export {
   getNodesInside,
   getMarkerId,
 } from './utils/graph'
+
 export { defaultEdgeTypes, defaultNodeTypes } from './store'
+export { VueFlow as VueFlowInjection, NodeId as NodeIdInjection } from './context'
+
 export { default as useZoomPanHelper } from './composables/useZoomPanHelper'
 export { default as useVueFlow } from './composables/useVueFlow'
 export { default as useHandle } from './composables/useHandle'
 
 export * from './additional-components'
 export * from './types'
-export * from './context'