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
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 99da369..27a0bde 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 @@
     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 f5ffe54..f5186c8 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.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 @@
   }
 
   @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 60d38e7..66555c7 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 @@
     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 @@
   }
 
   @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 d740565..f6fcc27 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 @@
     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 @@
   }
 
   @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 2e19807..a1cc7a7 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 @@
     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 @@
   }
 
   @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 2732b68..b290713 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 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 0000000..f04ed1e
--- /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 46d4b94..434c56f 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 @@
   }
 
   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 24d0334..6cd72dd 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 @@
  * <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 41aa928..1369d8a 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 @@
   }
 
   @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 32eca66..e7d8c00 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 @@
  *
  * <pre>expr ::= NAME | WORD</pre>
  */
-final class TargetLiteral extends QueryExpression {
+public final class TargetLiteral extends QueryExpression {
 
   private final String pattern;
 
@@ -36,6 +36,10 @@
     this.pattern = Preconditions.checkNotNull(pattern);
   }
 
+  public String getPattern() {
+    return pattern;
+  }
+
   public boolean isVariableReference() {
     return LetExpression.isValidVarReference(pattern);
   }
@@ -63,6 +67,11 @@
   }
 
   @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 45a76e5..32159f4 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 @@
         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 @@
           null, "Error while parsing '" + query + "': " + e.getMessage()));
       return ExitCode.COMMAND_LINE_ERROR;
     }
+    expr = queryEnv.transformParsedQuery(expr);
 
     QueryEvalResult result;
     PrintStream output = null;