Change WalkableGraphFactory#prepareAndGet to take multiple SkyKeys as graph roots

It also changes a few accessors of utility methods in Skyframe library. It
refactors the QueryExpressionMapper to use a general QueryExpressionVisitor.

RELNOTES: None
PiperOrigin-RevId: 165534908

17 files changed, 252 insertions, 90 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java b/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java
index 43f77c8089..92a6dbb1e5 100644
--- a/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java
+++ b/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java
@@ -293,6 +293,11 @@ public abstract class TargetPattern implements Serializable {
     throw new IllegalStateException();
   }
 
+  /** For patterns of type {@link Type#SINGLE_TARGET}, returns the target path. */
+  public String getSingleTargetPath() {
+    throw new IllegalStateException();
+  }
+
   /**
    * For patterns of type {@link Type#SINGLE_TARGET} and {@link Type#TARGETS_IN_PACKAGE}, returns
    * the {@link PackageIdentifier} corresponding to the package that would contain the target(s)
@@ -349,6 +354,11 @@ public abstract class TargetPattern implements Serializable {
     }
 
     @Override
+    public String getSingleTargetPath() {
+      return targetName;
+    }
+
+    @Override
     public boolean equals(Object o) {
       if (this == o) {
         return true;
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 0121e00875..383a28ffb8 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
@@ -163,7 +163,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   private GraphBackedRecursivePackageProvider graphBackedRecursivePackageProvider;
   private ListeningExecutorService executor;
   private RecursivePackageProviderBackedTargetPatternResolver resolver;
-  private final SkyKey universeKey;
+  protected final SkyKey universeKey;
   private final ImmutableList<TargetPatternKey> universeTargetPatternKeys;
 
   public SkyQueryEnvironment(
@@ -239,17 +239,27 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     }
   }
 
-  private void beforeEvaluateQuery() throws InterruptedException {
-    if (graph == null || !graphFactory.isUpToDate(universeKey)) {
+  /** Gets roots of graph which contains all nodes needed to evaluate {@code expr}. */
+  protected Set<SkyKey> getGraphRootsFromExpression(QueryExpression expr)
+      throws QueryException, InterruptedException {
+    return ImmutableSet.of(universeKey);
+  }
+
+  private void beforeEvaluateQuery(QueryExpression expr)
+      throws QueryException, InterruptedException {
+    Set<SkyKey> roots = getGraphRootsFromExpression(expr);
+    if (graph == null || !graphFactory.isUpToDate(roots)) {
       // If this environment is uninitialized or the graph factory needs to evaluate, do so. We
       // assume here that this environment cannot be initialized-but-stale if the factory is up
       // to date.
       EvaluationResult<SkyValue> result;
       try (AutoProfiler p = AutoProfiler.logged("evaluation and walkable graph", LOG)) {
-        result =
-            graphFactory.prepareAndGet(universeKey, loadingPhaseThreads, universeEvalEventHandler);
+        result = graphFactory.prepareAndGet(roots, loadingPhaseThreads, universeEvalEventHandler);
+      }
+
+      if (roots.size() == 1 && Iterables.getOnlyElement(roots).equals(universeKey)) {
+        checkEvaluationResult(result);
       }
-      checkEvaluationResult(result);
 
       packageSemaphore = makeFreshPackageMultisetSemaphore();
       graph = result.getWalkableGraph();
@@ -335,7 +345,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     }
 
     @Override
-    public QueryExpression map(FunctionExpression functionExpression) {
+    public QueryExpression visit(FunctionExpression functionExpression) {
       if (functionExpression.getFunction().getName().equals(new RdepsFunction().getName())) {
         List<Argument> args = functionExpression.getArgs();
         QueryExpression universeExpression = args.get(0).getExpression();
@@ -349,14 +359,14 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
           }
         }
       }
-      return super.map(functionExpression);
+      return super.visit(functionExpression);
     }
   }
 
   @Override
   public final QueryExpression transformParsedQuery(QueryExpression queryExpression) {
     QueryExpressionMapper mapper = getQueryExpressionMapper();
-    QueryExpression transformedQueryExpression = queryExpression.getMapped(mapper);
+    QueryExpression transformedQueryExpression = queryExpression.accept(mapper);
     LOG.info(String.format(
         "transformed query [%s] to [%s]",
         Ascii.truncate(
@@ -406,7 +416,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   public QueryEvalResult evaluateQuery(
       QueryExpression expr, ThreadSafeOutputFormatterCallback<Target> callback)
           throws QueryException, InterruptedException, IOException {
-    beforeEvaluateQuery();
+    beforeEvaluateQuery(expr);
 
     // SkyQueryEnvironment batches callback invocations using a BatchStreamedCallback, created here
     // so that there's one per top-level evaluateQuery call. The batch size is large enough that
@@ -1153,7 +1163,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
       implements InterruptibleSupplier<ImmutableSet<PathFragment>> {
     private final WalkableGraph graph;
 
-    BlacklistSupplier(WalkableGraph graph) {
+    private BlacklistSupplier(WalkableGraph graph) {
       this.graph = graph;
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/AggregatingQueryExpressionVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/engine/AggregatingQueryExpressionVisitor.java
new file mode 100644
index 0000000000..cde80e072a
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/AggregatingQueryExpressionVisitor.java
@@ -0,0 +1,114 @@
+// Copyright 2017 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.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.primitives.Booleans;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.ArgumentType;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
+
+/**
+ * An implementation of {@link QueryExpressionVisitor} which recursively visits all nested {@link
+ * QueryExpression}s
+ */
+public abstract class AggregatingQueryExpressionVisitor<T> implements QueryExpressionVisitor<T> {
+
+  @Override
+  public T visit(BinaryOperatorExpression binaryOperatorExpression) {
+    ImmutableMap.Builder<QueryExpression, T> builder = ImmutableMap.builder();
+    for (QueryExpression expr : binaryOperatorExpression.getOperands()) {
+      builder.put(expr, expr.accept(this));
+    }
+
+    return aggregate(builder.build());
+  }
+
+  @Override
+  public T visit(FunctionExpression functionExpression) {
+    ImmutableMap.Builder<QueryExpression, T> builder = ImmutableMap.builder();
+    for (Argument argument : functionExpression.getArgs()) {
+      if (argument.getType() == ArgumentType.EXPRESSION) {
+        builder.put(argument.getExpression(), argument.getExpression().accept(this));
+      }
+    }
+
+    return aggregate(builder.build());
+  }
+
+  @Override
+  public T visit(LetExpression letExpression) {
+    return aggregate(
+        ImmutableMap.of(
+            letExpression.getVarExpr(), letExpression.getVarExpr().accept(this),
+            letExpression.getBodyExpr(), letExpression.getBodyExpr().accept(this)));
+  }
+
+  @Override
+  public T visit(SetExpression setExpression) {
+    ImmutableMap.Builder<QueryExpression, T> builder = ImmutableMap.builder();
+    for (TargetLiteral targetLiteral : setExpression.getWords()) {
+      builder.put(targetLiteral, targetLiteral.accept(this));
+    }
+
+    return aggregate(builder.build());
+  }
+
+  /**
+   * Aggregates results from all sub-expression visitations. The map is guaranteed to include
+   * results for all direct sub-expressions.
+   */
+  protected abstract T aggregate(ImmutableMap<QueryExpression, T> resultMap);
+
+  /**
+   * Returns {@code true} when the query expression contains at least one {@link QueryFunction}
+   * whose name is in the set of {@code functionName}.
+   */
+  public static class ContainsFunctionQueryExpressionVisitor
+      extends AggregatingQueryExpressionVisitor<Boolean>
+      implements QueryExpressionVisitor<Boolean> {
+
+    private final ImmutableSet<String> functionNames;
+
+    public ContainsFunctionQueryExpressionVisitor(Iterable<String> functionNames) {
+      this.functionNames = ImmutableSet.copyOf(functionNames);
+    }
+
+    @Override
+    public Boolean visit(TargetLiteral targetLiteral) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(SetExpression setExpression) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(FunctionExpression functionExpression) {
+      QueryFunction function = functionExpression.getFunction();
+      if (functionNames.contains(function.getName())) {
+        return true;
+      } else {
+        return super.visit(functionExpression);
+      }
+    }
+
+    @Override
+    protected Boolean aggregate(ImmutableMap<QueryExpression, Boolean> resultMap) {
+      return Booleans.contains(Booleans.toArray(resultMap.values()), true);
+    }
+  }
+}
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 31fa78577d..fa5a4d8343 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
@@ -183,8 +183,8 @@ public class BinaryOperatorExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
-    return mapper.map(this);
+  public <T> T accept(QueryExpressionVisitor<T> visitor) {
+    return visitor.visit(this);
   }
 
   @Override
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 a8c21d61c4..bd78b2071c 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
@@ -60,8 +60,8 @@ public class FunctionExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
-    return mapper.map(this);
+  public <T> T accept(QueryExpressionVisitor<T> visitor) {
+    return visitor.visit(this);
   }
 
   @Override
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 fd10135e57..eabf67ccf1 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
@@ -91,8 +91,8 @@ class LetExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
-    return mapper.map(this);
+  public <T> T accept(QueryExpressionVisitor<T> visitor) {
+    return visitor.visit(this);
   }
 
   @Override
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 920722db4b..5398f4dc0d 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
@@ -81,8 +81,8 @@ 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);
+  /* Implementations should just be {@code return visitor.visit(this)}. */
+  public abstract <T> T accept(QueryExpressionVisitor<T> visitor);
 
   /**
    * 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
index 8b4e3befe4..10e05ebc44 100644
--- 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
@@ -19,22 +19,25 @@ 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.
+ * <p>For each subclass of {@link QueryExpression}, there's a corresponding {@link #visit} 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 abstract class QueryExpressionMapper {
-  public QueryExpression map(TargetLiteral targetLiteral) {
+public abstract class QueryExpressionMapper implements QueryExpressionVisitor<QueryExpression> {
+
+  @Override
+  public QueryExpression visit(TargetLiteral targetLiteral) {
     return targetLiteral;
   }
 
-  public QueryExpression map(BinaryOperatorExpression binaryOperatorExpression) {
+  @Override
+  public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression) {
     boolean changed = false;
     ImmutableList.Builder<QueryExpression> mappedOperandsBuilder = ImmutableList.builder();
     for (QueryExpression operand : binaryOperatorExpression.getOperands()) {
-      QueryExpression mappedOperand = operand.getMapped(this);
+      QueryExpression mappedOperand = operand.accept(this);
       if (mappedOperand != operand) {
         changed = true;
       }
@@ -46,20 +49,22 @@ public abstract class QueryExpressionMapper {
         : binaryOperatorExpression;
   }
 
-  public QueryExpression map(FunctionExpression functionExpression) {
+  @Override
+  public QueryExpression visit(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;
+        case EXPRESSION:
+          {
+            QueryExpression expr = argument.getExpression();
+            QueryExpression mappedExpression = expr.accept(this);
+            mappedArgumentBuilder.add(Argument.of(mappedExpression));
+            if (expr != mappedExpression) {
+              changed = true;
+            }
+            break;
           }
-          break;
-        }
         default:
           mappedArgumentBuilder.add(argument);
           break;
@@ -70,13 +75,14 @@ public abstract class QueryExpressionMapper {
         : functionExpression;
   }
 
-  public QueryExpression map(LetExpression letExpression) {
+  @Override
+  public QueryExpression visit(LetExpression letExpression) {
     boolean changed = false;
-    QueryExpression mappedVarExpr = letExpression.getVarExpr().getMapped(this);
+    QueryExpression mappedVarExpr = letExpression.getVarExpr().accept(this);
     if (mappedVarExpr != letExpression.getVarExpr()) {
       changed = true;
     }
-    QueryExpression mappedBodyExpr = letExpression.getBodyExpr().getMapped(this);
+    QueryExpression mappedBodyExpr = letExpression.getBodyExpr().accept(this);
     if (mappedBodyExpr != letExpression.getBodyExpr()) {
       changed = true;
     }
@@ -85,7 +91,8 @@ public abstract class QueryExpressionMapper {
         : letExpression;
   }
 
-  public QueryExpression map(SetExpression setExpression) {
+  @Override
+  public QueryExpression visit(SetExpression setExpression) {
     return setExpression;
   }
 
@@ -109,27 +116,27 @@ public abstract class QueryExpressionMapper {
     }
 
     @Override
-    public QueryExpression map(TargetLiteral targetLiteral) {
+    public QueryExpression visit(TargetLiteral targetLiteral) {
       return mapAll(targetLiteral, mappers);
     }
 
     @Override
-    public QueryExpression map(BinaryOperatorExpression binaryOperatorExpression) {
+    public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression) {
       return mapAll(binaryOperatorExpression, mappers);
     }
 
     @Override
-    public QueryExpression map(FunctionExpression functionExpression) {
+    public QueryExpression visit(FunctionExpression functionExpression) {
       return mapAll(functionExpression, mappers);
     }
 
     @Override
-    public QueryExpression map(LetExpression letExpression) {
+    public QueryExpression visit(LetExpression letExpression) {
       return mapAll(letExpression, mappers);
     }
 
     @Override
-    public QueryExpression map(SetExpression setExpression) {
+    public QueryExpression visit(SetExpression setExpression) {
       return mapAll(setExpression, mappers);
     }
 
@@ -137,7 +144,7 @@ public abstract class QueryExpressionMapper {
         QueryExpression expression, QueryExpressionMapper[] mappers) {
       QueryExpression expr = expression;
       for (int i = mappers.length - 1; i >= 0; i--) {
-        expr = expr.getMapped(mappers[i]);
+        expr = expr.accept(mappers[i]);
       }
 
       return expr;
@@ -151,27 +158,27 @@ public abstract class QueryExpressionMapper {
     }
 
     @Override
-    public QueryExpression map(TargetLiteral targetLiteral) {
+    public QueryExpression visit(TargetLiteral targetLiteral) {
       return targetLiteral;
     }
 
     @Override
-    public QueryExpression map(BinaryOperatorExpression binaryOperatorExpression) {
+    public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression) {
       return binaryOperatorExpression;
     }
 
     @Override
-    public QueryExpression map(FunctionExpression functionExpression) {
+    public QueryExpression visit(FunctionExpression functionExpression) {
       return functionExpression;
     }
 
     @Override
-    public QueryExpression map(LetExpression letExpression) {
+    public QueryExpression visit(LetExpression letExpression) {
       return letExpression;
     }
 
     @Override
-    public QueryExpression map(SetExpression setExpression) {
+    public QueryExpression visit(SetExpression setExpression) {
       return setExpression;
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionVisitor.java
new file mode 100644
index 0000000000..3f233eeb4c
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionVisitor.java
@@ -0,0 +1,27 @@
+// Copyright 2017 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;
+
+/** Provides interfaces to visit all {@link QueryExpression}s. */
+public interface QueryExpressionVisitor<T> {
+  T visit(TargetLiteral targetLiteral);
+
+  T visit(BinaryOperatorExpression binaryOperatorExpression);
+
+  T visit(FunctionExpression functionExpression);
+
+  T visit(LetExpression letExpression);
+
+  T visit(SetExpression setExpression);
+}
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 57db015f9e..eb129711f3 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
@@ -64,8 +64,13 @@ public class SetExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
-    return mapper.map(this);
+  public <T> T accept(QueryExpressionVisitor<T> visitor) {
+    return visitor.visit(this);
+  }
+
+  /** Gets the list of {@link TargetLiteral}s contained in the expression. */
+  List<TargetLiteral> getWords() {
+    return words;
   }
 
   @Override
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 733bffb065..8cb1ab8656 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
@@ -80,8 +80,8 @@ public final class TargetLiteral extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
-    return mapper.map(this);
+  public <T> T accept(QueryExpressionVisitor<T> visitor) {
+    return visitor.visit(this);
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
index acc5a54644..951a2027e7 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
@@ -38,15 +38,12 @@ import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.packages.TriState;
 import com.google.devtools.build.lib.query2.AbstractBlazeQueryEnvironment;
-import com.google.devtools.build.lib.query2.engine.BuildFilesFunction;
-import com.google.devtools.build.lib.query2.engine.FunctionExpression;
-import com.google.devtools.build.lib.query2.engine.LoadFilesFunction;
+import com.google.devtools.build.lib.query2.engine.AggregatingQueryExpressionVisitor.ContainsFunctionQueryExpressionVisitor;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
-import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
 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.QueryExpressionVisitor;
 import com.google.devtools.build.lib.query2.engine.SynchronizedDelegatingOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.ThreadSafeOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.output.QueryOptions.OrderOutput;
@@ -67,7 +64,6 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
 import javax.annotation.Nullable;
 
 /**
@@ -380,20 +376,11 @@ public abstract class OutputFormatter implements Serializable {
       if (!(env instanceof AbstractBlazeQueryEnvironment)) {
         return;
       }
-      final AtomicBoolean found = new AtomicBoolean(false);
-      QueryExpressionMapper noteBuildFilesAndLoadLilesMapper = new QueryExpressionMapper() {
-        @Override
-        public QueryExpression map(FunctionExpression functionExpression) {
-          QueryFunction queryFunction = functionExpression.getFunction();
-          if (queryFunction instanceof LoadFilesFunction
-              || queryFunction instanceof BuildFilesFunction) {
-            found.set(true);
-          }
-          return super.map(functionExpression);
-        }
-      };
-      expr.getMapped(noteBuildFilesAndLoadLilesMapper);
-      if (found.get()) {
+
+      QueryExpressionVisitor<Boolean> noteBuildFilesAndLoadLilesVisitor =
+          new ContainsFunctionQueryExpressionVisitor(ImmutableList.of("loadfiles", "buildfiles"));
+
+      if (expr.accept(noteBuildFilesAndLoadLilesVisitor)) {
         throw new QueryException(
             "Query expressions involving 'buildfiles' or 'loadfiles' cannot be used with "
             + "--output=location");
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CollectPackagesUnderDirectoryValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/CollectPackagesUnderDirectoryValue.java
index 5f819e58ae..9e2f280382 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/CollectPackagesUnderDirectoryValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/CollectPackagesUnderDirectoryValue.java
@@ -225,7 +225,7 @@ public abstract class CollectPackagesUnderDirectoryValue implements SkyValue {
 
   /** Create a collect packages under directory request. */
   @ThreadSafe
-  static SkyKey key(
+  public static SkyKey key(
       RepositoryName repository, RootedPath rootedPath, ImmutableSet<PathFragment> excludedPaths) {
     return key(new RecursivePkgKey(repository, rootedPath, excludedPaths));
   }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CollectTargetsInPackageValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/CollectTargetsInPackageValue.java
index 700d32226a..632a79427b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/CollectTargetsInPackageValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/CollectTargetsInPackageValue.java
@@ -19,6 +19,7 @@ import com.google.devtools.build.lib.pkgcache.FilteringPolicy;
 import com.google.devtools.build.skyframe.LegacySkyKey;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.io.Serializable;
 
 /** Singleton result of {@link CollectTargetsInPackageFunction}. */
 public class CollectTargetsInPackageValue implements SkyValue {
@@ -30,7 +31,7 @@ public class CollectTargetsInPackageValue implements SkyValue {
    * Creates a key for evaluation of {@link CollectTargetsInPackageFunction}. See that class's
    * comment for what callers should have done beforehand.
    */
-  static SkyKey key(PackageIdentifier packageId, FilteringPolicy filteringPolicy) {
+  public static SkyKey key(PackageIdentifier packageId, FilteringPolicy filteringPolicy) {
     return LegacySkyKey.create(
         SkyFunctions.COLLECT_TARGETS_IN_PACKAGE,
         CollectTargetsInPackageKey.create(packageId, filteringPolicy));
@@ -38,7 +39,7 @@ public class CollectTargetsInPackageValue implements SkyValue {
 
   /** {@link SkyKey} argument. */
   @AutoValue
-  public abstract static class CollectTargetsInPackageKey {
+  public abstract static class CollectTargetsInPackageKey implements Serializable {
     public static CollectTargetsInPackageKey create(
         PackageIdentifier packageId, FilteringPolicy filteringPolicy) {
       return new AutoValue_CollectTargetsInPackageValue_CollectTargetsInPackageKey(
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index 96e89fa3d1..774944350e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -1746,17 +1746,16 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
    */
   @Override
   public EvaluationResult<SkyValue> prepareAndGet(
-      SkyKey universeKey, int numThreads, ExtendedEventHandler eventHandler)
+      Set<SkyKey> roots, int numThreads, ExtendedEventHandler eventHandler)
       throws InterruptedException {
     EvaluationResult<SkyValue> evaluationResult =
-        buildDriver.evaluate(ImmutableList.of(universeKey), true, numThreads, eventHandler);
-    Preconditions.checkNotNull(evaluationResult.getWalkableGraph(), universeKey);
+        buildDriver.evaluate(roots, true, numThreads, eventHandler);
     return evaluationResult;
   }
 
   @Override
-  public boolean isUpToDate(SkyKey universeKey) {
-    return buildDriver.alreadyEvaluated(ImmutableList.of(universeKey));
+  public boolean isUpToDate(Set<SkyKey> roots) {
+    return buildDriver.alreadyEvaluated(roots);
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java b/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
index 5b980792c1..b5b744f6af 100644
--- a/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
+++ b/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
@@ -18,6 +18,7 @@ import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.skyframe.QueryableGraph.Reason;
 import java.util.Collection;
 import java.util.Map;
+import java.util.Set;
 import javax.annotation.Nullable;
 
 /**
@@ -99,15 +100,15 @@ public interface WalkableGraph {
   /** Provides a WalkableGraph on demand after preparing it. */
   interface WalkableGraphFactory {
     EvaluationResult<SkyValue> prepareAndGet(
-        SkyKey universeKey, int numThreads, ExtendedEventHandler eventHandler)
+        Set<SkyKey> roots, int numThreads, ExtendedEventHandler eventHandler)
         throws InterruptedException;
 
     /**
-     * Returns true if this instance has already been used to {@link #prepareAndGet} {@code
-     * universeKey}. If so, cached results from {@link #prepareAndGet} can be re-used safely,
-     * potentially saving some processing time.
+     * Returns true if this instance has already been used to {@link #prepareAndGet} {@code roots}.
+     * If so, cached results from {@link #prepareAndGet} can be re-used safely, potentially saving
+     * some processing time.
      */
-    boolean isUpToDate(SkyKey universeKey);
+    boolean isUpToDate(Set<SkyKey> roots);
 
     /** Returns the {@link SkyKey} that defines this universe. */
     SkyKey getUniverseKey(Collection<String> roots, String offset);
diff --git a/src/test/java/com/google/devtools/build/lib/pkgcache/IOExceptionsTest.java b/src/test/java/com/google/devtools/build/lib/pkgcache/IOExceptionsTest.java
index 3cbcf24a36..dcbff513b4 100644
--- a/src/test/java/com/google/devtools/build/lib/pkgcache/IOExceptionsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/pkgcache/IOExceptionsTest.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.pkgcache;
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.base.Function;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.packages.ConstantRuleVisibility;
 import com.google.devtools.build.lib.packages.util.PackageLoadingTestCase;
@@ -64,7 +65,7 @@ public class IOExceptionsTest extends PackageLoadingTestCase {
   private boolean visitTransitively(Label label) throws InterruptedException {
     SkyKey key = TransitiveTargetValue.key(label);
     EvaluationResult<SkyValue> result =
-        skyframeExecutor.prepareAndGet(key, /*numThreads=*/5, reporter);
+        skyframeExecutor.prepareAndGet(ImmutableSet.of(key), /*numThreads=*/ 5, reporter);
     TransitiveTargetValue value = (TransitiveTargetValue) result.get(key);
     System.out.println(value);
     boolean hasTransitiveError = (value == null) || value.getTransitiveRootCauses() != null;