From a09cf42c8f0fa7d7bc57560d5217516d09030e9f Mon Sep 17 00:00:00 2001
From: Daniel Wagner-Hall <dwagnerhall@apple.com>
Date: Thu, 10 Aug 2023 11:50:50 +0100
Subject: [PATCH] cquery: Include toolchains as ruleInputs

These were already returned in a text output for a deps query - this
adds them also as ruleInputs.
---
 .../cquery/ProtoOutputFormatterCallback.java  | 17 +++-
 .../query/output/ProtoOutputFormatter.java    | 15 ++--
 .../StreamedJSONProtoOutputFormatter.java     |  3 +-
 .../output/StreamedProtoOutputFormatter.java  |  3 +-
 .../integration/configured_query_test.sh      | 81 +++++++++++++++++++
 5 files changed, 111 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java b/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java
index 0289a3b6874ca6..f83c094a46c995 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java
@@ -17,6 +17,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Ordering;
 import com.google.devtools.build.lib.actions.BuildConfigurationEvent;
@@ -27,6 +28,7 @@
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.ConfigMatchingProvider;
+import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
@@ -243,7 +245,20 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult)
       // we will want to add relevant tests.
       currentTarget = keyedConfiguredTarget;
       Target target = accessor.getTarget(keyedConfiguredTarget);
-      Build.Target.Builder targetBuilder = formatter.toTargetProtoBuffer(target).toBuilder();
+
+      // Toolchain dependencies appear as implicit dependencies, but don't appear on any attribute,
+      // which means they won't get detected by ruleInputs discovery which traverses attributes to get inputs.
+      // But we know here from our ConfiguredTarget that they are dependencies.
+      // Artificially provide these implicit dependencies as inputs.
+      ImmutableSortedSet.Builder<Label> implicitDeps = ImmutableSortedSet.naturalOrder();
+      if (keyedConfiguredTarget instanceof RuleConfiguredTarget) {
+        RuleConfiguredTarget ruleConfiguredTarget = (RuleConfiguredTarget) keyedConfiguredTarget;
+        for (ConfiguredTargetKey implicitDep : ruleConfiguredTarget.getImplicitDeps()) {
+          implicitDeps.add(implicitDep.getLabel());
+        }
+      }
+
+      Build.Target.Builder targetBuilder = formatter.toTargetProtoBuffer(target, implicitDeps.build()).toBuilder();
       if (target instanceof Rule && !Transitions.NONE.equals(options.transitions)) {
         try {
           for (CqueryTransitionResolver.ResolvedTransition resolvedTransition :
diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoOutputFormatter.java
index c716080d62dd17..2c89ef39a89a52 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/output/ProtoOutputFormatter.java
@@ -26,6 +26,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.hash.HashFunction;
@@ -170,12 +171,12 @@ public ThreadSafeOutputFormatterCallback<Target> createStreamCallback(
   }
 
   /** Converts a logical {@link Target} object into a {@link Build.Target} protobuffer. */
-  public Build.Target toTargetProtoBuffer(Target target) throws InterruptedException {
-    return toTargetProtoBuffer(target, /*extraDataForAttrHash=*/ "");
+  public Build.Target toTargetProtoBuffer(Target target, ImmutableSortedSet<Label> extraImplicitDeps) throws InterruptedException {
+    return toTargetProtoBuffer(target, extraImplicitDeps, /*extraDataForAttrHash=*/ "");
   }
 
   /** Converts a logical {@link Target} object into a {@link Build.Target} protobuffer. */
-  public Build.Target toTargetProtoBuffer(Target target, Object extraDataForAttrHash)
+  private Build.Target toTargetProtoBuffer(Target target, ImmutableSortedSet<Label> extraImplicitDeps, Object extraDataForAttrHash)
       throws InterruptedException {
     Build.Target.Builder targetPb = Build.Target.newBuilder();
 
@@ -243,7 +244,11 @@ public Build.Target toTargetProtoBuffer(Target target, Object extraDataForAttrHa
         // Include explicit elements for all direct inputs and outputs of a rule; this goes beyond
         // what is available from the attributes above, since it may also (depending on options)
         // include implicit outputs, exec-configuration outputs, and default values.
-        rule.getSortedLabels(dependencyFilter)
+        ImmutableSortedSet<Label> ruleInputs = ImmutableSortedSet.<Label>naturalOrder()
+            .addAll(rule.getSortedLabels(dependencyFilter))
+            .addAll(extraImplicitDeps)
+            .build();
+        ruleInputs
             .forEach(input -> rulePb.addRuleInput(input.toString()));
         rule.getOutputFiles().stream()
             .distinct()
@@ -565,7 +570,7 @@ public void processOutput(Iterable<Target> partialResult)
       // constructing and serializing a QueryResult proto are protected by test coverage and proto
       // best practices.
       for (Target target : partialResult) {
-        codedOut.writeMessage(QueryResult.TARGET_FIELD_NUMBER, toTargetProtoBuffer(target));
+        codedOut.writeMessage(QueryResult.TARGET_FIELD_NUMBER, toTargetProtoBuffer(target, ImmutableSortedSet.of()));
       }
     }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedJSONProtoOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedJSONProtoOutputFormatter.java
index d8509ded692fdf..025948c47fc05f 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedJSONProtoOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedJSONProtoOutputFormatter.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.query2.query.output;
 
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
@@ -44,7 +45,7 @@ public void processOutput(Iterable<Target> partialResult)
           out.write(
               jsonPrinter
                   .omittingInsignificantWhitespace()
-                  .print(toTargetProtoBuffer(target))
+                  .print(toTargetProtoBuffer(target, ImmutableSortedSet.of()))
                   .getBytes(StandardCharsets.UTF_8));
           out.write(System.lineSeparator().getBytes(StandardCharsets.UTF_8));
         }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java
index 4adb1a840778fd..9b630f21e89be4 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.query2.query.output;
 
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
@@ -38,7 +39,7 @@ public OutputFormatterCallback<Target> createPostFactoStreamCallback(
       public void processOutput(Iterable<Target> partialResult)
           throws IOException, InterruptedException {
         for (Target target : partialResult) {
-          toTargetProtoBuffer(target).writeDelimitedTo(out);
+          toTargetProtoBuffer(target, ImmutableSortedSet.of()).writeDelimitedTo(out);
         }
       }
     };
diff --git a/src/test/shell/integration/configured_query_test.sh b/src/test/shell/integration/configured_query_test.sh
index 4ddac2aefca188..72d5ae23ad5aed 100755
--- a/src/test/shell/integration/configured_query_test.sh
+++ b/src/test/shell/integration/configured_query_test.sh
@@ -1521,4 +1521,85 @@ EOF
   assert_contains "$pkg/single_file" output
 }
 
+function test_toolchain_rule_inputs() {
+  local -r pkg=$FUNCNAME
+  mkdir -p $pkg
+
+  cat >$pkg/BUILD <<EOF
+load("//$pkg:rules.bzl", "my_toolchain", "rule_using_toolchain")
+
+toolchain_type(
+    name = "toolchain_type",
+    visibility = ["//visibility:public"],
+)
+
+my_toolchain(
+    name = "my_toolchain",
+    bin = "echo_hello_to_argv0.sh",
+)
+
+toolchain(
+    name = "toolchain",
+    toolchain = ":my_toolchain",
+    toolchain_type = ":toolchain_type",
+)
+
+rule_using_toolchain(
+    name = "rule",
+)
+EOF
+
+  cat >$pkg/rules.bzl <<EOF
+def _rule_using_toolchain_impl(ctx):
+    output_file = ctx.actions.declare_file("output_file")
+
+    bin = ctx.toolchains["//$pkg:toolchain_type"].bin
+
+    args = ctx.actions.args()
+    args.add(output_file.path)
+
+    ctx.actions.run(
+        outputs = [output_file],
+        executable = bin.path,
+        arguments = [output_file.path],
+        tools = [bin],
+        use_default_shell_env = True,
+    )
+
+    return [
+        DefaultInfo(
+            files = depset([output_file]),
+        ),
+    ]
+
+rule_using_toolchain = rule(
+    implementation = _rule_using_toolchain_impl,
+    toolchains = ["//$pkg:toolchain_type"],
+)
+
+def _toolchain_impl(ctx):
+    return [
+        platform_common.ToolchainInfo(
+            bin = ctx.file.bin,
+        ),
+    ]
+
+my_toolchain = rule(
+    implementation = _toolchain_impl,
+    attrs = {
+        "bin": attr.label(
+            allow_single_file = True,
+            executable = True,
+            cfg = "exec",
+        ),
+    },
+)
+EOF
+  touch $pkg/echo_hello_to_argv0.sh
+  chmod 0755 $pkg/echo_hello_to_argv0.sh
+
+  bazel cquery --output=textproto --extra_toolchains=//$pkg:toolchain //$pkg:rule > output 2>"$TEST_log" || fail "Unexpected failure"
+  assert_contains "rule_input: \"//$pkg:my_toolchain\"" output
+}
+
 run_suite "${PRODUCT_NAME} configured query tests"