From 64843736578cb4a77eb0bb6b2d9473f3675818b4 Mon Sep 17 00:00:00 2001
From: John Cater <jcater@google.com>
Date: Mon, 28 Feb 2022 15:58:45 -0800
Subject: [PATCH] Update RuleClass to use ToolchainTypeRequirement.

A future change will update rules.

Part of Optional Toolchains (#14726).

Closes #14862.

PiperOrigin-RevId: 431546477
---
 .../build/lib/packages/RuleClass.java         | 47 ++++++++---
 .../google/devtools/build/lib/query2/BUILD    |  1 +
 .../cquery/ConfiguredTargetAccessor.java      |  6 +-
 .../skyframe/ConfiguredTargetFunction.java    |  9 +-
 .../lib/skyframe/ToolchainContextKey.java     | 15 ----
 .../devtools/build/lib/analysis/testing/BUILD |  2 +
 .../analysis/testing/RuleClassSubject.java    | 83 +++++++++++++++++++
 .../lib/packages/RuleClassBuilderTest.java    | 17 ++--
 .../build/lib/packages/RuleClassTest.java     | 17 ++--
 .../devtools/build/lib/query2/testutil/BUILD  |  1 +
 .../testutil/PostAnalysisQueryTest.java       |  6 +-
 .../StarlarkRuleClassFunctionsTest.java       |  6 +-
 .../google/devtools/build/lib/testutil/BUILD  |  1 +
 .../lib/testutil/TestRuleClassProvider.java   |  7 +-
 14 files changed, 163 insertions(+), 55 deletions(-)
 create mode 100644 src/test/java/com/google/devtools/build/lib/analysis/testing/RuleClassSubject.java

diff --git a/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java b/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
index f1e9f524f72750..a21517a2421b4a 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
@@ -35,6 +35,7 @@
 import com.google.common.collect.Maps;
 import com.google.common.collect.Ordering;
 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.analysis.config.transitions.PatchTransition;
 import com.google.devtools.build.lib.analysis.config.transitions.SplitTransition;
@@ -74,6 +75,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.Immutable;
 import net.starlark.java.eval.EvalException;
@@ -831,7 +833,7 @@ public enum ThirdPartyLicenseExistencePolicy {
     private ThirdPartyLicenseExistencePolicy thirdPartyLicenseExistencePolicy;
 
     private final Map<String, Attribute> attributes = new LinkedHashMap<>();
-    private final Set<Label> requiredToolchains = new HashSet<>();
+    private final Set<ToolchainTypeRequirement> toolchainTypes = new HashSet<>();
     private ToolchainResolutionMode useToolchainResolution = ToolchainResolutionMode.INHERIT;
     private ToolchainTransitionMode useToolchainTransition = ToolchainTransitionMode.INHERIT;
     private final Set<Label> executionPlatformConstraints = new HashSet<>();
@@ -868,7 +870,7 @@ public Builder(String name, RuleClassType type, boolean starlark, RuleClass... p
             .includeConfigurationFragmentsFrom(parent.getConfigurationFragmentPolicy());
         supportsConstraintChecking = parent.supportsConstraintChecking;
 
-        addRequiredToolchains(parent.getRequiredToolchains());
+        addToolchainTypes(parent.getToolchainTypes());
         this.useToolchainResolution =
             this.useToolchainResolution.apply(name, parent.useToolchainResolution);
         this.useToolchainTransition =
@@ -997,7 +999,7 @@ public RuleClass build(String name, String key) {
           configurationFragmentPolicy.build(),
           supportsConstraintChecking,
           thirdPartyLicenseExistencePolicy,
-          requiredToolchains,
+          toolchainTypes,
           useToolchainResolution,
           useToolchainTransition,
           executionPlatformConstraints,
@@ -1528,20 +1530,41 @@ public Builder setOptionReferenceFunctionForConfigSettingOnly(
     }
 
     /**
-     * Causes rules of this type to require the specified toolchains be available via toolchain
+     * Cause rules of this type to request the specified toolchains be available via toolchain
      * resolution when a target is configured.
      */
-    public Builder addRequiredToolchains(Iterable<Label> toolchainLabels) {
-      Iterables.addAll(this.requiredToolchains, toolchainLabels);
+    public Builder addToolchainTypes(Iterable<ToolchainTypeRequirement> toolchainTypes) {
+      Iterables.addAll(this.toolchainTypes, toolchainTypes);
       return this;
     }
 
+    /**
+     * Cause rules of this type to request the specified toolchains be available via toolchain
+     * resolution when a target is configured.
+     */
+    public Builder addToolchainTypes(ToolchainTypeRequirement... toolchainTypes) {
+      return addToolchainTypes(ImmutableList.copyOf(toolchainTypes));
+    }
+
     /**
      * Causes rules of this type to require the specified toolchains be available via toolchain
      * resolution when a target is configured.
      */
+    // TODO(katre): Remove this once all callers use addToolchainType.
     public Builder addRequiredToolchains(Label... toolchainLabels) {
-      return this.addRequiredToolchains(Lists.newArrayList(toolchainLabels));
+      return this.addRequiredToolchains(ImmutableList.copyOf(toolchainLabels));
+    }
+
+    /**
+     * Causes rules of this type to require the specified toolchains be available via toolchain
+     * resolution when a target is configured.
+     */
+    // TODO(katre): Remove this once all callers use addToolchainType.
+    public Builder addRequiredToolchains(Collection<Label> toolchainLabels) {
+      return this.addToolchainTypes(
+          toolchainLabels.stream()
+              .map(label -> ToolchainTypeRequirement.create(label))
+              .collect(Collectors.toList()));
     }
 
     /**
@@ -1756,7 +1779,7 @@ public Attribute.Builder<?> copy(String name) {
 
   private final ThirdPartyLicenseExistencePolicy thirdPartyLicenseExistencePolicy;
 
-  private final ImmutableSet<Label> requiredToolchains;
+  private final ImmutableSet<ToolchainTypeRequirement> toolchainTypes;
   private final ToolchainResolutionMode useToolchainResolution;
   private final ToolchainTransitionMode useToolchainTransition;
   private final ImmutableSet<Label> executionPlatformConstraints;
@@ -1813,7 +1836,7 @@ public Attribute.Builder<?> copy(String name) {
       ConfigurationFragmentPolicy configurationFragmentPolicy,
       boolean supportsConstraintChecking,
       ThirdPartyLicenseExistencePolicy thirdPartyLicenseExistencePolicy,
-      Set<Label> requiredToolchains,
+      Set<ToolchainTypeRequirement> toolchainTypes,
       ToolchainResolutionMode useToolchainResolution,
       ToolchainTransitionMode useToolchainTransition,
       Set<Label> executionPlatformConstraints,
@@ -1854,7 +1877,7 @@ public Attribute.Builder<?> copy(String name) {
     this.configurationFragmentPolicy = configurationFragmentPolicy;
     this.supportsConstraintChecking = supportsConstraintChecking;
     this.thirdPartyLicenseExistencePolicy = thirdPartyLicenseExistencePolicy;
-    this.requiredToolchains = ImmutableSet.copyOf(requiredToolchains);
+    this.toolchainTypes = ImmutableSet.copyOf(toolchainTypes);
     this.useToolchainResolution = useToolchainResolution;
     this.useToolchainTransition = useToolchainTransition;
     this.executionPlatformConstraints = ImmutableSet.copyOf(executionPlatformConstraints);
@@ -2768,8 +2791,8 @@ public boolean ignoreLicenses() {
     return ignoreLicenses;
   }
 
-  public ImmutableSet<Label> getRequiredToolchains() {
-    return requiredToolchains;
+  public ImmutableSet<ToolchainTypeRequirement> getToolchainTypes() {
+    return toolchainTypes;
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/query2/BUILD b/src/main/java/com/google/devtools/build/lib/query2/BUILD
index 52f8d619328e50..fb3ae2783cc881 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/query2/BUILD
@@ -37,6 +37,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:config/core_options",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment_options",
         "//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/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/null_transition",
diff --git a/src/main/java/com/google/devtools/build/lib/query2/cquery/ConfiguredTargetAccessor.java b/src/main/java/com/google/devtools/build/lib/query2/cquery/ConfiguredTargetAccessor.java
index 14fd45674e6d97..6d004daafb05ba 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/cquery/ConfiguredTargetAccessor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/cquery/ConfiguredTargetAccessor.java
@@ -28,6 +28,7 @@
 import com.google.devtools.build.lib.analysis.ToolchainContext;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationValue;
 import com.google.devtools.build.lib.analysis.config.ConfigMatchingProvider;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.configuredtargets.OutputFileConfiguredTarget;
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -215,7 +216,8 @@ private static ToolchainCollection<ToolchainContext> getToolchainContexts(
       return null;
     }
 
-    ImmutableSet<Label> requiredToolchains = rule.getRuleClassObject().getRequiredToolchains();
+    ImmutableSet<ToolchainTypeRequirement> toolchainTypes =
+        rule.getRuleClassObject().getToolchainTypes();
     // Collect local (target, rule) constraints for filtering out execution platforms.
     ImmutableSet<Label> execConstraintLabels =
         getExecutionPlatformConstraints(rule, config.getFragment(PlatformConfiguration.class));
@@ -249,7 +251,7 @@ private static ToolchainCollection<ToolchainContext> getToolchainContexts(
               walkableGraph.getValue(
                   ToolchainContextKey.key()
                       .configurationKey(configurationKey)
-                      .requiredToolchainTypeLabels(requiredToolchains)
+                      .toolchainTypes(toolchainTypes)
                       .execConstraintLabels(execConstraintLabels)
                       .debugTarget(debugTarget)
                       .build());
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java
index cf35827e001bfb..a647bf57ecc016 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java
@@ -52,6 +52,7 @@
 import com.google.devtools.build.lib.analysis.config.ConfigMatchingProvider;
 import com.google.devtools.build.lib.analysis.config.ConfigurationResolver;
 import com.google.devtools.build.lib.analysis.config.DependencyEvaluationException;
+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.configuredtargets.RuleConfiguredTarget;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
@@ -550,8 +551,8 @@ public static ComputedToolchainContexts computeUnloadedToolchainContexts(
     Rule rule = ((Rule) targetAndConfig.getTarget());
     BuildConfigurationValue configuration = targetAndConfig.getConfiguration();
 
-    ImmutableSet<Label> requiredDefaultToolchains =
-        rule.getRuleClassObject().getRequiredToolchains();
+    ImmutableSet<ToolchainTypeRequirement> toolchainTypes =
+        rule.getRuleClassObject().getToolchainTypes();
     // Collect local (target, rule) constraints for filtering out execution platforms.
     ImmutableSet<Label> defaultExecConstraintLabels =
         getExecutionPlatformConstraints(
@@ -560,7 +561,7 @@ public static ComputedToolchainContexts computeUnloadedToolchainContexts(
     // Create a merged version of the exec groups that handles exec group inheritance properly.
     ExecGroup defaultExecGroup =
         ExecGroup.builder()
-            .requiredToolchains(requiredDefaultToolchains)
+            .toolchainTypes(toolchainTypes)
             .execCompatibleWith(defaultExecConstraintLabels)
             .copyFrom(null)
             .build();
@@ -619,7 +620,7 @@ public static ComputedToolchainContexts computeUnloadedToolchainContexts(
     ToolchainContextKey.Builder toolchainContextKeyBuilder =
         ToolchainContextKey.key()
             .configurationKey(toolchainConfig)
-            .requiredToolchainTypeLabels(requiredDefaultToolchains)
+            .toolchainTypes(toolchainTypes)
             .execConstraintLabels(defaultExecConstraintLabels)
             .debugTarget(debugTarget);
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainContextKey.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainContextKey.java
index 9a0438512f964f..6e043f2b0425b2 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainContextKey.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainContextKey.java
@@ -13,7 +13,6 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
-import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
@@ -60,20 +59,6 @@ public SkyFunctionName functionName() {
   public interface Builder {
     Builder configurationKey(BuildConfigurationKey key);
 
-    // TODO(katre): Remove this once all callers use toolchainTypes.
-    default Builder requiredToolchainTypeLabels(Label... requiredToolchainTypeLabels) {
-      return this.requiredToolchainTypeLabels(ImmutableSet.copyOf(requiredToolchainTypeLabels));
-    }
-
-    // TODO(katre): Remove this once all callers use toolchainTypes.
-    default Builder requiredToolchainTypeLabels(ImmutableSet<Label> requiredToolchainTypeLabels) {
-      ImmutableSet<ToolchainTypeRequirement> toolchainTypeRequirements =
-          requiredToolchainTypeLabels.stream()
-              .map(label -> ToolchainTypeRequirement.create(label))
-              .collect(toImmutableSet());
-      return this.toolchainTypes(toolchainTypeRequirements);
-    }
-
     Builder toolchainTypes(ImmutableSet<ToolchainTypeRequirement> toolchainTypes);
 
     default Builder toolchainTypes(ToolchainTypeRequirement... toolchainTypes) {
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/testing/BUILD b/src/test/java/com/google/devtools/build/lib/analysis/testing/BUILD
index a34d123621e594..720ed205b5844c 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/testing/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/analysis/testing/BUILD
@@ -19,6 +19,7 @@ java_library(
         "ExecGroupCollectionSubject.java",
         "ExecGroupSubject.java",
         "ResolvedToolchainContextSubject.java",
+        "RuleClassSubject.java",
         "ToolchainCollectionSubject.java",
         "ToolchainContextSubject.java",
         "ToolchainInfoSubject.java",
@@ -33,6 +34,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis/platform",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/cmdline:cmdline-primitives",
+        "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/packages:exec_group",
         "//src/main/java/net/starlark/java/eval",
         "//third_party:guava",
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/testing/RuleClassSubject.java b/src/test/java/com/google/devtools/build/lib/analysis/testing/RuleClassSubject.java
new file mode 100644
index 00000000000000..837a7585fd2c98
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/analysis/testing/RuleClassSubject.java
@@ -0,0 +1,83 @@
+// Copyright 2022 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.analysis.testing;
+
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.truth.Truth.assertAbout;
+
+import com.google.common.base.Functions;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.truth.FailureMetadata;
+import com.google.common.truth.MapSubject;
+import com.google.common.truth.Subject;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.packages.RuleClass;
+import java.util.Map;
+
+/** A Truth {@link Subject} for {@link RuleClass}. */
+public class RuleClassSubject extends Subject {
+  // Static data.
+
+  /** Entry point for test assertions related to {@link RuleClass}. */
+  public static RuleClassSubject assertThat(RuleClass ruleClass) {
+    return assertAbout(RuleClassSubject::new).that(ruleClass);
+  }
+
+  /** Static method for getting the subject factory (for use with assertAbout()). */
+  public static Subject.Factory<RuleClassSubject, RuleClass> ruleClasses() {
+    return RuleClassSubject::new;
+  }
+
+  // Instance fields.
+
+  private final RuleClass actual;
+  private final Map<Label, ToolchainTypeRequirement> toolchainTypesMap;
+
+  protected RuleClassSubject(FailureMetadata failureMetadata, RuleClass subject) {
+    super(failureMetadata, subject);
+    this.actual = subject;
+    this.toolchainTypesMap = makeToolchainTypesMap(subject);
+  }
+
+  private static ImmutableMap<Label, ToolchainTypeRequirement> makeToolchainTypesMap(
+      RuleClass subject) {
+    return subject.getToolchainTypes().stream()
+        .collect(toImmutableMap(ToolchainTypeRequirement::toolchainType, Functions.identity()));
+  }
+
+  public MapSubject toolchainTypes() {
+    return check("toolchainTypes()").that(toolchainTypesMap);
+  }
+
+  public ToolchainTypeRequirementSubject toolchainType(String toolchainTypeLabel) {
+    return toolchainType(Label.parseAbsoluteUnchecked(toolchainTypeLabel));
+  }
+
+  public ToolchainTypeRequirementSubject toolchainType(Label toolchainType) {
+    return check("toolchainType(%s)", toolchainType)
+        .about(ToolchainTypeRequirementSubject.toolchainTypeRequirements())
+        .that(toolchainTypesMap.get(toolchainType));
+  }
+
+  public void hasToolchainType(String toolchainTypeLabel) {
+    toolchainType(toolchainTypeLabel).isNotNull();
+  }
+
+  public void hasToolchainType(Label toolchainType) {
+    toolchainType(toolchainType).isNotNull();
+  }
+
+  // TODO(blaze-team): Add more useful methods.
+}
diff --git a/src/test/java/com/google/devtools/build/lib/packages/RuleClassBuilderTest.java b/src/test/java/com/google/devtools/build/lib/packages/RuleClassBuilderTest.java
index de44d4cf098ef3..985c22349d80a2 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/RuleClassBuilderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/RuleClassBuilderTest.java
@@ -15,6 +15,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.analysis.testing.ExecGroupSubject.assertThat;
+import static com.google.devtools.build.lib.analysis.testing.RuleClassSubject.assertThat;
 import static com.google.devtools.build.lib.packages.Attribute.attr;
 import static com.google.devtools.build.lib.packages.Type.BOOLEAN;
 import static com.google.devtools.build.lib.packages.Type.INTEGER;
@@ -185,14 +186,15 @@ public void testRequiredToolchainsAreInherited() throws Exception {
     RuleClass parent =
         new RuleClass.Builder("$parent", RuleClassType.ABSTRACT, false)
             .add(attr("tags", STRING_LIST))
-            .addRequiredToolchains(ImmutableList.of(mockToolchainType))
+            .addToolchainTypes(ToolchainTypeRequirement.create(mockToolchainType))
             .build();
     RuleClass child =
         new RuleClass.Builder("child", RuleClassType.NORMAL, false, parent)
             .factory(DUMMY_CONFIGURED_TARGET_FACTORY)
             .add(attr("attr", STRING))
             .build();
-    assertThat(child.getRequiredToolchains()).contains(mockToolchainType);
+
+    assertThat(child).hasToolchainType(mockToolchainType);
   }
 
   @Test
@@ -233,18 +235,23 @@ public void testDuplicateExecGroupsThatInheritFromRuleIsOk() throws Exception {
             .factory(DUMMY_CONFIGURED_TARGET_FACTORY)
             .addExecGroups(ImmutableMap.of("blueberry", ExecGroup.copyFromDefault()))
             .add(attr("tags", STRING_LIST))
-            .addRequiredToolchains(Label.parseAbsoluteUnchecked("//some/toolchain"))
+            .addToolchainTypes(
+                ToolchainTypeRequirement.create(Label.parseAbsoluteUnchecked("//some/toolchain")))
             .build();
     RuleClass b =
         new RuleClass.Builder("ruleB", RuleClassType.NORMAL, false)
             .factory(DUMMY_CONFIGURED_TARGET_FACTORY)
             .addExecGroups(ImmutableMap.of("blueberry", ExecGroup.copyFromDefault()))
             .add(attr("tags", STRING_LIST))
-            .addRequiredToolchains(Label.parseAbsoluteUnchecked("//some/other/toolchain"))
+            .addToolchainTypes(
+                ToolchainTypeRequirement.create(
+                    Label.parseAbsoluteUnchecked("//some/other/toolchain")))
             .build();
     RuleClass c =
         new RuleClass.Builder("$ruleC", RuleClassType.ABSTRACT, false, a, b)
-            .addRequiredToolchains(Label.parseAbsoluteUnchecked("//actual/toolchain/we/care/about"))
+            .addToolchainTypes(
+                ToolchainTypeRequirement.create(
+                    Label.parseAbsoluteUnchecked("//actual/toolchain/we/care/about")))
             .build();
     assertThat(c.getExecGroups()).containsKey("blueberry");
     ExecGroup blueberry = c.getExecGroups().get("blueberry");
diff --git a/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java b/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java
index 68845ce2780a3a..25341794ad3db3 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java
@@ -16,6 +16,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.devtools.build.lib.analysis.testing.ExecGroupSubject.assertThat;
+import static com.google.devtools.build.lib.analysis.testing.RuleClassSubject.assertThat;
 import static com.google.devtools.build.lib.packages.Attribute.attr;
 import static com.google.devtools.build.lib.packages.BuildType.LABEL;
 import static com.google.devtools.build.lib.packages.BuildType.LABEL_LIST;
@@ -861,7 +862,7 @@ private static RuleClass newRuleClass(
             .build(),
         supportsConstraintChecking,
         ThirdPartyLicenseExistencePolicy.USER_CONTROLLABLE,
-        /*requiredToolchains=*/ ImmutableSet.of(),
+        /*toolchainTypes=*/ ImmutableSet.of(),
         /*useToolchainResolution=*/ ToolchainResolutionMode.ENABLED,
         /*useToolchainTransition=*/ ToolchainTransitionMode.ENABLED,
         /* executionPlatformConstraints= */ ImmutableSet.of(),
@@ -991,22 +992,20 @@ private void expectError(RuleClassType type, String name) {
   }
 
   @Test
-  public void testRequiredToolchains() throws Exception {
+  public void testToolchainTypes() throws Exception {
     RuleClass.Builder ruleClassBuilder =
         new RuleClass.Builder("ruleClass", RuleClassType.NORMAL, false)
             .factory(DUMMY_CONFIGURED_TARGET_FACTORY)
             .add(attr("tags", STRING_LIST));
 
-    ruleClassBuilder.addRequiredToolchains(
-        Label.parseAbsolute("//toolchain:tc1", ImmutableMap.of()),
-        Label.parseAbsolute("//toolchain:tc2", ImmutableMap.of()));
+    ruleClassBuilder.addToolchainTypes(
+        ToolchainTypeRequirement.create(Label.parseAbsolute("//toolchain:tc1", ImmutableMap.of())),
+        ToolchainTypeRequirement.create(Label.parseAbsolute("//toolchain:tc2", ImmutableMap.of())));
 
     RuleClass ruleClass = ruleClassBuilder.build();
 
-    assertThat(ruleClass.getRequiredToolchains())
-        .containsExactly(
-            Label.parseAbsolute("//toolchain:tc1", ImmutableMap.of()),
-            Label.parseAbsolute("//toolchain:tc2", ImmutableMap.of()));
+    assertThat(ruleClass).hasToolchainType("//toolchain:tc1");
+    assertThat(ruleClass).hasToolchainType("//toolchain:tc2");
   }
 
   @Test
diff --git a/src/test/java/com/google/devtools/build/lib/query2/testutil/BUILD b/src/test/java/com/google/devtools/build/lib/query2/testutil/BUILD
index 9284a3ce1ea3ee..b2fed83bc9ccac 100644
--- a/src/test/java/com/google/devtools/build/lib/query2/testutil/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/query2/testutil/BUILD
@@ -22,6 +22,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_configuration",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_options",
         "//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:server_directories",
         "//src/main/java/com/google/devtools/build/lib/clock",
diff --git a/src/test/java/com/google/devtools/build/lib/query2/testutil/PostAnalysisQueryTest.java b/src/test/java/com/google/devtools/build/lib/query2/testutil/PostAnalysisQueryTest.java
index ee73acdd4a9e24..0c1c4c01060d72 100644
--- a/src/test/java/com/google/devtools/build/lib/query2/testutil/PostAnalysisQueryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/query2/testutil/PostAnalysisQueryTest.java
@@ -26,6 +26,7 @@
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.BuildOptionsView;
 import com.google.devtools.build.lib.analysis.config.FragmentOptions;
+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.util.DummyTestFragment.DummyTestOptions;
 import com.google.devtools.build.lib.analysis.util.MockRule;
@@ -225,8 +226,9 @@ public void testNoImplicitDeps_toolchains() throws Exception {
             MockRule.define(
                 "implicit_toolchain_deps_rule",
                 (builder, env) ->
-                    builder.addRequiredToolchains(
-                        Label.parseAbsoluteUnchecked("//test:toolchain_type")));
+                    builder.addToolchainTypes(
+                        ToolchainTypeRequirement.create(
+                            Label.parseAbsoluteUnchecked("//test:toolchain_type"))));
     helper.useRuleClassProvider(setRuleClassProviders(ruleWithImplicitDeps).build());
 
     writeFile(
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 89e4f667f3585a..b43e359c62530c 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
@@ -16,6 +16,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.analysis.testing.ExecGroupSubject.assertThat;
+import static com.google.devtools.build.lib.analysis.testing.RuleClassSubject.assertThat;
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.base.Joiner;
@@ -2299,8 +2300,7 @@ public void testRuleAddToolchain() throws Exception {
         "def impl(ctx): return None",
         "r1 = rule(impl, toolchains=['//test:my_toolchain_type'])");
     RuleClass c = ((StarlarkRuleFunction) ev.lookup("r1")).getRuleClass();
-    assertThat(c.getRequiredToolchains())
-        .containsExactly(Label.parseAbsoluteUnchecked("//test:my_toolchain_type"));
+    assertThat(c).hasToolchainType("//test:my_toolchain_type");
   }
 
   @Test
@@ -2337,7 +2337,7 @@ public void testRuleAddExecGroup() throws Exception {
         "  },",
         ")");
     RuleClass plum = ((StarlarkRuleFunction) ev.lookup("plum")).getRuleClass();
-    assertThat(plum.getRequiredToolchains()).isEmpty();
+    assertThat(plum.getToolchainTypes()).isEmpty();
     // TODO(https://github.com/bazelbuild/bazel/issues/14726): Add tests of optional toolchains.
     assertThat(plum.getExecGroups().get("group")).hasToolchainType("//test:my_toolchain_type");
     assertThat(plum.getExecGroups().get("group"))
diff --git a/src/test/java/com/google/devtools/build/lib/testutil/BUILD b/src/test/java/com/google/devtools/build/lib/testutil/BUILD
index 860e02238a38a5..1822c419883a85 100644
--- a/src/test/java/com/google/devtools/build/lib/testutil/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/testutil/BUILD
@@ -49,6 +49,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:common_prerequisite_validator",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/core_options",
+        "//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:platform_configuration",
         "//src/main/java/com/google/devtools/build/lib/analysis:rule_definition_environment",
diff --git a/src/test/java/com/google/devtools/build/lib/testutil/TestRuleClassProvider.java b/src/test/java/com/google/devtools/build/lib/testutil/TestRuleClassProvider.java
index b0d8b4dfd9b5d6..a475a1cb47c0cd 100644
--- a/src/test/java/com/google/devtools/build/lib/testutil/TestRuleClassProvider.java
+++ b/src/test/java/com/google/devtools/build/lib/testutil/TestRuleClassProvider.java
@@ -20,7 +20,6 @@
 import static com.google.devtools.build.lib.packages.Type.INTEGER;
 import static com.google.devtools.build.lib.packages.Type.STRING_LIST;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
 import com.google.devtools.build.lib.analysis.BaseRuleClasses;
@@ -36,6 +35,7 @@
 import com.google.devtools.build.lib.analysis.RunfilesProvider;
 import com.google.devtools.build.lib.analysis.TemplateVariableInfo;
 import com.google.devtools.build.lib.analysis.config.CoreOptions;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
@@ -216,8 +216,9 @@ public static class MockToolchainRule implements RuleDefinition {
     public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env) {
       return builder
           .requiresConfigurationFragments(PlatformConfiguration.class)
-          .addRequiredToolchains(
-              ImmutableList.of(Label.parseAbsoluteUnchecked("//toolchain:test_toolchain")))
+          .addToolchainTypes(
+              ToolchainTypeRequirement.create(
+                  Label.parseAbsoluteUnchecked("//toolchain:test_toolchain")))
           .build();
     }