diff --git a/web/extensions/core/rerouteNode.js b/web/extensions/core/rerouteNode.js index bdc7306b..8ce349ff 100644 --- a/web/extensions/core/rerouteNode.js +++ b/web/extensions/core/rerouteNode.js @@ -14,11 +14,92 @@ app.registerExtension({ this.addInput("", "*"); this.addOutput(this.properties.showOutputText ? "*" : "", "*"); - this.onConnectInput = function (_, type) { - if (type !== this.outputs[0].type) { - this.removeOutput(0); - this.addOutput(this.properties.showOutputText ? type : "", type); - this.size = this.computeSize(); + + this.onConnectionsChange = function (type, index, connected, link_info) { + // Prevent multiple connections to different types when we have no input + if (connected && type === LiteGraph.OUTPUT) { + // Ignore wildcard nodes as these will be updated to real types + const types = new Set(this.outputs[0].links.map((l) => app.graph.links[l].type).filter((t) => t !== "*")); + if (types.size > 1) { + for (let i = 0; i < this.outputs[0].links.length - 1; i++) { + const linkId = this.outputs[0].links[i]; + const link = app.graph.links[linkId]; + const node = app.graph.getNodeById(link.target_id); + node.disconnectInput(link.target_slot); + } + } + } + + // Find root input + let currentNode = this; + let updateNodes = []; + let inputType = null; + while (currentNode) { + updateNodes.unshift(currentNode); + const linkId = currentNode.inputs[0].link; + if (linkId !== null) { + const link = app.graph.links[linkId]; + const node = app.graph.getNodeById(link.origin_id); + const type = node.constructor.type; + if (type === "Reroute") { + // Move the previous node + currentNode = node; + } else { + // We've found the end + inputType = node.outputs[link.origin_slot].type; + break; + } + } else { + // This path has no input node + currentNode = null; + break; + } + } + + // Find all outputs + const nodes = [this]; + let outputType = null; + while (nodes.length) { + currentNode = nodes.pop(); + const outputs = (currentNode.outputs ? currentNode.outputs[0].links : []) || []; + if (outputs.length) { + for (const linkId of outputs) { + const link = app.graph.links[linkId]; + + // When disconnecting sometimes the link is still registered + if (!link) continue; + + const node = app.graph.getNodeById(link.target_id); + const type = node.constructor.type; + + if (type === "Reroute") { + // Follow reroute nodes + nodes.push(node); + updateNodes.push(node); + } else { + // We've found an output + const nodeOutType = node.inputs[link.target_slot].type; + if (inputType && nodeOutType !== inputType) { + // The output doesnt match our input so disconnect it + node.disconnectInput(link.target_slot); + } else { + outputType = nodeOutType; + } + } + } + } else { + // No more outputs for this path + } + } + + // Update the types of each node + for (const node of updateNodes) { + // If we dont have an input type we are always wildcard but we'll show the output type + // This lets you change the output link to a different type and all nodes will update + node.outputs[0].type = inputType || "*"; + node.__outputType = inputType || outputType || "*"; + node.outputs[0].name = node.properties.showOutputText ? inputType || outputType || "*" : ""; + node.size = node.computeSize(); } }; @@ -41,12 +122,12 @@ app.registerExtension({ callback: () => { this.properties.showOutputText = !this.properties.showOutputText; if (this.properties.showOutputText) { - this.outputs[0].name = this.outputs[0].type; + this.outputs[0].name = this.__outputType || this.outputs[0].type; } else { this.outputs[0].name = ""; } this.size = this.computeSize(); - app.graph.setDirtyCanvas(true); + app.graph.setDirtyCanvas(true, true); }, }, {