From dff88b54e6c753fefd4e9456d5d245b1806ff34c Mon Sep 17 00:00:00 2001
From: Ben Sherman <bentshermann@gmail.com>
Date: Fri, 17 Jan 2025 16:53:04 -0600
Subject: [PATCH] Fix detection of tap assignment (close #89)

Signed-off-by: Ben Sherman <bentshermann@gmail.com>
---
 .../script/control/VariableScopeVisitor.java  | 48 +++++++++----------
 .../src/main/java/script/types/Channel.java   |  8 ++++
 2 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/modules/compiler/src/main/java/script/control/VariableScopeVisitor.java b/modules/compiler/src/main/java/script/control/VariableScopeVisitor.java
index 6619241..9e83d5f 100644
--- a/modules/compiler/src/main/java/script/control/VariableScopeVisitor.java
+++ b/modules/compiler/src/main/java/script/control/VariableScopeVisitor.java
@@ -456,15 +456,6 @@ public void visitExpressionStatement(ExpressionStatement node) {
             }
             return;
         }
-        if( exp instanceof MethodCallExpression mce ) {
-            var source = mce.getObjectExpression();
-            var target = checkSetAssignment(mce);
-            if( target != null ) {
-                visit(source);
-                declareAssignedVariable(target);
-                return;
-            }
-        }
         super.visitExpressionStatement(node);
     }
 
@@ -562,21 +553,6 @@ private void checkExternalWriteInAsyncClosure(VariableExpression target, Variabl
             addFutureWarning("Mutating an external variable in an operator closure may lead to a race condition", target, "External variable declared here", (ASTNode) variable);
     }
 
-    /**
-     * Treat `set` operator as an assignment.
-     */
-    private VariableExpression checkSetAssignment(MethodCallExpression node) {
-        if( !(currentDefinition instanceof WorkflowNode) )
-            return null;
-        var name = node.getMethodAsString();
-        if( !"set".equals(name) )
-            return null;
-        var code = asDslBlock(node, 1);
-        if( code == null || code.getStatements().size() != 1 )
-            return null;
-        return asVarX(code.getStatements().get(0));
-    }
-
     // expressions
 
     private static final List<String> KEYWORDS = List.of(
@@ -588,6 +564,15 @@ private VariableExpression checkSetAssignment(MethodCallExpression node) {
 
     @Override
     public void visitMethodCallExpression(MethodCallExpression node) {
+        if( !node.isImplicitThis() ) {
+            var source = node.getObjectExpression();
+            var target = checkSetAssignment(node);
+            if( target != null ) {
+                visit(source);
+                declareAssignedVariable(target);
+                return;
+            }
+        }
         if( node.isImplicitThis() && node.getMethod() instanceof ConstantExpression ) {
             var name = node.getMethodAsString();
             var variable = findVariableDeclaration(name, node);
@@ -599,6 +584,21 @@ public void visitMethodCallExpression(MethodCallExpression node) {
         super.visitMethodCallExpression(node);
     }
 
+    /**
+     * Treat `set` and `tap` operators as assignments.
+     */
+    private VariableExpression checkSetAssignment(MethodCallExpression node) {
+        if( !(currentDefinition instanceof WorkflowNode) )
+            return null;
+        var name = node.getMethodAsString();
+        if( !"set".equals(name) && !"tap".equals(name) )
+            return null;
+        var code = asDslBlock(node, 1);
+        if( code == null || code.getStatements().size() != 1 )
+            return null;
+        return asVarX(code.getStatements().get(0));
+    }
+
     @Override
     public void visitDeclarationExpression(DeclarationExpression node) {
         visit(node.getRightExpression());
diff --git a/modules/compiler/src/main/java/script/types/Channel.java b/modules/compiler/src/main/java/script/types/Channel.java
index 67624c3..bcbcc6b 100644
--- a/modules/compiler/src/main/java/script/types/Channel.java
+++ b/modules/compiler/src/main/java/script/types/Channel.java
@@ -419,6 +419,14 @@ public static Channel<Path> watchPath(String filePattern, String events) {
     """)
     public abstract Channel take(int n);
 
+    @Operator
+    @Description("""
+        The `tap` operator assigns a source channel to a variable, whose name is specified in a closure.
+
+        [Read more](https://nextflow.io/docs/latest/reference/operator.html#tap)
+    """)
+    public abstract Channel tap(Closure holder);
+
     @Operator
     @Description("""
         The `toList` operator collects all the values from a source channel into a list and emits the list as a single value.