From c1218b6f5ef4e70555e1fd12c621c25d479dad2e Mon Sep 17 00:00:00 2001
From: huchenlei <huchenlei@proton.me>
Date: Sun, 9 Feb 2025 11:17:28 -0500
Subject: [PATCH] [Refactor] NodeSlot.isValidTarget

---
 src/LGraphCanvas.ts | 13 ++-----------
 src/NodeSlot.ts     | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/src/LGraphCanvas.ts b/src/LGraphCanvas.ts
index 733d779..236638d 100644
--- a/src/LGraphCanvas.ts
+++ b/src/LGraphCanvas.ts
@@ -4703,9 +4703,6 @@ export class LGraphCanvas implements ConnectionColorContext {
       LiteGraph.NODE_SELECTED_TITLE_COLOR ??
       LiteGraph.NODE_TEXT_COLOR
 
-    const out_slot = this.connecting_links?.[0]?.output
-    const in_slot = this.connecting_links?.[0]?.input
-
     let max_y = 0
     const slot_pos = new Float32Array(2) // to reuse
 
@@ -4716,9 +4713,7 @@ export class LGraphCanvas implements ConnectionColorContext {
         const slot = toClass(NodeInputSlot, input)
 
         // change opacity of incompatible slots when dragging a connection
-        const isValid =
-          !this.connecting_links ||
-          (out_slot && LiteGraph.isValidConnection(slot.type, out_slot.type))
+        const isValid = slot.isValidTarget(this.connecting_links?.[0])
         const highlight = isValid && node.mouseOver?.inputId === i
         const label_color = highlight
           ? highlightColour
@@ -4747,12 +4742,8 @@ export class LGraphCanvas implements ConnectionColorContext {
       for (const [i, output] of (node.outputs ?? []).entries()) {
         const slot = toClass(NodeOutputSlot, output)
 
-        const slot_type = slot.type
-
         // change opacity of incompatible slots when dragging a connection
-        const isValid =
-          !this.connecting_links ||
-          (in_slot && LiteGraph.isValidConnection(slot_type, in_slot.type))
+        const isValid = slot.isValidTarget(this.connecting_links?.[0])
         const highlight = isValid && node.mouseOver?.outputId === i
         const label_color = highlight
           ? highlightColour
diff --git a/src/NodeSlot.ts b/src/NodeSlot.ts
index 27cebc6..8d58546 100644
--- a/src/NodeSlot.ts
+++ b/src/NodeSlot.ts
@@ -1,8 +1,9 @@
-import type { CanvasColour, Dictionary, INodeInputSlot, INodeOutputSlot, INodeSlot, ISlotType, Point } from "./interfaces"
+import type { CanvasColour, ConnectingLink, Dictionary, INodeInputSlot, INodeOutputSlot, INodeSlot, ISlotType, Point } from "./interfaces"
 import type { IWidget } from "./types/widgets"
 import type { LinkId } from "./LLink"
 import { LinkDirection, RenderShape } from "./types/globalEnums"
 import { LabelPosition, SlotShape, SlotType } from "./draw"
+import { LiteGraph } from "./litegraph"
 
 export interface ConnectionColorContext {
   default_connection_color: {
@@ -48,6 +49,12 @@ export abstract class NodeSlot implements INodeSlot {
     this.type = slot.type
   }
 
+  /**
+   * Whether this slot is a valid target for a dragging link.
+   * @param link - The link to check against.
+   */
+  abstract isValidTarget(link: ConnectingLink | null): boolean
+
   /**
    * The label to display in the UI.
    */
@@ -233,6 +240,12 @@ export class NodeInputSlot extends NodeSlot implements INodeInputSlot {
     return this.link != null
   }
 
+  override isValidTarget(link: ConnectingLink | null): boolean {
+    if (!link) return true
+
+    return !!link.output && LiteGraph.isValidConnection(this.type, link.output.type)
+  }
+
   override draw(ctx: CanvasRenderingContext2D, options: Omit<IDrawOptions, "doStroke" | "labelPosition">) {
     const originalTextAlign = ctx.textAlign
     ctx.textAlign = options.horizontal ? "center" : "left"
@@ -259,6 +272,12 @@ export class NodeOutputSlot extends NodeSlot implements INodeOutputSlot {
     this.slot_index = slot.slot_index
   }
 
+  override isValidTarget(link: ConnectingLink | null): boolean {
+    if (!link) return true
+
+    return !!link.input && LiteGraph.isValidConnection(this.type, link.input.type)
+  }
+
   override isConnected(): boolean {
     return this.links != null && this.links.length > 0
   }