Have QueryExpressionMapper#compose take arbitrary number of mappers

AttrFunction and LabelsFunction are also made public.

RELNOTES: None
PiperOrigin-RevId: 155108260

3 files changed, 36 insertions, 28 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/AttrFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/AttrFunction.java
index 6c0baa9cbd..381f205a7a 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/AttrFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/AttrFunction.java
@@ -16,28 +16,26 @@ package com.google.devtools.build.lib.query2.engine;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.ArgumentType;
-
 import java.util.List;
 
 /**
- * An attr(attribute, pattern, argument) filter expression, which computes
- * the set of subset of nodes in 'argument' which correspond to rules with
- * defined attribute 'attribute' with attribute value matching the unanchored
- * regexp 'pattern'. For list attributes, the attribute value will be defined as
- * a usual List.toString() representation (using '[' as first character, ']' as
- * last character and ", " as a delimiter between multiple values). Also, all
- * label-based attributes will use fully-qualified label names instead of
- * original value specified in the BUILD file.
+ * An attr(attribute, pattern, argument) filter expression, which computes the set of subset of
+ * nodes in 'argument' which correspond to rules with defined attribute 'attribute' with attribute
+ * value matching the unanchored regexp 'pattern'. For list attributes, the attribute value will be
+ * defined as a usual List.toString() representation (using '[' as first character, ']' as last
+ * character and ", " as a delimiter between multiple values). Also, all label-based attributes will
+ * use fully-qualified label names instead of original value specified in the BUILD file.
  *
  * <pre>expr ::= ATTR '(' ATTRNAME ',' WORD ',' expr ')'</pre>
  *
  * Examples
+ *
  * <pre>
  * attr(linkshared,1,//project/...)    find all rules under in the //project/... that
  *                                 have attribute linkshared set to 1.
  * </pre>
  */
-class AttrFunction extends RegexFilterExpression {
+public class AttrFunction extends RegexFilterExpression {
   AttrFunction() {
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/LabelsFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/LabelsFunction.java
index 1d68573578..edd0935fc7 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/LabelsFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/LabelsFunction.java
@@ -22,17 +22,18 @@ import java.util.ArrayList;
 import java.util.List;
 
 /**
- * A label(attr_name, argument) expression, which computes the set of targets
- * whose labels appear in the specified attribute of some rule in 'argument'.
+ * A label(attr_name, argument) expression, which computes the set of targets whose labels appear in
+ * the specified attribute of some rule in 'argument'.
  *
  * <pre>expr ::= LABELS '(' WORD ',' expr ')'</pre>
  *
  * Example:
+ *
  * <pre>
  *  labels(srcs, //foo)      The 'srcs' source files to the //foo rule.
  * </pre>
  */
-class LabelsFunction implements QueryFunction {
+public class LabelsFunction implements QueryFunction {
   LabelsFunction() {
   }
 
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 0a0336c8af..8b4e3befe4 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
@@ -93,45 +93,54 @@ public abstract class QueryExpressionMapper {
     return IdentityMapper.INSTANCE;
   }
 
-  public static QueryExpressionMapper compose(
-      QueryExpressionMapper outerMapper, QueryExpressionMapper innerMapper) {
-    return new ComposedQueryExpressionMapper(outerMapper, innerMapper);
+  /**
+   * 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) {
+    return new ComposedQueryExpressionMapper(mappers);
   }
 
   private static class ComposedQueryExpressionMapper extends QueryExpressionMapper {
-    private QueryExpressionMapper outerMapper;
-    private QueryExpressionMapper innerMapper;
+    private final QueryExpressionMapper[] mappers;
 
-    private ComposedQueryExpressionMapper(
-        QueryExpressionMapper outerMapper,
-        QueryExpressionMapper innerMapper) {
-      this.outerMapper = outerMapper;
-      this.innerMapper = innerMapper;
+    private ComposedQueryExpressionMapper(QueryExpressionMapper... mappers) {
+      this.mappers = mappers;
     }
 
     @Override
     public QueryExpression map(TargetLiteral targetLiteral) {
-      return innerMapper.map(targetLiteral).getMapped(outerMapper);
+      return mapAll(targetLiteral, mappers);
     }
 
     @Override
     public QueryExpression map(BinaryOperatorExpression binaryOperatorExpression) {
-      return innerMapper.map(binaryOperatorExpression).getMapped(outerMapper);
+      return mapAll(binaryOperatorExpression, mappers);
     }
 
     @Override
     public QueryExpression map(FunctionExpression functionExpression) {
-      return innerMapper.map(functionExpression).getMapped(outerMapper);
+      return mapAll(functionExpression, mappers);
     }
 
     @Override
     public QueryExpression map(LetExpression letExpression) {
-      return innerMapper.map(letExpression).getMapped(outerMapper);
+      return mapAll(letExpression, mappers);
     }
 
     @Override
     public QueryExpression map(SetExpression setExpression) {
-      return innerMapper.map(setExpression).getMapped(outerMapper);
+      return mapAll(setExpression, mappers);
+    }
+
+    private static QueryExpression mapAll(
+        QueryExpression expression, QueryExpressionMapper[] mappers) {
+      QueryExpression expr = expression;
+      for (int i = mappers.length - 1; i >= 0; i--) {
+        expr = expr.getMapped(mappers[i]);
+      }
+
+      return expr;
     }
   }