Fix lineage checks in UI

src/common/nodes/checkNodeValidity.ts

@@ -2,8 +2,10 @@ import { Input, InputData, InputId, NodeSchema } from '../common-types';
 import { FunctionInstance } from '../types/function';
 import { generateAssignmentErrorTrace, printErrorTrace, simpleError } from '../types/mismatch';
 import { withoutNull } from '../types/util';
+import { assertNever } from '../util';
 import { VALID, Validity, invalid } from '../Validity';
 import { testInputCondition } from './condition';
+import { ChainLineage, Lineage } from './lineage';
 import { getRequireConditions } from './required';
 
 const formatMissingInputs = (missingInputs: Input[]) => {
@@ -15,12 +17,16 @@ export interface CheckNodeValidityOptions {
     inputData: InputData;
     connectedInputs: ReadonlySet<InputId>;
     functionInstance: FunctionInstance | undefined;
+    chainLineage: ChainLineage | undefined;
+    nodeId: string;
 }
 export const checkNodeValidity = ({
     schema,
     inputData,
     connectedInputs,
     functionInstance,
+    chainLineage,
+    nodeId,
 }: CheckNodeValidityOptions): Validity => {
     const isOptional = (input: Input): boolean => {
         if (input.kind !== 'generic' || !input.optional) {
@@ -60,6 +66,7 @@ export const checkNodeValidity = ({
         return invalid(formatMissingInputs(missingInputs));
     }
 
+    // Type check
     if (functionInstance) {
         for (const { inputId, assignedType, inputType } of functionInstance.inputErrors) {
             const input = schema.inputs.find((i) => i.id === inputId)!;
@@ -89,6 +96,29 @@ export const checkNodeValidity = ({
         }
     }
 
+    // Lineage check
+    if (chainLineage) {
+        for (const input of schema.inputs) {
+            const sourceLineage = chainLineage.getConnectedOutputLineage({
+                nodeId,
+                inputId: input.id,
+            });
+            if (sourceLineage !== undefined) {
+                // eslint-disable-next-line @typescript-eslint/no-use-before-define
+                const lineageValid = checkAssignedLineage(
+                    sourceLineage,
+                    nodeId,
+                    input.id,
+                    schema,
+                    chainLineage
+                );
+                if (!lineageValid.isValid) {
+                    return lineageValid;
+                }
+            }
+        }
+    }
+
     return VALID;
 };
 
@@ -118,3 +148,62 @@ export const checkRequiredInputs = (schema: NodeSchema, inputData: InputData): V
         reason: formatMissingInputs(missingInputs),
     };
 };
+
+export const checkAssignedLineage = (
+    sourceLineage: Lineage | null,
+    nodeId: string,
+    inputId: InputId,
+    schema: NodeSchema,
+    chainLineage: ChainLineage
+): Validity => {
+    const input = schema.inputs.find((i) => i.id === inputId)!;
+
+    const differentLineage = () =>
+        invalid('Cannot connect node to 2 different sequence of different origin.');
+    const nonIteratorInput = () => invalid(`Input ${input.label} cannot be connect to a sequence.`);
+
+    switch (schema.kind) {
+        case 'regularNode': {
+            // regular is auto-iterated, so it has to be treated separately
+            if (sourceLineage) {
+                const targetLineage = chainLineage.getInputLineage(nodeId, {
+                    exclude: new Set([inputId]),
+                });
+                if (targetLineage && !targetLineage.equals(sourceLineage)) {
+                    return differentLineage();
+                }
+            } else {
+                // it's always valid connect a node with no lineage
+            }
+            break;
+        }
+        case 'newIterator': {
+            if (sourceLineage) {
+                return nonIteratorInput();
+            }
+            break;
+        }
+        case 'collector': {
+            const isIterated = schema.iteratorInputs[0].inputs.includes(inputId);
+            if (isIterated) {
+                if (!sourceLineage) {
+                    return invalid(`Input ${input.label} expects a sequence.`);
+                }
+
+                const targetLineage = chainLineage.getInputLineage(nodeId, {
+                    exclude: new Set([inputId]),
+                });
+                if (targetLineage && !targetLineage.equals(sourceLineage)) {
+                    return differentLineage();
+                }
+            } else if (sourceLineage) {
+                return nonIteratorInput();
+            }
+            break;
+        }
+        default:
+            return assertNever(schema.kind);
+    }
+
+    return VALID;
+};

src/common/nodes/lineage.ts

@@ -1,13 +1,15 @@
 import { Edge, Node } from 'reactflow';
-import { EdgeData, NodeData, NodeSchema } from '../common-types';
+import { EdgeData, InputId, NodeData, NodeSchema } from '../common-types';
 import { SchemaMap } from '../SchemaMap';
 import {
     EMPTY_ARRAY,
+    EMPTY_SET,
     ParsedSourceHandle,
     ParsedTargetHandle,
     assertNever,
     groupBy,
     parseSourceHandle,
+    parseTargetHandle,
     stringifyTargetHandle,
 } from '../util';
 
@@ -57,82 +59,107 @@ export class ChainLineage {
 
     static readonly EMPTY: ChainLineage = new ChainLineage(SchemaMap.EMPTY, [], []);
 
-    getEdgeByTarget(handle: ParsedTargetHandle): Edge<EdgeData> | undefined {
+    private getEdgeByTarget(handle: ParsedTargetHandle): Edge<EdgeData> | undefined {
         return this.byTargetHandle.get(stringifyTargetHandle(handle));
     }
 
-    /**
-     * Returns the single lineage (if any) of all iterated inputs of the given node.
-     *
-     * Note: regular nodes are auto-iterated, so their lineage is that of the first iterated input (if any).
-     *
-     * Note: the input lineage of collector nodes is `null` if there are no connected iterated inputs (invalid chain).
-     */
-    getInputLineage(nodeId: string): Lineage | null {
-        const schema = this.nodeSchemata.get(nodeId);
-        if (!schema) return null;
-
+    private getInputLineageImpl(
+        nodeId: string,
+        schema: NodeSchema,
+        exclude: ReadonlySet<InputId>
+    ): Lineage | null {
         switch (schema.kind) {
             case 'newIterator': {
                 // iterator source nodes do not support iterated inputs
                 return null;
             }
             case 'regularNode': {
                 // regular nodes are auto-iterated, so their lineage is that of the first iterated input
-                let lineage = this.nodeLineageCache.get(nodeId);
-                if (lineage === undefined) {
-                    lineage = null;
-
-                    const edges = this.byTargetNode.get(nodeId) ?? EMPTY_ARRAY;
-                    for (const edge of edges) {
-                        const inputLineage = this.getOutputLineage(
-                            parseSourceHandle(edge.sourceHandle!)
-                        );
-                        if (inputLineage !== null) {
-                            lineage = inputLineage;
-                            break;
-                        }
+                let lineage: Lineage | null = null;
+
+                const edges = this.byTargetNode.get(nodeId) ?? EMPTY_ARRAY;
+                for (const edge of edges) {
+                    if (
+                        exclude.size > 0 &&
+                        exclude.has(parseTargetHandle(edge.targetHandle!).inputId)
+                    ) {
+                        // eslint-disable-next-line no-continue
+                        continue;
                     }
 
-                    this.nodeLineageCache.set(nodeId, lineage);
+                    const inputLineage = this.getOutputLineage(
+                        parseSourceHandle(edge.sourceHandle!)
+                    );
+                    if (inputLineage !== null) {
+                        lineage = inputLineage;
+                        break;
+                    }
                 }
+
                 return lineage;
             }
             case 'collector': {
                 // collectors already return non-iterator outputs
-                let lineage = this.nodeLineageCache.get(nodeId);
-                if (lineage === undefined) {
-                    lineage = null;
-
-                    if (schema.iteratorInputs.length !== 1) {
-                        throw new Error(
-                            `Collector nodes should have exactly 1 iterator input info (${schema.schemaId})`
-                        );
-                    }
-                    const info = schema.iteratorInputs[0];
+                let lineage: Lineage | null = null;
 
-                    for (const inputId of info.inputs) {
-                        const edge = this.getEdgeByTarget({ nodeId, inputId });
-                        // eslint-disable-next-line no-continue
-                        if (!edge) continue;
-
-                        const handle = parseSourceHandle(edge.sourceHandle!);
-                        const inputLineage = this.getOutputLineage(handle);
-                        if (inputLineage !== null) {
-                            lineage = inputLineage;
-                            break;
-                        }
-                    }
+                if (schema.iteratorInputs.length !== 1) {
+                    throw new Error(
+                        `Collector nodes should have exactly 1 iterator input info (${schema.schemaId})`
+                    );
+                }
+                const info = schema.iteratorInputs[0];
+
+                for (const inputId of info.inputs) {
+                    // eslint-disable-next-line no-continue
+                    if (exclude.has(inputId)) continue;
 
-                    this.nodeLineageCache.set(nodeId, lineage);
+                    const edge = this.getEdgeByTarget({ nodeId, inputId });
+                    // eslint-disable-next-line no-continue
+                    if (!edge) continue;
+
+                    const handle = parseSourceHandle(edge.sourceHandle!);
+                    const inputLineage = this.getOutputLineage(handle);
+                    if (inputLineage !== null) {
+                        lineage = inputLineage;
+                        break;
+                    }
                 }
+
                 return lineage;
             }
             default:
                 return assertNever(schema.kind);
         }
     }
 
+    /**
+     * Returns the single lineage (if any) of all iterated inputs of the given node.
+     *
+     * Note: regular nodes are auto-iterated, so their lineage is that of the first iterated input (if any).
+     *
+     * Note: the input lineage of collector nodes is `null` if there are no connected iterated inputs (invalid chain).
+     */
+    getInputLineage(
+        nodeId: string,
+        { exclude }: { exclude?: ReadonlySet<InputId> } = {}
+    ): Lineage | null {
+        const schema = this.nodeSchemata.get(nodeId);
+        if (!schema) return null;
+
+        const useCache = exclude === undefined || exclude.size === 0;
+
+        if (!useCache) {
+            return this.getInputLineageImpl(nodeId, schema, exclude);
+        }
+
+        let lineage = this.nodeLineageCache.get(nodeId);
+        if (lineage === undefined) {
+            lineage = this.getInputLineageImpl(nodeId, schema, EMPTY_SET);
+            this.nodeLineageCache.set(nodeId, lineage);
+        }
+        return lineage;
+    }
+
     /**
      * Returns the lineage of the given specific output.
      */
@@ -164,4 +191,15 @@ export class ChainLineage {
                 return assertNever(schema.kind);
         }
     }
+
+    /**
+     * Returns `getOutputLineage` for the source handle connected to the given target handle.
+     * If no such connection exists, returns `undefined`.
+     *
+     */
+    getConnectedOutputLineage(target: ParsedTargetHandle): Lineage | null | undefined {
+        const edge = this.getEdgeByTarget(target);
+        if (!edge) return undefined;
+        return this.getOutputLineage(parseSourceHandle(edge.sourceHandle!));
+    }
 }

src/main/cli/run.ts

@@ -14,6 +14,7 @@ import { log } from '../../common/log';
 import { checkNodeValidity } from '../../common/nodes/checkNodeValidity';
 import { getConnectedInputs } from '../../common/nodes/connectedInputs';
 import { getEffectivelyDisabledNodes } from '../../common/nodes/disabled';
+import { ChainLineage } from '../../common/nodes/lineage';
 import { parseFunctionDefinitions } from '../../common/nodes/parseFunctionDefinitions';
 import { getNodesWithSideEffects } from '../../common/nodes/sideEffect';
 import { toBackendJson } from '../../common/nodes/toBackendJson';
@@ -154,6 +155,7 @@ const ensureStaticCorrectness = (
 
     const byId = new Map(nodes.map((n) => [n.id, n]));
     const typeState = TypeState.create(byId, edges, new Map(), functionDefinitions);
+    const chainLineage = new ChainLineage(schemata, nodes, edges);
 
     const invalidNodes = nodes.flatMap((node) => {
         const functionInstance = typeState.functions.get(node.data.id);
@@ -164,6 +166,8 @@ const ensureStaticCorrectness = (
             connectedInputs: getConnectedInputs(node.id, edges),
             schema,
             functionInstance,
+            chainLineage,
+            nodeId: node.id,
         });
         if (validity.isValid) return [];
 

src/renderer/components/NodeDocumentation/NodeExample.tsx

@@ -140,6 +140,8 @@ export const NodeExample = memo(({ accentColor, selectedSchema }: NodeExamplePro
         connectedInputs: requiredGenericInputs,
         inputData,
         functionInstance: typeInfo.instance,
+        chainLineage: undefined,
+        nodeId,
     });
 
     const { iteratedInputs, iteratedOutputs } = useMemo(() => {

src/renderer/contexts/ExecutionContext.tsx

@@ -133,6 +133,7 @@ const useRegisterNodeEvents = (
 export const ExecutionProvider = memo(({ children }: React.PropsWithChildren<{}>) => {
     const {
         typeStateRef,
+        chainLineageRef,
         outputDataActions,
         getInputHash,
         setManualOutputType,
@@ -384,6 +385,8 @@ export const ExecutionProvider = memo(({ children }: React.PropsWithChildren<{}>
                     connectedInputs: getConnectedInputs(node.id, edges),
                     schema,
                     functionInstance,
+                    chainLineage: chainLineageRef.current,
+                    nodeId: node.id,
                 })
             );
             if (validity.isValid) return [];
@@ -448,6 +451,7 @@ export const ExecutionProvider = memo(({ children }: React.PropsWithChildren<{}>
         schemata,
         sendAlert,
         typeStateRef,
+        chainLineageRef,
         features,
         featureStates,
         backend,