Add collected for-loop nodes

2026-06-24 19:05:28 +02:00
parent 6a5e71719e
commit a469e6cef2
4 changed files with 603 additions and 0 deletions
@@ -0,0 +1,130 @@
+import { app } from "../../scripts/app.js";
+
+const EXTENSION = "ethanfel.prompt_builder.loop_slots";
+const LOOP_NODES = new Set(["SxCPForLoopStart", "SxCPForLoopEnd", "SxCPWhileLoopStart", "SxCPWhileLoopEnd"]);
+const MAX_CARRY = 18;
+
+function isCarryInput(input) {
+    return /^initial_value\d+$/.test(input?.name || "");
+}
+
+function isCarryOutput(output) {
+    return /^value\d+$/.test(output?.name || "");
+}
+
+function carryNumber(slot) {
+    const match = String(slot?.name || "").match(/\d+$/);
+    return match ? Number(match[0]) : -1;
+}
+
+function resizeNode(node) {
+    const size = node.computeSize?.();
+    if (size) node.setSize?.(size);
+    app.graph?.setDirtyCanvas(true, true);
+}
+
+function getCarryLimit(nodeName) {
+    return nodeName === "SxCPWhileLoopStart" || nodeName === "SxCPWhileLoopEnd" ? 19 : MAX_CARRY;
+}
+
+function getFirstCarry(nodeName) {
+    return nodeName === "SxCPWhileLoopStart" || nodeName === "SxCPWhileLoopEnd" ? 0 : 1;
+}
+
+function addCarryPair(node, nodeName, number) {
+    if (number > getCarryLimit(nodeName)) return;
+    const inputName = `initial_value${number}`;
+    const outputName = `value${number}`;
+    if (!node.inputs?.some((input) => input.name === inputName)) node.addInput(inputName, "*");
+    if (!node.outputs?.some((output) => output.name === outputName)) node.addOutput(outputName, "*");
+}
+
+function removeCarryPair(node, number) {
+    const inputIndex = node.inputs?.findIndex((input) => input.name === `initial_value${number}`) ?? -1;
+    if (inputIndex >= 0 && !node.inputs[inputIndex]?.link) node.removeInput(inputIndex);
+    const outputIndex = node.outputs?.findIndex((output) => output.name === `value${number}`) ?? -1;
+    if (outputIndex >= 0 && !(node.outputs[outputIndex]?.links?.length)) node.removeOutput(outputIndex);
+}
+
+function trimCarryTail(node, nodeName) {
+    const first = getFirstCarry(nodeName);
+    for (let number = getCarryLimit(nodeName); number > first; number--) {
+        const input = node.inputs?.find((slot) => slot.name === `initial_value${number}`);
+        const output = node.outputs?.find((slot) => slot.name === `value${number}`);
+        const previousInput = node.inputs?.find((slot) => slot.name === `initial_value${number - 1}`);
+        const previousOutput = node.outputs?.find((slot) => slot.name === `value${number - 1}`);
+        const currentUsed = Boolean(input?.link || output?.links?.length);
+        const previousUsed = Boolean(previousInput?.link || previousOutput?.links?.length);
+        if (!currentUsed && !previousUsed) removeCarryPair(node, number);
+    }
+}
+
+function setupNodeSlots(node, nodeName) {
+    const first = getFirstCarry(nodeName);
+    const limit = getCarryLimit(nodeName);
+    addCarryPair(node, nodeName, first);
+    for (let number = first + 1; number <= limit; number++) {
+        const input = node.inputs?.find((slot) => slot.name === `initial_value${number}`);
+        const output = node.outputs?.find((slot) => slot.name === `value${number}`);
+        if (!input?.link && !output?.links?.length) removeCarryPair(node, number);
+    }
+    for (const output of node.outputs || []) {
+        if (output.name === "flow") output.shape = 5;
+    }
+    for (const input of node.inputs || []) {
+        if (input.name === "flow") input.shape = 5;
+    }
+    trimCarryTail(node, nodeName);
+    resizeNode(node);
+}
+
+function maybeGrow(node, nodeName) {
+    const carryInputs = (node.inputs || []).filter(isCarryInput);
+    const carryOutputs = (node.outputs || []).filter(isCarryOutput);
+    const lastInput = carryInputs.reduce((max, input) => Math.max(max, carryNumber(input)), -1);
+    const lastOutput = carryOutputs.reduce((max, output) => Math.max(max, carryNumber(output)), -1);
+    const last = Math.max(lastInput, lastOutput, getFirstCarry(nodeName));
+    const input = node.inputs?.find((slot) => slot.name === `initial_value${last}`);
+    const output = node.outputs?.find((slot) => slot.name === `value${last}`);
+    if ((input?.link || output?.links?.length) && last < getCarryLimit(nodeName)) {
+        addCarryPair(node, nodeName, last + 1);
+        resizeNode(node);
+    }
+}
+
+app.registerExtension({
+    name: EXTENSION,
+
+    async beforeRegisterNodeDef(nodeType, nodeData) {
+        if (!LOOP_NODES.has(nodeData.name)) return;
+
+        const onNodeCreated = nodeType.prototype.onNodeCreated;
+        nodeType.prototype.onNodeCreated = function () {
+            const result = onNodeCreated?.apply(this, arguments);
+            queueMicrotask(() => setupNodeSlots(this, nodeData.name));
+            return result;
+        };
+
+        const onConfigure = nodeType.prototype.onConfigure;
+        nodeType.prototype.onConfigure = function () {
+            const result = onConfigure?.apply(this, arguments);
+            queueMicrotask(() => setupNodeSlots(this, nodeData.name));
+            return result;
+        };
+
+        const onConnectionsChange = nodeType.prototype.onConnectionsChange;
+        nodeType.prototype.onConnectionsChange = function (type, index, connected, linkInfo) {
+            const result = onConnectionsChange?.apply(this, arguments);
+            if (!linkInfo) return result;
+            const slot = type === LiteGraph.INPUT ? this.inputs?.[index] : this.outputs?.[index];
+            if (isCarryInput(slot) || isCarryOutput(slot)) {
+                if (connected) maybeGrow(this, nodeData.name);
+                else {
+                    trimCarryTail(this, nodeData.name);
+                    resizeNode(this);
+                }
+            }
+            return result;
+        };
+    },
+});