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

1 files changed, 7 insertions, 20 deletions

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");