src/main/java/com/google/devtools/build/lib/syntax/AbstractComprehension.java

// 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 com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
import java.io.IOException;
import java.io.Serializable;
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.
 */
public abstract class AbstractComprehension extends Expression {

  /**
   * The interface implemented by ForClause and (later) IfClause.
   * A comprehension consists of one or many Clause.
   */
  public interface Clause extends Serializable {

    /** 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 #getLValue} 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 LValue defined in Clause, i.e. the loop variables for ForClause and null for
     * IfClause. This is needed for SyntaxTreeVisitor.
     */
    @Nullable  // for the IfClause
    LValue getLValue();

    /**
     * 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. */
  @AutoCodec
  public static final class ForClause implements Clause {
    private final LValue lvalue;
    private final Expression iterable;

    @Override
    public Kind getKind() {
      return Kind.FOR;
    }

    public ForClause(LValue lvalue, Expression iterable) {
      this.lvalue = lvalue;
      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) {
          lvalue.assign(listElement, env, loc);
          evalStep(env, collector, step);
        }
      } finally {
        EvalUtils.unlock(iterableObject, loc);
      }
    }

    @Override
    public LValue getLValue() {
      return lvalue;
    }

    @Override
    public Expression getExpression() {
      return iterable;
    }

    @Override
    public void prettyPrint(Appendable buffer) throws IOException {
      buffer.append("for ");
      lvalue.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. */
  @AutoCodec
  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 LValue getLValue() {
      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;

  public 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(LValue lvalue, Expression iterable) {
      Clause forClause = new ForClause(lvalue, 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.
      LValue lvalue = clause.getLValue();
      if (lvalue != null) {
        for (Identifier ident : lvalue.boundIdentifiers()) {
          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;
  }
}