bazel syntax: simplify comprehensions

This change merges {List,Dict,Abstract}Comprehension into a
single concrete class. The "k: v" body in a dict comprehension
{k: v for vars in iterable} is now represented as
DictionaryLiteral.Entry.

"Clause" was a poor abstraction: clients need to know that
there are only two instances, If and For, and handle them
specifically.

The doEval method is now a single recursive
function that fits easily in a single page.

We delete the following classes:

  AbstractBuilder
  OutputCollector
  ListComprehension
    Builder
    ListOutputCollector
  Clause.Kind
  DictComprehension
    Builder
    DictOutputCollector

PiperOrigin-RevId: 269185401
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
deleted file mode 100644
index f9c3076..0000000
--- a/src/main/java/com/google/devtools/build/lib/syntax/AbstractComprehension.java
+++ /dev/null
@@ -1,337 +0,0 @@
-// Copyright 2014 The Bazel Authors. All rights reserved.
-//
-// 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 com.google.devtools.build.lib.syntax;
-
-import com.google.common.collect.ImmutableList;
-import com.google.devtools.build.lib.events.Location;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import javax.annotation.Nullable;
-
-/**
- * Base class for list and dict comprehension expressions.
- *
- * <p>A comprehension contains one or more clauses, e.g. [a+d for a in b if c for d in e] contains
- * three clauses: "for a in b", "if c", "for d in e". For and If clauses can happen in any order,
- * except that the first one has to be a For.
- *
- * <p>The code above can be expanded as:
- *
- * <pre>
- *   for a in b:
- *     if c:
- *       for d in e:
- *         result.append(a+d)
- * </pre>
- *
- * result is initialized to [] (list) or {} (dict) and is the return value of the whole expression.
- */
-// TODO(adonovan): replace {Abstract,List,Dict}Comprehension by a single class.
-public abstract class AbstractComprehension extends Expression {
-
-  /**
-   * The interface implemented by ForClause and (later) IfClause. A comprehension consists of one or
-   * more Clauses.
-   */
-  public interface Clause {
-
-    /** Enum for distinguishing clause types. */
-    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 #getLHS} is null.
-     */
-    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
-     * call it zero or one time) which will evaluate the next clause. To know which clause
-     * is the next one, we pass a step argument (it represents the index in the clauses
-     * list). Results are aggregated in the result argument, and are populated by
-     * evalStep.
-     *
-     * @param env environment in which we do the evaluation.
-     * @param collector the aggregated results of the comprehension.
-     * @param step the index of the next clause to evaluate.
-     */
-    void eval(Environment env, OutputCollector collector, int step)
-        throws EvalException, InterruptedException;
-
-    /** The LHS variables defined in a ForClause, or null for an IfClause. */
-    @Nullable
-    Expression getLHS();
-
-    /**
-     * The Expression defined in Clause, i.e. the collection for ForClause and the
-     * condition for IfClause. This is needed for SyntaxTreeVisitor.
-     */
-    Expression getExpression();
-
-    /** Pretty print to a buffer. */
-    void prettyPrint(Appendable buffer) throws IOException;
-  }
-
-  /** A for clause in a comprehension, e.g. "for a in b" in the example above. */
-  public static final class ForClause implements Clause {
-    private final Expression lhs;
-    private final Expression iterable;
-
-    @Override
-    public Kind getKind() {
-      return Kind.FOR;
-    }
-
-    public ForClause(Expression lhs, Expression iterable) {
-      this.lhs = lhs;
-      this.iterable = iterable;
-    }
-
-    @Override
-    public void eval(Environment env, OutputCollector collector, int step)
-        throws EvalException, InterruptedException {
-      Object iterableObject = iterable.eval(env);
-      Location loc = collector.getLocation();
-      Iterable<?> listValue = EvalUtils.toIterable(iterableObject, loc, env);
-      EvalUtils.lock(iterableObject, loc);
-      try {
-        for (Object listElement : listValue) {
-          Eval.assign(lhs, listElement, env, loc);
-          evalStep(env, collector, step);
-        }
-      } finally {
-        EvalUtils.unlock(iterableObject, loc);
-      }
-    }
-
-    @Override
-    public Expression getLHS() {
-      return lhs;
-    }
-
-    @Override
-    public Expression getExpression() {
-      return iterable;
-    }
-
-    @Override
-    public void prettyPrint(Appendable buffer) throws IOException {
-      buffer.append("for ");
-      lhs.prettyPrint(buffer);
-      buffer.append(" in ");
-      iterable.prettyPrint(buffer);
-    }
-
-    @Override
-    public String toString() {
-      StringBuilder builder = new StringBuilder();
-      try {
-        prettyPrint(builder);
-      } catch (IOException e) {
-        // Not possible for StringBuilder.
-        throw new AssertionError(e);
-      }
-      return builder.toString();
-    }
-  }
-
-  /** A if clause in a comprehension, e.g. "if c" in the example above. */
-  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;
-    }
-
-    @Override
-    public void eval(Environment env, OutputCollector collector, int step)
-        throws EvalException, InterruptedException {
-      if (EvalUtils.toBoolean(condition.eval(env))) {
-        evalStep(env, collector, step);
-      }
-    }
-
-    @Override
-    public Expression getLHS() {
-      return null;
-    }
-
-    @Override
-    public Expression getExpression() {
-      return condition;
-    }
-
-    @Override
-    public void prettyPrint(Appendable buffer) throws IOException {
-      buffer.append("if ");
-      condition.prettyPrint(buffer);
-    }
-
-    @Override
-    public String toString() {
-      StringBuilder builder = new StringBuilder();
-      try {
-        prettyPrint(builder);
-      } catch (IOException e) {
-        // Not possible for StringBuilder.
-        throw new AssertionError(e);
-      }
-      return builder.toString();
-    }
-  }
-
-  /**
-   * The output expressions, e.g. "a+d" in the example above. This list has either one (list) or two
-   * (dict) items.
-   */
-  private final ImmutableList<Expression> outputExpressions;
-
-  private final ImmutableList<Clause> clauses;
-
-  AbstractComprehension(List<Clause> clauses, Expression... outputExpressions) {
-    this.clauses = ImmutableList.copyOf(clauses);
-    this.outputExpressions = ImmutableList.copyOf(outputExpressions);
-  }
-
-  protected abstract char openingBracket();
-
-  protected abstract char closingBracket();
-
-  public ImmutableList<Expression> getOutputExpressions() {
-    return outputExpressions;
-  }
-
-  @Override
-  public void prettyPrint(Appendable buffer) throws IOException {
-    buffer.append(openingBracket());
-    printExpressions(buffer);
-    for (Clause clause : clauses) {
-      buffer.append(' ');
-      clause.prettyPrint(buffer);
-    }
-    buffer.append(closingBracket());
-  }
-
-  /** Base class for comprehension builders. */
-  public abstract static class AbstractBuilder {
-
-    protected final List<Clause> clauses = new ArrayList<>();
-
-    public void addFor(Expression lhs, Expression iterable) {
-      Clause forClause = new ForClause(lhs, iterable);
-      clauses.add(forClause);
-    }
-
-    public void addIf(Expression condition) {
-      clauses.add(new IfClause(condition));
-    }
-
-    public abstract AbstractComprehension build();
-  }
-
-  public ImmutableList<Clause> getClauses() {
-    return clauses;
-  }
-
-  @Override
-  public void accept(SyntaxTreeVisitor visitor) {
-    visitor.visit(this);
-  }
-
-  @Override
-  public Kind kind() {
-    return Kind.COMPREHENSION;
-  }
-
-  @Override
-  Object doEval(Environment env) throws EvalException, InterruptedException {
-    OutputCollector collector = createCollector(env);
-    evalStep(env, collector, 0);
-    Object result = collector.getResult(env);
-
-    // Undefine loop variables (remove them from the environment).
-    // This code is useful for the transition, to make sure no one relies on the old behavior
-    // (where loop variables were leaking).
-    // TODO(laurentlb): Instead of removing variables, we should create them in a nested scope.
-    for (Clause clause : clauses) {
-      // Check if a loop variable conflicts with another local variable.
-      Expression lhs = clause.getLHS();
-      if (lhs != null) {
-        for (Identifier ident : Identifier.boundIdentifiers(lhs)) {
-          env.removeLocalBinding(ident.getName());
-        }
-      }
-    }
-    return result;
-  }
-
-  /**
-   * Evaluate the clause indexed by step, or elementExpression. When we evaluate the
-   * comprehension, step is 0 and we evaluate the first clause. Each clause may
-   * recursively call evalStep any number of times. After the last clause,
-   * the output expression(s) is/are evaluated and added to the results.
-   *
-   * <p> In the expanded example above, you can consider that evalStep is equivalent to
-   * evaluating the line number 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 {
-      clauses.get(step).eval(env, collector, step + 1);
-    }
-  }
-
-  /** Pretty-prints the output expression(s). */
-  protected abstract void printExpressions(Appendable buffer) throws IOException;
-
-  abstract OutputCollector createCollector(Environment env);
-
-  /**
-   * Interface for collecting the intermediate output of an {@code AbstractComprehension} and for
-   * 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.
-     */
-    void evaluateAndCollect(Environment env) throws EvalException, InterruptedException;
-
-    /**
-     * Returns the final result of the comprehension.
-     */
-    Object getResult(Environment env) throws EvalException;
-  }
-}
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Comprehension.java b/src/main/java/com/google/devtools/build/lib/syntax/Comprehension.java
new file mode 100644
index 0000000..a6d9af0
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Comprehension.java
@@ -0,0 +1,202 @@
+// Copyright 2014 The Bazel Authors. All rights reserved.
+//
+// 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 com.google.devtools.build.lib.syntax;
+
+import com.google.common.collect.ImmutableList;
+import com.google.devtools.build.lib.events.Location;
+import java.io.IOException;
+import java.util.ArrayList;
+
+/**
+ * Syntax node for list and dict comprehensions.
+ *
+ * <p>A comprehension contains one or more clauses, e.g. [a+d for a in b if c for d in e] contains
+ * three clauses: "for a in b", "if c", "for d in e". For and If clauses can happen in any order,
+ * except that the first one has to be a For.
+ *
+ * <p>The code above can be expanded as:
+ *
+ * <pre>
+ *   for a in b:
+ *     if c:
+ *       for d in e:
+ *         result.append(a+d)
+ * </pre>
+ *
+ * result is initialized to [] (list) or {} (dict) and is the return value of the whole expression.
+ */
+public final class Comprehension extends Expression {
+
+  /** For or If */
+  public abstract static class Clause {}
+
+  /** A for clause in a comprehension, e.g. "for a in b" in the example above. */
+  public static final class For extends Clause {
+    private final Expression vars;
+    private final Expression iterable;
+
+    For(Expression vars, Expression iterable) {
+      this.vars = vars;
+      this.iterable = iterable;
+    }
+
+    public Expression getVars() {
+      return vars;
+    }
+
+    public Expression getIterable() {
+      return iterable;
+    }
+  }
+
+  /** A if clause in a comprehension, e.g. "if c" in the example above. */
+  public static final class If extends Clause {
+    private final Expression condition;
+
+    If(Expression condition) {
+      this.condition = condition;
+    }
+
+    public Expression getCondition() {
+      return condition;
+    }
+  }
+
+  private final boolean isDict; // {k: v for vars in iterable}
+  private final ASTNode body; // Expression or DictionaryLiteral.Entry
+  private final ImmutableList<Clause> clauses;
+
+  Comprehension(boolean isDict, ASTNode body, ImmutableList<Clause> clauses) {
+    this.isDict = isDict;
+    this.body = body;
+    this.clauses = clauses;
+  }
+
+  public boolean isDict() {
+    return isDict;
+  }
+
+  /**
+   * Returns the loop body: an expression for a list comprehension, or a DictionaryLiteral.Entry for
+   * a dict comprehension.
+   */
+  public ASTNode getBody() {
+    return body;
+  }
+
+  public ImmutableList<Clause> getClauses() {
+    return clauses;
+  }
+
+  @Override
+  public void prettyPrint(Appendable buffer) throws IOException {
+    buffer.append(isDict ? '{' : '[');
+    body.prettyPrint(buffer);
+    for (Clause clause : clauses) {
+      buffer.append(' ');
+      if (clause instanceof For) {
+        For forClause = (For) clause;
+        buffer.append("for ");
+        forClause.vars.prettyPrint(buffer);
+        buffer.append(" in ");
+        forClause.iterable.prettyPrint(buffer);
+      } else {
+        If ifClause = (If) clause;
+        buffer.append("if ");
+        ifClause.condition.prettyPrint(buffer);
+      }
+    }
+    buffer.append(isDict ? '}' : ']');
+  }
+
+  @Override
+  public void accept(SyntaxTreeVisitor visitor) {
+    visitor.visit(this);
+  }
+
+  @Override
+  public Kind kind() {
+    return Kind.COMPREHENSION;
+  }
+
+  @Override
+  Object doEval(Environment env) throws EvalException, InterruptedException {
+    final SkylarkDict<Object, Object> dict = isDict ? SkylarkDict.of(env) : null;
+    final ArrayList<Object> list = isDict ? null : new ArrayList<>();
+
+    // The Lambda class serves as a recursive lambda closure.
+    class Lambda {
+      // execClauses(index) recursively executes the clauses starting at index,
+      // and finally evaluates the body and adds its value to the result.
+      void execClauses(int index) throws EvalException, InterruptedException {
+        // recursive case: one or more clauses
+        if (index < clauses.size()) {
+          Clause clause = clauses.get(index);
+          if (clause instanceof For) {
+            For forClause = (For) clause;
+
+            Object iterable = forClause.getIterable().eval(env);
+            Location loc = getLocation();
+            Iterable<?> listValue = EvalUtils.toIterable(iterable, loc, env);
+            EvalUtils.lock(iterable, loc);
+            try {
+              for (Object elem : listValue) {
+                Eval.assign(forClause.getVars(), elem, env, loc);
+                execClauses(index + 1);
+              }
+            } finally {
+              EvalUtils.unlock(iterable, loc);
+            }
+
+          } else {
+            If ifClause = (If) clause;
+            if (EvalUtils.toBoolean(ifClause.condition.eval(env))) {
+              execClauses(index + 1);
+            }
+          }
+          return;
+        }
+
+        // base case: evaluate body and add to result.
+        if (dict != null) {
+          DictionaryLiteral.Entry body = (DictionaryLiteral.Entry) Comprehension.this.body;
+          Object k = body.getKey().eval(env);
+          EvalUtils.checkValidDictKey(k, env);
+          Object v = body.getValue().eval(env);
+          dict.put(k, v, getLocation(), env);
+        } else {
+          list.add(((Expression) body).eval(env));
+        }
+      }
+    }
+    new Lambda().execClauses(0);
+
+    // Undefine loop variables (remove them from the environment).
+    // This code is useful for the transition, to make sure no one relies on the old behavior
+    // (where loop variables were leaking).
+    // TODO(laurentlb): Instead of removing variables, we should create them in a nested scope.
+    for (Clause clause : clauses) {
+      // Check if a loop variable conflicts with another local variable.
+      if (clause instanceof For) {
+        for (Identifier ident : Identifier.boundIdentifiers(((For) clause).getVars())) {
+          env.removeLocalBinding(ident.getName());
+        }
+      }
+    }
+
+    return isDict ? dict : SkylarkList.MutableList.copyOf(env, list);
+  }
+
+}
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
deleted file mode 100644
index a8094f3..0000000
--- a/src/main/java/com/google/devtools/build/lib/syntax/DictComprehension.java
+++ /dev/null
@@ -1,123 +0,0 @@
-// Copyright 2014 The Bazel Authors. All rights reserved.
-//
-// 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 com.google.devtools.build.lib.syntax;
-
-import com.google.common.base.Preconditions;
-import com.google.devtools.build.lib.events.Location;
-import java.io.IOException;
-import java.util.List;
-import java.util.Map;
-
-/** Syntax node for dictionary comprehension expressions. */
-public final class DictComprehension extends AbstractComprehension {
-  private final Expression keyExpression;
-  private final Expression valueExpression;
-
-  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
-  protected void printExpressions(Appendable buffer) throws IOException {
-    keyExpression.prettyPrint(buffer);
-    buffer.append(": ");
-    valueExpression.prettyPrint(buffer);
-  }
-
-  /** 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);
-  }
-
-  /**
-   * Helper class that collects the intermediate results of the {@link DictComprehension} and
-   * provides access to the resulting {@link Map}.
-   */
-  private final class DictOutputCollector implements OutputCollector {
-    private final SkylarkDict<Object, Object> result;
-
-    DictOutputCollector(Environment env) {
-      // We want to keep the iteration order
-      result = SkylarkDict.of(env);
-    }
-
-    @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, env);
-      result.put(key, valueExpression.eval(env), getLocation(), env);
-    }
-
-    @Override
-    public Object getResult(Environment env) throws EvalException {
-      return result;
-    }
-  }
-}
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 3dbc643..c7df54f 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
@@ -19,15 +19,16 @@
 import java.util.List;
 
 /** Syntax node for dictionary literals. */
+// TODO(adonovan): rename to DictExpression; it's not a literal.
 public final class DictionaryLiteral extends Expression {
 
-  /** Node for an individual key-value pair in a dictionary literal. */
-  public static final class DictionaryEntryLiteral extends ASTNode {
+  /** A key/value pair in a dict expression or comprehension. */
+  public static final class Entry extends ASTNode {
 
     private final Expression key;
     private final Expression value;
 
-    public DictionaryEntryLiteral(Expression key, Expression value) {
+    Entry(Expression key, Expression value) {
       this.key = key;
       this.value = value;
     }
@@ -53,9 +54,9 @@
     }
   }
 
-  private final ImmutableList<DictionaryEntryLiteral> entries;
+  private final ImmutableList<Entry> entries;
 
-  DictionaryLiteral(List<DictionaryEntryLiteral> entries) {
+  DictionaryLiteral(List<Entry> entries) {
     this.entries = ImmutableList.copyOf(entries);
   }
 
@@ -63,7 +64,7 @@
   Object doEval(Environment env) throws EvalException, InterruptedException {
     SkylarkDict<Object, Object> dict = SkylarkDict.of(env);
     Location loc = getLocation();
-    for (DictionaryEntryLiteral entry : entries) {
+    for (Entry entry : entries) {
       Object key = entry.key.eval(env);
       Object val = entry.value.eval(env);
       if (dict.containsKey(key)) {
@@ -79,7 +80,7 @@
   public void prettyPrint(Appendable buffer) throws IOException {
     buffer.append("{");
     String sep = "";
-    for (DictionaryEntryLiteral e : entries) {
+    for (Entry e : entries) {
       buffer.append(sep);
       e.prettyPrint(buffer);
       sep = ", ";
@@ -97,7 +98,7 @@
     return Kind.DICTIONARY_LITERAL;
   }
 
-  public ImmutableList<DictionaryEntryLiteral> getEntries() {
+  public ImmutableList<Entry> getEntries() {
     return entries;
   }
 }
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
deleted file mode 100644
index 12d31e6..0000000
--- a/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java
+++ /dev/null
@@ -1,104 +0,0 @@
-// Copyright 2014 The Bazel Authors. All rights reserved.
-//
-// 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 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.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-
-/** Syntax node for lists comprehension expressions. */
-public final class ListComprehension extends AbstractComprehension {
-  private final Expression outputExpression;
-
-  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
-  protected void printExpressions(Appendable buffer) throws IOException {
-    outputExpression.prettyPrint(buffer);
-  }
-
-  /** 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();
-  }
-
-  /**
-   * Helper class that collects the intermediate results of the {@code ListComprehension} and
-   * provides access to the resulting {@code List}.
-   */
-  private final class ListOutputCollector implements OutputCollector {
-    private final List<Object> result;
-
-    ListOutputCollector() {
-      result = new ArrayList<>();
-    }
-
-    @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 MutableList.copyOf(env, result);
-    }
-  }
-}
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 e2a6b88..f1ed111 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
@@ -26,7 +26,6 @@
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
-import com.google.devtools.build.lib.syntax.DictionaryLiteral.DictionaryEntryLiteral;
 import com.google.devtools.build.lib.syntax.IfStatement.ConditionalStatements;
 import java.util.ArrayList;
 import java.util.EnumSet;
@@ -568,8 +567,8 @@
   }
 
   // dict_entry_list ::= ( (dict_entry ',')* dict_entry ','? )?
-  private List<DictionaryEntryLiteral> parseDictEntryList() {
-    List<DictionaryEntryLiteral> list = new ArrayList<>();
+  private List<DictionaryLiteral.Entry> parseDictEntryList() {
+    List<DictionaryLiteral.Entry> list = new ArrayList<>();
     // the terminating token for a dict entry list
     while (token.kind != TokenKind.RBRACE) {
       list.add(parseDictEntry());
@@ -583,12 +582,12 @@
   }
 
   // dict_entry ::= nontupleexpr ':' nontupleexpr
-  private DictionaryEntryLiteral parseDictEntry() {
+  private DictionaryLiteral.Entry parseDictEntry() {
     int start = token.left;
     Expression key = parseNonTupleExpression();
     expect(TokenKind.COLON);
     Expression value = parseNonTupleExpression();
-    return setLocation(new DictionaryEntryLiteral(key, value), start, value);
+    return setLocation(new DictionaryLiteral.Entry(key, value), start, value);
   }
 
   /**
@@ -774,36 +773,37 @@
 
   // comprehension_suffix ::= 'FOR' loop_variables 'IN' expr comprehension_suffix
   //                        | 'IF' expr comprehension_suffix
-  //                        | ']'
-  private Expression parseComprehensionSuffix(
-      AbstractComprehension.AbstractBuilder comprehensionBuilder,
-      TokenKind closingBracket,
-      int comprehensionStartOffset) {
+  //                        | ']' | '}'
+  private Expression parseComprehensionSuffix(ASTNode body, TokenKind closingBracket, int offset) {
+    ImmutableList.Builder<Comprehension.Clause> clauses = ImmutableList.builder();
     while (true) {
       if (token.kind == TokenKind.FOR) {
         nextToken();
-        Expression lhs = parseForLoopVariables();
+        Expression vars = parseForLoopVariables();
         expect(TokenKind.IN);
         // 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);
-        comprehensionBuilder.addFor(lhs, listExpression);
+        Expression seq = parseNonTupleExpression(0);
+        clauses.add(new Comprehension.For(vars, seq));
       } else if (token.kind == TokenKind.IF) {
         nextToken();
         // [x for x in li if 1, 2]  # parse error
         // [x for x in li if (1, 2)]  # ok
-        comprehensionBuilder.addIf(parseNonTupleExpression(0));
+        Expression cond = parseNonTupleExpression(0);
+        clauses.add(new Comprehension.If(cond));
       } else if (token.kind == closingBracket) {
-        Expression expr = comprehensionBuilder.build();
-        setLocation(expr, comprehensionStartOffset, token.right);
-        nextToken();
-        return expr;
+        break;
       } else {
         syntaxError("expected '" + closingBracket + "', 'for' or 'if'");
         syncPast(LIST_TERMINATOR_SET);
-        return makeErrorExpression(comprehensionStartOffset, token.right);
+        return makeErrorExpression(offset, token.right);
       }
     }
+    boolean isDict = closingBracket == TokenKind.RBRACE;
+    Comprehension comp = new Comprehension(isDict, body, clauses.build());
+    setLocation(comp, offset, token.right);
+    nextToken();
+    return comp;
   }
 
   // list_maker ::= '[' ']'
@@ -832,10 +832,7 @@
         }
       case FOR:
         { // list comprehension
-          return parseComprehensionSuffix(
-              new ListComprehension.Builder().setOutputExpression(expression),
-              TokenKind.RBRACKET,
-              start);
+          return parseComprehensionSuffix(expression, TokenKind.RBRACKET, start);
         }
       case COMMA:
         {
@@ -877,17 +874,12 @@
       nextToken();
       return literal;
     }
-    DictionaryEntryLiteral entry = parseDictEntry();
+    DictionaryLiteral.Entry entry = parseDictEntry();
     if (token.kind == TokenKind.FOR) {
       // Dict comprehension
-      return parseComprehensionSuffix(
-          new DictComprehension.Builder()
-              .setKeyExpression(entry.getKey())
-              .setValueExpression(entry.getValue()),
-          TokenKind.RBRACE,
-          start);
+      return parseComprehensionSuffix(entry, TokenKind.RBRACE, start);
     }
-    List<DictionaryEntryLiteral> entries = new ArrayList<>();
+    List<DictionaryLiteral.Entry> entries = new ArrayList<>();
     entries.add(entry);
     if (token.kind == TokenKind.COMMA) {
       expect(TokenKind.COMMA);
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 348191f..25ed6fa 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
@@ -13,7 +13,6 @@
 // limitations under the License.
 package com.google.devtools.build.lib.syntax;
 
-import com.google.devtools.build.lib.syntax.DictionaryLiteral.DictionaryEntryLiteral;
 import com.google.devtools.build.lib.syntax.IfStatement.ConditionalStatements;
 import java.util.List;
 
@@ -73,14 +72,18 @@
 
   public void visit(@SuppressWarnings("unused") Identifier node) {}
 
-  public void visit(AbstractComprehension node) {
-    for (AbstractComprehension.Clause clause : node.getClauses()) {
-      visit(clause.getExpression());
-      if (clause.getLHS() != null) {
-        visit(clause.getLHS());
+  public void visit(Comprehension node) {
+    for (Comprehension.Clause clause : node.getClauses()) {
+      if (clause instanceof Comprehension.For) {
+        Comprehension.For forClause = (Comprehension.For) clause;
+        visit(forClause.getVars());
+        visit(forClause.getIterable());
+      } else {
+        Comprehension.If ifClause = (Comprehension.If) clause;
+        visit(ifClause.getCondition());
       }
     }
-    visitAll(node.getOutputExpressions());
+    visit(node.getBody());
   }
 
   public void visit(ForStatement node) {
@@ -149,7 +152,7 @@
     visitAll(node.getEntries());
   }
 
-  public void visit(DictionaryEntryLiteral node) {
+  public void visit(DictionaryLiteral.Entry node) {
     visit(node.getKey());
     visit(node.getValue());
   }
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/ValidationEnvironment.java b/src/main/java/com/google/devtools/build/lib/syntax/ValidationEnvironment.java
index 8c6662f..e9a7240 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/ValidationEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/ValidationEnvironment.java
@@ -244,21 +244,26 @@
   }
 
   @Override
-  public void visit(AbstractComprehension node) {
+  public void visit(Comprehension node) {
     openBlock(Scope.Local);
-    for (AbstractComprehension.Clause clause : node.getClauses()) {
-      if (clause.getLHS() != null) {
-        collectDefinitions(clause.getLHS());
+    for (Comprehension.Clause clause : node.getClauses()) {
+      if (clause instanceof Comprehension.For) {
+        Comprehension.For forClause = (Comprehension.For) clause;
+        collectDefinitions(forClause.getVars());
       }
     }
     // TODO(adonovan): opt: combine loops
-    for (AbstractComprehension.Clause clause : node.getClauses()) {
-      visit(clause.getExpression());
-      if (clause.getLHS() != null) {
-        assign(clause.getLHS());
+    for (Comprehension.Clause clause : node.getClauses()) {
+      if (clause instanceof Comprehension.For) {
+        Comprehension.For forClause = (Comprehension.For) clause;
+        visit(forClause.getIterable());
+        assign(forClause.getVars());
+      } else {
+        Comprehension.If ifClause = (Comprehension.If) clause;
+        visit(ifClause.getCondition());
       }
     }
-    visitAll(node.getOutputExpressions());
+    visit(node.getBody());
     closeBlock();
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
index b00c886..c406499 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
@@ -22,7 +22,6 @@
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.events.Location.LineAndColumn;
-import com.google.devtools.build.lib.syntax.DictionaryLiteral.DictionaryEntryLiteral;
 import com.google.devtools.build.lib.syntax.SkylarkImport.SkylarkImportSyntaxException;
 import com.google.devtools.build.lib.syntax.util.EvaluationTestCase;
 import java.util.LinkedList;
@@ -78,7 +77,7 @@
   }
 
   // helper func for testListLiterals:
-  private static int getIntElem(DictionaryEntryLiteral entry, boolean key) {
+  private static int getIntElem(DictionaryLiteral.Entry entry, boolean key) {
     return ((IntegerLiteral) (key ? entry.getKey() : entry.getValue())).getValue();
   }
 
@@ -88,7 +87,7 @@
   }
 
   // helper func for testListLiterals:
-  private static DictionaryEntryLiteral getElem(DictionaryLiteral list, int index) {
+  private static DictionaryLiteral.Entry getElem(DictionaryLiteral list, int index) {
     return list.getEntries().get(index);
   }
 
@@ -745,7 +744,7 @@
     DictionaryLiteral dictionaryList =
       (DictionaryLiteral) parseExpression("{1:42}"); // a singleton dictionary
     assertThat(dictionaryList.getEntries()).hasSize(1);
-    DictionaryEntryLiteral tuple = getElem(dictionaryList, 0);
+    DictionaryLiteral.Entry tuple = getElem(dictionaryList, 0);
     assertThat(getIntElem(tuple, true)).isEqualTo(1);
     assertThat(getIntElem(tuple, false)).isEqualTo(42);
   }
@@ -762,7 +761,7 @@
     DictionaryLiteral dictionaryList =
       (DictionaryLiteral) parseExpression("{1:42,}"); // a singleton dictionary
     assertThat(dictionaryList.getEntries()).hasSize(1);
-    DictionaryEntryLiteral tuple = getElem(dictionaryList, 0);
+    DictionaryLiteral.Entry tuple = getElem(dictionaryList, 0);
     assertThat(getIntElem(tuple, true)).isEqualTo(1);
     assertThat(getIntElem(tuple, false)).isEqualTo(42);
   }
@@ -772,7 +771,7 @@
     DictionaryLiteral dictionaryList = (DictionaryLiteral) parseExpression("{1:42,2:43,3:44}");
     assertThat(dictionaryList.getEntries()).hasSize(3);
     for (int i = 0; i < 3; i++) {
-      DictionaryEntryLiteral tuple = getElem(dictionaryList, i);
+      DictionaryLiteral.Entry tuple = getElem(dictionaryList, i);
       assertThat(getIntElem(tuple, true)).isEqualTo(i + 1);
       assertThat(getIntElem(tuple, false)).isEqualTo(i + 42);
     }
@@ -817,31 +816,38 @@
 
   @Test
   public void testListComprehensionEmptyList() throws Exception {
-    List<ListComprehension.Clause> clauses = ((ListComprehension) parseExpression(
-        "['foo/%s.java' % x for x in []]")).getClauses();
+    List<Comprehension.Clause> clauses =
+        ((Comprehension) parseExpression("['foo/%s.java' % x for x in []]")).getClauses();
     assertThat(clauses).hasSize(1);
-    assertThat(clauses.get(0).getExpression().toString()).isEqualTo("[]");
-    assertThat(clauses.get(0).getLHS().toString()).isEqualTo("x");
+    Comprehension.For for0 = (Comprehension.For) clauses.get(0);
+    assertThat(for0.getIterable().toString()).isEqualTo("[]");
+    assertThat(for0.getVars().toString()).isEqualTo("x");
   }
 
   @Test
   public void testListComprehension() throws Exception {
-    List<ListComprehension.Clause> clauses = ((ListComprehension) parseExpression(
-        "['foo/%s.java' % x for x in ['bar', 'wiz', 'quux']]")).getClauses();
+    List<Comprehension.Clause> clauses =
+        ((Comprehension) parseExpression("['foo/%s.java' % x for x in ['bar', 'wiz', 'quux']]"))
+            .getClauses();
     assertThat(clauses).hasSize(1);
-    assertThat(clauses.get(0).getLHS().toString()).isEqualTo("x");
-    assertThat(clauses.get(0).getExpression()).isInstanceOf(ListLiteral.class);
+    Comprehension.For for0 = (Comprehension.For) clauses.get(0);
+    assertThat(for0.getVars().toString()).isEqualTo("x");
+    assertThat(for0.getIterable()).isInstanceOf(ListLiteral.class);
   }
 
   @Test
   public void testForForListComprehension() throws Exception {
-    List<ListComprehension.Clause> clauses = ((ListComprehension) parseExpression(
-        "['%s/%s.java' % (x, y) for x in ['foo', 'bar'] for y in list]")).getClauses();
+    List<Comprehension.Clause> clauses =
+        ((Comprehension)
+                parseExpression("['%s/%s.java' % (x, y) for x in ['foo', 'bar'] for y in list]"))
+            .getClauses();
     assertThat(clauses).hasSize(2);
-    assertThat(clauses.get(0).getLHS().toString()).isEqualTo("x");
-    assertThat(clauses.get(0).getExpression()).isInstanceOf(ListLiteral.class);
-    assertThat(clauses.get(1).getLHS().toString()).isEqualTo("y");
-    assertThat(clauses.get(1).getExpression()).isInstanceOf(Identifier.class);
+    Comprehension.For for0 = (Comprehension.For) clauses.get(0);
+    assertThat(for0.getVars().toString()).isEqualTo("x");
+    assertThat(for0.getIterable()).isInstanceOf(ListLiteral.class);
+    Comprehension.For for1 = (Comprehension.For) clauses.get(1);
+    assertThat(for1.getVars().toString()).isEqualTo("y");
+    assertThat(for1.getIterable()).isInstanceOf(Identifier.class);
   }
 
   @Test