From fb16270b56409e2c5feb3ccae24fb781015d0b17 Mon Sep 17 00:00:00 2001
From: John Cater <jcater@google.com>
Date: Wed, 16 Feb 2022 10:09:34 -0800
Subject: [PATCH] Update AspectDefinition to use ToolchainTypeRequirement.

Part of Optional Toolchains (#14726).

Closes #14823.

PiperOrigin-RevId: 429077498
---
 .../build/lib/packages/AspectDefinition.java  | 34 ++++++++++++-------
 .../google/devtools/build/lib/packages/BUILD  |  1 +
 .../devtools/build/lib/rules/android/BUILD    |  1 +
 .../lib/rules/android/DexArchiveAspect.java   |  6 ++--
 .../devtools/build/lib/rules/cpp/proto/BUILD  |  1 +
 .../lib/rules/cpp/proto/CcProtoAspect.java    |  8 ++++-
 .../devtools/build/lib/rules/objc/BUILD       |  2 +-
 .../build/lib/rules/objc/J2ObjcAspect.java    |  8 ++++-
 .../build/lib/skyframe/AspectFunction.java    |  3 +-
 9 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java b/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
index 556b498eb849c9..1b1a4a0088fc84 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
@@ -19,11 +19,10 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
 import com.google.common.collect.Multimap;
 import com.google.devtools.build.lib.analysis.TransitiveInfoProvider;
 import com.google.devtools.build.lib.analysis.config.Fragment;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.config.transitions.ConfigurationTransition;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
@@ -40,6 +39,7 @@
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.BiPredicate;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -67,7 +67,7 @@ public final class AspectDefinition {
   private final RequiredProviders requiredProviders;
   private final RequiredProviders requiredProvidersForAspects;
   private final ImmutableMap<String, Attribute> attributes;
-  private final ImmutableSet<Label> requiredToolchains;
+  private final ImmutableSet<ToolchainTypeRequirement> toolchainTypes;
   private final boolean useToolchainTransition;
 
   /**
@@ -103,7 +103,7 @@ private AspectDefinition(
       RequiredProviders requiredProviders,
       RequiredProviders requiredProvidersForAspects,
       ImmutableMap<String, Attribute> attributes,
-      ImmutableSet<Label> requiredToolchains,
+      ImmutableSet<ToolchainTypeRequirement> toolchainTypes,
       boolean useToolchainTransition,
       @Nullable ImmutableSet<String> restrictToAttributes,
       @Nullable ConfigurationFragmentPolicy configurationFragmentPolicy,
@@ -117,7 +117,7 @@ private AspectDefinition(
     this.requiredProvidersForAspects = requiredProvidersForAspects;
 
     this.attributes = attributes;
-    this.requiredToolchains = requiredToolchains;
+    this.toolchainTypes = toolchainTypes;
     this.useToolchainTransition = useToolchainTransition;
     this.restrictToAttributes = restrictToAttributes;
     this.configurationFragmentPolicy = configurationFragmentPolicy;
@@ -141,8 +141,8 @@ public ImmutableMap<String, Attribute> getAttributes() {
   }
 
   /** Returns the required toolchains declared by this aspect. */
-  public ImmutableSet<Label> getRequiredToolchains() {
-    return requiredToolchains;
+  public ImmutableSet<ToolchainTypeRequirement> getToolchainTypes() {
+    return toolchainTypes;
   }
 
   public boolean useToolchainTransition() {
@@ -281,7 +281,7 @@ public static final class Builder {
         new ConfigurationFragmentPolicy.Builder();
     private boolean applyToFiles = false;
     private boolean applyToGeneratingRules = false;
-    private final List<Label> requiredToolchains = new ArrayList<>();
+    private final List<ToolchainTypeRequirement> toolchainTypes = new ArrayList<>();
     private boolean useToolchainTransition = false;
     private ImmutableSet<AspectClass> requiredAspectClasses = ImmutableSet.of();
 
@@ -587,15 +587,23 @@ public Builder applyToGeneratingRules(boolean applyToGeneratingRules) {
     }
 
     /** Adds the given toolchains as requirements for this aspect. */
-    public Builder addRequiredToolchains(Label... toolchainLabels) {
-      Iterables.addAll(this.requiredToolchains, Lists.newArrayList(toolchainLabels));
+    public Builder addToolchainType(ToolchainTypeRequirement... toolchainTypes) {
+      return this.addToolchainType(ImmutableList.copyOf(toolchainTypes));
+    }
+
+    /** Adds the given toolchains as requirements for this aspect. */
+    public Builder addToolchainType(List<ToolchainTypeRequirement> toolchainTypes) {
+      this.toolchainTypes.addAll(toolchainTypes);
       return this;
     }
 
     /** Adds the given toolchains as requirements for this aspect. */
+    // TODO(katre): Remove this once all callers use addToolchainType.
     public Builder addRequiredToolchains(List<Label> requiredToolchains) {
-      this.requiredToolchains.addAll(requiredToolchains);
-      return this;
+      return addToolchainType(
+          requiredToolchains.stream()
+              .map(label -> ToolchainTypeRequirement.create(label))
+              .collect(Collectors.toList()));
     }
 
     public Builder useToolchainTransition(boolean useToolchainTransition) {
@@ -622,7 +630,7 @@ public AspectDefinition build() {
           requiredProviders,
           requiredAspectProviders.build(),
           ImmutableMap.copyOf(attributes),
-          ImmutableSet.copyOf(requiredToolchains),
+          ImmutableSet.copyOf(toolchainTypes),
           useToolchainTransition,
           propagateAlongAttributes == null ? null : ImmutableSet.copyOf(propagateAlongAttributes),
           configurationFragmentPolicy.build(),
diff --git a/src/main/java/com/google/devtools/build/lib/packages/BUILD b/src/main/java/com/google/devtools/build/lib/packages/BUILD
index 15776790292c7c..0b496c032b0edf 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/packages/BUILD
@@ -59,6 +59,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:thread_state_receiver",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment_class_set",
+        "//src/main/java/com/google/devtools/build/lib/analysis:config/toolchain_type_requirement",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/configuration_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/no_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/patch_transition",
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/BUILD b/src/main/java/com/google/devtools/build/lib/rules/android/BUILD
index 626a113ca3b2c9..e413de358c911e 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/BUILD
@@ -53,6 +53,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment_options",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/invalid_configuration_exception",
+        "//src/main/java/com/google/devtools/build/lib/analysis:config/toolchain_type_requirement",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transition_factories",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/patch_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/split_transition",
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
index f758ea5ec7b8f6..d83b02264b150a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
@@ -45,6 +45,7 @@
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.analysis.config.ExecutionTransitionFactory;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.ToolchainInfo;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
@@ -182,8 +183,9 @@ public AspectDefinition getDefinition(AspectParameters params) {
                     // For proto_lang_toolchain rules, where we just want to get at their runtime
                     // deps.
                     ImmutableSet.of(ProtoLangToolchainProvider.class)))
-            .addRequiredToolchains(
-                Label.parseAbsoluteUnchecked(toolsRepository + sdkToolchainLabel))
+            .addToolchainType(
+                ToolchainTypeRequirement.create(
+                    Label.parseAbsoluteUnchecked(toolsRepository + sdkToolchainLabel)))
             // Parse labels since we don't have RuleDefinitionEnvironment.getLabel like in a rule
             .add(
                 attr(ASPECT_DESUGAR_PREREQ, LABEL)
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/BUILD b/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/BUILD
index 59bf9494d00e7c..77de621ae078d6 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/BUILD
@@ -21,6 +21,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//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:config/toolchain_type_requirement",
         "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
         "//src/main/java/com/google/devtools/build/lib/analysis:rule_definition_environment",
         "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java
index 7e467a0467343b..9fd1da34a0fc8b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/proto/CcProtoAspect.java
@@ -33,6 +33,7 @@
 import com.google.devtools.build.lib.analysis.TransitiveInfoCollection;
 import com.google.devtools.build.lib.analysis.TransitiveInfoProviderMap;
 import com.google.devtools.build.lib.analysis.TransitiveInfoProviderMapBuilder;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
@@ -124,7 +125,12 @@ public AspectDefinition getDefinition(AspectParameters aspectParameters) {
             .propagateAlongAttribute("deps")
             .requiresConfigurationFragments(CppConfiguration.class, ProtoConfiguration.class)
             .requireStarlarkProviders(ProtoInfo.PROVIDER.id())
-            .addRequiredToolchains(ccToolchainType)
+            .addToolchainType(
+                ToolchainTypeRequirement.builder(ccToolchainType)
+                    // TODO(https://github.com/bazelbuild/bazel/issues/14727): Evaluate whether this
+                    // can be optional.
+                    .mandatory(true)
+                    .build())
             .useToolchainTransition(true)
             .add(
                 attr(PROTO_TOOLCHAIN_ATTR, LABEL)
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/BUILD b/src/main/java/com/google/devtools/build/lib/rules/objc/BUILD
index 6669ae664fa531..0e56b7b68b0dd1 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/BUILD
@@ -19,7 +19,6 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/analysis:actions/custom_command_line",
         "//src/main/java/com/google/devtools/build/lib/analysis:actions/parameter_file_write_action",
-        "//src/main/java/com/google/devtools/build/lib/analysis:actions/symlink_action",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_configuration",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_options",
@@ -30,6 +29,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:config/execution_transition_factory",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment_options",
+        "//src/main/java/com/google/devtools/build/lib/analysis:config/toolchain_type_requirement",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/patch_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/split_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/starlark_exposed_rule_transition_factory",
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java b/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
index f0822145978480..2581e65c986f4a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/J2ObjcAspect.java
@@ -39,6 +39,7 @@
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.config.ConfigAwareAspectBuilder;
 import com.google.devtools.build.lib.analysis.config.ExecutionTransitionFactory;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.ToolchainInfo;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
@@ -152,7 +153,12 @@ public AspectDefinition getDefinition(AspectParameters aspectParameters) {
             J2ObjcConfiguration.class,
             ObjcConfiguration.class,
             ProtoConfiguration.class)
-        .addRequiredToolchains(ccToolchainType)
+        .addToolchainType(
+            ToolchainTypeRequirement.builder(ccToolchainType)
+                // TODO(https://github.com/bazelbuild/bazel/issues/14727): Evaluate whether this can
+                // be optional.
+                .mandatory(true)
+                .build())
         .useToolchainTransition(true)
         .add(
             attr("$grep_includes", LABEL)
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java
index e1cec1e4115e8c..1bb4214211b66c 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java
@@ -567,13 +567,12 @@ private static UnloadedToolchainContext getUnloadedToolchainContext(
       // Configuration can be null in the case of aspects applied to input files. In this case,
       // there are no chances of toolchains being used, so skip it.
       try {
-        ImmutableSet<Label> requiredToolchains = aspect.getDefinition().getRequiredToolchains();
         unloadedToolchainContext =
             (UnloadedToolchainContext)
                 env.getValueOrThrow(
                     ToolchainContextKey.key()
                         .configurationKey(configuration.getKey())
-                        .requiredToolchainTypeLabels(requiredToolchains)
+                        .toolchainTypes(aspect.getDefinition().getToolchainTypes())
                         .build(),
                     ToolchainException.class);
       } catch (ToolchainException e) {