Add basic type inference for tuple channel operators

bentsherman · bentsherman · commit f78ff0d84355 · 2026-01-27T17:50:09.000-06:00
diff --git a/src/main/java/nextflow/script/control/TupleOpResolver.java b/src/main/java/nextflow/script/control/TupleOpResolver.java
@@ -0,0 +1,204 @@
+/*
+ * Copyright 2024-2025, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package nextflow.script.control;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import nextflow.script.types.Bag;
+import nextflow.script.types.Channel;
+import nextflow.script.types.Tuple;
+import nextflow.script.types.Value;
+import org.codehaus.groovy.ast.ClassHelper;
+import org.codehaus.groovy.ast.ClassNode;
+import org.codehaus.groovy.ast.GenericsType;
+import org.codehaus.groovy.ast.MethodNode;
+import org.codehaus.groovy.ast.expr.Expression;
+import org.codehaus.groovy.ast.expr.NamedArgumentListExpression;
+
+import static nextflow.script.ast.ASTUtils.*;
+import static nextflow.script.types.TypeCheckingUtils.*;
+
+/**
+ *
+ * @author Ben Sherman <bentshermann@gmail.com>
+ */
+class TupleOpResolver {
+
+    private static final ClassNode BAG_TYPE = ClassHelper.makeCached(Bag.class);
+    private static final ClassNode CHANNEL_TYPE = ClassHelper.makeCached(Channel.class);
+    private static final ClassNode TUPLE_TYPE = ClassHelper.makeCached(Tuple.class);
+    private static final ClassNode VALUE_TYPE = ClassHelper.makeCached(Value.class);
+
+    /**
+     * Resolve the return type of dataflow operators that tranform
+     * tuples, such as `combine`, `groupTuple`, and `join`.
+     *
+     * @param lhsType
+     * @param method
+     * @param arguments
+     */
+    public ClassNode apply(ClassNode lhsType, MethodNode method, List<Expression> arguments) {
+        var name = method.getName();
+
+        if( "combine".equals(name) )
+            return applyCombine(lhsType, arguments);
+
+        if( "groupTuple".equals(name) )
+            return applyGroupBy(lhsType, arguments);
+
+        if( "join".equals(name) )
+            return applyJoin(lhsType, arguments);
+
+        return ClassHelper.dynamicType();
+    }
+
+    /**
+     * Resolve the result type of a `combine` operation in terms of the left
+     * and right operands.
+     *
+     * Given arguments of type `(L1, L2, ..., Lm)` and `R`, `combine`
+     * produces a tuple of type `(L1, L2, ..., Lm, R).
+     *
+     * When the `by` option is specified, `combine` produces the same result
+     * type as `join`.
+     *
+     * @param lhsType
+     * @param arguments
+     */
+    private ClassNode applyCombine(ClassNode lhsType, List<Expression> arguments) {
+        if( !TUPLE_TYPE.equals(lhsType) )
+            return ClassHelper.dynamicType();
+
+        var namedArgs = namedArgs(arguments);
+        if( namedArgs.containsKey("by") )
+            return applyJoin(lhsType, arguments);
+
+        var argType = getType(arguments.get(arguments.size() - 1));
+        var rhsType = dataflowElementType(argType);
+
+        var lgts = lhsType.getGenericsTypes();
+        if( lgts == null || lgts.length == 0 )
+            return ClassHelper.dynamicType();
+
+        var gts = new GenericsType[lgts.length + 1];
+        for( int i = 0; i < lgts.length; i++ )
+            gts[i] = lgts[i];
+        gts[lgts.length] = new GenericsType(rhsType);
+
+        return channelTupleType(gts);
+    }
+
+    /**
+     * Resolve the result type of a `groupTuple` operation.
+     *
+     * Given source tuples of type `(K, V1, V2, ..., Vn)`,
+     * `groupTuple` produces a tuple of type `(K, Bag<V1>, Bag<V2>, ..., Bag<Vn>)`.
+     *
+     * @param lhsType
+     * @param arguments
+     */
+    private ClassNode applyGroupBy(ClassNode lhsType, List<Expression> arguments) {
+        if( !TUPLE_TYPE.equals(lhsType) )
+            return ClassHelper.dynamicType();
+
+        var namedArgs = namedArgs(arguments);
+        if( namedArgs.containsKey("by") )
+            return ClassHelper.dynamicType();
+
+        var lgts = lhsType.getGenericsTypes();
+        if( lgts == null || lgts.length == 0 )
+            return ClassHelper.dynamicType();
+
+        // TODO: group on index specified by `by` option
+        // TODO: skip if `by` option isn't a single integer
+        var gts = new GenericsType[lgts.length];
+        gts[0] = lgts[0];
+        for( int i = 1; i < lgts.length; i++ ) {
+            var groupType = makeType(BAG_TYPE, lgts[i].getType());
+            gts[i] = new GenericsType(groupType);
+        }
+
+        return channelTupleType(gts);
+    }
+
+    /**
+     * Resolve the result type of a `join` operation in terms of the left
+     * and right operands.
+     *
+     * Given tuples of type `(K, L1, L2, ..., Lm)` and `(K, R1, R2, ..., Rn)`,
+     * `join` produces a tuple of type `(K, L1, L2, ..., Lm, R1, R2, ..., Rn).
+     *
+     * @param lhsType
+     * @param arguments
+     */
+    private ClassNode applyJoin(ClassNode lhsType, List<Expression> arguments) {
+        if( !TUPLE_TYPE.equals(lhsType) )
+            return ClassHelper.dynamicType();
+
+        var namedArgs = namedArgs(arguments);
+        if( namedArgs.containsKey("by") )
+            return ClassHelper.dynamicType();
+
+        var argType = getType(arguments.get(arguments.size() - 1));
+        var rhsType = dataflowElementType(argType);
+        if( !TUPLE_TYPE.equals(rhsType) )
+            return ClassHelper.dynamicType();
+
+        var lgts = lhsType.getGenericsTypes();
+        var rgts = rhsType.getGenericsTypes();
+        if( lgts == null || lgts.length == 0 || rgts == null || rgts.length == 0 )
+            return ClassHelper.dynamicType();
+
+        // TODO: join on index specified by `by` option
+        // TODO: skip if `by` option isn't a single integer
+        var gts = new GenericsType[lgts.length + rgts.length - 1];
+        for( int i = 0; i < lgts.length; i++ )
+            gts[i] = lgts[i];
+        for( int i = 1; i < rgts.length; i++ )
+            gts[lgts.length + i - 1] = rgts[i];
+
+        return channelTupleType(gts);
+    }
+
+    private static Map<String,Expression> namedArgs(List<Expression> args) {
+        return args.size() > 0 && args.get(0) instanceof NamedArgumentListExpression nale
+            ? Map.ofEntries(
+                nale.getMapEntryExpressions().stream()
+                    .map((entry) -> {
+                        var name = entry.getKeyExpression().getText();
+                        var value = entry.getValueExpression();
+                        return Map.entry(name, value);
+                    })
+                    .toArray(Map.Entry[]::new)
+            )
+            : Collections.emptyMap();
+    }
+
+    private static ClassNode dataflowElementType(ClassNode type) {
+        if( CHANNEL_TYPE.equals(type) || VALUE_TYPE.equals(type) )
+            return elementType(type);
+        return ClassHelper.dynamicType();
+    }
+
+    private static ClassNode channelTupleType(GenericsType[] gts) {
+        var tupleType = TUPLE_TYPE.getPlainNodeReference();
+        tupleType.setGenericsTypes(gts);
+        return makeType(CHANNEL_TYPE, tupleType);
+    }
+
+}
diff --git a/src/main/java/nextflow/script/control/TypeCheckingVisitorEx.java b/src/main/java/nextflow/script/control/TypeCheckingVisitorEx.java
@@ -23,7 +23,6 @@
 import java.util.Map;
 import java.util.stream.IntStream;
 
-import nextflow.lsp.ast.ASTNodeCache;
 import nextflow.script.ast.ASTNodeMarker;
 import nextflow.script.ast.AssignmentExpression;
 import nextflow.script.ast.FeatureFlagNode;
@@ -117,8 +116,18 @@ protected SourceUnit getSourceUnit() {
 
     public void visit() {
         var moduleNode = sourceUnit.getAST();
-        if( moduleNode instanceof ScriptNode sn )
-            visit(sn);
+        if( moduleNode instanceof ScriptNode sn ) {
+            for( var featureFlag : sn.getFeatureFlags() )
+                visitFeatureFlag(featureFlag);
+            if( sn.getParams() != null )
+                visitParams(sn.getParams());
+            for( var functionNode : sn.getFunctions() )
+                visitFunction(functionNode);
+            for( var processNode : sn.getProcesses() )
+                visitProcess(processNode);
+            for( var workflowNode : sn.getWorkflows() )
+                visitWorkflow(workflowNode);
+        }
     }
 
     // script declarations
@@ -478,6 +487,8 @@ public void visitMethodCallExpression(MethodCallExpression node) {
             var dummyMethod = resolveGenericReturnType(receiverType, target, arguments);
             node.putNodeMetaData(ASTNodeMarker.METHOD_TARGET, dummyMethod);
             node.putNodeMetaData(ASTNodeMarker.INFERRED_TYPE, dummyMethod.getReturnType());
+
+            checkOperatorCall(node);
         }
         else if( node.getNodeMetaData(ASTNodeMarker.METHOD_TARGET) instanceof MethodNode mn ) {
             var parameters = mn.getParameters();
@@ -550,6 +561,30 @@ private void checkSpreadMethodCall(MethodCallExpression node) {
         }
     }
 
+    /**
+     * Resolve the return type of operators that transform tuples.
+     * such as `combine`, `groupTuple`, and `join`.
+     *
+     * @param node
+     */
+    private void checkOperatorCall(MethodCallExpression node) {
+        if( node.isImplicitThis() )
+            return;
+
+        var receiverType = getType(node.getObjectExpression());
+        if( !CHANNEL_TYPE.equals(receiverType) )
+            return;
+
+        var lhsType = elementType(receiverType);
+        var method = (MethodNode) node.getNodeMetaData(ASTNodeMarker.METHOD_TARGET);
+        var arguments = asMethodCallArguments(node);
+        var resultType = new TupleOpResolver().apply(lhsType, method, arguments);
+        if( ClassHelper.isDynamicTyped(resultType) )
+            return;
+
+        node.putNodeMetaData(ASTNodeMarker.INFERRED_TYPE, resultType);
+    }
+
     /**
      * Check the arguments of an invalid method call and report appropriate
      * errors for each invalid argument.
@@ -691,13 +726,6 @@ private static ClassNode dataflowElementType(ClassNode type) {
         return type;
     }
 
-    private static ClassNode elementType(ClassNode type) {
-        var gts = type.getGenericsTypes();
-        if( gts == null || gts.length != 1 )
-            return ClassHelper.dynamicType();
-        return gts[0].getType();
-    }
-
     private static ClassNode processOutputType(ClassNode dataflowType, Statement block) {
         var outputs = asBlockStatements(block);
         if( outputs.size() == 1 ) {
diff --git a/src/main/java/nextflow/script/types/TypeCheckingUtils.java b/src/main/java/nextflow/script/types/TypeCheckingUtils.java
@@ -859,4 +859,16 @@ public static ClassNode makeType(ClassNode type, ClassNode... typeArguments) {
         cn.setGenericsTypes(gts);
         return cn;
     }
+
+    /**
+     * Return the element type of a type with one type parameter.
+     *
+     * @param type
+     */
+    public static ClassNode elementType(ClassNode type) {
+        var gts = type.getGenericsTypes();
+        if( gts == null || gts.length != 1 )
+            return ClassHelper.dynamicType();
+        return gts[0].getType();
+    }
 }
diff --git a/src/test/groovy/nextflow/script/types/TypeCheckingTest.groovy b/src/test/groovy/nextflow/script/types/TypeCheckingTest.groovy
@@ -643,6 +643,31 @@ class TypeCheckingTest extends Specification {
         type = getType(exp)
         then:
         TypesEx.getName(type) == 'Value<String>'
+
+        when:
+        exp = parseExpression(
+            '''\
+            nextflow.preview.types = true
+
+            process hello {
+                input:
+                target: String
+
+                output:
+                tuple(target, "Hello, $target!")
+
+                exec:
+                true
+            }
+
+            workflow {
+                hello( 'World' )
+            }
+            '''
+        )
+        type = getType(exp)
+        then:
+        TypesEx.getName(type) == 'Value<Tuple<String, String>>'
     }
 
     def 'should resolve tuple type' () {
@@ -705,4 +730,45 @@ class TypeCheckingTest extends Specification {
         "files('*.txt')*.toUriString()" | null
     }
 
+    def 'should resolve a `combine` operation' () {
+        when:
+        def exp = parseExpression(
+            '''\
+            left  = channel.of( tuple(42, 'hello') )
+            right = channel.of( true )
+            left.combine(right)
+            '''
+        )
+        def type = getType(exp)
+        then:
+        TypesEx.getName(type) == 'Channel<Tuple<Integer, String, Boolean>>'
+    }
+
+    def 'should resolve a `groupTuple` operation' () {
+        when:
+        def exp = parseExpression(
+            '''\
+            left = channel.of( tuple(42, 'hello'), tuple(42, 'goodbye') )
+            left.groupTuple()
+            '''
+        )
+        def type = getType(exp)
+        then:
+        TypesEx.getName(type) == 'Channel<Tuple<Integer, Bag<String>>>'
+    }
+
+    def 'should resolve a `join` operation' () {
+        when:
+        def exp = parseExpression(
+            '''\
+            left  = channel.of( tuple(42, 'hello') )
+            right = channel.of( tuple(42, true) )
+            left.join(right)
+            '''
+        )
+        def type = getType(exp)
+        then:
+        TypesEx.getName(type) == 'Channel<Tuple<Integer, String, Boolean>>'
+    }
+
 }
