Add rbuildfiles function to queries, only for use when using SkyQueryEnvironment.

rbuildfiles returns the packages (in the form of BUILD files) that depend on the given source files as BUILD files or subincludes.

--
MOS_MIGRATED_REVID=101157700

4 files changed, 190 insertions, 28 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java b/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java
new file mode 100644
index 0000000000..299b03443c
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java
@@ -0,0 +1,82 @@
+// Copyright 2015 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.query2;
+
+import com.google.common.base.Function;
+import com.google.common.collect.Collections2;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.ArgumentType;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
+import com.google.devtools.build.lib.query2.engine.QueryException;
+import com.google.devtools.build.lib.query2.engine.QueryExpression;
+import com.google.devtools.build.lib.vfs.PathFragment;
+
+import java.util.List;
+import java.util.Set;
+
+/**
+ * An "rbuildfiles" query expression, which computes the set of packages (as represented by their
+ * BUILD source file targets) that depend on the given set of files, either as BUILD files directly
+ * or as subincludes. Is morally the inverse of the "buildfiles" operator, although that operator
+ * takes targets and returns subinclude targets, while this takes files and returns BUILD file
+ * targets.
+ *
+ * <pre>expr ::= RBUILDFILES '(' WORD, ... ')'</pre>
+ *
+ * <p>This expression can only be used with SkyQueryEnvironment.
+ */
+public class RBuildFilesFunction implements QueryFunction {
+
+  @Override
+  public String getName() {
+    return "rbuildfiles";
+  }
+
+  @Override
+  public int getMandatoryArguments() {
+    return 1;
+  }
+
+  @Override
+  public Iterable<ArgumentType> getArgumentTypes() {
+    return Iterables.cycle(ArgumentType.WORD);
+  }
+
+  private static final Function<Argument, PathFragment> ARGUMENT_TO_PATH_FRAGMENT =
+      new Function<Argument, PathFragment>() {
+        @Override
+        public PathFragment apply(Argument argument) {
+          return new PathFragment(argument.getWord());
+        }
+      };
+
+  @Override
+  @SuppressWarnings("unchecked") // Cast from <Target> to <T>. This will only be used with <Target>.
+  public <T> Set<T> eval(QueryEnvironment<T> env, QueryExpression expression, List<Argument> args)
+      throws QueryException {
+    if (!(env instanceof SkyQueryEnvironment)) {
+      throw new QueryException("rbuildfiles can only be used with SkyQueryEnvironment");
+    }
+    ImmutableList.Builder<String> fileNames = ImmutableList.builder();
+    for (Argument arg : args) {
+      fileNames.add(arg.getWord());
+    }
+    return (Set<T>)
+        ((SkyQueryEnvironment) env)
+            .getRBuildFiles(Collections2.transform(args, ARGUMENT_TO_PATH_FRAGMENT));
+  }
+}
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 17ecb747da..1910762914 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
@@ -20,12 +20,14 @@ import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.Collections2;
+import com.google.common.collect.CompactHashSet;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
+import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.cmdline.ResolvedTargets;
 import com.google.devtools.build.lib.cmdline.TargetParsingException;
 import com.google.devtools.build.lib.cmdline.TargetPattern;
@@ -35,6 +37,7 @@ import com.google.devtools.build.lib.graph.Digraph;
 import com.google.devtools.build.lib.packages.NoSuchTargetException;
 import com.google.devtools.build.lib.packages.NoSuchThingException;
 import com.google.devtools.build.lib.packages.Package;
+import com.google.devtools.build.lib.packages.PackageIdentifier;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
@@ -44,7 +47,9 @@ import com.google.devtools.build.lib.query2.engine.AllRdepsFunction;
 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.skyframe.FileValue;
 import com.google.devtools.build.lib.skyframe.GraphBackedRecursivePackageProvider;
+import com.google.devtools.build.lib.skyframe.PackageLookupValue;
 import com.google.devtools.build.lib.skyframe.PackageValue;
 import com.google.devtools.build.lib.skyframe.PrepareDepsOfPatternsValue;
 import com.google.devtools.build.lib.skyframe.RecursivePackageProviderBackedTargetPatternResolver;
@@ -53,6 +58,8 @@ import com.google.devtools.build.lib.skyframe.TargetPatternValue;
 import com.google.devtools.build.lib.skyframe.TargetPatternValue.TargetPatternKey;
 import com.google.devtools.build.lib.skyframe.TransitiveTraversalValue;
 import com.google.devtools.build.lib.syntax.Label;
+import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.EvaluationResult;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -499,7 +506,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
         }
       };
 
-  private Iterable<SkyKey> makeKeys(Iterable<Target> targets) {
+  private static Iterable<SkyKey> makeKeys(Iterable<Target> targets) {
     return Iterables.transform(targets, TARGET_TO_SKY_KEY);
   }
 
@@ -508,9 +515,89 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
     return target;
   }
 
+  /**
+   * Get SkyKeys for the FileValues for the given {@param pathFragments}. To do this, we look for a
+   * package lookup node for each path fragment, since package lookup nodes contain the "root" of a
+   * package. The returned SkyKeys correspond to FileValues that may not exist in the graph.
+   */
+  private Collection<SkyKey> getSkyKeysForFileFragments(Iterable<PathFragment> pathFragments) {
+    Set<SkyKey> result = new HashSet<>();
+    Multimap<PathFragment, PathFragment> currentToOriginal = ArrayListMultimap.create();
+    for (PathFragment pathFragment : pathFragments) {
+      currentToOriginal.put(pathFragment, pathFragment);
+    }
+    while (!currentToOriginal.isEmpty()) {
+      Map<SkyKey, PathFragment> keys = new HashMap<>();
+      for (PathFragment pathFragment : currentToOriginal.keySet()) {
+        keys.put(
+            PackageLookupValue.key(PackageIdentifier.createInDefaultRepo(pathFragment)),
+            pathFragment);
+      }
+      Map<SkyKey, SkyValue> lookupValues = graph.getDoneValues(keys.keySet());
+      for (Map.Entry<SkyKey, SkyValue> entry : lookupValues.entrySet()) {
+        PackageLookupValue packageLookupValue = (PackageLookupValue) entry.getValue();
+        PathFragment dir = keys.get(entry.getKey());
+        if (packageLookupValue.packageExists()) {
+          Collection<PathFragment> originalFiles = currentToOriginal.get(dir);
+          Preconditions.checkState(!originalFiles.isEmpty(), entry);
+          for (PathFragment fileName : originalFiles) {
+            result.add(
+                FileValue.key(RootedPath.toRootedPath(packageLookupValue.getRoot(), fileName)));
+          }
+        }
+        currentToOriginal.removeAll(dir);
+      }
+      Multimap<PathFragment, PathFragment> newCurrentToOriginal = ArrayListMultimap.create();
+      for (PathFragment pathFragment : currentToOriginal.keySet()) {
+        PathFragment parent = pathFragment.getParentDirectory();
+        if (parent != null) {
+          newCurrentToOriginal.putAll(parent, currentToOriginal.get(pathFragment));
+        }
+      }
+      currentToOriginal = newCurrentToOriginal;
+    }
+    return result;
+  }
+
+  /**
+   * Calculates the set of {@link Package} objects, represented as source file targets, that depend
+   * on the given list of BUILD files and subincludes (other files are filtered out).
+   */
+  @Nullable
+  Set<Target> getRBuildFiles(Collection<PathFragment> fileIdentifiers) {
+    Collection<SkyKey> files = getSkyKeysForFileFragments(fileIdentifiers);
+    Collection<SkyKey> current = graph.getDoneValues(files).keySet();
+    Set<SkyKey> resultKeys = CompactHashSet.create();
+    while (!current.isEmpty()) {
+      Collection<Iterable<SkyKey>> reverseDeps = graph.getReverseDeps(current).values();
+      current = new HashSet<>();
+      for (SkyKey rdep : Iterables.concat(reverseDeps)) {
+        if (rdep.functionName().equals(SkyFunctions.PACKAGE)) {
+          resultKeys.add(rdep);
+        } else if (!rdep.functionName().equals(SkyFunctions.PACKAGE_LOOKUP)) {
+          // Packages may depend on subpackages for existence, but we don't report them as rdeps.
+          current.add(rdep);
+        }
+      }
+    }
+    Map<SkyKey, SkyValue> packageValues = graph.getDoneValues(resultKeys);
+    if (packageValues.size() != resultKeys.size()) {
+      throw new IllegalStateException(
+          "Missing values: " + Sets.difference(resultKeys, packageValues.keySet()));
+    }
+    ImmutableSet.Builder<Target> result = ImmutableSet.builder();
+    for (SkyValue value : packageValues.values()) {
+      result.add(((PackageValue) value).getPackage().getBuildFile());
+    }
+    return result.build();
+  }
+
   @Override
   public Iterable<QueryFunction> getFunctions() {
     return ImmutableList.<QueryFunction>builder()
-        .addAll(super.getFunctions()).add(new AllRdepsFunction()).build();
+        .addAll(super.getFunctions())
+        .add(new AllRdepsFunction())
+        .add(new RBuildFilesFunction())
+        .build();
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryEnvironment.java
index 6dfabfd818..78603c9e0d 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryEnvironment.java
@@ -30,17 +30,15 @@ import javax.annotation.Nonnull;
  * @param <T> the node type of the dependency graph
  */
 public interface QueryEnvironment<T> {
-  /**
-   * Type of an argument of a user-defined query function.
-   */
-  public enum ArgumentType {
-    EXPRESSION, WORD, INTEGER;
+  /** Type of an argument of a user-defined query function. */
+  enum ArgumentType {
+    EXPRESSION,
+    WORD,
+    INTEGER;
   }
 
-  /**
-   * Value of an argument of a user-defined query function.
-   */
-  public static class Argument {
+  /** Value of an argument of a user-defined query function. */
+  class Argument {
     private final ArgumentType type;
     private final QueryExpression expression;
     private final String word;
@@ -92,10 +90,8 @@ public interface QueryEnvironment<T> {
     }
   }
 
-  /**
-   * A user-defined query function.
-   */
-  public interface QueryFunction {
+  /** A user-defined query function. */
+  interface QueryFunction {
     /**
      * Name of the function as it appears in the query language.
      */
@@ -109,10 +105,8 @@ public interface QueryEnvironment<T> {
      */
     int getMandatoryArguments();
 
-    /**
-     * The types of the arguments of the function.
-     */
-    List<ArgumentType> getArgumentTypes();
+    /** The types of the arguments of the function. */
+    Iterable<ArgumentType> getArgumentTypes();
 
     /**
      * Called when a user-defined function is to be evaluated.
@@ -130,7 +124,7 @@ public interface QueryEnvironment<T> {
    * Exception type for the case where a target cannot be found. It's basically a wrapper for
    * whatever exception is internally thrown.
    */
-  public static final class TargetNotFoundException extends Exception {
+  final class TargetNotFoundException extends Exception {
     public TargetNotFoundException(String msg) {
       super(msg);
     }
@@ -227,7 +221,7 @@ public interface QueryEnvironment<T> {
   /**
    * Settings for the query engine. See {@link QueryEnvironment#isSettingEnabled}.
    */
-  public static enum Setting {
+  enum Setting {
 
     /**
      * Whether to evaluate tests() expressions in strict mode. If {@link #isSettingEnabled} returns
@@ -257,7 +251,7 @@ public interface QueryEnvironment<T> {
    * An adapter interface giving access to properties of T. There are four types of targets: rules,
    * package groups, source files, and generated files. Of these, only rules can have attributes.
    */
-  public static interface TargetAccessor<T> {
+  interface TargetAccessor<T> {
     /**
      * Returns the target type represented as a string of the form {@code &lt;type&gt; rule} or
      * {@code package group} or {@code source file} or {@code generated file}. This is widely used
@@ -345,8 +339,8 @@ public interface QueryEnvironment<T> {
   }
 
   /** List of the default query functions. */
-  public static final List<QueryFunction> DEFAULT_QUERY_FUNCTIONS =
-      ImmutableList.<QueryFunction>of(
+  List<QueryFunction> DEFAULT_QUERY_FUNCTIONS =
+      ImmutableList.of(
           new AllPathsFunction(),
           new BuildFilesFunction(),
           new AttrFunction(),
@@ -358,6 +352,5 @@ public interface QueryEnvironment<T> {
           new TestsFunction(),
           new DepsFunction(),
           new RdepsFunction(),
-          new VisibleFunction()
-          );
+          new VisibleFunction());
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PackageLookupValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/PackageLookupValue.java
index c26252a198..3bf53b6ea4 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/PackageLookupValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PackageLookupValue.java
@@ -34,7 +34,7 @@ import java.util.Objects;
  *
  * <p>Implementation detail: we use inheritance here to optimize for memory usage.
  */
-abstract class PackageLookupValue implements SkyValue {
+public abstract class PackageLookupValue implements SkyValue {
 
   enum ErrorReason {
     // There is no BUILD file.
@@ -115,7 +115,7 @@ abstract class PackageLookupValue implements SkyValue {
     return key(PackageIdentifier.createInDefaultRepo(directory));
   }
 
-  static SkyKey key(PackageIdentifier pkgIdentifier) {
+  public static SkyKey key(PackageIdentifier pkgIdentifier) {
     return new SkyKey(SkyFunctions.PACKAGE_LOOKUP, pkgIdentifier);
   }