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

1 files changed, 18 insertions, 17 deletions

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