From 9e941966a4ae639e0d4729e2bf0eb6edeb3f6601 Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Thu, 25 Nov 2021 13:25:28 +0100
Subject: [PATCH] feat(nodes): add option to pass a valid target pos / valid
 source pos func to nodes

* edge updater cannot discern a valid connection by itself so we allow an option to pass a valid connection func to nodes which will then be used to determine if updating connection is valid
* it is optional, by default all connections are valid

Signed-off-by: Braks <78412429+bcakmakoglu@users.noreply.github.com>
---
 examples/Validation/CustomInput.vue       |  9 ++++++---
 examples/Validation/CustomNode.vue        |  8 +++-----
 examples/Validation/ValidationExample.vue | 22 +++++++++++++++-------
 src/components/Edges/EdgeWrapper.vue      |  6 ++----
 src/components/Handle/Handle.vue          |  2 +-
 src/components/Nodes/DefaultNode.vue      | 18 +++++++++++++++---
 src/components/Nodes/InputNode.vue        | 10 ++++++++--
 src/components/Nodes/NodeWrapper.vue      |  4 ++++
 src/components/Nodes/OutputNode.vue       | 10 ++++++++--
 src/composables/useHandle.ts              | 17 ++++++++++-------
 src/types/connection.ts                   | 10 +++++-----
 src/types/hooks.ts                        |  3 +--
 src/types/node.ts                         |  6 +++++-
 13 files changed, 83 insertions(+), 42 deletions(-)

diff --git a/examples/Validation/CustomInput.vue b/examples/Validation/CustomInput.vue
index 32c30fa9..3d24fd6b 100644
--- a/examples/Validation/CustomInput.vue
+++ b/examples/Validation/CustomInput.vue
@@ -1,9 +1,12 @@
 <script lang="ts" setup>
-import { Position, Handle, Connection } from '~/index'
+import { Position, Handle, ValidConnectionFunc } from '~/index'
 
-const isValidConnection = (connection: Connection) => connection.target === 'B'
+interface CustomInputProps {
+  isValidTargetPos: ValidConnectionFunc
+}
+const props = defineProps<CustomInputProps>()
 </script>
 <template>
   <div>Only connectable with B</div>
-  <Handle type="source" :position="Position.Right" :is-valid-connection="isValidConnection" />
+  <Handle type="source" :position="Position.Right" :is-valid-connection="props.isValidTargetPos" />
 </template>
diff --git a/examples/Validation/CustomNode.vue b/examples/Validation/CustomNode.vue
index b455ca79..0c754302 100644
--- a/examples/Validation/CustomNode.vue
+++ b/examples/Validation/CustomNode.vue
@@ -1,16 +1,14 @@
 <script lang="ts" setup>
-import { Position, Handle, Connection, NodeProps } from '~/index'
+import { Position, Handle, NodeProps, ValidConnectionFunc } from '~/index'
 
 interface CustomNodeProps extends NodeProps {
   id: string
+  isValidSourcePos: ValidConnectionFunc
 }
 
 const props = defineProps<CustomNodeProps>()
-
-const isValidConnection = (connection: Connection) => connection.target === 'B'
 </script>
 <template>
-  <Handle type="target" :position="Position.Left" :is-valid-connection="isValidConnection" />
+  <Handle type="target" :position="Position.Left" :is-valid-connection="props.isValidSourcePos" />
   <div>{{ props.id }}</div>
-  <Handle type="source" :position="Position.Right" :is-valid-connection="isValidConnection" />
 </template>
diff --git a/examples/Validation/ValidationExample.vue b/examples/Validation/ValidationExample.vue
index 87f8972f..8888deb2 100644
--- a/examples/Validation/ValidationExample.vue
+++ b/examples/Validation/ValidationExample.vue
@@ -1,24 +1,32 @@
 <script lang="ts" setup>
 import CustomInput from './CustomInput.vue'
 import CustomNode from './CustomNode.vue'
-import { VueFlow, addEdge, Connection, Elements, Edge, OnConnectStartParams, FlowInstance } from '~/index'
+import { VueFlow, addEdge, Connection, Elements, OnConnectStartParams, FlowInstance } from '~/index'
 
 import './validation.css'
 
 const initialElements: Elements = [
-  { id: '0', type: 'custominput', position: { x: 0, y: 150 } },
-  { id: 'A', type: 'customnode', position: { x: 250, y: 0 } },
-  { id: 'B', type: 'customnode', position: { x: 250, y: 150 } },
-  { id: 'C', type: 'customnode', position: { x: 250, y: 300 } },
+  { id: '0', type: 'custominput', position: { x: 0, y: 150 }, isValidTargetPos: (connection) => connection.target === 'B' },
+  {
+    id: 'A',
+    type: 'customnode',
+    position: { x: 250, y: 0 },
+    isValidSourcePos: (connection) => {
+      console.log(connection)
+      return false
+    },
+  },
+  { id: 'B', type: 'customnode', position: { x: 250, y: 150 }, isValidSourcePos: (connection) => connection.target === 'B' },
+  { id: 'C', type: 'customnode', position: { x: 250, y: 300 }, isValidSourcePos: (connection) => connection.target === 'B' },
 ]
 
-const onLoad = (reactFlowInstance: FlowInstance) => reactFlowInstance.fitView()
+const onLoad = (flowInstance: FlowInstance) => flowInstance.fitView()
 const onConnectStart = ({ nodeId, handleType }: OnConnectStartParams) => console.log('on connect start', { nodeId, handleType })
 const onConnectStop = (event: MouseEvent) => console.log('on connect stop', event)
 const onConnectEnd = (event: MouseEvent) => console.log('on connect end', event)
 
 const elements = ref<Elements>(initialElements)
-const onConnect = (params: Connection | Edge) => {
+const onConnect = (params: Connection) => {
   console.log('on connect', params)
   elements.value = addEdge(params, elements.value)
 }
diff --git a/src/components/Edges/EdgeWrapper.vue b/src/components/Edges/EdgeWrapper.vue
index 9875b2be..94707361 100644
--- a/src/components/Edges/EdgeWrapper.vue
+++ b/src/components/Edges/EdgeWrapper.vue
@@ -43,16 +43,14 @@ const onEdgeUpdaterTargetMouseDown = (event: MouseEvent) => handleEdgeUpdater(ev
 const handleEdgeUpdater = (event: MouseEvent, isSourceHandle: boolean) => {
   const nodeId = isSourceHandle ? props.edge.target : props.edge.source
   const handleId = (isSourceHandle ? props.edge.targetHandle : props.edge.sourceHandle) ?? ''
-  const isValidConnection = () => true
-  const isTarget = isSourceHandle
 
   store.hooks.edgeUpdateStart.trigger({ event, edge: props.edge })
   handler(
     event,
     handleId,
     nodeId,
-    isTarget,
-    isValidConnection,
+    isSourceHandle,
+    undefined,
     isSourceHandle ? 'target' : 'source',
     (connection) => store.hooks.edgeUpdate.trigger({ edge: props.edge, connection }),
     () => store.hooks.edgeUpdateEnd.trigger({ event, edge: props.edge }),
diff --git a/src/components/Handle/Handle.vue b/src/components/Handle/Handle.vue
index 8b5901f0..a328e2dc 100644
--- a/src/components/Handle/Handle.vue
+++ b/src/components/Handle/Handle.vue
@@ -22,7 +22,7 @@ const nodeId = inject(NodeId, '')
 
 const handler = useHandle()
 const onMouseDownHandler = (event: MouseEvent) =>
-  handler(event, props.id, nodeId, props.type === 'target', props.isValidConnection ?? (() => true))
+  handler(event, props.id, nodeId, props.type === 'target', props.isValidConnection)
 </script>
 <script lang="ts">
 export default {
diff --git a/src/components/Nodes/DefaultNode.vue b/src/components/Nodes/DefaultNode.vue
index 5b1018bf..cfcf6fe2 100644
--- a/src/components/Nodes/DefaultNode.vue
+++ b/src/components/Nodes/DefaultNode.vue
@@ -1,12 +1,14 @@
 <script lang="ts" setup>
 import Handle from '../Handle/Handle.vue'
-import { NodeProps, Position } from '../../types'
+import { NodeProps, Position, ValidConnectionFunc } from '../../types'
 
 interface DefaultNodeProps extends NodeProps {
   data?: NodeProps['data']
   connectable?: NodeProps['connectable']
   targetPosition?: NodeProps['targetPosition']
   sourcePosition?: NodeProps['sourcePosition']
+  isValidTargetPos?: ValidConnectionFunc
+  isValidSourcePos?: ValidConnectionFunc
 }
 
 const props = withDefaults(defineProps<DefaultNodeProps>(), {
@@ -23,10 +25,20 @@ export default {
 }
 </script>
 <template>
-  <Handle type="target" :position="props.targetPosition" :is-connectable="props.connectable" />
+  <Handle
+    type="target"
+    :position="props.targetPosition"
+    :is-connectable="props.connectable"
+    :is-valid-connection="props.isValidTargetPos"
+  />
   <slot v-bind="props">
     <component :is="props.data?.label" v-if="typeof props.data?.label !== 'string'" />
     <span v-else v-html="props.data?.label"></span>
   </slot>
-  <Handle type="source" :position="props.sourcePosition" :is-connectable="props.connectable" />
+  <Handle
+    type="source"
+    :position="props.sourcePosition"
+    :is-connectable="props.connectable"
+    :is-valid-connection="props.isValidSourcePos"
+  />
 </template>
diff --git a/src/components/Nodes/InputNode.vue b/src/components/Nodes/InputNode.vue
index e376bd15..b2bea4be 100644
--- a/src/components/Nodes/InputNode.vue
+++ b/src/components/Nodes/InputNode.vue
@@ -1,11 +1,12 @@
 <script lang="ts" setup>
 import Handle from '../Handle/Handle.vue'
-import { NodeProps, Position } from '../../types'
+import { NodeProps, Position, ValidConnectionFunc } from '../../types'
 
 interface InputNodeProps extends NodeProps {
   data?: NodeProps['data']
   connectable?: NodeProps['connectable']
   sourcePosition?: NodeProps['sourcePosition']
+  isValidSourcePos?: ValidConnectionFunc
 }
 
 const props = withDefaults(defineProps<InputNodeProps>(), {
@@ -25,5 +26,10 @@ export default {
     <component :is="props.data?.label" v-if="typeof props.data?.label !== 'string'" />
     <span v-else v-html="props.data?.label"></span>
   </slot>
-  <Handle type="source" :position="props.sourcePosition" :connectable="props.connectable" />
+  <Handle
+    type="source"
+    :position="props.sourcePosition"
+    :connectable="props.connectable"
+    :is-valid-connection="props.isValidSourcePos"
+  />
 </template>
diff --git a/src/components/Nodes/NodeWrapper.vue b/src/components/Nodes/NodeWrapper.vue
index 67e0aed6..3668e39b 100644
--- a/src/components/Nodes/NodeWrapper.vue
+++ b/src/components/Nodes/NodeWrapper.vue
@@ -180,6 +180,8 @@ export default {
           sourcePosition: props.node.sourcePosition,
           targetPosition: props.node.targetPosition,
           dragging: props.node.__vf.isDragging,
+          isValidTargetPos: props.node.isValidTargetPos,
+          isValidSourcePos: props.node.isValidSourcePos,
         }"
       >
         <component
@@ -195,6 +197,8 @@ export default {
             sourcePosition: props.node.sourcePosition,
             targetPosition: props.node.targetPosition,
             dragging: props.node.__vf.isDragging,
+            isValidTargetPos: props.node.isValidTargetPos,
+            isValidSourcePos: props.node.isValidSourcePos,
           }"
         />
       </slot>
diff --git a/src/components/Nodes/OutputNode.vue b/src/components/Nodes/OutputNode.vue
index b96f2b03..5b32e093 100644
--- a/src/components/Nodes/OutputNode.vue
+++ b/src/components/Nodes/OutputNode.vue
@@ -1,11 +1,12 @@
 <script lang="ts" setup>
 import Handle from '../Handle/Handle.vue'
-import { NodeProps, Position } from '../../types'
+import { NodeProps, Position, ValidConnectionFunc } from '../../types'
 
 interface OutputNodeProps extends NodeProps {
   data?: NodeProps['data']
   connectable?: NodeProps['connectable']
   targetPosition?: NodeProps['targetPosition']
+  isValidTargetPos?: ValidConnectionFunc
 }
 
 const props = withDefaults(defineProps<OutputNodeProps>(), {
@@ -25,5 +26,10 @@ export default {
     <component :is="props.data?.label" v-if="typeof props.data?.label !== 'string'" />
     <span v-else v-html="props.data?.label"></span>
   </slot>
-  <Handle type="source" :position="props.targetPosition" :is-connectable="props.connectable" />
+  <Handle
+    type="source"
+    :position="props.targetPosition"
+    :is-connectable="props.connectable"
+    :is-valid-connection="props.isValidTargetPos"
+  />
 </template>
diff --git a/src/composables/useHandle.ts b/src/composables/useHandle.ts
index d71a29d9..3ad9d7ce 100644
--- a/src/composables/useHandle.ts
+++ b/src/composables/useHandle.ts
@@ -73,9 +73,7 @@ export default (store = useStore()) =>
     handleId: ElementId,
     nodeId: ElementId,
     isTarget: boolean,
-    isValidConnection: ValidConnectionFunc = () => {
-      return true
-    },
+    isValidConnection?: ValidConnectionFunc,
     elementEdgeUpdaterType?: HandleType,
     onEdgeUpdate?: (connection: Connection) => void,
     onEdgeUpdateEnd?: () => void,
@@ -85,7 +83,11 @@ export default (store = useStore()) =>
     const doc = getHostForElement(event.target as HTMLElement)
 
     if (!doc) return
-
+    let validConnectFunc: ValidConnectionFunc = isValidConnection ?? (() => true)
+    if (!isValidConnection) {
+      const node = store.getNodes.find((n) => n.id === nodeId)
+      if (node) validConnectFunc = (!isTarget ? node.isValidTargetPos : node.isValidSourcePos) ?? (() => true)
+    }
     const elementBelow = doc.elementFromPoint(event.clientX, event.clientY)
     const elementBelowIsTarget = elementBelow?.classList.contains('target')
     const elementBelowIsSource = elementBelow?.classList.contains('source')
@@ -104,7 +106,7 @@ export default (store = useStore()) =>
       connectionHandleId: handleId,
       connectionHandleType: handleType,
     })
-    store.hooks.connectStart.trigger({ event, params: { nodeId, handleId, handleType } })
+    store.hooks.connectStart.trigger({ event, nodeId, handleId, handleType })
 
     function onMouseMove(event: MouseEvent) {
       store.connectionPosition.x = event.clientX - containerBounds.left
@@ -116,7 +118,7 @@ export default (store = useStore()) =>
         isTarget,
         nodeId,
         handleId,
-        isValidConnection,
+        validConnectFunc,
         doc,
       )
 
@@ -140,9 +142,10 @@ export default (store = useStore()) =>
         isTarget,
         nodeId,
         handleId,
-        isValidConnection,
+        validConnectFunc,
         doc,
       )
+      console.log(isValidConnection)
 
       store.hooks.connectStop.trigger(event)
 
diff --git a/src/types/connection.ts b/src/types/connection.ts
index 0b613e3c..70651f97 100644
--- a/src/types/connection.ts
+++ b/src/types/connection.ts
@@ -11,8 +11,8 @@ export enum ConnectionLineType {
 export interface Connection {
   source: ElementId
   target: ElementId
-  sourceHandle: ElementId
-  targetHandle: ElementId
+  sourceHandle?: ElementId
+  targetHandle?: ElementId
 }
 
 export type ConnectionLineProps = {
@@ -28,9 +28,9 @@ export type ConnectionLineProps = {
 
 export type OnConnectFunc = (connection: Connection) => void
 export type OnConnectStartParams = {
-  nodeId: ElementId | undefined
-  handleId: ElementId | undefined
-  handleType: HandleType | undefined
+  nodeId?: ElementId
+  handleId?: ElementId
+  handleType?: HandleType
 }
 export type OnConnectStartFunc = (event: MouseEvent, params: OnConnectStartParams) => void
 export type OnConnectStopFunc = (event: MouseEvent) => void
diff --git a/src/types/hooks.ts b/src/types/hooks.ts
index 7e40a3cc..108c5314 100644
--- a/src/types/hooks.ts
+++ b/src/types/hooks.ts
@@ -21,8 +21,7 @@ export interface FlowEvents {
   connect: Connection
   connectStart: {
     event: MouseEvent
-    params: OnConnectStartParams
-  }
+  } & { [key in keyof OnConnectStartParams]: OnConnectStartParams[key] }
   connectStop: MouseEvent
   connectEnd: MouseEvent
   load: FlowInstance
diff --git a/src/types/node.ts b/src/types/node.ts
index e556e920..b5fad056 100644
--- a/src/types/node.ts
+++ b/src/types/node.ts
@@ -1,7 +1,7 @@
 import { Component, CSSProperties, DefineComponent } from 'vue'
 import { DraggableOptions } from '@braks/revue-draggable'
 import { XYPosition, ElementId, Position, SnapGrid } from './flow'
-import { HandleElement } from './components'
+import { HandleElement, ValidConnectionFunc } from './components'
 
 export interface VFInternals {
   isDragging?: boolean
@@ -24,6 +24,8 @@ export interface NodeProps<T = any> {
   targetPosition?: Position
   sourcePosition?: Position
   dragging?: boolean
+  isValidTargetPos?: ValidConnectionFunc
+  isValidSourcePos?: ValidConnectionFunc
 }
 
 export type NodeComponent = Component<NodeProps> | DefineComponent<NodeProps, any, any, any, any> | string
@@ -46,6 +48,8 @@ export interface Node<T = any> {
   connectable?: boolean
   dragHandle?: string
   snapGrid?: SnapGrid
+  isValidTargetPos?: ValidConnectionFunc
+  isValidSourcePos?: ValidConnectionFunc
 }
 
 export interface GraphNode<T = any> extends Node<T> {