Improve rendering of DAG preview

build.gradle

@@ -27,7 +27,7 @@ test {
 }
 
 dependencies {
-  implementation 'io.nextflow:nf-lang:25.05.0-edge'
+  implementation 'io.nextflow:nf-lang:25.06.0-edge'
   implementation 'org.apache.groovy:groovy:4.0.27'
   implementation 'org.apache.groovy:groovy-json:4.0.27'
   implementation 'org.eclipse.lsp4j:org.eclipse.lsp4j:0.23.0'

src/main/java/nextflow/lsp/services/script/ScriptCodeLensProvider.java

@@ -73,7 +73,7 @@ public List<CodeLens> codeLens(TextDocumentIdentifier textDocument) {
     public Map<String,String> previewDag(String documentUri, String name) {
         var uri = URI.create(documentUri);
         if( !ast.hasAST(uri) || ast.hasErrors(uri) )
-            return Map.ofEntries(Map.entry("error", "DAG preview cannot be shown because the script has errors."));
+            return Map.of("error", "DAG preview cannot be shown because the script has errors.");
 
         var sourceUnit = ast.getSourceUnit(uri);
         return ast.getWorkflowNodes(uri).stream()
@@ -86,7 +86,7 @@ public Map<String,String> previewDag(String documentUri, String name) {
                 var graph = visitor.getGraph(wn.isEntry() ? "<entry>" : wn.getName());
                 var result = new MermaidRenderer().render(wn.getName(), graph);
                 log.debug(result);
-                return Map.ofEntries(Map.entry("result", result));
+                return Map.of("result", result);
             })
             .orElse(null);
     }

src/main/java/nextflow/lsp/services/script/dag/DataflowVisitor.java

@@ -15,21 +15,17 @@
  */
 package nextflow.lsp.services.script.dag;
 
-import java.net.URI;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.Stack;
 import java.util.stream.Collectors;
 
 import groovy.lang.Tuple3;
-import nextflow.lsp.ast.LanguageServerASTUtils;
 import nextflow.lsp.services.script.ScriptAstCache;
 import nextflow.script.ast.ASTNodeMarker;
 import nextflow.script.ast.AssignmentExpression;
@@ -47,10 +43,10 @@
 import static nextflow.script.ast.ASTUtils.*;
 import static org.codehaus.groovy.ast.tools.GeneralUtils.*;
 
-
 /**
  *
  * @author Ben Sherman <bentshermann@gmail.com>
+ * @author Erik Danielsson <danielsson.erik.0@gmail.com>
  */
 public class DataflowVisitor extends ScriptVisitorSupport {
 
@@ -62,6 +58,8 @@ public class DataflowVisitor extends ScriptVisitorSupport {
 
     private Stack<Set<Node>> stackPreds = new Stack<>();
 
+    private VariableContext vc = new VariableContext();
+
     public DataflowVisitor(SourceUnit sourceUnit, ScriptAstCache ast) {
         this.sourceUnit = sourceUnit;
         this.ast = ast;
@@ -114,7 +112,7 @@ private void visitWorkflowTakes(WorkflowNode node, Map<String,Node> result) {
         for( var stmt : asBlockStatements(node.takes) ) {
             var name = asVarX(stmt).getName();
             var dn = addNode(name, Node.Type.NAME, stmt);
-            current.putSymbol(name, dn);
+            vc.putSymbol(name, dn);
             result.put(name, dn);
         }
     }
@@ -135,7 +133,7 @@ else if( emit instanceof AssignmentExpression assign ) {
                 name = "$out";
                 visit(new AssignmentExpression(varX(name), emit));
             }
-            var dn = current.getSymbol(name);
+            var dn = getSymbol(name);
             if( dn == null )
                 System.err.println("missing emit: " + name);
             result.put(name, dn);
@@ -150,10 +148,50 @@ private void visitWorkflowPublishers(WorkflowNode node, Map<String,Node> result)
             var source = publisher.getRightExpression();
             visit(new AssignmentExpression(target, source));
             var name = target.getName();
-            result.put(name, current.getSymbol(name));
+            var dn = getSymbol(name);
+            if( dn == null )
+                System.err.println("missing publisher: " + name);
+            result.put(name, dn);
         }
     }
 
+    // statements
+
+    @Override
+    public void visitIfElse(IfStatement node) {
+        // visit the conditional expression
+        var preds = visitWithPreds(node.getBooleanExpression());
+        var controlDn = addNode("", Node.Type.CONTROL, null, preds);
+
+        // visit the if branch
+        vc.pushScope();
+        current.pushSubgraph(controlDn);
+        visitWithPreds(node.getIfBlock());
+
+        var ifScope = vc.popScope();
+        var ifSubgraph = current.popSubgraph();
+
+        // visit the else branch
+        Map<String,Variable> elseScope;
+
+        if( !node.getElseBlock().isEmpty() ) {
+            vc.pushScope();
+            current.pushSubgraph(controlDn);
+            visitWithPreds(node.getElseBlock());
+
+            elseScope = vc.popScope();
+            current.popSubgraph();
+        }
+        else {
+            // if there is no else branch, then the set of active symbols
+            // after the if statement is the union of the active symbols
+            // from before the if and the active symbols in the if
+            elseScope = vc.peekScope();
+        }
+
+        vc.mergeConditionalScopes(ifScope, elseScope);
+    }
+
     // expressions
 
     @Override
@@ -173,7 +211,8 @@ public void visitMethodCallExpression(MethodCallExpression node) {
         var defNode = (MethodNode) node.getNodeMetaData(ASTNodeMarker.METHOD_TARGET);
         if( defNode instanceof WorkflowNode || defNode instanceof ProcessNode ) {
             var preds = visitWithPreds(node.getArguments());
-            current.putSymbol(name, addNode(name, Node.Type.OPERATOR, defNode, preds));
+            var dn = addNode(name, Node.Type.OPERATOR, defNode, preds);
+            vc.putSymbol(name, dn);
             return;
         }
 
@@ -183,7 +222,7 @@ public void visitMethodCallExpression(MethodCallExpression node) {
     @Override
     public void visitBinaryExpression(BinaryExpression node) {
         if( node instanceof AssignmentExpression ) {
-            visitAssignment(node);
+            visitAssignment(node, false);
             return;
         }
         if( node.getOperation().getType() == Types.PIPE ) {
@@ -194,17 +233,17 @@ public void visitBinaryExpression(BinaryExpression node) {
         super.visitBinaryExpression(node);
     }
 
-    private void visitAssignment(BinaryExpression node) {
+    private void visitAssignment(BinaryExpression node, boolean isLocal) {
         var preds = visitWithPreds(node.getRightExpression());
         var targets = getAssignmentTargets(node.getLeftExpression());
         for( var name : targets ) {
             var dn = addNode(name, Node.Type.NAME, null, preds);
-            current.putSymbol(name, dn);
+            vc.putSymbol(name, dn, isLocal);
         }
     }
 
     private Set<String> getAssignmentTargets(Expression node) {
-        // e.g. (x, y, z) = [1, 2, 3]
+        // e.g. (x, y, z) = xyz
         if( node instanceof TupleExpression te ) {
             return te.getExpressions().stream()
                 .map(el -> getAssignmentTarget(el).getName())
@@ -238,7 +277,8 @@ private void visitPipeline(BinaryExpression node) {
             if( defNode instanceof WorkflowNode || defNode instanceof ProcessNode ) {
                 var label = defNode.getName();
                 var preds = visitWithPreds(lhs);
-                current.putSymbol(label, addNode(label, Node.Type.OPERATOR, defNode, preds));
+                var dn = addNode(label, Node.Type.OPERATOR, defNode, preds);
+                vc.putSymbol(label, dn);
                 return;
             }
         }
@@ -248,11 +288,12 @@ private void visitPipeline(BinaryExpression node) {
 
     @Override
     public void visitDeclarationExpression(DeclarationExpression node) {
-        visitAssignment(node);
+        visitAssignment(node, true);
     }
 
     @Override
     public void visitClosureExpression(ClosureExpression node) {
+        // skip closures since they can't contain dataflow logic
     }
 
     @Override
@@ -365,25 +406,34 @@ else if( emit instanceof AssignmentExpression assign ) {
     }
 
     private void addOperatorPred(String label, ASTNode an) {
-        var dn = current.getSymbol(label);
+        var dn = getSymbol(label);
         if( dn != null )
             currentPreds().add(dn);
         else
-            current.putSymbol(label, addNode(label, Node.Type.OPERATOR, an));
+            vc.putSymbol(label, addNode(label, Node.Type.OPERATOR, an));
     }
 
     @Override
     public void visitVariableExpression(VariableExpression node) {
         var name = node.getName();
-        var dn = current.getSymbol(name);
+        var dn = getSymbol(name);
         if( dn != null )
             currentPreds().add(dn);
     }
 
     // helpers
 
+    private Node getSymbol(String name) {
+        var preds = vc.getSymbol(name);
+        if( preds.isEmpty() )
+            return null;
+        if( preds.size() == 1 )
+            return preds.iterator().next();
+        return addNode(name, Node.Type.NAME, null, preds);
+    }
+
     private Set<Node> currentPreds() {
-        return stackPreds.lastElement();
+        return stackPreds.peek();
     }
 
     private Set<Node> visitWithPreds(ASTNode... nodes) {
@@ -392,7 +442,7 @@ private Set<Node> visitWithPreds(ASTNode... nodes) {
 
     private Set<Node> visitWithPreds(Collection<? extends ASTNode> nodes) {
         // traverse a set of nodes and extract predecessor nodes
-        stackPreds.add(new HashSet<>());
+        stackPreds.push(new HashSet<>());
 
         for( var node : nodes ) {
             if( node != null )
@@ -402,10 +452,6 @@ private Set<Node> visitWithPreds(Collection<? extends ASTNode> nodes) {
         return stackPreds.pop();
     }
 
-    private Node visitWithPred(ASTNode node) {
-        return visitWithPreds(node).stream().findFirst().orElse(null);
-    }
-
     private Node addNode(String label, Node.Type type, ASTNode an, Set<Node> preds) {
         var uri = ast.getURI(an);
         var dn = current.addNode(label, type, uri, preds);
@@ -418,96 +464,3 @@ private Node addNode(String label, Node.Type type, ASTNode an) {
     }
 
 }
-
-
-class Graph {
-
-    public final Map<String,Node> inputs = new HashMap<>();
-
-    public final Map<Integer,Node> nodes = new HashMap<>();
-
-    public final Map<String,Node> outputs = new HashMap<>();
-
-    private List<Map<String,Node>> scopes = new ArrayList<>();
-
-    public Graph() {
-        pushScope();
-    }
-
-    public void pushScope() {
-        scopes.add(0, new HashMap<>());
-    }
-
-    public void popScope() {
-        scopes.remove(0);
-    }
-
-    public Node getSymbol(String name) {
-        // get a variable node from the name table
-        for( var scope : scopes )
-            if( scope.containsKey(name) )
-                return scope.get(name);
-
-        return null;
-    }
-
-    public void putSymbol(String name, Node dn) {
-        // put a variable node into the name table
-        for( var scope : scopes ) {
-            if( scope.containsKey(name) ) {
-                scope.put(name, dn);
-                return;
-            }
-        }
-
-        scopes.get(0).put(name, dn);
-    }
-
-    public Node addNode(String label, Node.Type type, URI uri, Set preds) {
-        var id = nodes.size();
-        var dn = new Node(id, label, type, uri, preds);
-        nodes.put(id, dn);
-        return dn;
-    }
-}
-
-
-class Node {
-    public enum Type {
-        NAME,
-        OPERATOR
-    }
-
-    public final int id;
-    public final String label;
-    public final Type type;
-    public final URI uri;
-    public final Set<Node> preds;
-
-    public Node(int id, String label, Type type, URI uri, Set<Node> preds) {
-        this.id = id;
-        this.label = label;
-        this.type = type;
-        this.uri = uri;
-        this.preds = preds;
-    }
-
-    public void addPredecessors(Set<Node> preds) {
-        this.preds.addAll(preds);
-    }
-
-    @Override
-    public boolean equals(Object other) {
-        return other instanceof Node n && this.id == n.id;
-    }
-
-    @Override
-    public int hashCode() {
-        return id;
-    }
-
-    @Override
-    public String toString() {
-        return String.format("id=%s,label='%s',type=%s", id, label, type);
-    }
-}