Rename VariableContext to the more general purpose QueryExpressionContext and thread it through to graph traversal functions. Some other light refactorings as well.

PiperOrigin-RevId: 200292556
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 382a8d2..a53f1d6 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
@@ -61,31 +61,26 @@
 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.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.QueryEvalResult;
 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.QueryExpressionContext;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionMapper;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.MinDepthUniquifierImpl;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.MutableKeyExtractorBackedMapImpl;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.ThreadSafeMutableKeyExtractorBackedSetImpl;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.UniquifierImpl;
-import com.google.devtools.build.lib.query2.engine.RdepsFunction;
 import com.google.devtools.build.lib.query2.engine.StreamableQueryEnvironment;
-import com.google.devtools.build.lib.query2.engine.TargetLiteral;
 import com.google.devtools.build.lib.query2.engine.ThreadSafeOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.Uniquifier;
-import com.google.devtools.build.lib.query2.engine.VariableContext;
 import com.google.devtools.build.lib.skyframe.BlacklistedPackagePrefixesValue;
 import com.google.devtools.build.lib.skyframe.ContainingPackageLookupFunction;
 import com.google.devtools.build.lib.skyframe.GraphBackedRecursivePackageProvider;
@@ -270,22 +265,13 @@
             new ThreadFactoryBuilder().setNameFormat("QueryEnvironment %d").build()));
     }
     resolver =
-        createTargetPatternResolver(
+        new RecursivePackageProviderBackedTargetPatternResolver(
             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();
   }
@@ -320,41 +306,6 @@
     }
   }
 
-  /**
-   * A {@link QueryExpressionMapper} that transforms each occurrence of an expression of the form
-   * {@literal 'rdeps(<universeScope>, <T>)'} to {@literal 'allrdeps(<T>)'}. The latter is more
-   * efficient.
-   */
-  protected static class RdepsToAllRdepsQueryExpressionMapper extends QueryExpressionMapper<Void> {
-    protected final TargetPattern.Parser targetPatternParser;
-    private final String absoluteUniverseScopePattern;
-
-    protected RdepsToAllRdepsQueryExpressionMapper(
-        TargetPattern.Parser targetPatternParser,
-        String universeScopePattern) {
-      this.targetPatternParser = targetPatternParser;
-      this.absoluteUniverseScopePattern = targetPatternParser.absolutize(universeScopePattern);
-    }
-
-    @Override
-    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();
-        if (universeExpression instanceof TargetLiteral) {
-          TargetLiteral literalUniverseExpression = (TargetLiteral) universeExpression;
-          String absolutizedUniverseExpression =
-              targetPatternParser.absolutize(literalUniverseExpression.getPattern());
-          if (absolutizedUniverseExpression.equals(absoluteUniverseScopePattern)) {
-            List<Argument> argsTail = args.subList(1, functionExpression.getArgs().size());
-            return new FunctionExpression(new AllRdepsFunction(), argsTail);
-          }
-        }
-      }
-      return super.visit(functionExpression, context);
-    }
-  }
-
   @Override
   public final QueryExpression transformParsedQuery(QueryExpression queryExpression) {
     QueryExpressionMapper<Void> mapper = getQueryExpressionMapper();
@@ -482,7 +433,8 @@
   }
 
   @Override
-  public ThreadSafeMutableSet<Target> getFwdDeps(Iterable<Target> targets)
+  public ThreadSafeMutableSet<Target> getFwdDeps(
+      Iterable<Target> targets, QueryExpressionContext<Target> context)
       throws InterruptedException {
     Map<SkyKey, Target> targetsByKey = Maps.newHashMapWithExpectedSize(Iterables.size(targets));
     for (Target target : targets) {
@@ -519,7 +471,9 @@
   }
 
   @Override
-  public Collection<Target> getReverseDeps(Iterable<Target> targets) throws InterruptedException {
+  public Collection<Target> getReverseDeps(
+      Iterable<Target> targets, QueryExpressionContext<Target> context)
+      throws InterruptedException {
     return getReverseDepsOfTransitiveTraversalKeys(Iterables.transform(targets, TARGET_TO_SKY_KEY));
   }
 
@@ -571,16 +525,18 @@
   }
 
   @Override
-  public ThreadSafeMutableSet<Target> getTransitiveClosure(ThreadSafeMutableSet<Target> targets)
+  public ThreadSafeMutableSet<Target> getTransitiveClosure(
+      ThreadSafeMutableSet<Target> targets, QueryExpressionContext<Target> context)
       throws InterruptedException {
     return SkyQueryUtils.getTransitiveClosure(
-        targets, this::getFwdDeps, createThreadSafeMutableSet());
+        targets, targets1 -> getFwdDeps(targets1, context), createThreadSafeMutableSet());
   }
 
   @Override
-  public ImmutableList<Target> getNodesOnPath(Target from, Target to)
-      throws InterruptedException {
-    return SkyQueryUtils.getNodesOnPath(from, to, this::getFwdDeps, Target::getLabel);
+  public ImmutableList<Target> getNodesOnPath(
+      Target from, Target to, QueryExpressionContext<Target> context) throws InterruptedException {
+    return SkyQueryUtils.getNodesOnPath(
+        from, to, targets -> getFwdDeps(targets, context), Target::getLabel);
   }
 
   private <R> ListenableFuture<R> safeSubmit(Callable<R> callable) {
@@ -603,7 +559,7 @@
   @Override
   public QueryTaskFuture<Void> eval(
       final QueryExpression expr,
-      final VariableContext<Target> context,
+      final QueryExpressionContext<Target> context,
       final Callback<Target> callback) {
     // TODO(bazel-team): As in here, use concurrency for the async #eval of other QueryEnvironment
     // implementations.
@@ -739,7 +695,11 @@
   @ThreadSafe
   @Override
   public ThreadSafeMutableSet<Target> getBuildFiles(
-      QueryExpression caller, ThreadSafeMutableSet<Target> nodes, boolean buildFiles, boolean loads)
+      QueryExpression caller,
+      ThreadSafeMutableSet<Target> nodes,
+      boolean buildFiles,
+      boolean loads,
+      QueryExpressionContext<Target> context)
       throws QueryException, InterruptedException {
     ThreadSafeMutableSet<Target> dependentFiles = createThreadSafeMutableSet();
     Set<PackageIdentifier> seenPackages = new HashSet<>();
@@ -1220,7 +1180,7 @@
   @Override
   public QueryTaskFuture<Void> getAllRdepsUnboundedParallel(
       QueryExpression expression,
-      VariableContext<Target> context,
+      QueryExpressionContext<Target> context,
       Callback<Target> callback) {
     return ParallelSkyQueryUtils.getAllRdepsUnboundedParallel(
         this, expression, context, callback, packageSemaphore);
@@ -1231,15 +1191,14 @@
   public QueryTaskFuture<Void> getAllRdepsBoundedParallel(
       QueryExpression expression,
       int depth,
-      VariableContext<Target> context,
+      QueryExpressionContext<Target> context,
       Callback<Target> callback) {
     return ParallelSkyQueryUtils.getAllRdepsBoundedParallel(
         this, expression, depth, context, callback, packageSemaphore);
   }
 
   protected QueryTaskFuture<Predicate<SkyKey>> getUniverseDTCSkyKeyPredicateFuture(
-      QueryExpression universe,
-      VariableContext<Target> context) {
+      QueryExpression universe, QueryExpressionContext<Target> context) {
     return ParallelSkyQueryUtils.getDTCSkyKeyPredicateFuture(
         this,
         universe,
@@ -1253,7 +1212,7 @@
   public QueryTaskFuture<Void> getRdepsUnboundedParallel(
       QueryExpression expression,
       QueryExpression universe,
-      VariableContext<Target> context,
+      QueryExpressionContext<Target> context,
       Callback<Target> callback) {
     return transformAsync(
         getUniverseDTCSkyKeyPredicateFuture(universe, context),
@@ -1264,7 +1223,7 @@
   @Override
   public QueryTaskFuture<Void> getDepsUnboundedParallel(
       QueryExpression expression,
-      VariableContext<Target> context,
+      QueryExpressionContext<Target> context,
       Callback<Target> callback,
       Callback<Target> errorReporter) {
     return ParallelSkyQueryUtils.getDepsUnboundedParallel(
@@ -1283,7 +1242,7 @@
       QueryExpression expression,
       int depth,
       QueryExpression universe,
-      VariableContext<Target> context,
+      QueryExpressionContext<Target> context,
       Callback<Target> callback) {
     return transformAsync(
         getUniverseDTCSkyKeyPredicateFuture(universe, context),
@@ -1318,32 +1277,4 @@
       }
     }
   }
-
-  /** Pair of a key and a depth, useful for driving usages of {@link MinDepthUniquifier}. */
-  public static class KeyAtDepth {
-    public final SkyKey key;
-    public final int depth;
-
-    public KeyAtDepth(SkyKey key, int depth) {
-      this.key = key;
-      this.depth = depth;
-    }
-
-    @Override
-    public int hashCode() {
-      // N.B. - We deliberately use a garbage-free hashCode implementation (rather than e.g.
-      // Objects#hash). This method is very hot during large visitations done by
-      // ParallelSkyQueryUtils.
-      return 31 * key.hashCode() + Integer.hashCode(depth);
-    }
-
-    @Override
-    public boolean equals(Object obj) {
-      if (!(obj instanceof KeyAtDepth)) {
-        return false;
-      }
-      KeyAtDepth other = (KeyAtDepth) obj;
-      return key.equals(other.key) && depth == other.depth;
-    }
-  }
 }