From e9c1372a249a061bfd56bc7a7109b40702352100 Mon Sep 17 00:00:00 2001
From: John Cater <jcater@google.com>
Date: Thu, 3 Mar 2022 14:32:11 -0500
Subject: [PATCH] Use ToolchainTypeRequirement in ToolchainContext and related
 classes.

Part of Optional Toolchains (#14726).
---
 .../google/devtools/build/lib/analysis/BUILD  |  3 +
 .../analysis/ResolvedToolchainContext.java    |  2 +-
 .../build/lib/analysis/ToolchainContext.java  |  4 +-
 .../config/ToolchainTypeRequirement.java      |  3 +
 .../starlark/StarlarkToolchainContext.java    |  5 +-
 .../google/devtools/build/lib/skyframe/BUILD  |  1 +
 .../RegisteredToolchainsCycleReporter.java    |  2 +-
 .../SingleToolchainResolutionFunction.java    | 48 ++++++-----
 .../SingleToolchainResolutionValue.java       | 19 ++--
 .../skyframe/ToolchainResolutionFunction.java | 86 +++++++++++++------
 .../UnloadedToolchainContextImpl.java         |  5 +-
 .../ResolvedToolchainContextTest.java         | 10 ++-
 .../testing/ToolchainContextSubject.java      | 41 ++++++---
 ...SingleToolchainResolutionFunctionTest.java |  9 +-
 .../ToolchainResolutionFunctionTest.java      |  4 +-
 .../skyframe/ToolchainsForTargetsTest.java    |  2 +-
 .../lib/starlark/StarlarkRuleContextTest.java |  4 +-
 17 files changed, 152 insertions(+), 96 deletions(-)

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 449d7067488075..2a3f68112033bc 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/analysis/BUILD
@@ -985,6 +985,7 @@ java_library(
     name = "resolved_toolchain_context",
     srcs = ["ResolvedToolchainContext.java"],
     deps = [
+        ":config/toolchain_type_requirement",
         ":template_variable_info",
         ":toolchain_context",
         "//src/main/java/com/google/devtools/build/lib/analysis/platform",
@@ -1091,6 +1092,7 @@ java_library(
     name = "toolchain_context",
     srcs = ["ToolchainContext.java"],
     deps = [
+        ":config/toolchain_type_requirement",
         "//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/skyframe:toolchain_context_key",
@@ -2296,6 +2298,7 @@ java_library(
     name = "starlark/starlark_toolchain_context",
     srcs = ["starlark/StarlarkToolchainContext.java"],
     deps = [
+        ":config/toolchain_type_requirement",
         ":resolved_toolchain_context",
         "//src/main/java/com/google/devtools/build/lib/analysis/platform",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContext.java b/src/main/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContext.java
index 0e004046af1007..849cb8e3fa1ffa 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContext.java
@@ -86,7 +86,7 @@ public static ResolvedToolchainContext load(
         unloadedToolchainContext.key(),
         unloadedToolchainContext.executionPlatform(),
         unloadedToolchainContext.targetPlatform(),
-        unloadedToolchainContext.requiredToolchainTypes(),
+        unloadedToolchainContext.toolchainTypes(),
         unloadedToolchainContext.resolvedToolchainLabels(),
         // this:
         targetDescription,
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ToolchainContext.java b/src/main/java/com/google/devtools/build/lib/analysis/ToolchainContext.java
index a14c8e1f1408fb..f1778e33a24158 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/ToolchainContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ToolchainContext.java
@@ -14,8 +14,8 @@
 package com.google.devtools.build.lib.analysis;
 
 import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
-import com.google.devtools.build.lib.analysis.platform.ToolchainTypeInfo;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.skyframe.ToolchainContextKey;
 
@@ -31,7 +31,7 @@ public interface ToolchainContext {
   PlatformInfo targetPlatform();
 
   /** Returns the toolchain types that were requested. */
-  ImmutableSet<ToolchainTypeInfo> requiredToolchainTypes();
+  ImmutableSet<ToolchainTypeRequirement> toolchainTypes();
 
   /** Returns the labels of the specific toolchains being used. */
   ImmutableSet<Label> resolvedToolchainLabels();
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/ToolchainTypeRequirement.java b/src/main/java/com/google/devtools/build/lib/analysis/config/ToolchainTypeRequirement.java
index 6b22dd6505466e..4111ee0eaa6779 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/config/ToolchainTypeRequirement.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/config/ToolchainTypeRequirement.java
@@ -62,6 +62,9 @@ public static ToolchainTypeRequirement strictest(
    */
   public abstract boolean mandatory();
 
+  /** Returns a new Builder to copy this ToolchainTypeRequirement. */
+  public abstract Builder toBuilder();
+
   /** A builder for a new {@link ToolchainTypeRequirement}. */
   @AutoValue.Builder
   public interface Builder {
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkToolchainContext.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkToolchainContext.java
index c19438a74be426..1e33cff784eb01 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkToolchainContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkToolchainContext.java
@@ -17,6 +17,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.devtools.build.lib.analysis.ResolvedToolchainContext;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.ToolchainInfo;
 import com.google.devtools.build.lib.analysis.platform.ToolchainTypeInfo;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -110,8 +111,8 @@ public ToolchainInfo getIndex(
           "In %s, toolchain type %s was requested but only types [%s] are configured",
           toolchainContext().targetDescription(),
           toolchainTypeLabel,
-          toolchainContext().requiredToolchainTypes().stream()
-              .map(ToolchainTypeInfo::typeLabel)
+          toolchainContext().toolchainTypes().stream()
+              .map(ToolchainTypeRequirement::toolchainType)
               .map(Label::toString)
               .collect(joining(", ")));
     }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BUILD b/src/main/java/com/google/devtools/build/lib/skyframe/BUILD
index 4b6bb085982b27..3ccedd4dc21815 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/BUILD
@@ -2726,6 +2726,7 @@ java_library(
     deps = [
         ":toolchain_context_key",
         ":unloaded_toolchain_context",
+        "//src/main/java/com/google/devtools/build/lib/analysis:config/toolchain_type_requirement",
         "//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/skyframe:skyframe-objects",
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredToolchainsCycleReporter.java b/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredToolchainsCycleReporter.java
index db31d4c77a1395..f75e6a06613a3f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredToolchainsCycleReporter.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/RegisteredToolchainsCycleReporter.java
@@ -85,7 +85,7 @@ public boolean maybeReportCycle(
           }
           if (input.argument() instanceof SingleToolchainResolutionKey) {
             Label toolchainType =
-                ((SingleToolchainResolutionKey) input.argument()).toolchainTypeLabel();
+                ((SingleToolchainResolutionKey) input.argument()).toolchainType().toolchainType();
             return String.format("toolchain type %s", toolchainType);
           }
           if (input.argument() instanceof ToolchainContextKey) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunction.java
index 5fb763688a8fdc..0b1b15d770b06d 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunction.java
@@ -23,6 +23,7 @@
 import com.google.devtools.build.lib.analysis.PlatformConfiguration;
 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.platform.ConstraintCollection;
 import com.google.devtools.build.lib.analysis.platform.ConstraintSettingInfo;
 import com.google.devtools.build.lib.analysis.platform.DeclaredToolchainInfo;
@@ -84,11 +85,11 @@ public SkyValue compute(SkyKey skyKey, Environment env)
         key.debugTarget()
             || configuration
                 .getFragment(PlatformConfiguration.class)
-                .debugToolchainResolution(key.toolchainTypeLabel());
+                .debugToolchainResolution(key.toolchainType().toolchainType());
 
     // Find the right one.
     return resolveConstraints(
-        key.toolchainTypeLabel(),
+        key.toolchainType(),
         key.availableExecutionPlatformKeys(),
         key.targetPlatformKey(),
         toolchains.registeredToolchains(),
@@ -103,7 +104,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
    */
   @Nullable
   private static SingleToolchainResolutionValue resolveConstraints(
-      Label toolchainTypeLabel,
+      ToolchainTypeRequirement toolchainType,
       List<ConfiguredTargetKey> availableExecutionPlatformKeys,
       ConfiguredTargetKey targetPlatformKey,
       ImmutableList<DeclaredToolchainInfo> toolchains,
@@ -135,13 +136,15 @@ private static SingleToolchainResolutionValue resolveConstraints(
     // check whether a platform has already been seen during processing.
     Set<ConfiguredTargetKey> platformKeysSeen = new HashSet<>();
     ImmutableMap.Builder<ConfiguredTargetKey, Label> builder = ImmutableMap.builder();
-    ToolchainTypeInfo toolchainType = null;
+    ToolchainTypeInfo toolchainTypeInfo = null;
 
     // Pre-filter for the correct toolchain type. This simplifies the loop and makes debugging
     // toolchain resolution much, much easier.
     ImmutableList<DeclaredToolchainInfo> filteredToolchains =
         toolchains.stream()
-            .filter(toolchain -> toolchain.toolchainType().typeLabel().equals(toolchainTypeLabel))
+            .filter(
+                toolchain ->
+                    toolchain.toolchainType().typeLabel().equals(toolchainType.toolchainType()))
             .collect(toImmutableList());
 
     for (DeclaredToolchainInfo toolchain : filteredToolchains) {
@@ -155,7 +158,7 @@ private static SingleToolchainResolutionValue resolveConstraints(
         debugMessage(
             eventHandler,
             "    Type %s: %s platform %s: Rejected toolchain %s; mismatching config settings: %s",
-            toolchainTypeLabel,
+            toolchainType.toolchainType(),
             "target",
             targetPlatform.label(),
             toolchain.toolchainLabel(),
@@ -169,7 +172,7 @@ private static SingleToolchainResolutionValue resolveConstraints(
           toolchain.targetConstraints(),
           "target",
           targetPlatform,
-          toolchainTypeLabel,
+          toolchainType,
           toolchain.toolchainLabel())) {
         continue;
       }
@@ -187,7 +190,7 @@ private static SingleToolchainResolutionValue resolveConstraints(
             toolchain.execConstraints(),
             "execution",
             executionPlatform,
-            toolchainTypeLabel,
+            toolchainType,
             toolchain.toolchainLabel())) {
           continue;
         }
@@ -195,11 +198,11 @@ private static SingleToolchainResolutionValue resolveConstraints(
         debugMessage(
             eventHandler,
             "  Type %s: target platform %s: execution %s: Selected toolchain %s",
-            toolchainTypeLabel,
+            toolchainType.toolchainType(),
             targetPlatform.label(),
             executionPlatformKey.getLabel(),
             toolchain.toolchainLabel());
-        toolchainType = toolchain.toolchainType();
+        toolchainTypeInfo = toolchain.toolchainType();
         builder.put(executionPlatformKey, toolchain.toolchainLabel());
         platformKeysSeen.add(executionPlatformKey);
       }
@@ -210,13 +213,12 @@ private static SingleToolchainResolutionValue resolveConstraints(
       debugMessage(
           eventHandler,
           "  Type %s: target platform %s: No toolchains found.",
-          toolchainTypeLabel,
+          toolchainType.toolchainType(),
           targetPlatform.label());
-      throw new ToolchainResolutionFunctionException(
-          new NoToolchainFoundException(toolchainTypeLabel));
+      throw new ToolchainResolutionFunctionException(new NoToolchainFoundException(toolchainType));
     }
 
-    return SingleToolchainResolutionValue.create(toolchainType, resolvedToolchainLabels);
+    return SingleToolchainResolutionValue.create(toolchainTypeInfo, resolvedToolchainLabels);
   }
 
   /**
@@ -242,7 +244,7 @@ private static boolean checkConstraints(
       ConstraintCollection toolchainConstraints,
       String platformType,
       PlatformInfo platform,
-      Label toolchainTypeLabel,
+      ToolchainTypeRequirement toolchainType,
       Label toolchainLabel) {
 
     // Check every constraint_setting in either the toolchain or the platform.
@@ -283,7 +285,7 @@ private static boolean checkConstraints(
       debugMessage(
           eventHandler,
           "    Type %s: %s platform %s: Rejected toolchain %s%s%s",
-          toolchainTypeLabel,
+          toolchainType.toolchainType(),
           platformType,
           platform.label(),
           toolchainLabel,
@@ -296,15 +298,17 @@ private static boolean checkConstraints(
 
   /** Used to indicate that a toolchain was not found for the current request. */
   public static final class NoToolchainFoundException extends NoSuchThingException {
-    private final Label missingToolchainTypeLabel;
+    private final ToolchainTypeRequirement missingToolchainType;
 
-    public NoToolchainFoundException(Label missingToolchainTypeLabel) {
-      super(String.format("no matching toolchain found for %s", missingToolchainTypeLabel));
-      this.missingToolchainTypeLabel = missingToolchainTypeLabel;
+    public NoToolchainFoundException(ToolchainTypeRequirement missingToolchainType) {
+      super(
+          String.format(
+              "no matching toolchain found for %s", missingToolchainType.toolchainType()));
+      this.missingToolchainType = missingToolchainType;
     }
 
-    public Label missingToolchainTypeLabel() {
-      return missingToolchainTypeLabel;
+    public ToolchainTypeRequirement missingToolchainType() {
+      return missingToolchainType;
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionValue.java
index ac6ab1a63e62ad..a19c27e81cf49e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionValue.java
@@ -18,6 +18,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.ToolchainTypeInfo;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.skyframe.SkyFunctionName;
@@ -35,26 +36,22 @@ public abstract class SingleToolchainResolutionValue implements SkyValue {
   // A key representing the input data.
   public static SingleToolchainResolutionKey key(
       BuildConfigurationKey configurationKey,
-      Label toolchainTypeLabel,
+      ToolchainTypeRequirement toolchainType,
       ConfiguredTargetKey targetPlatformKey,
       List<ConfiguredTargetKey> availableExecutionPlatformKeys) {
     return key(
-        configurationKey,
-        toolchainTypeLabel,
-        targetPlatformKey,
-        availableExecutionPlatformKeys,
-        false);
+        configurationKey, toolchainType, targetPlatformKey, availableExecutionPlatformKeys, false);
   }
 
   public static SingleToolchainResolutionKey key(
       BuildConfigurationKey configurationKey,
-      Label toolchainTypeLabel,
+      ToolchainTypeRequirement toolchainType,
       ConfiguredTargetKey targetPlatformKey,
       List<ConfiguredTargetKey> availableExecutionPlatformKeys,
       boolean debugTarget) {
     return SingleToolchainResolutionKey.create(
         configurationKey,
-        toolchainTypeLabel,
+        toolchainType,
         targetPlatformKey,
         availableExecutionPlatformKeys,
         debugTarget);
@@ -71,7 +68,7 @@ public SkyFunctionName functionName() {
 
     abstract BuildConfigurationKey configurationKey();
 
-    public abstract Label toolchainTypeLabel();
+    public abstract ToolchainTypeRequirement toolchainType();
 
     abstract ConfiguredTargetKey targetPlatformKey();
 
@@ -81,13 +78,13 @@ public SkyFunctionName functionName() {
 
     static SingleToolchainResolutionKey create(
         BuildConfigurationKey configurationKey,
-        Label toolchainTypeLabel,
+        ToolchainTypeRequirement toolchainType,
         ConfiguredTargetKey targetPlatformKey,
         List<ConfiguredTargetKey> availableExecutionPlatformKeys,
         boolean debugTarget) {
       return new AutoValue_SingleToolchainResolutionValue_SingleToolchainResolutionKey(
           configurationKey,
-          toolchainTypeLabel,
+          toolchainType,
           targetPlatformKey,
           ImmutableList.copyOf(availableExecutionPlatformKeys),
           debugTarget);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java
index ad60ed47f00e4c..ab1edb29567567 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import static com.google.common.collect.ImmutableList.builder;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.util.stream.Collectors.joining;
@@ -90,14 +91,16 @@ public UnloadedToolchainContext compute(SkyKey skyKey, Environment env)
 
       // Load the configured target for the toolchain types to ensure that they are valid and
       // resolve aliases.
-      ImmutableMap<Label, ToolchainTypeInfo> resolvedToolchainTypes =
-          loadToolchainTypes(env, configuration, key.toolchainTypes());
-      builder.setRequestedLabelToToolchainType(resolvedToolchainTypes);
-      // TODO(katre): Need to correct for possible aliases, map to mandatory/optional.
-      ImmutableSet<Label> resolvedToolchainTypeLabels =
-          resolvedToolchainTypes.values().stream()
-              .map(ToolchainTypeInfo::typeLabel)
-              .collect(toImmutableSet());
+      ImmutableMap<Label, ToolchainTypeInfo> resolvedToolchainTypeInfos =
+          loadToolchainTypes(
+              env,
+              configuration,
+              key.toolchainTypes().stream()
+                  .map(ToolchainTypeRequirement::toolchainType)
+                  .collect(toImmutableSet()));
+      builder.setRequestedLabelToToolchainType(resolvedToolchainTypeInfos);
+      ImmutableSet<ToolchainTypeRequirement> resolvedToolchainTypes =
+          loadToolchainTypeRequirements(resolvedToolchainTypeInfos, key.toolchainTypes());
 
       // Create keys for all platforms that will be used, and validate them early.
       PlatformKeys platformKeys =
@@ -116,7 +119,7 @@ public UnloadedToolchainContext compute(SkyKey skyKey, Environment env)
       determineToolchainImplementations(
           env,
           key.configurationKey(),
-          resolvedToolchainTypeLabels,
+          resolvedToolchainTypes,
           key.forceExecutionPlatform().map(platformKeys::find),
           builder,
           platformKeys,
@@ -149,15 +152,17 @@ public UnloadedToolchainContext compute(SkyKey skyKey, Environment env)
     }
   }
 
-  /** Returns a map from the requested toolchain type to the {@link ToolchainTypeInfo} provider. */
+  /**
+   * Returns a map from the requested toolchain type Label (after any alias chains) to the {@link
+   * ToolchainTypeInfo} provider.
+   */
   private static ImmutableMap<Label, ToolchainTypeInfo> loadToolchainTypes(
       Environment environment,
       BuildConfigurationValue configuration,
-      ImmutableSet<ToolchainTypeRequirement> toolchainTypeRequirements)
+      ImmutableSet<Label> toolchainTypeLabels)
       throws InvalidToolchainTypeException, InterruptedException, ValueMissingException {
     ImmutableSet<ConfiguredTargetKey> toolchainTypeKeys =
-        toolchainTypeRequirements.stream()
-            .map(ToolchainTypeRequirement::toolchainType)
+        toolchainTypeLabels.stream()
             .map(
                 label ->
                     ConfiguredTargetKey.builder()
@@ -174,6 +179,32 @@ private static ImmutableMap<Label, ToolchainTypeInfo> loadToolchainTypes(
     return resolvedToolchainTypes;
   }
 
+  /**
+   * Returns a map from the actual post-alias Label to the ToolchainTypeRequirement for that type.
+   */
+  private ImmutableSet<ToolchainTypeRequirement> loadToolchainTypeRequirements(
+      ImmutableMap<Label, ToolchainTypeInfo> resolvedToolchainTypeInfos,
+      ImmutableSet<ToolchainTypeRequirement> toolchainTypes) {
+    ImmutableSet.Builder<ToolchainTypeRequirement> resolved = new ImmutableSet.Builder<>();
+
+    for (ToolchainTypeRequirement toolchainTypeRequirement : toolchainTypes) {
+      // Find the actual Label.
+      Label toolchainTypeLabel = toolchainTypeRequirement.toolchainType();
+      if (resolvedToolchainTypeInfos.containsKey(toolchainTypeLabel)) {
+        toolchainTypeLabel = resolvedToolchainTypeInfos.get(toolchainTypeLabel).typeLabel();
+      }
+
+      // If the labels don't match, re-build the TTR.
+      if (!toolchainTypeLabel.equals(toolchainTypeRequirement.toolchainType())) {
+        toolchainTypeRequirement =
+            toolchainTypeRequirement.toBuilder().toolchainType(toolchainTypeLabel).build();
+      }
+
+      resolved.add(toolchainTypeRequirement);
+    }
+    return resolved.build();
+  }
+
   @AutoValue
   abstract static class PlatformKeys {
     abstract ConfiguredTargetKey hostPlatformKey();
@@ -357,7 +388,7 @@ private static boolean filterPlatform(
   private static void determineToolchainImplementations(
       Environment environment,
       BuildConfigurationKey configurationKey,
-      ImmutableSet<Label> requiredToolchainTypeLabels,
+      ImmutableSet<ToolchainTypeRequirement> toolchainTypes,
       Optional<ConfiguredTargetKey> forcedExecutionPlatform,
       UnloadedToolchainContextImpl.Builder builder,
       PlatformKeys platformKeys,
@@ -368,11 +399,11 @@ private static void determineToolchainImplementations(
 
     // Find the toolchains for the requested toolchain types.
     List<SingleToolchainResolutionKey> registeredToolchainKeys = new ArrayList<>();
-    for (Label toolchainTypeLabel : requiredToolchainTypeLabels) {
+    for (ToolchainTypeRequirement toolchainType : toolchainTypes) {
       registeredToolchainKeys.add(
           SingleToolchainResolutionValue.key(
               configurationKey,
-              toolchainTypeLabel,
+              toolchainType,
               platformKeys.targetPlatformKey(),
               platformKeys.executionPlatformKeys(),
               debugTarget));
@@ -404,7 +435,7 @@ private static void determineToolchainImplementations(
             findPlatformsAndLabels(requiredToolchainType, singleToolchainResolutionValue));
       } catch (NoToolchainFoundException e) {
         // Save the missing type and continue looping to check for more.
-        missingToolchains.add(e.missingToolchainTypeLabel());
+        missingToolchains.add(e.missingToolchainType().toolchainType());
       }
     }
 
@@ -429,9 +460,7 @@ private static void determineToolchainImplementations(
 
     if (selectedExecutionPlatformKey.isEmpty()) {
       throw new NoMatchingPlatformException(
-          requiredToolchainTypeLabels,
-          platformKeys.executionPlatformKeys(),
-          platformKeys.targetPlatformKey());
+          toolchainTypes, platformKeys.executionPlatformKeys(), platformKeys.targetPlatformKey());
     }
 
     Map<ConfiguredTargetKey, PlatformInfo> platforms =
@@ -442,7 +471,7 @@ private static void determineToolchainImplementations(
       throw new ValueMissingException();
     }
 
-    builder.setRequiredToolchainTypes(requiredToolchainTypes);
+    builder.setToolchainTypes(toolchainTypes);
     builder.setExecutionPlatform(platforms.get(selectedExecutionPlatformKey.get()));
     builder.setTargetPlatform(platforms.get(platformKeys.targetPlatformKey()));
 
@@ -521,19 +550,17 @@ private ValueMissingException() {
   /** Exception used when no execution platform can be found. */
   static final class NoMatchingPlatformException extends ToolchainException {
     NoMatchingPlatformException(
-        Set<Label> requiredToolchainTypeLabels,
+        Set<ToolchainTypeRequirement> toolchainTypes,
         ImmutableList<ConfiguredTargetKey> availableExecutionPlatformKeys,
         ConfiguredTargetKey targetPlatformKey) {
-      super(
-          formatError(
-              requiredToolchainTypeLabels, availableExecutionPlatformKeys, targetPlatformKey));
+      super(formatError(toolchainTypes, availableExecutionPlatformKeys, targetPlatformKey));
     }
 
     private static String formatError(
-        Set<Label> requiredToolchainTypeLabels,
+        Set<ToolchainTypeRequirement> toolchainTypes,
         ImmutableList<ConfiguredTargetKey> availableExecutionPlatformKeys,
         ConfiguredTargetKey targetPlatformKey) {
-      if (requiredToolchainTypeLabels.isEmpty()) {
+      if (toolchainTypes.isEmpty()) {
         return String.format(
             "Unable to find an execution platform for target platform %s"
                 + " from available execution platforms [%s]",
@@ -545,7 +572,10 @@ private static String formatError(
       return String.format(
           "Unable to find an execution platform for toolchains [%s] and target platform %s"
               + " from available execution platforms [%s]",
-          requiredToolchainTypeLabels.stream().map(Label::toString).collect(joining(", ")),
+          toolchainTypes.stream()
+              .map(ToolchainTypeRequirement::toolchainType)
+              .map(Label::toString)
+              .collect(joining(", ")),
           targetPlatformKey.getLabel(),
           availableExecutionPlatformKeys.stream()
               .map(key -> key.getLabel().toString())
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/UnloadedToolchainContextImpl.java b/src/main/java/com/google/devtools/build/lib/skyframe/UnloadedToolchainContextImpl.java
index 56787549bc7b66..b3eef1e2ef39d8 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/UnloadedToolchainContextImpl.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/UnloadedToolchainContextImpl.java
@@ -17,6 +17,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
 import com.google.devtools.build.lib.analysis.platform.ToolchainTypeInfo;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -33,7 +34,7 @@ public abstract class UnloadedToolchainContextImpl implements SkyValue, Unloaded
   public static Builder builder(ToolchainContextKey key) {
     return new AutoValue_UnloadedToolchainContextImpl.Builder()
         .setKey(key)
-        .setRequiredToolchainTypes(ImmutableSet.of())
+        .setToolchainTypes(ImmutableSet.of())
         .setRequestedLabelToToolchainType(ImmutableMap.of())
         .setToolchainTypeToResolved(
             ImmutableSetMultimap.<ToolchainTypeInfo, Label>builder().build());
@@ -52,7 +53,7 @@ public interface Builder {
     Builder setTargetPlatform(PlatformInfo targetPlatform);
 
     /** Sets the toolchain types that were requested. */
-    Builder setRequiredToolchainTypes(Set<ToolchainTypeInfo> requiredToolchainTypes);
+    Builder setToolchainTypes(Set<ToolchainTypeRequirement> toolchainTypes);
 
     /**
      * Maps from the actual toolchain type to the resolved toolchain implementation that should be
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContextTest.java b/src/test/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContextTest.java
index cbdfeac0fd8f79..3412155c0dc5be 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContextTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/ResolvedToolchainContextTest.java
@@ -38,6 +38,8 @@
 @RunWith(JUnit4.class)
 public class ResolvedToolchainContextTest extends ToolchainTestCase {
 
+  // TODO(https://github.com/bazelbuild/bazel/issues/14726): Add tests for optional toolchain types.
+
   @Test
   public void load() throws Exception {
     addToolchain(
@@ -58,7 +60,7 @@ public void load() throws Exception {
         UnloadedToolchainContextImpl.builder(toolchainContextKey)
             .setExecutionPlatform(linuxPlatform)
             .setTargetPlatform(linuxPlatform)
-            .setRequiredToolchainTypes(ImmutableSet.of(testToolchainTypeInfo))
+            .setToolchainTypes(ImmutableSet.of(testToolchainType))
             .setRequestedLabelToToolchainType(
                 ImmutableMap.of(testToolchainTypeLabel, testToolchainTypeInfo))
             .setToolchainTypeToResolved(
@@ -110,7 +112,7 @@ public void load_aliasedToolchain() throws Exception {
         UnloadedToolchainContextImpl.builder(toolchainContextKey)
             .setExecutionPlatform(linuxPlatform)
             .setTargetPlatform(linuxPlatform)
-            .setRequiredToolchainTypes(ImmutableSet.of(testToolchainTypeInfo))
+            .setToolchainTypes(ImmutableSet.of(testToolchainType))
             .setRequestedLabelToToolchainType(
                 ImmutableMap.of(testToolchainTypeLabel, testToolchainTypeInfo))
             .setToolchainTypeToResolved(
@@ -152,7 +154,7 @@ public void load_notToolchain() throws Exception {
         UnloadedToolchainContextImpl.builder(toolchainContextKey)
             .setExecutionPlatform(linuxPlatform)
             .setTargetPlatform(linuxPlatform)
-            .setRequiredToolchainTypes(ImmutableSet.of(testToolchainTypeInfo))
+            .setToolchainTypes(ImmutableSet.of(testToolchainType))
             .setRequestedLabelToToolchainType(
                 ImmutableMap.of(testToolchainTypeLabel, testToolchainTypeInfo))
             .setToolchainTypeToResolved(
@@ -216,7 +218,7 @@ public void load_withTemplateVariables() throws Exception {
         UnloadedToolchainContextImpl.builder(toolchainContextKey)
             .setExecutionPlatform(linuxPlatform)
             .setTargetPlatform(linuxPlatform)
-            .setRequiredToolchainTypes(ImmutableSet.of(variableToolchainTypeInfo))
+            .setToolchainTypes(ImmutableSet.of(variableToolchainType))
             .setRequestedLabelToToolchainType(
                 ImmutableMap.of(variableToolchainTypeLabel, variableToolchainTypeInfo))
             .setToolchainTypeToResolved(
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/testing/ToolchainContextSubject.java b/src/test/java/com/google/devtools/build/lib/analysis/testing/ToolchainContextSubject.java
index 5e4751d321719d..55c97e6e2b3445 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/testing/ToolchainContextSubject.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/testing/ToolchainContextSubject.java
@@ -13,17 +13,20 @@
 // 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.IterableSubject;
+import com.google.common.truth.MapSubject;
 import com.google.common.truth.Subject;
 import com.google.devtools.build.lib.analysis.ToolchainContext;
-import com.google.devtools.build.lib.analysis.platform.ToolchainTypeInfo;
+import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
-import java.util.stream.Collectors;
+import java.util.Map;
 
 /** A Truth {@link Subject} for {@link ToolchainContext}. */
 public class ToolchainContextSubject extends Subject {
@@ -42,10 +45,18 @@ public static Subject.Factory<ToolchainContextSubject, ToolchainContext> toolcha
   // Instance fields.
 
   private final ToolchainContext actual;
+  private final Map<Label, ToolchainTypeRequirement> toolchainTypesMap;
 
   protected ToolchainContextSubject(FailureMetadata failureMetadata, ToolchainContext subject) {
     super(failureMetadata, subject);
     this.actual = subject;
+    this.toolchainTypesMap = makeToolchainTypesMap(subject);
+  }
+
+  private static ImmutableMap<Label, ToolchainTypeRequirement> makeToolchainTypesMap(
+      ToolchainContext subject) {
+    return subject.toolchainTypes().stream()
+        .collect(toImmutableMap(ToolchainTypeRequirement::toolchainType, Functions.identity()));
   }
 
   public void hasExecutionPlatform(String platformLabel) throws LabelSyntaxException {
@@ -66,20 +77,26 @@ public void hasTargetPlatform(Label platform) {
     check("targetPlatform()").that(actual.targetPlatform().label()).isEqualTo(platform);
   }
 
-  public void hasToolchainType(String toolchainTypeLabel) throws LabelSyntaxException {
-    hasToolchainType(Label.parseAbsolute(toolchainTypeLabel, ImmutableMap.of()));
+  public MapSubject toolchainTypes() {
+    return check("toolchainTypes()").that(toolchainTypesMap);
   }
 
-  public void hasToolchainType(Label toolchainType) {
-    toolchainTypeLabels().contains(toolchainType);
+  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 IterableSubject toolchainTypeLabels() {
-    return check("requiredToolchainTypes()")
-        .that(
-            actual.requiredToolchainTypes().stream()
-                .map(ToolchainTypeInfo::typeLabel)
-                .collect(Collectors.toList()));
+  public void hasToolchainType(Label toolchainType) {
+    toolchainType(toolchainType).isNotNull();
   }
 
   public void hasResolvedToolchain(String resolvedToolchainLabel) throws LabelSyntaxException {
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunctionTest.java
index 78cb21ef11de1d..8547bcfbd28c91 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/SingleToolchainResolutionFunctionTest.java
@@ -78,7 +78,7 @@ private EvaluationResult<SingleToolchainResolutionValue> invokeToolchainResoluti
   public void testResolution_singleExecutionPlatform() throws Exception {
     SkyKey key =
         SingleToolchainResolutionValue.key(
-            targetConfigKey, testToolchainTypeLabel, linuxCtkey, ImmutableList.of(macCtkey));
+            targetConfigKey, testToolchainType, linuxCtkey, ImmutableList.of(macCtkey));
     EvaluationResult<SingleToolchainResolutionValue> result = invokeToolchainResolution(key);
 
     assertThatEvaluationResult(result).hasNoError();
@@ -104,10 +104,7 @@ public void testResolution_multipleExecutionPlatforms() throws Exception {
 
     SkyKey key =
         SingleToolchainResolutionValue.key(
-            targetConfigKey,
-            testToolchainTypeLabel,
-            linuxCtkey,
-            ImmutableList.of(linuxCtkey, macCtkey));
+            targetConfigKey, testToolchainType, linuxCtkey, ImmutableList.of(linuxCtkey, macCtkey));
     EvaluationResult<SingleToolchainResolutionValue> result = invokeToolchainResolution(key);
 
     assertThatEvaluationResult(result).hasNoError();
@@ -128,7 +125,7 @@ public void testResolution_noneFound() throws Exception {
 
     SkyKey key =
         SingleToolchainResolutionValue.key(
-            targetConfigKey, testToolchainTypeLabel, linuxCtkey, ImmutableList.of(macCtkey));
+            targetConfigKey, testToolchainType, linuxCtkey, ImmutableList.of(macCtkey));
     EvaluationResult<SingleToolchainResolutionValue> result = invokeToolchainResolution(key);
 
     assertThatEvaluationResult(result)
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunctionTest.java
index 7cc35cd6ba7942..1d75a0b4432908 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunctionTest.java
@@ -136,7 +136,7 @@ public void resolve_noToolchainType() throws Exception {
     UnloadedToolchainContext unloadedToolchainContext = result.get(key);
     assertThat(unloadedToolchainContext).isNotNull();
 
-    assertThat(unloadedToolchainContext.requiredToolchainTypes()).isEmpty();
+    assertThat(unloadedToolchainContext.toolchainTypes()).isEmpty();
     // Even with no toolchains requested, should still select the first execution platform.
     assertThat(unloadedToolchainContext).hasExecutionPlatform("//platforms:mac");
     assertThat(unloadedToolchainContext).hasTargetPlatform("//platforms:linux");
@@ -173,7 +173,7 @@ public void resolve_noToolchainType_hostNotAvailable() throws Exception {
     UnloadedToolchainContext unloadedToolchainContext = result.get(key);
     assertThat(unloadedToolchainContext).isNotNull();
 
-    assertThat(unloadedToolchainContext.requiredToolchainTypes()).isEmpty();
+    assertThat(unloadedToolchainContext.toolchainTypes()).isEmpty();
     assertThat(unloadedToolchainContext).hasExecutionPlatform("//sample:sample_b");
     assertThat(unloadedToolchainContext).hasTargetPlatform("//platforms:linux");
   }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainsForTargetsTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainsForTargetsTest.java
index 8f09572732a80b..63f1271e6fb255 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainsForTargetsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ToolchainsForTargetsTest.java
@@ -362,7 +362,7 @@ public void execGroups_named() throws Exception {
         getToolchainCollection("//a");
     assertThat(toolchainCollection).isNotNull();
     assertThat(toolchainCollection).hasDefaultExecGroup();
-    assertThat(toolchainCollection).defaultToolchainContext().toolchainTypeLabels().isEmpty();
+    assertThat(toolchainCollection).defaultToolchainContext().toolchainTypes().isEmpty();
     assertThat(toolchainCollection).defaultToolchainContext().resolvedToolchainLabels().isEmpty();
 
     assertThat(toolchainCollection).hasExecGroup("temp");
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
index 2d42c602cf0a23..075fd3fcaad81a 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
+++ b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
@@ -3261,7 +3261,7 @@ public void testExecGroup_toolchain() throws Exception {
         ((StarlarkExecGroupCollection) info.getValue("exec_groups"))
             .getToolchainCollectionForTesting();
     assertThat(toolchainContexts.keySet()).containsExactly(DEFAULT_EXEC_GROUP_NAME, "dragonfruit");
-    assertThat(toolchainContexts.get(DEFAULT_EXEC_GROUP_NAME).requiredToolchainTypes()).isEmpty();
+    assertThat(toolchainContexts.get(DEFAULT_EXEC_GROUP_NAME).toolchainTypes()).isEmpty();
     assertThat(toolchainContexts.get("dragonfruit").resolvedToolchainLabels())
         .containsExactly(Label.parseAbsoluteUnchecked("//toolchain:foo"));
   }
@@ -3318,7 +3318,7 @@ public void testExecGroup_duplicateToolchainType() throws Exception {
             .getToolchainCollectionForTesting();
     assertThat(toolchainContexts.keySet())
         .containsExactly(DEFAULT_EXEC_GROUP_NAME, "dragonfruit", "passionfruit");
-    assertThat(toolchainContexts.get(DEFAULT_EXEC_GROUP_NAME).requiredToolchainTypes()).isEmpty();
+    assertThat(toolchainContexts.get(DEFAULT_EXEC_GROUP_NAME).toolchainTypes()).isEmpty();
     assertThat(toolchainContexts.get("dragonfruit").resolvedToolchainLabels())
         .containsExactly(Label.parseAbsoluteUnchecked("//toolchain:foo"));
     assertThat(toolchainContexts.get("passionfruit").resolvedToolchainLabels())