RELNOTES: add loadfiles() query operator, to find skylark files loaded by targets.

--
MOS_MIGRATED_REVID=110275427

5 files changed, 153 insertions, 29 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
index 1b3cc8adad..d4614f0cbf 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
@@ -48,10 +48,12 @@ import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -348,7 +350,12 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   // TODO(bazel-team): rename this to getDependentFiles when all implementations
   // of QueryEnvironment is fixed.
   @Override
-  public Set<Target> getBuildFiles(final QueryExpression caller, Set<Target> nodes)
+  public Set<Target> getBuildFiles(
+      final QueryExpression caller,
+      Set<Target> nodes,
+      boolean buildFiles,
+      boolean subincludes,
+      boolean loads)
       throws QueryException {
     Set<Target> dependentFiles = new LinkedHashSet<>();
     Set<Package> seenPackages = new HashSet<>();
@@ -361,17 +368,32 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     for (Target x : nodes) {
       Package pkg = x.getPackage();
       if (seenPackages.add(pkg)) {
-        addIfUniqueLabel(getNode(pkg.getBuildFile()), seenLabels, dependentFiles);
-        for (Label subinclude
-            : Iterables.concat(pkg.getSubincludeLabels(), pkg.getSkylarkFileDependencies())) {
+        if (buildFiles) {
+          addIfUniqueLabel(getNode(pkg.getBuildFile()), seenLabels, dependentFiles);
+        }
+
+        List<Label> extensions = new ArrayList<>();
+        if (subincludes) {
+          extensions.addAll(pkg.getSubincludeLabels());
+        }
+        if (loads) {
+          extensions.addAll(pkg.getSkylarkFileDependencies());
+        }
+
+        for (Label subinclude : extensions) {
           addIfUniqueLabel(getSubincludeTarget(subinclude, pkg), seenLabels, dependentFiles);
 
           // Also add the BUILD file of the subinclude.
-          try {
-            addIfUniqueLabel(getSubincludeTarget(
-                subinclude.getLocalTargetLabel("BUILD"), pkg), seenLabels, dependentFiles);
-          } catch (LabelSyntaxException e) {
-            throw new AssertionError("BUILD should always parse as a target name", e);
+          if (buildFiles) {
+            try {
+              addIfUniqueLabel(
+                  getSubincludeTarget(subinclude.getLocalTargetLabel("BUILD"), pkg),
+                  seenLabels,
+                  dependentFiles);
+
+            } catch (LabelSyntaxException e) {
+              throw new AssertionError("BUILD should always parse as a target name", e);
+            }
           }
         }
       }
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 4ca2ffac32..d304382b18 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
@@ -339,7 +339,12 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
   }
 
   @Override
-  public Set<Target> getBuildFiles(QueryExpression caller, Set<Target> nodes)
+  public Set<Target> getBuildFiles(
+      QueryExpression caller,
+      Set<Target> nodes,
+      boolean buildFiles,
+      boolean subincludes,
+      boolean loads)
       throws QueryException {
     Set<Target> dependentFiles = new LinkedHashSet<>();
     Set<Package> seenPackages = new HashSet<>();
@@ -352,17 +357,32 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
     for (Target x : nodes) {
       Package pkg = x.getPackage();
       if (seenPackages.add(pkg)) {
-        addIfUniqueLabel(pkg.getBuildFile(), seenLabels, dependentFiles);
-        for (Label subinclude
-            : Iterables.concat(pkg.getSubincludeLabels(), pkg.getSkylarkFileDependencies())) {
+        if (buildFiles) {
+          addIfUniqueLabel(pkg.getBuildFile(), seenLabels, dependentFiles);
+        }
+
+        List<Label> extensions = new ArrayList<>();
+        if (subincludes) {
+          extensions.addAll(pkg.getSubincludeLabels());
+        }
+        if (loads) {
+          extensions.addAll(pkg.getSkylarkFileDependencies());
+        }
+
+        for (Label subinclude : extensions) {
           addIfUniqueLabel(getSubincludeTarget(subinclude, pkg), seenLabels, dependentFiles);
 
-          // Also add the BUILD file of the subinclude.
-          try {
-            addIfUniqueLabel(getSubincludeTarget(
-                subinclude.getLocalTargetLabel("BUILD"), pkg), seenLabels, dependentFiles);
-          } catch (LabelSyntaxException e) {
-            throw new AssertionError("BUILD should always parse as a target name", e);
+          if (buildFiles) {
+            // Also add the BUILD file of the subinclude.
+            try {
+              addIfUniqueLabel(
+                  getSubincludeTarget(subinclude.getLocalTargetLabel("BUILD"), pkg),
+                  seenLabels,
+                  dependentFiles);
+
+            } catch (LabelSyntaxException e) {
+              throw new AssertionError("BUILD should always parse as a target name", e);
+            }
           }
         }
       }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/BuildFilesFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/BuildFilesFunction.java
index c36f98a29f..ea391e9627 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/BuildFilesFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/BuildFilesFunction.java
@@ -44,14 +44,19 @@ class BuildFilesFunction implements QueryFunction {
   public <T> void eval(final QueryEnvironment<T> env, final QueryExpression expression,
       List<Argument> args, final Callback<T> callback)
       throws QueryException, InterruptedException {
-    env.eval(args.get(0).getExpression(), new Callback<T>() {
-      @Override
-      public void process(Iterable<T> partialResult) throws QueryException, InterruptedException {
-        Set<T> result = CompactHashSet.create();
-        Iterables.addAll(result, partialResult);
-        callback.process(env.getBuildFiles(expression, result));
-      }
-    });
+    env.eval(
+        args.get(0).getExpression(),
+        new Callback<T>() {
+          @Override
+          public void process(Iterable<T> partialResult)
+              throws QueryException, InterruptedException {
+            Set<T> result = CompactHashSet.create();
+            Iterables.addAll(result, partialResult);
+            callback.process(
+                env.getBuildFiles(
+                    expression, result, /* BUILD */ true, /* subinclude */ true, /* load */ true));
+          }
+        });
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/LoadFilesFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/LoadFilesFunction.java
new file mode 100644
index 0000000000..1566f2166c
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/LoadFilesFunction.java
@@ -0,0 +1,74 @@
+// Copyright 2014 The Bazel Authors. 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.engine;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.devtools.build.lib.collect.CompactHashSet;
+
+import java.util.List;
+import java.util.Set;
+
+/**
+ * A loadfiles(x) query expression, which computes the set of .bzl files
+ * for each target in set x.  The result is unordered.  This
+ * operator is typically used for determinining what files or packages to check
+ * out.
+ *
+ * <pre>expr ::= LOADFILES '(' expr ')'</pre>
+ */
+class LoadFilesFunction implements QueryEnvironment.QueryFunction {
+  LoadFilesFunction() {}
+
+  @Override
+  public String getName() {
+    return "loadfiles";
+  }
+
+  @Override
+  public <T> void eval(
+      final QueryEnvironment<T> env,
+      final QueryExpression expression,
+      List<QueryEnvironment.Argument> args,
+      final Callback<T> callback)
+      throws QueryException, InterruptedException {
+    env.eval(
+        args.get(0).getExpression(),
+        new Callback<T>() {
+          @Override
+          public void process(Iterable<T> partialResult)
+              throws QueryException, InterruptedException {
+            Set<T> result = CompactHashSet.create();
+            Iterables.addAll(result, partialResult);
+            callback.process(
+                env.getBuildFiles(
+                    expression,
+                    result,
+                    /* BUILD */ false,
+                    /* subinclude */ false,
+                    /* load */ true));
+          }
+        });
+  }
+
+  @Override
+  public int getMandatoryArguments() {
+    return 1;
+  }
+
+  @Override
+  public List<QueryEnvironment.ArgumentType> getArgumentTypes() {
+    return ImmutableList.of(QueryEnvironment.ArgumentType.EXPRESSION);
+  }
+}
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 b8c065f936..13d7f4ca12 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
@@ -207,10 +207,12 @@ public interface QueryEnvironment<T> {
   void reportBuildFileError(QueryExpression expression, String msg) throws QueryException;
 
   /**
-   * Returns the set of BUILD, included, sub-included and Skylark files that define the given set of
+   * Returns the set of BUILD, and optionally sub-included and Skylark files that define the given set of
    * targets. Each such file is itself represented as a target in the result.
    */
-  Set<T> getBuildFiles(QueryExpression caller, Set<T> nodes) throws QueryException;
+  Set<T> getBuildFiles(
+      QueryExpression caller, Set<T> nodes, boolean buildFiles, boolean subincludes, boolean loads)
+      throws QueryException;
 
   /**
    * Returns an object that can be used to query information about targets. Implementations should
@@ -359,6 +361,7 @@ public interface QueryEnvironment<T> {
       ImmutableList.of(
           new AllPathsFunction(),
           new BuildFilesFunction(),
+          new LoadFilesFunction(),
           new AttrFunction(),
           new FilterFunction(),
           new LabelsFunction(),