From 0e876b1794d4db58b72949014401d22cc65d94a1 Mon Sep 17 00:00:00 2001
From: Googler <jcater@google.com>
Date: Thu, 10 Oct 2024 09:33:23 -0700
Subject: [PATCH] Allow genrule.toolchains to accept toolchain_type targets.

PiperOrigin-RevId: 684473797
Change-Id: If7d3f2bdcfb9813853af5cb50496d974ff985d81
---
 .../devtools/build/lib/rules/genrule/BUILD    |   3 +
 .../build/lib/rules/genrule/GenRuleBase.java  |   7 +
 .../lib/rules/genrule/GenRuleBaseRule.java    |  29 +++
 .../toolchains/ToolchainContextUtil.java      |  37 +++-
 .../build/lib/bazel/rules/genrule/BUILD       |   1 +
 .../rules/genrule/GenRuleIntegrationTest.java | 169 ++++++++++++++++++
 6 files changed, 244 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/rules/genrule/BUILD b/src/main/java/com/google/devtools/build/lib/rules/genrule/BUILD
index 4ae0f8cfcababf..2d4ed50ac7e7d5 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/genrule/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/genrule/BUILD
@@ -26,11 +26,14 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:file_provider",
         "//src/main/java/com/google/devtools/build/lib/analysis:make_variable_supplier",
         "//src/main/java/com/google/devtools/build/lib/analysis:rule_definition_environment",
+        "//src/main/java/com/google/devtools/build/lib/analysis:template_variable_info",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
+        "//src/main/java/com/google/devtools/build/lib/analysis/platform",
         "//src/main/java/com/google/devtools/build/lib/analysis/stringtemplate",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
         "//src/main/java/com/google/devtools/build/lib/packages",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:configured_target_and_data",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:filetype",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
diff --git a/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBase.java b/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBase.java
index 9bb899f4a05232..9ed01ef86b6969 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBase.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBase.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.rules.genrule;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 
 import com.google.common.collect.ImmutableList;
@@ -47,6 +48,7 @@
 import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.packages.TriState;
 import com.google.devtools.build.lib.packages.Type;
+import com.google.devtools.build.lib.skyframe.ConfiguredTargetAndData;
 import com.google.devtools.build.lib.util.FileTypeSet;
 import com.google.devtools.build.lib.util.OnDemandString;
 import com.google.devtools.build.lib.util.Pair;
@@ -90,12 +92,17 @@ public final ConfiguredTarget create(RuleContext ruleContext)
     labelMap.values().forEach(resolvedSrcsBuilder::addTransitive);
     NestedSet<Artifact> resolvedSrcs = resolvedSrcsBuilder.build();
 
+    ImmutableList<ConfiguredTarget> toolchainPrerequisites =
+        ruleContext.getToolchainContext().prerequisiteTargets().stream()
+            .map(ConfiguredTargetAndData::getConfiguredTarget)
+            .collect(toImmutableList());
     // The CommandHelper class makes an explicit copy of this in the constructor, so flattening
     // here should be benign.
     CommandHelper commandHelper =
         CommandHelper.builder(ruleContext)
             .addToolDependencies("tools")
             .addToolDependencies("toolchains")
+            .addToolDependencies(toolchainPrerequisites)
             .addLabelMap(
                 labelMap.entrySet().stream()
                     .collect(toImmutableMap(Map.Entry::getKey, e -> e.getValue().toList())))
diff --git a/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBaseRule.java b/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBaseRule.java
index 1da01f19d2115c..ebf77bf84dd76b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBaseRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/genrule/GenRuleBaseRule.java
@@ -20,10 +20,13 @@
 import static com.google.devtools.build.lib.packages.Type.BOOLEAN;
 import static com.google.devtools.build.lib.packages.Type.STRING;
 
+import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.analysis.BaseRuleClasses;
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.analysis.RuleDefinitionEnvironment;
+import com.google.devtools.build.lib.analysis.TemplateVariableInfo;
 import com.google.devtools.build.lib.analysis.config.ExecutionTransitionFactory;
+import com.google.devtools.build.lib.analysis.platform.ToolchainTypeInfo;
 import com.google.devtools.build.lib.packages.Attribute;
 import com.google.devtools.build.lib.packages.AttributeMap;
 import com.google.devtools.build.lib.packages.BuildType;
@@ -302,6 +305,32 @@ public Object getDefault(AttributeMap rule) {
         // This is a misfeature, so don't document it. We would like to get rid of it, but that
         // would require a cleanup of existing rules.
         .add(attr("heuristic_label_expansion", BOOLEAN).value(false))
+
+        /* <!-- #BLAZE_RULE(genrule).ATTRIBUTE(toolchains) -->
+        <p>
+          The set of targets whose <a href="${link make-variables}">Make variables</a> this genrule
+          is allowed to access, or the <a href="${link toolchains}"><code>toolchain_type</code></a>
+          targets that this genrule will access.
+        </p>
+
+        <p>
+          Toolchains accessed via <code>toolchain_type</code> must also provide a
+          <code>TemplateVariableInfo</code> provider, which the target can use to access toolchain
+          details.
+        </p>
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        // Override `toolchains` to allow make variables and toolchain types. Toolchain types are
+        // handled specially in ToolchainContextUtil.
+        .override(
+            attr("toolchains", LABEL_LIST)
+                .allowedFileTypes(FileTypeSet.NO_FILE)
+                // Accept TemplateVariableInfo and/or ToolchainTypeInfo.
+                .mandatoryProvidersList(
+                    ImmutableList.of(
+                        ImmutableList.of(TemplateVariableInfo.PROVIDER.id()),
+                        ImmutableList.of(ToolchainTypeInfo.PROVIDER.id())))
+                .dontCheckConstraints()
+                .nonconfigurable("Don't allow dynamic toolchain types"))
         .removeAttribute("data")
         .removeAttribute("deps")
         .build();
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/toolchains/ToolchainContextUtil.java b/src/main/java/com/google/devtools/build/lib/skyframe/toolchains/ToolchainContextUtil.java
index 8ae7594a675b57..2ed4ed00c65008 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/toolchains/ToolchainContextUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/toolchains/ToolchainContextUtil.java
@@ -27,6 +27,7 @@
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.skyframe.config.BuildConfigurationKey;
+import java.util.List;
 import javax.annotation.Nullable;
 
 /** Utility methods for toolchain context creation. */
@@ -49,11 +50,18 @@ public static UnloadedToolchainContextsInputs getUnloadedToolchainContextsInputs
             .getAutoExecGroupsMode()
             .isEnabled(RawAttributeMapper.of(rule), coreOptions.useAutoExecGroups);
     var defaultExecConstraintLabels = getExecutionPlatformConstraints(rule, platformConfig);
+    ImmutableSet<ToolchainTypeRequirement> toolchainTypes = ruleClass.getToolchainTypes();
+
+    if (!ruleClass.isStarlark() && ruleClass.getName().equals("genrule")) {
+      // Override the toolchain types based on the target-level "toolchains" attribute.
+      toolchainTypes = updateToolchainTypesFromAttribute(rule, toolchainTypes);
+    }
+
     var processedExecGroups =
         ExecGroupCollection.process(
             ruleClass.getExecGroups(),
             defaultExecConstraintLabels,
-            ruleClass.getToolchainTypes(),
+            toolchainTypes,
             useAutoExecGroups);
 
     if (platformConfig == null || !rule.useToolchainResolution()) {
@@ -68,10 +76,35 @@ public static UnloadedToolchainContextsInputs getUnloadedToolchainContextsInputs
             defaultExecConstraintLabels,
             /* debugTarget= */ platformConfig.debugToolchainResolution(rule.getLabel()),
             /* useAutoExecGroups= */ useAutoExecGroups,
-            ruleClass.getToolchainTypes(),
+            toolchainTypes,
             parentExecutionPlatformLabel));
   }
 
+  private static ImmutableSet<ToolchainTypeRequirement> updateToolchainTypesFromAttribute(
+      Rule rule, ImmutableSet<ToolchainTypeRequirement> toolchainTypes) {
+    NonconfigurableAttributeMapper attributes = NonconfigurableAttributeMapper.of(rule);
+    List<Label> targetToolchainTypes = attributes.get("toolchains", BuildType.LABEL_LIST);
+    if (targetToolchainTypes.isEmpty()) {
+      // No need to update.
+      return toolchainTypes;
+    }
+
+    ImmutableSet.Builder<ToolchainTypeRequirement> updatedToolchainTypes =
+        new ImmutableSet.Builder<>();
+    updatedToolchainTypes.addAll(toolchainTypes);
+    targetToolchainTypes.stream()
+        .map(
+            toolchainTypeLabel ->
+                ToolchainTypeRequirement.builder(toolchainTypeLabel)
+                    .mandatory(false)
+                    // Some of these may be template variables, not toolchain types, and so should
+                    // be silently ignored.
+                    .ignoreIfInvalid(true)
+                    .build())
+        .forEach(updatedToolchainTypes::add);
+    return updatedToolchainTypes.build();
+  }
+
   public static ToolchainContextKey createDefaultToolchainContextKey(
       BuildConfigurationKey configurationKey,
       ImmutableSet<Label> defaultExecConstraintLabels,
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/BUILD b/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/BUILD
index 8044205c998c7c..f4f48f073582f1 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/BUILD
@@ -21,6 +21,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
+        "//src/main/java/com/google/devtools/build/lib/analysis:view_creation_failed_exception",
         "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/util/io:out-err",
         "//src/main/java/com/google/devtools/build/lib/vfs",
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleIntegrationTest.java b/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleIntegrationTest.java
index 26294a532aef1b..94aae6c112a3e5 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleIntegrationTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/rules/genrule/GenRuleIntegrationTest.java
@@ -15,8 +15,10 @@
 package com.google.devtools.build.lib.bazel.rules.genrule;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
 
 import com.google.devtools.build.lib.analysis.CommandHelper;
+import com.google.devtools.build.lib.analysis.ViewCreationFailedException;
 import com.google.devtools.build.lib.analysis.configuredtargets.OutputFileConfiguredTarget;
 import com.google.devtools.build.lib.buildtool.util.BuildIntegrationTestCase;
 import com.google.devtools.build.lib.util.io.RecordingOutErr;
@@ -182,4 +184,171 @@ public void testRealisticIndirectExecution() throws Exception {
       }
     }
   }
+
+  @Test
+  public void testToolchains_fromTemplateVariableInfo() throws Exception {
+    // Write a rule that generates templated data.
+    write(
+        "test/template_rule.bzl",
+        """
+        def _impl(ctx):
+            vars = ctx.attr.vars
+            return [platform_common.TemplateVariableInfo(vars)]
+
+        template_rule = rule(
+            _impl,
+            attrs = {
+                "vars": attr.string_dict(),
+            },
+        )
+        """);
+
+    // Write a BUILD file that uses the data.
+    write(
+        "test/BUILD",
+        """
+        load(":template_rule.bzl", "template_rule")
+
+        template_rule(
+            name = "data",
+            vars = {
+                "foo": "bar",
+            },
+        )
+
+        genrule(
+            name = "g",
+            srcs = [],
+            outs = ["g.out"],
+            cmd = "echo foo: $(foo) > $@",
+            toolchains = [":data"],
+        )
+        """);
+
+    buildTarget("//test:g");
+    OutputFileConfiguredTarget output =
+        (OutputFileConfiguredTarget) getConfiguredTarget("//test:g.out");
+    assertThat(readContentAsLatin1String(output.getArtifact())).isEqualTo("foo: bar\n");
+  }
+
+  @Test
+  public void testToolchains_fromToolchain() throws Exception {
+    // Write a toolchain rule that generates templated data.
+    write(
+        "test/toolchain/template_toolchain.bzl",
+        """
+        def _impl(ctx):
+            vars = ctx.attr.vars
+            return [
+                platform_common.TemplateVariableInfo(vars),
+                platform_common.ToolchainInfo(data = "from " + ctx.label.name),
+            ]
+
+        template_toolchain = rule(
+            _impl,
+            attrs = {
+                "vars": attr.string_dict(),
+            },
+        )
+        """);
+    write(
+        "test/toolchain/BUILD",
+        """
+        load(":template_toolchain.bzl", "template_toolchain")
+
+        toolchain_type(name = "toolchain_type")
+
+        template_toolchain(
+            name = "data",
+            vars = {
+                "foo": "bar",
+            },
+        )
+
+        toolchain(
+            name = "data_impl",
+            toolchain_type = ":toolchain_type",
+            toolchain = ":data",
+        )
+        """);
+
+    // Write a BUILD file that uses the toolchain type.
+    write(
+        "test/BUILD",
+        """
+
+genrule(
+    name = "g",
+    srcs = [],
+    outs = ["g.out"],
+    cmd = "echo foo: $(foo) > $@",
+    toolchains = ["//test/toolchain:toolchain_type"],
+)
+""");
+
+    // Make sure the toolchain is available.
+    addOptions("--extra_toolchains=//test/toolchain:data_impl");
+    buildTarget("//test:g");
+    OutputFileConfiguredTarget output =
+        (OutputFileConfiguredTarget) getConfiguredTarget("//test:g.out");
+    assertThat(readContentAsLatin1String(output.getArtifact())).isEqualTo("foo: bar\n");
+  }
+
+  @Test
+  public void testToolchains_fromToolchain_noToolchainFound() throws Exception {
+    // Define a toolchain type.
+    write(
+        "test/toolchain/BUILD",
+        """
+        toolchain_type(name = "toolchain_type")
+        """);
+
+    // Write a BUILD file that uses the toolchain type.
+    write(
+        "test/BUILD",
+        """
+
+genrule(
+    name = "g",
+    srcs = [],
+    outs = ["g.out"],
+    cmd = "echo foo: $(foo) > $@",
+    toolchains = ["//test/toolchain:toolchain_type"],
+)
+""");
+
+    assertThrows(ViewCreationFailedException.class, () -> buildTarget("//test:g"));
+    assertContainsError("$(foo) not defined");
+  }
+
+  @Test
+  public void testToolchains_fromToolchain_noToolchainFound_unused() throws Exception {
+    // Define a toolchain type.
+    write(
+        "test/toolchain/BUILD",
+        """
+        toolchain_type(name = "toolchain_type")
+        """);
+
+    // Write a BUILD file that uses the toolchain type.
+    write(
+        "test/BUILD",
+        """
+
+genrule(
+    name = "g",
+    srcs = [],
+    outs = ["g.out"],
+    cmd = "echo no template variables used > $@",
+    toolchains = ["//test/toolchain:toolchain_type"],
+)
+""");
+
+    // Invoke the target, even though the toolchain isn't resolved.
+    buildTarget("//test:g");
+    OutputFileConfiguredTarget output =
+        (OutputFileConfiguredTarget) getConfiguredTarget("//test:g.out");
+    assertThat(readContentAsLatin1String(output.getArtifact()))
+        .isEqualTo("no template variables used\n");
+  }
 }