Compile list and dict comprehensions to byte code.

--
MOS_MIGRATED_REVID=107394025

3 files changed, 209 insertions, 1 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 f37a2a4ce2..962f2fcb31 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
@@ -14,13 +14,29 @@
 
 package com.google.devtools.build.lib.syntax;
 
+import static com.google.devtools.build.lib.syntax.compiler.ByteCodeUtils.append;
+
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.events.Location;
+import com.google.devtools.build.lib.syntax.compiler.ByteCodeMethodCalls;
+import com.google.devtools.build.lib.syntax.compiler.ByteCodeUtils;
+import com.google.devtools.build.lib.syntax.compiler.DebugInfo;
+import com.google.devtools.build.lib.syntax.compiler.DebugInfo.AstAccessors;
+import com.google.devtools.build.lib.syntax.compiler.Jump;
+import com.google.devtools.build.lib.syntax.compiler.Jump.PrimitiveComparison;
+import com.google.devtools.build.lib.syntax.compiler.LabelAdder;
+import com.google.devtools.build.lib.syntax.compiler.Variable.InternalVariable;
+import com.google.devtools.build.lib.syntax.compiler.VariableScope;
+
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
 
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
+import java.util.ListIterator;
 
 import javax.annotation.Nullable;
 
@@ -63,6 +79,13 @@ public abstract class AbstractComprehension extends Expression {
     abstract void eval(Environment env, OutputCollector collector, int step)
         throws EvalException, InterruptedException;
 
+    ByteCodeAppender compile(
+        ByteCodeAppender inner,
+        VariableScope scope,
+        DebugInfo debugInfo,
+        ASTNode node,
+        AstAccessors debugAccessors) throws EvalException;
+
     abstract void validate(ValidationEnvironment env) throws EvalException;
 
     /**
@@ -123,6 +146,48 @@ public abstract class AbstractComprehension extends Expression {
     public String toString() {
       return Printer.format("for %s in %r", variables.toString(), list);
     }
+
+    @Override
+    public ByteCodeAppender compile(
+        @Nullable ByteCodeAppender inner,
+        VariableScope scope,
+        DebugInfo debugInfo,
+        ASTNode node,
+        AstAccessors debugAccessors)
+        throws EvalException {
+      List<ByteCodeAppender> code = new ArrayList<>();
+      InternalVariable iterator =
+          scope.freshVariable(new TypeDescription.ForLoadedType(Iterator.class));
+      // compute the collection and get it on the stack and transform it to the right type
+      code.add(list.compile(scope, debugInfo));
+      append(
+          code,
+          debugAccessors.loadLocation,
+          EvalUtils.toIterable,
+          ByteCodeMethodCalls.BCImmutableList.copyOf,
+          ByteCodeMethodCalls.BCImmutableList.iterator);
+      code.add(iterator.store());
+      LabelAdder loopHeader = new LabelAdder();
+      LabelAdder loopBody = new LabelAdder();
+      append(
+          code,
+          Jump.to(loopHeader),
+          loopBody,
+          iterator.load(),
+          ByteCodeMethodCalls.BCIterator.next);
+      // store current element into l-values
+      code.add(variables.compileAssignment(node, debugAccessors, scope));
+      code.add(inner);
+      // compile code for the loop header
+      append(
+          code,
+          loopHeader,
+          iterator.load(),
+          ByteCodeMethodCalls.BCIterator.hasNext,
+          // falls through to end of loop if hasNext() was false, otherwise jumps back
+          Jump.ifIntOperandToZero(PrimitiveComparison.NOT_EQUAL).to(loopBody));
+      return ByteCodeUtils.compoundAppender(code);
+    }
   }
 
   /**
@@ -162,6 +227,27 @@ public abstract class AbstractComprehension extends Expression {
     public String toString() {
       return String.format("if %s", condition);
     }
+
+    @Override
+    public ByteCodeAppender compile(
+        ByteCodeAppender inner,
+        VariableScope scope,
+        DebugInfo debugInfo,
+        ASTNode node,
+        AstAccessors debugAccessors) throws EvalException {
+      List<ByteCodeAppender> code = new ArrayList<>();
+      LabelAdder nopeLabel = new LabelAdder();
+      // compile condition and convert to boolean
+      code.add(condition.compile(scope, debugInfo));
+      append(
+          code,
+          EvalUtils.toBoolean,
+          // don't execute inner if false
+          Jump.ifIntOperandToZero(PrimitiveComparison.EQUAL).to(nopeLabel));
+      code.add(inner);
+      append(code, nopeLabel);
+      return ByteCodeUtils.compoundAppender(code);
+    }
   }
 
   /**
@@ -243,6 +329,27 @@ public abstract class AbstractComprehension extends Expression {
     }
   }
 
+  @Override
+  ByteCodeAppender compile(VariableScope scope, DebugInfo debugInfo) throws EvalException {
+    // We use a new scope for all comprehensions, as in Python 3 semantics
+    // In Python 2, list comprehensions are in the same scope as the function and the backported
+    // dict comprehensions (2.7) use a new scope
+    VariableScope ourScope = scope.createSubScope();
+    List<ByteCodeAppender> code = new ArrayList<>();
+    InternalVariable collection = compileInitialization(ourScope, code);
+    AstAccessors debugAccessors = debugInfo.add(this);
+    ByteCodeAppender collector = compileCollector(ourScope, collection, debugInfo, debugAccessors);
+    for (ListIterator<Clause> clauseIterator = clauses.listIterator(clauses.size());
+        clauseIterator.hasPrevious();
+        ) {
+      Clause clause = clauseIterator.previous();
+      collector = clause.compile(collector, ourScope, debugInfo, this, debugAccessors);
+    }
+    code.add(collector);
+    code.add(compileBuilding(ourScope, collection));
+    return ByteCodeUtils.compoundAppender(code);
+  }
+
   /**
    * 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
@@ -269,6 +376,25 @@ public abstract class AbstractComprehension extends Expression {
   abstract OutputCollector createCollector();
 
   /**
+   * Add byte code which initializes the collection and returns the variable it is saved in.
+   */
+  abstract InternalVariable compileInitialization(VariableScope scope, List<ByteCodeAppender> code);
+
+  /**
+   * Add byte code which adds a value to the collection.
+   */
+  abstract ByteCodeAppender compileCollector(
+      VariableScope scope,
+      InternalVariable collection,
+      DebugInfo debugInfo,
+      AstAccessors debugAccessors) throws EvalException;
+
+  /**
+   * Add byte code which finalizes and loads the collection on the stack.
+   */
+  abstract ByteCodeAppender compileBuilding(VariableScope scope, InternalVariable collection);
+
+  /**
    * Interface for collecting the intermediate output of an {@code AbstractComprehension} and for
    * providing access to the final results.
    */
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 2d6668a72d..51995bf8ed 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
@@ -13,9 +13,23 @@
 // limitations under the License.
 package com.google.devtools.build.lib.syntax;
 
+import static com.google.devtools.build.lib.syntax.compiler.ByteCodeUtils.append;
+
 import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.syntax.compiler.ByteCodeMethodCalls;
+import com.google.devtools.build.lib.syntax.compiler.ByteCodeUtils;
+import com.google.devtools.build.lib.syntax.compiler.DebugInfo;
+import com.google.devtools.build.lib.syntax.compiler.DebugInfo.AstAccessors;
+import com.google.devtools.build.lib.syntax.compiler.Variable.InternalVariable;
+import com.google.devtools.build.lib.syntax.compiler.VariableScope;
+
+import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
+import net.bytebuddy.implementation.bytecode.Duplication;
+import net.bytebuddy.implementation.bytecode.Removal;
 
+import java.util.ArrayList;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -41,6 +55,35 @@ public class DictComprehension extends AbstractComprehension {
     return new DictOutputCollector();
   }
 
+  @Override
+  InternalVariable compileInitialization(VariableScope scope, List<ByteCodeAppender> code) {
+    InternalVariable dict = scope.freshVariable(ImmutableMap.class);
+    append(code, ByteCodeMethodCalls.BCImmutableMap.builder);
+    code.add(dict.store());
+    return dict;
+  }
+
+  @Override
+  ByteCodeAppender compileCollector(
+      VariableScope scope,
+      InternalVariable collection,
+      DebugInfo debugInfo,
+      AstAccessors debugAccessors) throws EvalException {
+    List<ByteCodeAppender> code = new ArrayList<>();
+    append(code, collection.load());
+    code.add(keyExpression.compile(scope, debugInfo));
+    append(code, Duplication.SINGLE, EvalUtils.checkValidDictKey);
+    code.add(valueExpression.compile(scope, debugInfo));
+    append(code, ByteCodeMethodCalls.BCImmutableMap.Builder.put, Removal.SINGLE);
+    return ByteCodeUtils.compoundAppender(code);
+  }
+
+  @Override
+  ByteCodeAppender compileBuilding(VariableScope scope, InternalVariable collection) {
+    return new ByteCodeAppender.Simple(
+        collection.load(), ByteCodeMethodCalls.BCImmutableMap.Builder.build);
+  }
+
   /**
    * Helper class that collects the intermediate results of the {@link DictComprehension} and
    * provides access to the resulting {@link Map}.
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 567fdaf194..9573076b99 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
@@ -11,10 +11,21 @@
 // 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 static com.google.devtools.build.lib.syntax.compiler.ByteCodeUtils.append;
+
 import com.google.devtools.build.lib.syntax.SkylarkList.MutableList;
+import com.google.devtools.build.lib.syntax.compiler.ByteCodeMethodCalls;
+import com.google.devtools.build.lib.syntax.compiler.ByteCodeUtils;
+import com.google.devtools.build.lib.syntax.compiler.DebugInfo;
+import com.google.devtools.build.lib.syntax.compiler.DebugInfo.AstAccessors;
+import com.google.devtools.build.lib.syntax.compiler.NewObject;
+import com.google.devtools.build.lib.syntax.compiler.Variable.InternalVariable;
+import com.google.devtools.build.lib.syntax.compiler.VariableScope;
+
+import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
+import net.bytebuddy.implementation.bytecode.Removal;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -41,6 +52,34 @@ public final class ListComprehension extends AbstractComprehension {
     return new ListOutputCollector();
   }
 
+  @Override
+  InternalVariable compileInitialization(VariableScope scope, List<ByteCodeAppender> code) {
+    InternalVariable list = scope.freshVariable(ArrayList.class);
+    append(code, NewObject.fromConstructor(ArrayList.class).arguments());
+    code.add(list.store());
+    return list;
+  }
+
+  @Override
+  ByteCodeAppender compileCollector(
+      VariableScope scope,
+      InternalVariable collection,
+      DebugInfo debugInfo,
+      AstAccessors debugAccessors) throws EvalException {
+    List<ByteCodeAppender> code = new ArrayList<>();
+    append(code, collection.load());
+    code.add(outputExpression.compile(scope, debugInfo));
+    append(code, ByteCodeMethodCalls.BCList.add, Removal.SINGLE);
+    return ByteCodeUtils.compoundAppender(code);
+  }
+
+  @Override
+  ByteCodeAppender compileBuilding(VariableScope scope, InternalVariable collection) {
+    return new ByteCodeAppender.Simple(
+        NewObject.fromConstructor(MutableList.class, Iterable.class, Environment.class)
+            .arguments(collection.load(), scope.loadEnvironment()));
+  }
+
   /**
    * Helper class that collects the intermediate results of the {@code ListComprehension} and
    * provides access to the resulting {@code List}.