In SkyQueryEnvironment, rewrite queries using the semantics-preserving transformation 'rdeps(<sky_query_environment_universe_scope>, T, depth)' -> 'allrdeps(T, depth)'.

SkyQueryEnvironment can evaluate such allrdeps queries much more efficiently since it doesn't need to bother filtering out targets outside of universe, meaning it doesn't need to have all targets in the universe in memory at the same time.

--
MOS_MIGRATED_REVID=116075008

12 files changed, 204 insertions, 4 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
index 99da369f4b..27a0bde0ab 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
@@ -184,6 +184,10 @@ public abstract class AbstractBlazeQueryEnvironment<T> implements QueryEnvironme
     return new QueryEvalResult(!eventHandler.hasErrors(), empty.get());
   }
 
+  public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
+    return queryExpression;
+  }
+
   public QueryEvalResult evaluateQuery(String query, Callback<T> callback)
       throws QueryException, InterruptedException {
     return evaluateQuery(QueryExpression.parse(query, this), callback);
diff --git a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
index f5ffe544d9..f5186c85d3 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
@@ -47,10 +47,14 @@ import com.google.devtools.build.lib.pkgcache.TargetPatternEvaluator;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.query2.engine.AllRdepsFunction;
 import com.google.devtools.build.lib.query2.engine.Callback;
+import com.google.devtools.build.lib.query2.engine.FunctionExpression;
 import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
+import com.google.devtools.build.lib.query2.engine.QueryExpressionMapper;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.AbstractUniquifier;
+import com.google.devtools.build.lib.query2.engine.RdepsFunction;
+import com.google.devtools.build.lib.query2.engine.TargetLiteral;
 import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.lib.skyframe.FileValue;
 import com.google.devtools.build.lib.skyframe.GraphBackedRecursivePackageProvider;
@@ -173,6 +177,42 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
   }
 
   @Override
+  public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
+    // Transform each occurrence of an expressions of the form 'rdeps(<universeScope>, <T>)' to
+    // 'allrdeps(<T>)'. The latter is more efficient.
+    if (universeScope.size() != 1) {
+      return queryExpression;
+    }
+    final TargetPattern.Parser targetPatternParser = new TargetPattern.Parser(parserPrefix);
+    String universeScopePattern = Iterables.getOnlyElement(universeScope);
+    final String absoluteUniverseScopePattern =
+        targetPatternParser.absolutize(universeScopePattern);
+    QueryExpressionMapper rdepsToAllRDepsMapper = new QueryExpressionMapper() {
+      @Override
+      public QueryExpression map(FunctionExpression functionExpression) {
+        if (functionExpression.getFunction().getName().equals(new RdepsFunction().getName())) {
+          List<Argument> args = functionExpression.getArgs();
+          QueryExpression universeExpression = args.get(0).getExpression();
+          if (universeExpression instanceof TargetLiteral) {
+            TargetLiteral literalUniverseExpression = (TargetLiteral) universeExpression;
+            String absolutizedUniverseExpression =
+                targetPatternParser.absolutize(literalUniverseExpression.getPattern());
+            if (absolutizedUniverseExpression.equals(absoluteUniverseScopePattern)) {
+              List<Argument> argsTail = args.subList(1, functionExpression.getArgs().size());
+              return new FunctionExpression(new AllRdepsFunction(), argsTail);
+            }
+          }
+        }
+        return super.map(functionExpression);
+      }
+    };
+    QueryExpression transformedQueryExpression = queryExpression.getMapped(rdepsToAllRDepsMapper);
+    LOG.info(String.format("transformed query [%s] to [%s]", queryExpression,
+        transformedQueryExpression));
+    return transformedQueryExpression;
+  }
+
+  @Override
   public QueryEvalResult evaluateQuery(QueryExpression expr, Callback<Target> callback)
       throws QueryException, InterruptedException {
     // Some errors are reported as QueryExceptions and others as ERROR events (if --keep_going). The
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java
index 60d38e7a35..66555c7bee 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java
@@ -45,6 +45,14 @@ class BinaryOperatorExpression extends QueryExpression {
     this.operands = ImmutableList.copyOf(operands);
   }
 
+  Lexer.TokenKind getOperator() {
+    return operator;
+  }
+
+  ImmutableList<QueryExpression> getOperands() {
+    return operands;
+  }
+
   @Override
   public <T> void eval(QueryEnvironment<T> env, Callback<T> callback)
       throws QueryException, InterruptedException {
@@ -85,6 +93,11 @@ class BinaryOperatorExpression extends QueryExpression {
   }
 
   @Override
+  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+    return mapper.map(this);
+  }
+
+  @Override
   public String toString() {
     StringBuilder result = new StringBuilder();
     for (int i = 1; i < operands.size(); i++) {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java
index d74056516f..f6fcc27ef5 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java
@@ -36,6 +36,14 @@ public class FunctionExpression extends QueryExpression {
     this.args = ImmutableList.copyOf(args);
   }
 
+  public QueryFunction getFunction() {
+    return function;
+  }
+
+  public List<Argument> getArgs() {
+    return args;
+  }
+
   @Override
   public <T> void eval(QueryEnvironment<T> env, Callback<T> callback)
       throws QueryException, InterruptedException {
@@ -52,6 +60,11 @@ public class FunctionExpression extends QueryExpression {
   }
 
   @Override
+  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+    return mapper.map(this);
+  }
+
+  @Override
   public String toString() {
     return function.getName() +
         "(" + Joiner.on(", ").join(Iterables.transform(args, Functions.toStringFunction())) + ")";
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java
index 2e198077bb..a1cc7a7cbf 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java
@@ -51,6 +51,18 @@ class LetExpression extends QueryExpression {
     this.bodyExpr = bodyExpr;
   }
 
+  String getVarName() {
+    return varName;
+  }
+
+  QueryExpression getVarExpr() {
+    return varExpr;
+  }
+
+  QueryExpression getBodyExpr() {
+    return bodyExpr;
+  }
+
   @Override
   public <T> void eval(QueryEnvironment<T> env, Callback<T> callback)
       throws QueryException, InterruptedException {
@@ -77,6 +89,11 @@ class LetExpression extends QueryExpression {
   }
 
   @Override
+  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+    return mapper.map(this);
+  }
+
+  @Override
   public String toString() {
     return "let " + varName + " = " + varExpr + " in " + bodyExpr;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java
index 2732b6886c..b29071329b 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java
@@ -76,6 +76,9 @@ public abstract class QueryExpression {
    */
   public abstract void collectTargetPatterns(Collection<String> literals);
 
+  /* Implementations should just be {@code return mapper.map(this)}. */
+  public abstract QueryExpression getMapped(QueryExpressionMapper mapper);
+
   /**
    * Returns this query expression pretty-printed.
    */
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionMapper.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionMapper.java
new file mode 100644
index 0000000000..f04ed1ec4b
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionMapper.java
@@ -0,0 +1,92 @@
+// Copyright 2016 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.query2.engine;
+
+import com.google.common.collect.ImmutableList;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
+
+/**
+ * Performs an arbitrary transformation of a {@link QueryExpression}.
+ *
+ * <p>For each subclass of {@link QueryExpression}, there's a corresponding {@link #map} overload
+ * that transforms a node of that type. By default, this method recursively applies this
+ * {@link QueryExpressionMapper}'s transformation in a structure-preserving manner (trying to
+ * maintain reference-equality, as an optimization). Subclasses of {@link QueryExpressionMapper} can
+ * override these methods in order to implement an arbitrary transformation.
+ */
+public class QueryExpressionMapper {
+  public QueryExpression map(TargetLiteral targetLiteral) {
+    return targetLiteral;
+  }
+
+  public QueryExpression map(BinaryOperatorExpression binaryOperatorExpression) {
+    boolean changed = false;
+    ImmutableList.Builder<QueryExpression> mappedOperandsBuilder = ImmutableList.builder();
+    for (QueryExpression operand : binaryOperatorExpression.getOperands()) {
+      QueryExpression mappedOperand = operand.getMapped(this);
+      if (mappedOperand != operand) {
+        changed = true;
+      }
+      mappedOperandsBuilder.add(mappedOperand);
+    }
+    return changed
+        ? new BinaryOperatorExpression(
+            binaryOperatorExpression.getOperator(), mappedOperandsBuilder.build())
+        : binaryOperatorExpression;
+  }
+
+  public QueryExpression map(FunctionExpression functionExpression) {
+    boolean changed = false;
+    ImmutableList.Builder<Argument> mappedArgumentBuilder = ImmutableList.builder();
+    for (Argument argument : functionExpression.getArgs()) {
+      switch (argument.getType()) {
+        case EXPRESSION: {
+          QueryExpression expr = argument.getExpression();
+          QueryExpression mappedExpression = expr.getMapped(this);
+          mappedArgumentBuilder.add(Argument.of(mappedExpression));
+          if (expr != mappedExpression) {
+            changed = true;
+          }
+          break;
+        }
+        default:
+          mappedArgumentBuilder.add(argument);
+          break;
+      }
+    }
+    return changed
+        ? new FunctionExpression(functionExpression.getFunction(), mappedArgumentBuilder.build())
+        : functionExpression;
+  }
+
+  public QueryExpression map(LetExpression letExpression) {
+    boolean changed = false;
+    QueryExpression mappedVarExpr = letExpression.getVarExpr().getMapped(this);
+    if (mappedVarExpr != letExpression.getVarExpr()) {
+      changed = true;
+    }
+    QueryExpression mappedBodyExpr = letExpression.getBodyExpr().getMapped(this);
+    if (mappedBodyExpr != letExpression.getBodyExpr()) {
+      changed = true;
+    }
+    return changed
+        ? new LetExpression(letExpression.getVarName(), mappedVarExpr, mappedBodyExpr)
+        : letExpression;
+  }
+
+  public QueryExpression map(SetExpression setExpression) {
+    return setExpression;
+  }
+}
+
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryParser.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryParser.java
index 46d4b94b17..434c56fce6 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryParser.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryParser.java
@@ -68,6 +68,9 @@ final class QueryParser {
   }
 
   private QueryParser(List<Lexer.Token> tokens, QueryEnvironment<?> env) {
+    // TODO(bazel-team): We only need QueryEnvironment#getFunctions, consider refactoring users of
+    // QueryParser#parse to instead just pass in the set of functions to make testing, among other
+    // things, simpler.
     this.functions = new HashMap<>();
     for (QueryFunction queryFunction : env.getFunctions()) {
       this.functions.put(queryFunction.getName(), queryFunction);
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/RdepsFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/RdepsFunction.java
index 24d0334df2..6cd72dd75a 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/RdepsFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/RdepsFunction.java
@@ -30,8 +30,8 @@ import java.util.Set;
  * <pre>expr ::= RDEPS '(' expr ',' expr ')'</pre>
  * <pre>       | RDEPS '(' expr ',' expr ',' WORD ')'</pre>
  */
-final class RdepsFunction extends AllRdepsFunction {
-  RdepsFunction() {}
+public final class RdepsFunction extends AllRdepsFunction {
+  public RdepsFunction() {}
 
   @Override
   public String getName() {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java
index 41aa9281dc..1369d8ace6 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java
@@ -61,6 +61,11 @@ class SetExpression extends QueryExpression {
   }
 
   @Override
+  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+    return mapper.map(this);
+  }
+
+  @Override
   public String toString() {
     return "set(" + Joiner.on(' ').join(words) + ")";
   }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java b/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java
index 32eca663eb..e7d8c00d73 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java
@@ -28,7 +28,7 @@ import java.util.Set;
  *
  * <pre>expr ::= NAME | WORD</pre>
  */
-final class TargetLiteral extends QueryExpression {
+public final class TargetLiteral extends QueryExpression {
 
   private final String pattern;
 
@@ -36,6 +36,10 @@ final class TargetLiteral extends QueryExpression {
     this.pattern = Preconditions.checkNotNull(pattern);
   }
 
+  public String getPattern() {
+    return pattern;
+  }
+
   public boolean isVariableReference() {
     return LetExpression.isValidVarReference(pattern);
   }
@@ -63,6 +67,11 @@ final class TargetLiteral extends QueryExpression {
   }
 
   @Override
+  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+    return mapper.map(this);
+  }
+
+  @Override
   public String toString() {
     // Keep predicate consistent with Lexer.scanWord!
     boolean needsQuoting = Lexer.isReservedWord(pattern)
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
index 45a76e5ed9..32159f413a 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
@@ -138,7 +138,7 @@ public final class QueryCommand implements BlazeCommand {
         queryOptions.universeScope, queryOptions.loadingPhaseThreads,
         settings);
 
-    // 1. Parse query:
+    // 1. Parse and transform query:
     QueryExpression expr;
     try {
       expr = QueryExpression.parse(query, queryEnv);
@@ -147,6 +147,7 @@ public final class QueryCommand implements BlazeCommand {
           null, "Error while parsing '" + query + "': " + e.getMessage()));
       return ExitCode.COMMAND_LINE_ERROR;
     }
+    expr = queryEnv.transformParsedQuery(expr);
 
     QueryEvalResult result;
     PrintStream output = null;