Update AspectDefinition and StarlarkDefinedAspect to work with

ToolchainTypeRequirement.

Part of Optional Toolchains (#14726).

Closes #14946.

PiperOrigin-RevId: 432197702
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/BUILD b/src/main/java/com/google/devtools/build/lib/analysis/BUILD
index 55c84c3..2b44f57 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/analysis/BUILD
@@ -298,6 +298,7 @@
         ":config/per_label_options",
         ":config/run_under",
         ":config/starlark_defined_config_transition",
+        ":config/toolchain_type_requirement",
         ":config/transition_factories",
         ":config/transitions/composing_transition",
         ":config/transitions/composing_transition_factory",
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
index a16fc34..01721f2 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.analysis.starlark;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.devtools.build.lib.analysis.BaseRuleClasses.RUN_UNDER;
 import static com.google.devtools.build.lib.analysis.BaseRuleClasses.TEST_RUNNER_EXEC_GROUP;
 import static com.google.devtools.build.lib.analysis.BaseRuleClasses.TIMEOUT_DEFAULT;
@@ -43,6 +44,7 @@
 import com.google.devtools.build.lib.analysis.config.ExecutionTransitionFactory;
 import com.google.devtools.build.lib.analysis.config.HostTransition;
 import com.google.devtools.build.lib.analysis.config.StarlarkDefinedConfigTransition;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.config.transitions.PatchTransition;
 import com.google.devtools.build.lib.analysis.config.transitions.StarlarkExposedRuleTransitionFactory;
 import com.google.devtools.build.lib.analysis.config.transitions.TransitionFactory;
@@ -696,6 +698,7 @@
           "An aspect cannot simultaneously have required providers and apply to generating rules.");
     }
 
+    ImmutableList<Label> toolchainTypes = parseToolchains(toolchains, thread);
     return new StarlarkDefinedAspect(
         implementation,
         attrAspects.build(),
@@ -709,7 +712,9 @@
         ImmutableSet.copyOf(Sequence.cast(fragments, String.class, "fragments")),
         HostTransition.INSTANCE,
         ImmutableSet.copyOf(Sequence.cast(hostFragments, String.class, "host_fragments")),
-        parseToolchains(toolchains, thread),
+        toolchainTypes.stream()
+            .map(tt -> ToolchainTypeRequirement.create(tt))
+            .collect(toImmutableSet()),
         useToolchainTransition,
         applyToGeneratingRules);
   }
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 1b1a4a0..d568f30 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
@@ -31,15 +31,14 @@
 import com.google.devtools.build.lib.packages.Type.LabelClass;
 import com.google.devtools.build.lib.packages.Type.LabelVisitor;
 import java.io.Serializable;
-import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
-import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.BiPredicate;
-import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -115,7 +114,6 @@
     this.advertisedProviders = advertisedProviders;
     this.requiredProviders = requiredProviders;
     this.requiredProvidersForAspects = requiredProvidersForAspects;
-
     this.attributes = attributes;
     this.toolchainTypes = toolchainTypes;
     this.useToolchainTransition = useToolchainTransition;
@@ -281,7 +279,7 @@
         new ConfigurationFragmentPolicy.Builder();
     private boolean applyToFiles = false;
     private boolean applyToGeneratingRules = false;
-    private final List<ToolchainTypeRequirement> toolchainTypes = new ArrayList<>();
+    private final Set<ToolchainTypeRequirement> toolchainTypes = new HashSet<>();
     private boolean useToolchainTransition = false;
     private ImmutableSet<AspectClass> requiredAspectClasses = ImmutableSet.of();
 
@@ -587,25 +585,16 @@
     }
 
     /** Adds the given toolchains as requirements for this aspect. */
-    public Builder addToolchainType(ToolchainTypeRequirement... toolchainTypes) {
-      return this.addToolchainType(ImmutableList.copyOf(toolchainTypes));
+    public Builder addToolchainTypes(ToolchainTypeRequirement... toolchainTypes) {
+      return this.addToolchainTypes(ImmutableSet.copyOf(toolchainTypes));
     }
 
     /** Adds the given toolchains as requirements for this aspect. */
-    public Builder addToolchainType(List<ToolchainTypeRequirement> toolchainTypes) {
+    public Builder addToolchainTypes(Collection<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) {
-      return addToolchainType(
-          requiredToolchains.stream()
-              .map(label -> ToolchainTypeRequirement.create(label))
-              .collect(Collectors.toList()));
-    }
-
     public Builder useToolchainTransition(boolean useToolchainTransition) {
       this.useToolchainTransition = useToolchainTransition;
       return this;
diff --git a/src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java b/src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java
index bb5454e..e9d6fa5 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java
@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+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.events.EventHandler;
@@ -46,7 +47,7 @@
   private final ImmutableSet<String> fragments;
   private final ConfigurationTransition hostTransition;
   private final ImmutableSet<String> hostFragments;
-  private final ImmutableList<Label> requiredToolchains;
+  private final ImmutableSet<ToolchainTypeRequirement> toolchainTypes;
   private final boolean useToolchainTransition;
   private final boolean applyToGeneratingRules;
 
@@ -71,7 +72,7 @@
       // The host transition is in lib.analysis, so we can't reference it directly here.
       ConfigurationTransition hostTransition,
       ImmutableSet<String> hostFragments,
-      ImmutableList<Label> requiredToolchains,
+      ImmutableSet<ToolchainTypeRequirement> toolchainTypes,
       boolean useToolchainTransition,
       boolean applyToGeneratingRules) {
     this.implementation = implementation;
@@ -85,7 +86,7 @@
     this.fragments = fragments;
     this.hostTransition = hostTransition;
     this.hostFragments = hostFragments;
-    this.requiredToolchains = requiredToolchains;
+    this.toolchainTypes = toolchainTypes;
     this.useToolchainTransition = useToolchainTransition;
     this.applyToGeneratingRules = applyToGeneratingRules;
   }
@@ -177,7 +178,7 @@
     builder.advertiseProvider(advertisedStarlarkProviders.build());
     builder.requiresConfigurationFragmentsByStarlarkBuiltinName(fragments);
     builder.requiresConfigurationFragmentsByStarlarkBuiltinName(hostTransition, hostFragments);
-    builder.addRequiredToolchains(requiredToolchains);
+    builder.addToolchainTypes(toolchainTypes);
     builder.useToolchainTransition(useToolchainTransition);
     builder.applyToGeneratingRules(applyToGeneratingRules);
     ImmutableSet.Builder<AspectClass> requiredAspectsClasses = ImmutableSet.builder();
@@ -343,8 +344,8 @@
         getName(), name, value);
   }
 
-  public ImmutableList<Label> getRequiredToolchains() {
-    return requiredToolchains;
+  public ImmutableSet<ToolchainTypeRequirement> getToolchainTypes() {
+    return toolchainTypes;
   }
 
   public boolean useToolchainTransition() {
@@ -395,7 +396,7 @@
         && Objects.equals(fragments, that.fragments)
         && Objects.equals(hostTransition, that.hostTransition)
         && Objects.equals(hostFragments, that.hostFragments)
-        && Objects.equals(requiredToolchains, that.requiredToolchains)
+        && Objects.equals(toolchainTypes, that.toolchainTypes)
         && useToolchainTransition == that.useToolchainTransition
         && Objects.equals(aspectClass, that.aspectClass);
   }
@@ -414,7 +415,7 @@
         fragments,
         hostTransition,
         hostFragments,
-        requiredToolchains,
+        toolchainTypes,
         useToolchainTransition,
         aspectClass);
   }
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 d83b022..da9721b 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
@@ -183,7 +183,7 @@
                     // For proto_lang_toolchain rules, where we just want to get at their runtime
                     // deps.
                     ImmutableSet.of(ProtoLangToolchainProvider.class)))
-            .addToolchainType(
+            .addToolchainTypes(
                 ToolchainTypeRequirement.create(
                     Label.parseAbsoluteUnchecked(toolsRepository + sdkToolchainLabel)))
             // Parse labels since we don't have RuleDefinitionEnvironment.getLabel like in a rule
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 9fd1da3..78ac790 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
@@ -125,7 +125,7 @@
             .propagateAlongAttribute("deps")
             .requiresConfigurationFragments(CppConfiguration.class, ProtoConfiguration.class)
             .requireStarlarkProviders(ProtoInfo.PROVIDER.id())
-            .addToolchainType(
+            .addToolchainTypes(
                 ToolchainTypeRequirement.builder(ccToolchainType)
                     // TODO(https://github.com/bazelbuild/bazel/issues/14727): Evaluate whether this
                     // can be optional.
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 2581e65..1460eb3 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
@@ -153,7 +153,7 @@
             J2ObjcConfiguration.class,
             ObjcConfiguration.class,
             ProtoConfiguration.class)
-        .addToolchainType(
+        .addToolchainTypes(
             ToolchainTypeRequirement.builder(ccToolchainType)
                 // TODO(https://github.com/bazelbuild/bazel/issues/14727): Evaluate whether this can
                 // be optional.
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/BUILD b/src/test/java/com/google/devtools/build/lib/starlark/BUILD
index fd83aa7..f6c6625 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/starlark/BUILD
@@ -41,6 +41,7 @@
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/host_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/starlark_defined_config_transition",
+        "//src/main/java/com/google/devtools/build/lib/analysis:config/toolchain_type_requirement",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/no_transition",
         "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
         "//src/main/java/com/google/devtools/build/lib/analysis:file_provider",
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java
index 9a16de0..b264bd1 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java
@@ -25,6 +25,7 @@
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
 import com.google.devtools.build.lib.analysis.RuleContext;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.config.transitions.NoTransition;
 import com.google.devtools.build.lib.analysis.starlark.StarlarkAttrModule;
 import com.google.devtools.build.lib.analysis.starlark.StarlarkRuleClassFunctions.StarlarkRuleFunction;
@@ -659,8 +660,11 @@
     evalAndExport(
         ev, "def _impl(ctx): pass", "a1 = aspect(_impl, toolchains=['//test:my_toolchain_type'])");
     StarlarkDefinedAspect a = (StarlarkDefinedAspect) ev.lookup("a1");
-    assertThat(a.getRequiredToolchains())
-        .containsExactly(Label.parseAbsoluteUnchecked("//test:my_toolchain_type"));
+    // TODO(https://github.com/bazelbuild/bazel/issues/14726): Add tests of optional toolchains.
+    assertThat(a.getToolchainTypes())
+        .containsExactly(
+            ToolchainTypeRequirement.create(
+                Label.parseAbsoluteUnchecked("//test:my_toolchain_type")));
   }
 
   @Test
@@ -2300,6 +2304,7 @@
         ev,
         "def impl(ctx): return None",
         "r1 = rule(impl, toolchains=['//test:my_toolchain_type'])");
+    // TODO(https://github.com/bazelbuild/bazel/issues/14726): Add tests of optional toolchains.
     RuleClass c = ((StarlarkRuleFunction) ev.lookup("r1")).getRuleClass();
     assertThat(c).hasToolchainType("//test:my_toolchain_type");
   }