Refactor SkyQueryEnvironment to allow preprocessing or postprocessing targets

Refactor SkyQueryEnvironment and a few other query helpers to make it easier to
work with targets.

RELNOTES: None
PiperOrigin-RevId: 160165398

6 files changed, 119 insertions, 83 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/pkgcache/TargetPatternResolverUtil.java b/src/main/java/com/google/devtools/build/lib/pkgcache/TargetPatternResolverUtil.java
index 84a1fc7e73..99436907e6 100644
--- a/src/main/java/com/google/devtools/build/lib/pkgcache/TargetPatternResolverUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/pkgcache/TargetPatternResolverUtil.java
@@ -13,6 +13,8 @@
 // limitations under the License.
 package com.google.devtools.build.lib.pkgcache;
 
+import com.google.common.base.Function;
+import com.google.common.base.Functions;
 import com.google.devtools.build.lib.cmdline.LabelValidator;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.ResolvedTargets;
@@ -26,6 +28,9 @@ import com.google.devtools.build.lib.vfs.PathFragment;
  * Common utility methods for target pattern resolution.
  */
 public final class TargetPatternResolverUtil {
+
+  private static final Function<Target, Target> KEEP_ORIGINAL_TARGET = Functions.identity();
+
   private TargetPatternResolverUtil() {
     // Utility class.
   }
@@ -38,12 +43,16 @@ public final class TargetPatternResolverUtil {
     }
   }
 
-  public static ResolvedTargets<Target> resolvePackageTargets(Package pkg,
-                                                              FilteringPolicy policy) {
+  public static ResolvedTargets<Target> resolvePackageTargets(Package pkg, FilteringPolicy policy) {
+    return resolvePackageTargets(pkg, policy, KEEP_ORIGINAL_TARGET);
+  }
+
+  public static ResolvedTargets<Target> resolvePackageTargets(
+      Package pkg, FilteringPolicy policy, Function<Target, Target> convertTarget) {
     ResolvedTargets.Builder<Target> builder = ResolvedTargets.builder();
     for (Target target : pkg.getTargets().values()) {
       if (policy.shouldRetain(target, false)) {
-        builder.add(target);
+        builder.add(convertTarget.apply(target));
       }
     }
     return builder.build();
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 a2cf884ffd..6bf3ceab52 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
@@ -59,17 +59,17 @@ import com.google.devtools.build.lib.packages.NoSuchThingException;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
+import com.google.devtools.build.lib.pkgcache.FilteringPolicy;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
+import com.google.devtools.build.lib.pkgcache.RecursivePackageProvider;
 import com.google.devtools.build.lib.pkgcache.TargetPatternEvaluator;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
-import com.google.devtools.build.lib.query2.AbstractBlazeQueryEnvironment.TargetKeyExtractor;
 import com.google.devtools.build.lib.query2.engine.AllRdepsFunction;
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.FunctionExpression;
 import com.google.devtools.build.lib.query2.engine.KeyExtractor;
 import com.google.devtools.build.lib.query2.engine.MinDepthUniquifier;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
-import com.google.devtools.build.lib.query2.engine.QueryEnvironment.MutableMap;
 import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
@@ -263,13 +263,22 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
             new ThreadFactoryBuilder().setNameFormat("QueryEnvironment %d").build()));
     }
     resolver =
-        new RecursivePackageProviderBackedTargetPatternResolver(
+        createTargetPatternResolver(
             graphBackedRecursivePackageProvider,
             eventHandler,
             TargetPatternEvaluator.DEFAULT_FILTERING_POLICY,
             packageSemaphore);
   }
 
+  protected RecursivePackageProviderBackedTargetPatternResolver createTargetPatternResolver(
+      RecursivePackageProvider graphBackedRecursivePackageProvider,
+      ExtendedEventHandler eventHandler,
+      FilteringPolicy policy,
+      MultisetSemaphore<PackageIdentifier> packageSemaphore) {
+    return new RecursivePackageProviderBackedTargetPatternResolver(
+        graphBackedRecursivePackageProvider, eventHandler, policy, packageSemaphore);
+  }
+
   protected MultisetSemaphore<PackageIdentifier> makeFreshPackageMultisetSemaphore() {
     return MultisetSemaphore.unbounded();
   }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/AbstractQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/engine/AbstractQueryEnvironment.java
index 84529ca3ed..c9e9a0ab74 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/AbstractQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/AbstractQueryEnvironment.java
@@ -22,8 +22,6 @@ import com.google.common.collect.Iterables;
 import com.google.common.util.concurrent.AsyncFunction;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
-import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskCallable;
-import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskFuture;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java
index 2024b8d4c7..becfe2453c 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java
@@ -61,6 +61,46 @@ public class AllRdepsFunction implements QueryFunction {
     return eval(env, context, args, callback, Optional.<Predicate<T>>absent());
   }
 
+  /** Evaluates rdeps query. */
+  public static <T> QueryTaskFuture<Void> eval(
+      final QueryEnvironment<T> env,
+      QueryExpression expression,
+      final Predicate<T> universe,
+      VariableContext<T> context,
+      final Callback<T> callback,
+      final int depth) {
+    final MinDepthUniquifier<T> minDepthUniquifier = env.createMinDepthUniquifier();
+    return env.eval(
+        expression,
+        context,
+        new Callback<T>() {
+          @Override
+          public void process(Iterable<T> partialResult)
+              throws QueryException, InterruptedException {
+            Iterable<T> current = partialResult;
+            // We need to iterate depthBound + 1 times.
+            for (int i = 0; i <= depth; i++) {
+              List<T> next = new ArrayList<>();
+              // Restrict to nodes satisfying the universe predicate.
+              Iterable<T> currentInUniverse = Iterables.filter(current, universe);
+              // Filter already visited nodes: if we see a node in a later round, then we don't
+              // need to visit it again, because the depth at which we see it must be greater
+              // than or equal to the last visit.
+              Iterables.addAll(
+                  next,
+                  env.getReverseDeps(
+                      minDepthUniquifier.uniqueAtDepthLessThanOrEqualTo(currentInUniverse, i)));
+              callback.process(currentInUniverse);
+              if (next.isEmpty()) {
+                // Exit when there are no more nodes to visit.
+                break;
+              }
+              current = next;
+            }
+          }
+        });
+  }
+
   protected <T> QueryTaskFuture<Void> eval(
       final QueryEnvironment<T> env,
       VariableContext<T> context,
@@ -78,36 +118,7 @@ public class AllRdepsFunction implements QueryFunction {
         : streamableEnv.getAllRdeps(
             args.get(0).getExpression(), universe, context, callback, depth);
     } else {
-      final MinDepthUniquifier<T> minDepthUniquifier = env.createMinDepthUniquifier();
-      return env.eval(
-          args.get(0).getExpression(),
-          context,
-          new Callback<T>() {
-            @Override
-            public void process(Iterable<T> partialResult)
-                throws QueryException, InterruptedException {
-              Iterable<T> current = partialResult;
-              // We need to iterate depthBound + 1 times.
-              for (int i = 0; i <= depth; i++) {
-                List<T> next = new ArrayList<>();
-                // Restrict to nodes satisfying the universe predicate.
-                Iterable<T> currentInUniverse = Iterables.filter(current, universe);
-                // Filter already visited nodes: if we see a node in a later round, then we don't
-                // need to visit it again, because the depth at which we see it must be greater
-                // than or equal to the last visit.
-                Iterables.addAll(
-                    next,
-                    env.getReverseDeps(
-                        minDepthUniquifier.uniqueAtDepthLessThanOrEqualTo(currentInUniverse, i)));
-                callback.process(currentInUniverse);
-                if (next.isEmpty()) {
-                  // Exit when there are no more nodes to visit.
-                  break;
-                }
-                current = next;
-              }
-            }
-          });
+      return eval(env, args.get(0).getExpression(), universe, context, callback, depth);
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/TestsFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/TestsFunction.java
index 15950d7114..0331d480c5 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/TestsFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/TestsFunction.java
@@ -30,19 +30,18 @@ import java.util.List;
 import java.util.Set;
 
 /**
- * A tests(x) filter expression, which returns all the tests in set x,
- * expanding test_suite rules into their constituents.
+ * A tests(x) filter expression, which returns all the tests in set x, expanding test_suite rules
+ * into their constituents.
  *
- * <p>Unfortunately this class reproduces a substantial amount of logic from
- * {@code TestSuiteConfiguredTarget}, albeit in a somewhat simplified form.
- * This is basically inevitable since the expansion of test_suites cannot be
- * done during the loading phase, because it involves inter-package references.
- * We make no attempt to validate the input, or report errors or warnings other
- * than missing target.
+ * <p>Unfortunately this class reproduces a substantial amount of logic from {@code
+ * TestSuiteConfiguredTarget}, albeit in a somewhat simplified form. This is basically inevitable
+ * since the expansion of test_suites cannot be done during the loading phase, because it involves
+ * inter-package references. We make no attempt to validate the input, or report errors or warnings
+ * other than missing target.
  *
  * <pre>expr ::= TESTS '(' expr ')'</pre>
  */
-class TestsFunction implements QueryFunction {
+public class TestsFunction implements QueryFunction {
   TestsFunction() {
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java b/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java
index 1ca21124d2..2fb74861a2 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java
@@ -46,6 +46,7 @@ import com.google.devtools.build.lib.pkgcache.FilteringPolicy;
 import com.google.devtools.build.lib.pkgcache.RecursivePackageProvider;
 import com.google.devtools.build.lib.pkgcache.TargetPatternResolverUtil;
 import com.google.devtools.build.lib.util.BatchCallback;
+import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.util.ThreadSafeBatchCallback;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.ArrayList;
@@ -65,9 +66,9 @@ public class RecursivePackageProviderBackedTargetPatternResolver
   // TODO(janakr): Move this to a more generic place and unify with SkyQueryEnvironment's value?
   private static final int MAX_PACKAGES_BULK_GET = 1000;
 
+  protected final FilteringPolicy policy;
   private final RecursivePackageProvider recursivePackageProvider;
   private final ExtendedEventHandler eventHandler;
-  private final FilteringPolicy policy;
   private final MultisetSemaphore<PackageIdentifier> packageSemaphore;
 
   public RecursivePackageProviderBackedTargetPatternResolver(
@@ -100,13 +101,19 @@ public class RecursivePackageProviderBackedTargetPatternResolver
     return recursivePackageProvider.bulkGetPackages(pkgIds);
   }
 
+  /** Optionally convert a {@link Target} before including it in returns. */
+  protected Target maybeConvertTargetBeforeReturn(Target target) {
+    return target;
+  }
+
   @Override
   public Target getTargetOrNull(Label label) throws InterruptedException {
     try {
       if (!isPackage(label.getPackageIdentifier())) {
         return null;
       }
-      return recursivePackageProvider.getTarget(eventHandler, label);
+      return maybeConvertTargetBeforeReturn(
+          recursivePackageProvider.getTarget(eventHandler, label));
     } catch (NoSuchThingException e) {
       return null;
     }
@@ -118,7 +125,7 @@ public class RecursivePackageProviderBackedTargetPatternResolver
     try {
       Target target = recursivePackageProvider.getTarget(eventHandler, label);
       return policy.shouldRetain(target, true)
-          ? ResolvedTargets.of(target)
+          ? ResolvedTargets.of(maybeConvertTargetBeforeReturn(target))
           : ResolvedTargets.<Target>empty();
     } catch (NoSuchThingException e) {
       throw new TargetParsingException(e.getMessage(), e);
@@ -132,15 +139,17 @@ public class RecursivePackageProviderBackedTargetPatternResolver
     FilteringPolicy actualPolicy = rulesOnly
         ? FilteringPolicies.and(FilteringPolicies.RULES_ONLY, policy)
         : policy;
-    return getTargetsInPackage(originalPattern, packageIdentifier, actualPolicy);
-  }
-
-  private ResolvedTargets<Target> getTargetsInPackage(String originalPattern,
-      PackageIdentifier packageIdentifier, FilteringPolicy policy)
-      throws TargetParsingException, InterruptedException {
     try {
       Package pkg = getPackage(packageIdentifier);
-      return TargetPatternResolverUtil.resolvePackageTargets(pkg, policy);
+      return TargetPatternResolverUtil.resolvePackageTargets(
+          pkg,
+          actualPolicy,
+          new Function<Target, Target>() {
+            @Override
+            public Target apply(Target target) {
+              return maybeConvertTargetBeforeReturn(Preconditions.checkNotNull(target));
+            }
+          });
     } catch (NoSuchThingException e) {
       String message = TargetPatternResolverUtil.getParsingErrorMessage(
           e.getMessage(), originalPattern);
@@ -286,33 +295,35 @@ public class RecursivePackageProviderBackedTargetPatternResolver
         ImmutableList.copyOf(Iterables.partition(pkgIds, MAX_PACKAGES_BULK_GET));
     ArrayList<ListenableFuture<Void>> futures = new ArrayList<>(partitions.size());
     for (final Iterable<PackageIdentifier> pkgIdBatch : partitions) {
-      futures.add(executor.submit(new Callable<Void>() {
-          @Override
-          public Void call() throws E, TargetParsingException, InterruptedException {
-            ImmutableSet<PackageIdentifier> pkgIdBatchSet = ImmutableSet.copyOf(pkgIdBatch);
-            packageSemaphore.acquireAll(pkgIdBatchSet);
-            try {
-              Iterable<ResolvedTargets<Target>> resolvedTargets =
-                  bulkGetTargetsInPackage(originalPattern, pkgIdBatch, NO_FILTER).values();
-              List<Target> filteredTargets = new ArrayList<>(calculateSize(resolvedTargets));
-              for (ResolvedTargets<Target> targets : resolvedTargets) {
-                for (Target target : targets.getTargets()) {
-                  // Perform the no-targets-found check before applying the filtering policy
-                  // so we only return the error if the input directory's subtree really
-                  // contains no targets.
-                  foundTarget.set(true);
-                  if (actualPolicy.shouldRetain(target, false)) {
-                    filteredTargets.add(target);
+      futures.add(
+          executor.submit(
+              new Callable<Void>() {
+                @Override
+                public Void call() throws E, TargetParsingException, InterruptedException {
+                  ImmutableSet<PackageIdentifier> pkgIdBatchSet = ImmutableSet.copyOf(pkgIdBatch);
+                  packageSemaphore.acquireAll(pkgIdBatchSet);
+                  try {
+                    Iterable<ResolvedTargets<Target>> resolvedTargets =
+                        bulkGetTargetsInPackage(originalPattern, pkgIdBatch, NO_FILTER).values();
+                    List<Target> filteredTargets = new ArrayList<>(calculateSize(resolvedTargets));
+                    for (ResolvedTargets<Target> targets : resolvedTargets) {
+                      for (Target target : targets.getTargets()) {
+                        // Perform the no-targets-found check before applying the filtering policy
+                        // so we only return the error if the input directory's subtree really
+                        // contains no targets.
+                        foundTarget.set(true);
+                        if (actualPolicy.shouldRetain(target, false)) {
+                          filteredTargets.add(maybeConvertTargetBeforeReturn(target));
+                        }
+                      }
+                    }
+                    callback.process(filteredTargets);
+                  } finally {
+                    packageSemaphore.releaseAll(pkgIdBatchSet);
                   }
+                  return null;
                 }
-              }
-              callback.process(filteredTargets);
-            } finally {
-              packageSemaphore.releaseAll(pkgIdBatchSet);
-            }
-            return null;
-          }
-        }));
+              }));
     }
     return Futures.whenAllSucceed(futures)
         .call(
@@ -336,6 +347,5 @@ public class RecursivePackageProviderBackedTargetPatternResolver
     }
     return size;
   }
-
 }