1 files changed, 88 insertions, 4 deletions

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;