Fix AST visitor for for loops and comprehensions.

RELNOTES: None.
PiperOrigin-RevId: 164893252

1 files changed, 2 insertions, 3 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/syntax/SyntaxTreeVisitor.java b/src/main/java/com/google/devtools/build/lib/syntax/SyntaxTreeVisitor.java
index c086fe3220..345b4bbe96 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/SyntaxTreeVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/SyntaxTreeVisitor.java
@@ -66,18 +66,17 @@ public class SyntaxTreeVisitor {
   public void visit(@SuppressWarnings("unused") Identifier node) {}
 
   public void visit(AbstractComprehension node) {
-    visitAll(node.getOutputExpressions());
-
     for (ListComprehension.Clause clause : node.getClauses()) {
       if (clause.getLValue() != null) {
         visit(clause.getLValue());
       }
       visit(clause.getExpression());
     }
+    visitAll(node.getOutputExpressions());
   }
 
   public void visit(ForStatement node) {
-    visit(node.getVariable().getExpression());
+    visit(node.getVariable());
     visit(node.getCollection());
     visitAll(node.getBlock());
   }