Generalize QueryExpressionVisitor's visitation by permitting an arbitrary "context" object to be passed along. Also make some query internals public, for use in fancy QueryExpressionVisitor implementations.

RELNOTES: None
PiperOrigin-RevId: 184014063

12 files changed, 92 insertions, 85 deletions

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 03341de629..0485f9dd95 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
@@ -335,7 +335,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
    * {@literal 'rdeps(<universeScope>, <T>)'} to {@literal 'allrdeps(<T>)'}. The latter is more
    * efficient.
    */
-  protected static class RdepsToAllRdepsQueryExpressionMapper extends QueryExpressionMapper {
+  protected static class RdepsToAllRdepsQueryExpressionMapper extends QueryExpressionMapper<Void> {
     protected final TargetPattern.Parser targetPatternParser;
     private final String absoluteUniverseScopePattern;
 
@@ -347,7 +347,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     }
 
     @Override
-    public QueryExpression visit(FunctionExpression functionExpression) {
+    public QueryExpression visit(FunctionExpression functionExpression, Void context) {
       if (functionExpression.getFunction().getName().equals(new RdepsFunction().getName())) {
         List<Argument> args = functionExpression.getArgs();
         QueryExpression universeExpression = args.get(0).getExpression();
@@ -361,13 +361,13 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
           }
         }
       }
-      return super.visit(functionExpression);
+      return super.visit(functionExpression, context);
     }
   }
 
   @Override
   public final QueryExpression transformParsedQuery(QueryExpression queryExpression) {
-    QueryExpressionMapper mapper = getQueryExpressionMapper();
+    QueryExpressionMapper<Void> mapper = getQueryExpressionMapper();
     QueryExpression transformedQueryExpression = queryExpression.accept(mapper);
     logger.info(
         String.format(
@@ -381,7 +381,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     return transformedQueryExpression;
   }
 
-  protected QueryExpressionMapper getQueryExpressionMapper() {
+  protected QueryExpressionMapper<Void> getQueryExpressionMapper() {
     if (universeScope.size() != 1) {
       return QueryExpressionMapper.identity();
     }
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
index cde80e072a..fef1baaf36 100644
--- 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
@@ -22,26 +22,27 @@ import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunctio
 
 /**
  * An implementation of {@link QueryExpressionVisitor} which recursively visits all nested {@link
- * QueryExpression}s
+ * QueryExpression}s.
  */
-public abstract class AggregatingQueryExpressionVisitor<T> implements QueryExpressionVisitor<T> {
+public abstract class AggregatingQueryExpressionVisitor<T, C>
+    implements QueryExpressionVisitor<T, C> {
 
   @Override
-  public T visit(BinaryOperatorExpression binaryOperatorExpression) {
+  public T visit(BinaryOperatorExpression binaryOperatorExpression, C context) {
     ImmutableMap.Builder<QueryExpression, T> builder = ImmutableMap.builder();
     for (QueryExpression expr : binaryOperatorExpression.getOperands()) {
-      builder.put(expr, expr.accept(this));
+      builder.put(expr, expr.accept(this, context));
     }
 
     return aggregate(builder.build());
   }
 
   @Override
-  public T visit(FunctionExpression functionExpression) {
+  public T visit(FunctionExpression functionExpression, C context) {
     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));
+        builder.put(argument.getExpression(), argument.getExpression().accept(this, context));
       }
     }
 
@@ -49,18 +50,18 @@ public abstract class AggregatingQueryExpressionVisitor<T> implements QueryExpre
   }
 
   @Override
-  public T visit(LetExpression letExpression) {
+  public T visit(LetExpression letExpression, C context) {
     return aggregate(
         ImmutableMap.of(
-            letExpression.getVarExpr(), letExpression.getVarExpr().accept(this),
-            letExpression.getBodyExpr(), letExpression.getBodyExpr().accept(this)));
+            letExpression.getVarExpr(), letExpression.getVarExpr().accept(this, context),
+            letExpression.getBodyExpr(), letExpression.getBodyExpr().accept(this, context)));
   }
 
   @Override
-  public T visit(SetExpression setExpression) {
+  public T visit(SetExpression setExpression, C context) {
     ImmutableMap.Builder<QueryExpression, T> builder = ImmutableMap.builder();
     for (TargetLiteral targetLiteral : setExpression.getWords()) {
-      builder.put(targetLiteral, targetLiteral.accept(this));
+      builder.put(targetLiteral, targetLiteral.accept(this, context));
     }
 
     return aggregate(builder.build());
@@ -77,8 +78,8 @@ public abstract class AggregatingQueryExpressionVisitor<T> implements QueryExpre
    * whose name is in the set of {@code functionName}.
    */
   public static class ContainsFunctionQueryExpressionVisitor
-      extends AggregatingQueryExpressionVisitor<Boolean>
-      implements QueryExpressionVisitor<Boolean> {
+      extends AggregatingQueryExpressionVisitor<Boolean, Void>
+      implements QueryExpressionVisitor<Boolean, Void> {
 
     private final ImmutableSet<String> functionNames;
 
@@ -87,22 +88,22 @@ public abstract class AggregatingQueryExpressionVisitor<T> implements QueryExpre
     }
 
     @Override
-    public Boolean visit(TargetLiteral targetLiteral) {
+    public Boolean visit(TargetLiteral targetLiteral, Void context) {
       return false;
     }
 
     @Override
-    public Boolean visit(SetExpression setExpression) {
+    public Boolean visit(SetExpression setExpression, Void context) {
       return false;
     }
 
     @Override
-    public Boolean visit(FunctionExpression functionExpression) {
+    public Boolean visit(FunctionExpression functionExpression, Void context) {
       QueryFunction function = functionExpression.getFunction();
       if (functionNames.contains(function.getName())) {
         return true;
       } else {
-        return super.visit(functionExpression);
+        return super.visit(functionExpression, context);
       }
     }
 
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 968c7a967f..3b4968806e 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 <T> T accept(QueryExpressionVisitor<T> visitor) {
-    return visitor.visit(this);
+  public <T, C> T accept(QueryExpressionVisitor<T, C> visitor, C context) {
+    return visitor.visit(this, context);
   }
 
   @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 bd78b2071c..04c4358b4e 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 <T> T accept(QueryExpressionVisitor<T> visitor) {
-    return visitor.visit(this);
+  public <T, C> T accept(QueryExpressionVisitor<T, C> visitor, C context) {
+    return visitor.visit(this, context);
   }
 
   @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 eabf67ccf1..336f94ff0c 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
@@ -24,7 +24,7 @@ import java.util.regex.Pattern;
  *
  * <pre>expr ::= LET WORD = expr IN expr</pre>
  */
-class LetExpression extends QueryExpression {
+public class LetExpression extends QueryExpression {
 
   private static final String VAR_NAME_PATTERN = "[a-zA-Z_][a-zA-Z0-9_]*$";
 
@@ -39,7 +39,7 @@ class LetExpression extends QueryExpression {
     return REF_PATTERN.matcher(varName).matches();
   }
 
-  static String getNameFromReference(String reference) {
+  public static String getNameFromReference(String reference) {
     return reference.substring(1);
   }
 
@@ -47,21 +47,21 @@ class LetExpression extends QueryExpression {
   private final QueryExpression varExpr;
   private final QueryExpression bodyExpr;
 
-  LetExpression(String varName, QueryExpression varExpr, QueryExpression bodyExpr) {
+  public LetExpression(String varName, QueryExpression varExpr, QueryExpression bodyExpr) {
     this.varName = varName;
     this.varExpr = varExpr;
     this.bodyExpr = bodyExpr;
   }
 
-  String getVarName() {
+  public String getVarName() {
     return varName;
   }
 
-  QueryExpression getVarExpr() {
+  public QueryExpression getVarExpr() {
     return varExpr;
   }
 
-  QueryExpression getBodyExpr() {
+  public QueryExpression getBodyExpr() {
     return bodyExpr;
   }
 
@@ -91,8 +91,8 @@ class LetExpression extends QueryExpression {
   }
 
   @Override
-  public <T> T accept(QueryExpressionVisitor<T> visitor) {
-    return visitor.visit(this);
+  public <T, C> T accept(QueryExpressionVisitor<T, C> visitor, C context) {
+    return visitor.visit(this, context);
   }
 
   @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 5398f4dc0d..3a9afe8303 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,12 @@ public abstract class QueryExpression {
    */
   public abstract void collectTargetPatterns(Collection<String> literals);
 
-  /* Implementations should just be {@code return visitor.visit(this)}. */
-  public abstract <T> T accept(QueryExpressionVisitor<T> visitor);
+  /* Implementations should just be {@code return visitor.visit(this, context)}. */
+  public abstract <T, C> T accept(QueryExpressionVisitor<T, C> visitor, C context);
+
+  public final <T> T accept(QueryExpressionVisitor<T, Void> visitor) {
+    return accept(visitor, /*context=*/ null);
+  }
 
   /**
    * 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 10e05ebc44..af3b711021 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
@@ -17,7 +17,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
 
 /**
- * Performs an arbitrary transformation of a {@link QueryExpression}.
+ * Performs an arbitrary contextual transformation of a {@link QueryExpression}.
  *
  * <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
@@ -25,19 +25,20 @@ import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
  * reference-equality, as an optimization). Subclasses of {@link QueryExpressionMapper} can override
  * these methods in order to implement an arbitrary transformation.
  */
-public abstract class QueryExpressionMapper implements QueryExpressionVisitor<QueryExpression> {
+public abstract class QueryExpressionMapper<C>
+    implements QueryExpressionVisitor<QueryExpression, C> {
 
   @Override
-  public QueryExpression visit(TargetLiteral targetLiteral) {
+  public QueryExpression visit(TargetLiteral targetLiteral, C context) {
     return targetLiteral;
   }
 
   @Override
-  public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression) {
+  public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression, C context) {
     boolean changed = false;
     ImmutableList.Builder<QueryExpression> mappedOperandsBuilder = ImmutableList.builder();
     for (QueryExpression operand : binaryOperatorExpression.getOperands()) {
-      QueryExpression mappedOperand = operand.accept(this);
+      QueryExpression mappedOperand = operand.accept(this, context);
       if (mappedOperand != operand) {
         changed = true;
       }
@@ -50,7 +51,7 @@ public abstract class QueryExpressionMapper implements QueryExpressionVisitor<Qu
   }
 
   @Override
-  public QueryExpression visit(FunctionExpression functionExpression) {
+  public QueryExpression visit(FunctionExpression functionExpression, C context) {
     boolean changed = false;
     ImmutableList.Builder<Argument> mappedArgumentBuilder = ImmutableList.builder();
     for (Argument argument : functionExpression.getArgs()) {
@@ -58,7 +59,7 @@ public abstract class QueryExpressionMapper implements QueryExpressionVisitor<Qu
         case EXPRESSION:
           {
             QueryExpression expr = argument.getExpression();
-            QueryExpression mappedExpression = expr.accept(this);
+            QueryExpression mappedExpression = expr.accept(this, context);
             mappedArgumentBuilder.add(Argument.of(mappedExpression));
             if (expr != mappedExpression) {
               changed = true;
@@ -76,13 +77,13 @@ public abstract class QueryExpressionMapper implements QueryExpressionVisitor<Qu
   }
 
   @Override
-  public QueryExpression visit(LetExpression letExpression) {
+  public QueryExpression visit(LetExpression letExpression, C context) {
     boolean changed = false;
-    QueryExpression mappedVarExpr = letExpression.getVarExpr().accept(this);
+    QueryExpression mappedVarExpr = letExpression.getVarExpr().accept(this, context);
     if (mappedVarExpr != letExpression.getVarExpr()) {
       changed = true;
     }
-    QueryExpression mappedBodyExpr = letExpression.getBodyExpr().accept(this);
+    QueryExpression mappedBodyExpr = letExpression.getBodyExpr().accept(this, context);
     if (mappedBodyExpr != letExpression.getBodyExpr()) {
       changed = true;
     }
@@ -92,7 +93,7 @@ public abstract class QueryExpressionMapper implements QueryExpressionVisitor<Qu
   }
 
   @Override
-  public QueryExpression visit(SetExpression setExpression) {
+  public QueryExpression visit(SetExpression setExpression, C context) {
     return setExpression;
   }
 
@@ -104,81 +105,83 @@ public abstract class QueryExpressionMapper implements QueryExpressionVisitor<Qu
    * Returns a {@link QueryExpressionMapper} which applies all the mappings provided by {@code
    * mappers}, in the reverse order of mapper array.
    */
-  public static QueryExpressionMapper compose(QueryExpressionMapper... mappers) {
+  public static <C> QueryExpressionMapper<C> compose(QueryExpressionMapper<C>... mappers) {
     return new ComposedQueryExpressionMapper(mappers);
   }
 
-  private static class ComposedQueryExpressionMapper extends QueryExpressionMapper {
-    private final QueryExpressionMapper[] mappers;
+  private static class ComposedQueryExpressionMapper<C> extends QueryExpressionMapper<C> {
+    private final QueryExpressionMapper<C>[] mappers;
 
     private ComposedQueryExpressionMapper(QueryExpressionMapper... mappers) {
       this.mappers = mappers;
     }
 
     @Override
-    public QueryExpression visit(TargetLiteral targetLiteral) {
-      return mapAll(targetLiteral, mappers);
+    public QueryExpression visit(TargetLiteral targetLiteral, C context) {
+      return mapAll(targetLiteral, mappers, context);
     }
 
     @Override
-    public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression) {
-      return mapAll(binaryOperatorExpression, mappers);
+    public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression, C context) {
+      return mapAll(binaryOperatorExpression, mappers, context);
     }
 
     @Override
-    public QueryExpression visit(FunctionExpression functionExpression) {
-      return mapAll(functionExpression, mappers);
+    public QueryExpression visit(FunctionExpression functionExpression, C context) {
+      return mapAll(functionExpression, mappers, context);
     }
 
     @Override
-    public QueryExpression visit(LetExpression letExpression) {
-      return mapAll(letExpression, mappers);
+    public QueryExpression visit(LetExpression letExpression, C context) {
+      return mapAll(letExpression, mappers, context);
     }
 
     @Override
-    public QueryExpression visit(SetExpression setExpression) {
-      return mapAll(setExpression, mappers);
+    public QueryExpression visit(SetExpression setExpression, C context) {
+      return mapAll(setExpression, mappers, context);
     }
 
-    private static QueryExpression mapAll(
-        QueryExpression expression, QueryExpressionMapper[] mappers) {
+    private static <C> QueryExpression mapAll(
+        QueryExpression expression,
+        QueryExpressionMapper<C>[] mappers,
+        C context) {
       QueryExpression expr = expression;
       for (int i = mappers.length - 1; i >= 0; i--) {
-        expr = expr.accept(mappers[i]);
+        expr = expr.accept(mappers[i], context);
       }
 
       return expr;
     }
   }
 
-  private static class IdentityMapper extends QueryExpressionMapper {
+  private static class IdentityMapper extends QueryExpressionMapper<Void> {
     private static final IdentityMapper INSTANCE = new IdentityMapper();
 
     private IdentityMapper() {
     }
 
     @Override
-    public QueryExpression visit(TargetLiteral targetLiteral) {
+    public QueryExpression visit(TargetLiteral targetLiteral, Void context) {
       return targetLiteral;
     }
 
     @Override
-    public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression) {
+    public QueryExpression visit(BinaryOperatorExpression binaryOperatorExpression, Void context) {
       return binaryOperatorExpression;
     }
 
     @Override
-    public QueryExpression visit(FunctionExpression functionExpression) {
+    public QueryExpression visit(FunctionExpression functionExpression, Void context) {
       return functionExpression;
     }
 
     @Override
-    public QueryExpression visit(LetExpression letExpression) {
+    public QueryExpression visit(LetExpression letExpression, Void context) {
       return letExpression;
     }
 
     @Override
-    public QueryExpression visit(SetExpression setExpression) {
+    public QueryExpression visit(SetExpression setExpression, Void context) {
       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
index 3f233eeb4c..51f5351294 100644
--- 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
@@ -13,15 +13,15 @@
 // 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);
+/** Provides interfaces to visit all {@link QueryExpression}s with a context object. */
+public interface QueryExpressionVisitor<T, C> {
+  T visit(TargetLiteral targetLiteral, C context);
 
-  T visit(BinaryOperatorExpression binaryOperatorExpression);
+  T visit(BinaryOperatorExpression binaryOperatorExpression, C context);
 
-  T visit(FunctionExpression functionExpression);
+  T visit(FunctionExpression functionExpression, C context);
 
-  T visit(LetExpression letExpression);
+  T visit(LetExpression letExpression, C context);
 
-  T visit(SetExpression setExpression);
+  T visit(SetExpression setExpression, C context);
 }
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 9ba285ee29..8d35aedcfa 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
@@ -63,7 +63,7 @@ public final class QueryParser {
     return parse(query, functions);
   }
 
-  public static QueryExpression parse(String query, HashMap<String, QueryFunction> functions)
+  public static QueryExpression parse(String query, Map<String, QueryFunction> functions)
       throws QueryException {
     QueryParser parser = new QueryParser(Lexer.scan(query), functions);
     QueryExpression expr = parser.parseExpression();
@@ -74,7 +74,7 @@ public final class QueryParser {
     return expr;
   }
 
-  public QueryParser(List<Lexer.Token> tokens, HashMap<String, QueryFunction> functions) {
+  public QueryParser(List<Lexer.Token> tokens, Map<String, QueryFunction> functions) {
     this.functions = functions;
     this.tokens = tokens;
     this.tokenIterator = tokens.iterator();
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 eb129711f3..3f883c144f 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,12 +64,12 @@ public class SetExpression extends QueryExpression {
   }
 
   @Override
-  public <T> T accept(QueryExpressionVisitor<T> visitor) {
-    return visitor.visit(this);
+  public <T, C> T accept(QueryExpressionVisitor<T, C> visitor, C context) {
+    return visitor.visit(this, context);
   }
 
   /** Gets the list of {@link TargetLiteral}s contained in the expression. */
-  List<TargetLiteral> getWords() {
+  public List<TargetLiteral> getWords() {
     return 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 951d5ce8cb..4e3a8e7eb2 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 <T> T accept(QueryExpressionVisitor<T> visitor) {
-    return visitor.visit(this);
+  public <T, C> T accept(QueryExpressionVisitor<T, C> visitor, C context) {
+    return visitor.visit(this, context);
   }
 
   @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 d13beac249..c794038a16 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
@@ -43,7 +43,6 @@ 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.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
-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;
@@ -387,7 +386,7 @@ public abstract class OutputFormatter implements Serializable {
         return;
       }
 
-      QueryExpressionVisitor<Boolean> noteBuildFilesAndLoadLilesVisitor =
+      ContainsFunctionQueryExpressionVisitor noteBuildFilesAndLoadLilesVisitor =
           new ContainsFunctionQueryExpressionVisitor(ImmutableList.of("loadfiles", "buildfiles"));
 
       if (expr.accept(noteBuildFilesAndLoadLilesVisitor)) {