Implement the config(expr, word) cquery function.

expr - the expression to be evaluated
word - the configuration (represented by the strings 'host', 'target', or 'null') to try to find the result(s) of 'expr' in. If some but not all results of expr can be found in the specified config, then the subset that can be is returned. If no results of expr can be found in the specified config, then an error is thrown.

PiperOrigin-RevId: 185572590

4 files changed, 185 insertions, 7 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
index 3530399a9c..adc6b18113 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
@@ -66,6 +66,7 @@ import com.google.devtools.build.lib.profiler.ProfilePhase;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.query2.CommonQueryOptions;
 import com.google.devtools.build.lib.query2.ConfiguredTargetQueryEnvironment;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
 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;
@@ -474,11 +475,16 @@ public final class BuildTool {
 
     WalkableGraph walkableGraph =
         SkyframeExecutorWrappingWalkableGraph.of(env.getSkyframeExecutor());
+    ImmutableList<QueryFunction> extraFunctions =
+        new ImmutableList.Builder<QueryFunction>()
+            .addAll(ConfiguredTargetQueryEnvironment.CQUERY_FUNCTIONS)
+            .addAll(env.getRuntime().getQueryFunctions())
+            .build();
     ConfiguredTargetQueryEnvironment configuredTargetQueryEnvironment =
         new ConfiguredTargetQueryEnvironment(
             request.getKeepGoing(),
             env.getReporter(),
-            env.getRuntime().getQueryFunctions(),
+            extraFunctions,
             targetConfig,
             hostConfiguration,
             env.newTargetPatternEvaluator().getOffset(),
diff --git a/src/main/java/com/google/devtools/build/lib/query2/ConfigFunction.java b/src/main/java/com/google/devtools/build/lib/query2/ConfigFunction.java
new file mode 100644
index 0000000000..7a3f6a0b63
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/ConfigFunction.java
@@ -0,0 +1,88 @@
+// Copyright 2018 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;
+
+import com.google.common.collect.ImmutableList;
+import com.google.devtools.build.lib.analysis.ConfiguredTarget;
+import com.google.devtools.build.lib.query2.engine.Callback;
+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.QueryEnvironment.QueryTaskFuture;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.ThreadSafeMutableSet;
+import com.google.devtools.build.lib.query2.engine.QueryExpression;
+import com.google.devtools.build.lib.query2.engine.QueryUtil;
+import com.google.devtools.build.lib.query2.engine.VariableContext;
+import java.util.List;
+
+/**
+ * A "config" query expression for cquery. The first argument is the expression to be evaluated. The
+ * second argument is either "host", "target", or "null" to specify which configuration the
+ * user is seeking to query in. If some but not all results of expr can be found in the specified
+ * config, then the subset that can be is returned. If no results of expr can be found in the
+ * specified config, then an error is thrown.
+ *
+ * <pre> expr ::= CONFIG '(' expr ',' word ')'</pre>
+ */
+public final class ConfigFunction implements QueryFunction {
+
+  public ConfigFunction() {
+  }
+
+  @Override
+  public String getName() {
+    return "config";
+  }
+
+  @Override
+  public int getMandatoryArguments() {
+    return 2;
+  }
+
+  @Override
+  public List<ArgumentType> getArgumentTypes() {
+    return ImmutableList.of(ArgumentType.EXPRESSION, ArgumentType.WORD);
+  }
+
+  /**
+   * This function is only viable with ConfiguredTargetQueryEnvironment which extends {@link
+   * AbstractBlazeQueryEnvironment<ConfiguredTarget>}
+   */
+  @Override
+  @SuppressWarnings("unchecked")
+  public <T> QueryTaskFuture<Void> eval(
+      QueryEnvironment<T> env,
+      VariableContext<T> context,
+      QueryExpression expression,
+      List<Argument> args,
+      final Callback<T> callback) {
+
+    Argument targetExpression = args.get(0);
+    String configuration = args.get(1).toString();
+
+
+    final QueryTaskFuture<ThreadSafeMutableSet<T>> targets =
+        QueryUtil.evalAll(env, context, targetExpression.getExpression());
+
+    return env.whenSucceedsCall(
+        targets,
+        ((ConfiguredTargetQueryEnvironment) env)
+            .getConfiguredTargets(
+                targetExpression.toString(),
+                (ThreadSafeMutableSet<ConfiguredTarget>) targets.getIfSuccessful(),
+                configuration,
+                (Callback<ConfiguredTarget>) callback));
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/query2/ConfiguredTargetQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/ConfiguredTargetQueryEnvironment.java
index c87583db36..a694ab3f8c 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/ConfiguredTargetQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/ConfiguredTargetQueryEnvironment.java
@@ -111,6 +111,11 @@ public class ConfiguredTargetQueryEnvironment
   private static final KeyExtractor<ConfiguredTarget, ConfiguredTargetKey>
       CONFIGURED_TARGET_KEY_EXTRACTOR = ConfiguredTargetKey::of;
 
+  /** Common query functions and cquery specific functions. */
+  public static final ImmutableList<QueryFunction> FUNCTIONS = populateFunctions();
+  /** Cquery specific functions. */
+  public static final ImmutableList<QueryFunction> CQUERY_FUNCTIONS = getCqueryFunctions();
+
   static {
     TargetPattern targetPattern;
     try {
@@ -158,6 +163,17 @@ public class ConfiguredTargetQueryEnvironment
     checkSettings(settings);
   }
 
+  private static ImmutableList<QueryFunction> populateFunctions() {
+    return new ImmutableList.Builder<QueryFunction>()
+        .addAll(QueryEnvironment.DEFAULT_QUERY_FUNCTIONS)
+        .addAll(getCqueryFunctions())
+        .build();
+  }
+
+  private static ImmutableList<QueryFunction> getCqueryFunctions() {
+    return ImmutableList.of(new ConfigFunction());
+  }
+
   // Check to make sure the settings requested are currently supported by this class
   private void checkSettings(Set<Setting> settings) throws QueryException {
     if (settings.contains(Setting.NO_NODEP_DEPS)
@@ -181,20 +197,33 @@ public class ConfiguredTargetQueryEnvironment
 
   private ConfiguredTarget getConfiguredTarget(Label label) throws InterruptedException {
     // Try with target configuration.
-    ConfiguredTarget configuredTarget =
-        getConfiguredTarget(ConfiguredTargetValue.key(label, defaultTargetConfiguration));
+    ConfiguredTarget configuredTarget = getTargetConfiguredTarget(label);
     if (configuredTarget != null) {
       return configuredTarget;
     }
     // Try with host configuration (even when --nohost_deps is set in the case that top-level
     // targets are configured in the host configuration so we are doing a host-configuration-only
     // query).
-    configuredTarget =
-        getConfiguredTarget(ConfiguredTargetValue.key(label, hostConfiguration));
+    configuredTarget = getHostConfiguredTarget(label);
     if (configuredTarget != null) {
       return configuredTarget;
     }
     // Last chance: source file.
+    return getNullConfiguredTarget(label);
+  }
+
+  @Nullable
+  private ConfiguredTarget getHostConfiguredTarget(Label label) throws InterruptedException {
+    return getConfiguredTarget(ConfiguredTargetValue.key(label, hostConfiguration));
+  }
+
+  @Nullable
+  private ConfiguredTarget getTargetConfiguredTarget(Label label) throws InterruptedException {
+    return getConfiguredTarget(ConfiguredTargetValue.key(label, defaultTargetConfiguration));
+  }
+
+  @Nullable
+  private ConfiguredTarget getNullConfiguredTarget(Label label) throws InterruptedException {
     return getConfiguredTarget(ConfiguredTargetValue.key(label, null));
   }
 
@@ -268,6 +297,61 @@ public class ConfiguredTargetQueryEnvironment
             MoreExecutors.directExecutor()));
   }
 
+  /**
+   * Processes the targets in {@code targets} with the requested {@code configuration}
+   *
+   * @param pattern the original pattern that {@code targets} were parsed from. Used for error
+   *     message.
+   * @param targets the set of {@link ConfiguredTarget}s whose labels represent the targets being
+   *     requested.
+   * @param configuration the configuration to request {@code targets} in.
+   * @param callback the callback to receive the results of this method.
+   * @return {@link QueryTaskCallable} that returns the correctly configured targets.
+   */
+  QueryTaskCallable<Void> getConfiguredTargets(
+      String pattern,
+      ThreadSafeMutableSet<ConfiguredTarget> targets,
+      String configuration,
+      Callback<ConfiguredTarget> callback) {
+    return new QueryTaskCallable<Void>() {
+      @Override
+      public Void call() throws QueryException, InterruptedException {
+        List<ConfiguredTarget> transformedResult = new ArrayList<>();
+        for (ConfiguredTarget target : targets) {
+          Label label = target.getLabel();
+          ConfiguredTarget configuredTarget;
+          switch (configuration) {
+            case "\'host\'":
+              configuredTarget = getHostConfiguredTarget(label);
+              break;
+            case "\'target\'":
+              configuredTarget = getTargetConfiguredTarget(label);
+              break;
+            case "\'null\'":
+              configuredTarget = getNullConfiguredTarget(label);
+              break;
+            default:
+              throw new QueryException(
+                  "the second argument of the config function must be 'target', 'host', or 'null'");
+          }
+          if (configuredTarget != null) {
+            transformedResult.add(configuredTarget);
+          }
+        }
+        if (transformedResult.isEmpty()) {
+          throw new QueryException(
+              "No target (in) "
+                  + pattern
+                  + " could be found in the "
+                  + configuration
+                  + " configuration");
+        }
+        callback.process(transformedResult);
+        return null;
+      }
+    };
+  }
+
   @Override
   public ConfiguredTarget getOrCreate(ConfiguredTarget target) {
     return target;
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/CqueryCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/CqueryCommand.java
index 4a87a0da27..0ec5b8a4b0 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/CqueryCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/CqueryCommand.java
@@ -19,7 +19,7 @@ import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.buildtool.BuildTool;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.query2.CommonQueryOptions;
-import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
+import com.google.devtools.build.lib.query2.ConfiguredTargetQueryEnvironment;
 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;
@@ -77,7 +77,7 @@ public final class CqueryCommand implements BlazeCommand {
     }
     String query = Joiner.on(' ').join(options.getResidue());
     HashMap<String, QueryFunction> functions = new HashMap<>();
-    for (QueryFunction queryFunction : QueryEnvironment.DEFAULT_QUERY_FUNCTIONS) {
+    for (QueryFunction queryFunction : ConfiguredTargetQueryEnvironment.FUNCTIONS) {
       functions.put(queryFunction.getName(), queryFunction);
     }
     for (QueryFunction queryFunction : env.getRuntime().getQueryFunctions()) {