Refactor comprehensions and other AST nodes

* Cleanup ForClause and IfClause to be static classes
* Add missing/non-visible accessors and constructors
* Make use of final modifiers more consistent
* Make LValue an ASTNode.

RELNOTES: None
PiperOrigin-RevId: 156050683

14 files changed, 219 insertions, 78 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/syntax/AbstractComprehension.java b/src/main/java/com/google/devtools/build/lib/syntax/AbstractComprehension.java
index 483d1eaab4..c3dd10bd8c 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/AbstractComprehension.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/AbstractComprehension.java
@@ -18,7 +18,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.events.Location;
 import java.io.Serializable;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import javax.annotation.Nullable;
 
@@ -46,6 +45,21 @@ public abstract class AbstractComprehension extends Expression {
    * A comprehension consists of one or many Clause.
    */
   public interface Clause extends Serializable {
+
+    /** Enum for distinguishing clause types. */
+    public enum Kind {
+      FOR,
+      IF
+    }
+
+    /**
+     * Returns whether this is a For or If clause.
+     *
+     * <p>This avoids having to rely on reflection, or on checking whether {@link #getLValue} is
+     * null.
+     */
+    public abstract Kind getKind();
+
     /**
      * The evaluation of the comprehension is based on recursion. Each clause may
      * call recursively evalStep (ForClause will call it multiple times, IfClause will
@@ -61,7 +75,7 @@ public abstract class AbstractComprehension extends Expression {
     abstract void eval(Environment env, OutputCollector collector, int step)
         throws EvalException, InterruptedException;
 
-    abstract void validate(ValidationEnvironment env) throws EvalException;
+    abstract void validate(ValidationEnvironment env, Location loc) throws EvalException;
 
     /**
      * The LValue defined in Clause, i.e. the loop variables for ForClause and null for
@@ -80,10 +94,15 @@ public abstract class AbstractComprehension extends Expression {
   /**
    * A for clause in a comprehension, e.g. "for a in b" in the example above.
    */
-  public final class ForClause implements Clause {
+  public static final class ForClause implements Clause {
     private final LValue variables;
     private final Expression list;
 
+    @Override
+    public Kind getKind() {
+      return Kind.FOR;
+    }
+
     public ForClause(LValue variables, Expression list) {
       this.variables = variables;
       this.list = list;
@@ -93,22 +112,22 @@ public abstract class AbstractComprehension extends Expression {
     public void eval(Environment env, OutputCollector collector, int step)
         throws EvalException, InterruptedException {
       Object listValueObject = list.eval(env);
-      Location loc = getLocation();
+      Location loc = collector.getLocation();
       Iterable<?> listValue = EvalUtils.toIterable(listValueObject, loc);
-      EvalUtils.lock(listValueObject, getLocation());
+      EvalUtils.lock(listValueObject, loc);
       try {
         for (Object listElement : listValue) {
           variables.assign(env, loc, listElement);
           evalStep(env, collector, step);
         }
       } finally {
-        EvalUtils.unlock(listValueObject, getLocation());
+        EvalUtils.unlock(listValueObject, loc);
       }
     }
 
     @Override
-    public void validate(ValidationEnvironment env) throws EvalException {
-      variables.validate(env, getLocation());
+    public void validate(ValidationEnvironment env, Location loc) throws EvalException {
+      variables.validate(env, loc);
       list.validate(env);
     }
 
@@ -131,9 +150,14 @@ public abstract class AbstractComprehension extends Expression {
   /**
    * A if clause in a comprehension, e.g. "if c" in the example above.
    */
-  public final class IfClause implements Clause {
+  public static final class IfClause implements Clause {
     private final Expression condition;
 
+    @Override
+    public Kind getKind() {
+      return Kind.IF;
+    }
+
     public IfClause(Expression condition) {
       this.condition = condition;
     }
@@ -147,7 +171,7 @@ public abstract class AbstractComprehension extends Expression {
     }
 
     @Override
-    public void validate(ValidationEnvironment env) throws EvalException {
+    public void validate(ValidationEnvironment env, Location loc) throws EvalException {
       condition.validate(env);
     }
 
@@ -173,18 +197,17 @@ public abstract class AbstractComprehension extends Expression {
    */
   private final ImmutableList<Expression> outputExpressions;
 
-  private final List<Clause> clauses;
-  private final char openingBracket;
-  private final char closingBracket;
+  private final ImmutableList<Clause> clauses;
 
-  public AbstractComprehension(
-      char openingBracket, char closingBracket, Expression... outputExpressions) {
-    clauses = new ArrayList<>();
+  public AbstractComprehension(List<Clause> clauses, Expression... outputExpressions) {
+    this.clauses = ImmutableList.copyOf(clauses);
     this.outputExpressions = ImmutableList.copyOf(outputExpressions);
-    this.openingBracket = openingBracket;
-    this.closingBracket = closingBracket;
   }
 
+  protected abstract char openingBracket();
+
+  protected abstract char closingBracket();
+
   public ImmutableList<Expression> getOutputExpressions() {
     return outputExpressions;
   }
@@ -192,35 +215,33 @@ public abstract class AbstractComprehension extends Expression {
   @Override
   public String toString() {
     StringBuilder sb = new StringBuilder();
-    sb.append(openingBracket).append(printExpressions());
+    sb.append(openingBracket()).append(printExpressions());
     for (Clause clause : clauses) {
       sb.append(' ').append(clause);
     }
-    sb.append(closingBracket);
+    sb.append(closingBracket());
     return sb.toString();
   }
 
-  /**
-   * Add a new ForClause to the comprehension. This is used only by the parser and must
-   * not be called once AST is complete.
-   * TODO(bazel-team): Remove this side-effect. Clauses should be passed to the constructor
-   * instead.
-   */
-  void addFor(Expression loopVar, Expression listExpression) {
-    Clause forClause = new ForClause(new LValue(loopVar), listExpression);
-    clauses.add(forClause);
-  }
+  /** Base class for comprehension builders. */
+  public abstract static class AbstractBuilder {
 
-  /**
-   * Add a new ForClause to the comprehension.
-   * TODO(bazel-team): Remove this side-effect.
-   */
-  void addIf(Expression condition) {
-    clauses.add(new IfClause(condition));
+    protected List<Clause> clauses = new ArrayList<>();
+
+    public void addFor(Expression loopVar, Expression listExpression) {
+      Clause forClause = new ForClause(new LValue(loopVar), listExpression);
+      clauses.add(forClause);
+    }
+
+    public void addIf(Expression condition) {
+      clauses.add(new IfClause(condition));
+    }
+
+    public abstract AbstractComprehension build();
   }
 
-  public List<Clause> getClauses() {
-    return Collections.unmodifiableList(clauses);
+  public ImmutableList<Clause> getClauses() {
+    return clauses;
   }
 
   @Override
@@ -238,7 +259,7 @@ public abstract class AbstractComprehension extends Expression {
   @Override
   void validate(ValidationEnvironment env) throws EvalException {
     for (Clause clause : clauses) {
-      clause.validate(env);
+      clause.validate(env, getLocation());
     }
     // Clauses have to be validated before expressions in order to introduce the variable names.
     for (Expression expr : outputExpressions) {
@@ -255,8 +276,9 @@ public abstract class AbstractComprehension extends Expression {
    * <p> In the expanded example above, you can consider that evalStep is equivalent to
    * evaluating the line number step.
    */
-  private void evalStep(Environment env, OutputCollector collector, int step)
+  private static void evalStep(Environment env, OutputCollector collector, int step)
       throws EvalException, InterruptedException {
+    List<Clause> clauses = collector.getClauses();
     if (step >= clauses.size()) {
       collector.evaluateAndCollect(env);
     } else {
@@ -276,6 +298,13 @@ public abstract class AbstractComprehension extends Expression {
    * providing access to the final results.
    */
   interface OutputCollector {
+
+    /** Returns the location for the comprehension we are evaluating. */
+    Location getLocation();
+
+    /** Returns the list of clauses for the comprehension we are evaluating. */
+    List<Clause> getClauses();
+
     /**
      * Evaluates the output expression(s) of the comprehension and collects the result.
      */
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Argument.java b/src/main/java/com/google/devtools/build/lib/syntax/Argument.java
index af05260337..ce1b5452a2 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Argument.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Argument.java
@@ -70,7 +70,7 @@ public abstract class Argument extends ASTNode {
   }
 
   /** positional argument: Expression */
-  public static class Positional extends Passed {
+  public static final class Positional extends Passed {
 
     public Positional(Expression value) {
       super(value);
@@ -86,7 +86,7 @@ public abstract class Argument extends ASTNode {
   }
 
   /** keyword argument: K = Expression */
-  public static class Keyword extends Passed {
+  public static final class Keyword extends Passed {
 
     final String name;
 
@@ -108,7 +108,7 @@ public abstract class Argument extends ASTNode {
   }
 
   /** positional rest (starred) argument: *Expression */
-  public static class Star extends Passed {
+  public static final class Star extends Passed {
 
     public Star(Expression value) {
       super(value);
@@ -124,7 +124,7 @@ public abstract class Argument extends ASTNode {
   }
 
   /** keyword rest (star_starred) parameter: **Expression */
-  public static class StarStar extends Passed {
+  public static final class StarStar extends Passed {
 
     public StarStar(Expression value) {
       super(value);
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java b/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java
index b1f079a14b..718952f41c 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/BuildFileAST.java
@@ -33,7 +33,7 @@ import javax.annotation.Nullable;
 /**
  * Abstract syntax node for an entire BUILD file.
  */
-public class BuildFileAST extends ASTNode {
+public final class BuildFileAST extends ASTNode {
 
   private final ImmutableList<Statement> stmts;
 
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/DictComprehension.java b/src/main/java/com/google/devtools/build/lib/syntax/DictComprehension.java
index 44fa0e6ec4..be9191a2fe 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/DictComprehension.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/DictComprehension.java
@@ -14,26 +14,74 @@
 package com.google.devtools.build.lib.syntax;
 
 
+import com.google.common.base.Preconditions;
+import com.google.devtools.build.lib.events.Location;
+import java.util.List;
 import java.util.Map;
 
 /**
  * Syntax node for dictionary comprehension expressions.
  */
-public class DictComprehension extends AbstractComprehension {
+public final class DictComprehension extends AbstractComprehension {
   private final Expression keyExpression;
   private final Expression valueExpression;
 
-  public DictComprehension(Expression keyExpression, Expression valueExpression) {
-    super('{', '}', keyExpression, valueExpression);
+  public DictComprehension(
+      List<Clause> clauses, Expression keyExpression, Expression valueExpression) {
+    super(clauses, keyExpression, valueExpression);
     this.keyExpression = keyExpression;
     this.valueExpression = valueExpression;
   }
 
   @Override
+  protected char openingBracket() {
+    return '{';
+  }
+
+  @Override
+  protected char closingBracket() {
+    return '}';
+  }
+
+  public Expression getKeyExpression() {
+    return keyExpression;
+  }
+
+  public Expression getValueExpression() {
+    return valueExpression;
+  }
+
+  @Override
   String printExpressions() {
     return String.format("%s: %s", keyExpression, valueExpression);
   }
 
+  /** Builder for {@link DictComprehension}. */
+  public static class Builder extends AbstractBuilder {
+
+    private Expression keyExpression;
+    private Expression valueExpression;
+
+    public Builder setKeyExpression(Expression keyExpression) {
+      this.keyExpression = keyExpression;
+      return this;
+    }
+
+    public Builder setValueExpression(Expression valueExpression) {
+      this.valueExpression = valueExpression;
+      return this;
+    }
+
+    @Override
+    public DictComprehension build() {
+      Preconditions.checkState(!clauses.isEmpty());
+      return new DictComprehension(
+          clauses,
+          Preconditions.checkNotNull(keyExpression),
+          Preconditions.checkNotNull(valueExpression));
+    }
+  }
+
   @Override
   OutputCollector createCollector(Environment env) {
     return new DictOutputCollector(env);
@@ -52,6 +100,16 @@ public class DictComprehension extends AbstractComprehension {
     }
 
     @Override
+    public Location getLocation() {
+      return DictComprehension.this.getLocation();
+    }
+
+    @Override
+    public List<Clause> getClauses() {
+      return DictComprehension.this.getClauses();
+    }
+
+    @Override
     public void evaluateAndCollect(Environment env) throws EvalException, InterruptedException {
       Object key = keyExpression.eval(env);
       EvalUtils.checkValidDictKey(key);
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/DictionaryLiteral.java b/src/main/java/com/google/devtools/build/lib/syntax/DictionaryLiteral.java
index 64dbfa6d9c..87756990ae 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/DictionaryLiteral.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/DictionaryLiteral.java
@@ -20,9 +20,10 @@ import java.util.List;
 /**
  * Syntax node for dictionary literals.
  */
-public class DictionaryLiteral extends Expression {
+public final class DictionaryLiteral extends Expression {
 
-  static final class DictionaryEntryLiteral extends ASTNode {
+  /** Node for an individual key-value pair in a dictionary literal. */
+  public static final class DictionaryEntryLiteral extends ASTNode {
 
     private final Expression key;
     private final Expression value;
@@ -32,11 +33,11 @@ public class DictionaryLiteral extends Expression {
       this.value = value;
     }
 
-    Expression getKey() {
+    public Expression getKey() {
       return key;
     }
 
-    Expression getValue() {
+    public Expression getValue() {
       return value;
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/FunctionDefStatement.java b/src/main/java/com/google/devtools/build/lib/syntax/FunctionDefStatement.java
index 201d084f40..3951b2b72e 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/FunctionDefStatement.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/FunctionDefStatement.java
@@ -20,7 +20,7 @@ import java.util.List;
 /**
  * Syntax node for a function definition.
  */
-public class FunctionDefStatement extends Statement {
+public final class FunctionDefStatement extends Statement {
 
   private final Identifier ident;
   private final FunctionSignature.WithValues<Expression, Expression> signature;
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/LValue.java b/src/main/java/com/google/devtools/build/lib/syntax/LValue.java
index 1b7278b6a7..52753cc16f 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/LValue.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/LValue.java
@@ -17,7 +17,6 @@ package com.google.devtools.build.lib.syntax;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.util.Preconditions;
-import java.io.Serializable;
 import java.util.Collection;
 
 /**
@@ -28,11 +27,12 @@ import java.util.Collection;
  *    for lvalue in exp: pass
  * An LValue can be a simple variable or something more complex like a tuple.
  */
-public class LValue implements Serializable {
+public class LValue extends ASTNode {
   private final Expression expr;
 
   public LValue(Expression expr) {
     this.expr = expr;
+    setLocation(expr.getLocation());
   }
 
   public Expression getExpression() {
@@ -180,6 +180,11 @@ public class LValue implements Serializable {
     env.update(ident.getName(), result);
   }
 
+  @Override
+  public void accept(SyntaxTreeVisitor visitor) {
+    visitor.visit(this);
+  }
+
   void validate(ValidationEnvironment env, Location loc) throws EvalException {
     validate(env, loc, expr);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java b/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java
index eebe7ab796..f5cef330ad 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java
@@ -13,27 +13,59 @@
 // limitations under the License.
 package com.google.devtools.build.lib.syntax;
 
+import com.google.common.base.Preconditions;
+import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.syntax.SkylarkList.MutableList;
 import java.util.ArrayList;
 import java.util.List;
 
 /**
  * Syntax node for lists comprehension expressions.
- *
  */
 public final class ListComprehension extends AbstractComprehension {
   private final Expression outputExpression;
 
-  public ListComprehension(Expression outputExpression) {
-    super('[', ']', outputExpression);
+  public ListComprehension(List<Clause> clauses, Expression outputExpression) {
+    super(clauses, outputExpression);
     this.outputExpression = outputExpression;
   }
 
   @Override
+  protected char openingBracket() {
+    return '[';
+  }
+
+  @Override
+  protected char closingBracket() {
+    return ']';
+  }
+
+  public Expression getOutputExpression() {
+    return outputExpression;
+  }
+
+  @Override
   String printExpressions() {
     return outputExpression.toString();
   }
 
+  /** Builder for {@link ListComprehension}. */
+  public static class Builder extends AbstractBuilder {
+
+    private Expression outputExpression;
+
+    public Builder setOutputExpression(Expression outputExpression) {
+      this.outputExpression = outputExpression;
+      return this;
+    }
+
+    @Override
+    public ListComprehension build() {
+      Preconditions.checkState(!clauses.isEmpty());
+      return new ListComprehension(clauses, Preconditions.checkNotNull(outputExpression));
+    }
+  }
+
   @Override
   OutputCollector createCollector(Environment env) {
     return new ListOutputCollector();
@@ -51,13 +83,23 @@ public final class ListComprehension extends AbstractComprehension {
     }
 
     @Override
+    public Location getLocation() {
+      return ListComprehension.this.getLocation();
+    }
+
+    @Override
+    public List<Clause> getClauses() {
+      return ListComprehension.this.getClauses();
+    }
+
+    @Override
     public void evaluateAndCollect(Environment env) throws EvalException, InterruptedException {
       result.add(outputExpression.eval(env));
     }
 
     @Override
     public Object getResult(Environment env) throws EvalException {
-      return new MutableList(result, env);
+      return new MutableList<>(result, env);
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/ListLiteral.java b/src/main/java/com/google/devtools/build/lib/syntax/ListLiteral.java
index 49b3ec5719..6f5ece7715 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/ListLiteral.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/ListLiteral.java
@@ -37,7 +37,7 @@ public final class ListLiteral extends Expression {
 
   private final List<Expression> exprs;
 
-  private ListLiteral(Kind kind, List<Expression> exprs) {
+  public ListLiteral(Kind kind, List<Expression> exprs) {
     this.kind = kind;
     this.exprs = exprs;
   }
@@ -55,9 +55,10 @@ public final class ListLiteral extends Expression {
     return makeList(Collections.<Expression>emptyList());
   }
 
-  /**
-   * Returns the list of expressions for each element of the tuple.
-   */
+  public Kind getKind() {
+    return kind;
+  }
+
   public List<Expression> getElements() {
     return exprs;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Literal.java b/src/main/java/com/google/devtools/build/lib/syntax/Literal.java
index 6c6ee5b5fd..d6842cbadf 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Literal.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Literal.java
@@ -15,7 +15,7 @@
 package com.google.devtools.build.lib.syntax;
 
 /**
- *  Generic base class for primitive literals.
+ * Generic base class for primitive literals.
  */
 public abstract class Literal<T> extends Expression {
 
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/NotExpression.java b/src/main/java/com/google/devtools/build/lib/syntax/NotExpression.java
index 72aca2ba9b..069e92dbcf 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/NotExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/NotExpression.java
@@ -14,9 +14,9 @@
 package com.google.devtools.build.lib.syntax;
 
 /**
- * As syntax node for the not boolean operation.
+ * Syntax node for the not boolean operation.
  */
-public class NotExpression extends Expression {
+public final class NotExpression extends Expression {
 
   private final Expression expression;
 
@@ -24,7 +24,7 @@ public class NotExpression extends Expression {
     this.expression = expression;
   }
 
-  Expression getExpression() {
+  public Expression getExpression() {
     return expression;
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Parameter.java b/src/main/java/com/google/devtools/build/lib/syntax/Parameter.java
index 52bf2997ee..1b74dcde5a 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Parameter.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Parameter.java
@@ -69,7 +69,7 @@ public abstract class Parameter<V, T> extends Argument {
   }
 
   /** mandatory parameter (positional or key-only depending on position): Ident */
-  public static class Mandatory<V, T> extends Parameter<V, T> {
+  public static final class Mandatory<V, T> extends Parameter<V, T> {
 
     public Mandatory(String name) {
       super(name);
@@ -90,7 +90,7 @@ public abstract class Parameter<V, T> extends Argument {
   }
 
   /** optional parameter (positional or key-only depending on position): Ident = Value */
-  public static class Optional<V, T> extends Parameter<V, T> {
+  public static final class Optional<V, T> extends Parameter<V, T> {
     public final V defaultValue;
 
     public Optional(String name, @Nullable V defaultValue) {
@@ -118,7 +118,7 @@ public abstract class Parameter<V, T> extends Argument {
   }
 
   /** extra positionals parameter (star): *identifier */
-  public static class Star<V, T> extends Parameter<V, T> {
+  public static final class Star<V, T> extends Parameter<V, T> {
     public Star(@Nullable String name, @Nullable T type) {
       super(name, type);
     }
@@ -146,7 +146,7 @@ public abstract class Parameter<V, T> extends Argument {
   }
 
   /** extra keywords parameter (star_star): **identifier */
-  public static class StarStar<V, T> extends Parameter<V, T> {
+  public static final class StarStar<V, T> extends Parameter<V, T> {
     public StarStar(String name, @Nullable T type) {
       super(name, type);
     }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
index fd9711737a..23bb0cd0c6 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
@@ -785,7 +785,7 @@ public class Parser {
   //                        | 'IF' expr comprehension_suffix
   //                        | ']'
   private Expression parseComprehensionSuffix(
-      AbstractComprehension comprehension, TokenKind closingBracket) {
+      AbstractComprehension.AbstractBuilder comprehensionBuilder, TokenKind closingBracket) {
     while (true) {
       if (token.kind == TokenKind.FOR) {
         nextToken();
@@ -794,13 +794,13 @@ public class Parser {
         // The expression cannot be a ternary expression ('x if y else z') due to
         // conflicts in Python grammar ('if' is used by the comprehension).
         Expression listExpression = parseNonTupleExpression(0);
-        comprehension.addFor(loopVar, listExpression);
+        comprehensionBuilder.addFor(loopVar, listExpression);
       } else if (token.kind == TokenKind.IF) {
         nextToken();
-        comprehension.addIf(parseExpression());
+        comprehensionBuilder.addIf(parseExpression());
       } else if (token.kind == closingBracket) {
         nextToken();
-        return comprehension;
+        return comprehensionBuilder.build();
       } else {
         syntaxError(token, "expected '" + closingBracket.getPrettyName() + "', 'for' or 'if'");
         syncPast(LIST_TERMINATOR_SET);
@@ -836,7 +836,9 @@ public class Parser {
       case FOR:
         { // list comprehension
           Expression result =
-              parseComprehensionSuffix(new ListComprehension(expression), TokenKind.RBRACKET);
+              parseComprehensionSuffix(
+                  new ListComprehension.Builder().setOutputExpression(expression),
+                  TokenKind.RBRACKET);
           return setLocation(result, start, token.right);
         }
       case COMMA:
@@ -883,7 +885,10 @@ public class Parser {
     if (token.kind == TokenKind.FOR) {
       // Dict comprehension
       Expression result = parseComprehensionSuffix(
-          new DictComprehension(entry.getKey(), entry.getValue()), TokenKind.RBRACE);
+          new DictComprehension.Builder()
+              .setKeyExpression(entry.getKey())
+              .setValueExpression(entry.getValue()),
+          TokenKind.RBRACE);
       return setLocation(result, start, token.right);
     }
     List<DictionaryEntryLiteral> entries = new ArrayList<>();
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/ReturnStatement.java b/src/main/java/com/google/devtools/build/lib/syntax/ReturnStatement.java
index 08ed9e0104..0976982dcf 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/ReturnStatement.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/ReturnStatement.java
@@ -18,7 +18,7 @@ import com.google.devtools.build.lib.events.Location;
 /**
  * A wrapper Statement class for return expressions.
  */
-public class ReturnStatement extends Statement {
+public final class ReturnStatement extends Statement {
 
   /**
    * Exception sent by the return statement, to be caught by the function body.