Add --toolchain_resolution_override to allow specifying the toolchain for a given type.

Fixes https://github.com/katre/bazel/issues/6.

Change-Id: I6c6e303384277b013bdc27eb80743aa51f2fb98a
PiperOrigin-RevId: 162618674

5 files changed, 111 insertions, 8 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfiguration.java b/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfiguration.java
index 83d622fdbf..715afb7f2d 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfiguration.java
@@ -15,6 +15,8 @@
 package com.google.devtools.build.lib.analysis;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.analysis.PlatformOptions.ToolchainResolutionOverride;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.ThreadSafety;
@@ -35,13 +37,28 @@ public class PlatformConfiguration extends BuildConfiguration.Fragment {
   private final Label executionPlatform;
   private final ImmutableList<Label> targetPlatforms;
   private final ImmutableList<Label> extraToolchains;
+  private final ImmutableMap<Label, Label> toolchainResolutionOverrides;
 
   public PlatformConfiguration(
-      Label executionPlatform, List<Label> targetPlatforms, List<Label> extraToolchains) {
+      Label executionPlatform,
+      List<Label> targetPlatforms,
+      List<Label> extraToolchains,
+      List<ToolchainResolutionOverride> overrides) {
 
     this.executionPlatform = executionPlatform;
     this.targetPlatforms = ImmutableList.copyOf(targetPlatforms);
     this.extraToolchains = ImmutableList.copyOf(extraToolchains);
+    this.toolchainResolutionOverrides = convertOverrides(overrides);
+  }
+
+  private static ImmutableMap<Label, Label> convertOverrides(
+      List<ToolchainResolutionOverride> overrides) {
+    ImmutableMap.Builder<Label, Label> builder = new ImmutableMap.Builder<>();
+    for (ToolchainResolutionOverride override : overrides) {
+      builder.put(override.toolchainType(), override.toolchainLabel());
+    }
+
+    return builder.build();
   }
 
   @SkylarkCallable(
@@ -62,4 +79,14 @@ public class PlatformConfiguration extends BuildConfiguration.Fragment {
   public ImmutableList<Label> getExtraToolchains() {
     return extraToolchains;
   }
+
+  /** Returns {@code true} if the given toolchain type has a manual override set. */
+  public boolean hasToolchainOverride(Label toolchainType) {
+    return toolchainResolutionOverrides.containsKey(toolchainType);
+  }
+
+  /** Returns the {@link Label} of the toolchain to use for the given toolchain type. */
+  public Label getToolchainOverride(Label toolchainType) {
+    return toolchainResolutionOverrides.get(toolchainType);
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfigurationLoader.java b/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfigurationLoader.java
index 1eedd1265b..2e6ced471f 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfigurationLoader.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/PlatformConfigurationLoader.java
@@ -47,6 +47,10 @@ public class PlatformConfigurationLoader implements ConfigurationFragmentFactory
     // TODO(katre): This will change with remote execution.
     Label executionPlatform = options.hostPlatform;
 
-    return new PlatformConfiguration(executionPlatform, options.platforms, options.extraToolchains);
+    return new PlatformConfiguration(
+        executionPlatform,
+        options.platforms,
+        options.extraToolchains,
+        options.toolchainResolutionOverrides);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/PlatformOptions.java b/src/main/java/com/google/devtools/build/lib/analysis/PlatformOptions.java
index 47d64eefa2..438d6f2471 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/PlatformOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/PlatformOptions.java
@@ -14,15 +14,20 @@
 
 package com.google.devtools.build.lib.analysis;
 
+import static com.google.devtools.build.lib.analysis.config.BuildConfiguration.convertOptionsLabel;
+
+import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration.LabelListConverter;
 import com.google.devtools.build.lib.analysis.config.FragmentOptions;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionDocumentationCategory;
 import com.google.devtools.common.options.OptionEffectTag;
 import com.google.devtools.common.options.OptionMetadataTag;
+import com.google.devtools.common.options.OptionsParsingException;
 import java.util.List;
 
 /** Command-line options for platform-related configuration. */
@@ -63,10 +68,65 @@ public class PlatformOptions extends FragmentOptions {
   )
   public List<Label> extraToolchains;
 
+  @Option(
+    name = "toolchain_resolution_override",
+    converter = ToolchainResolutionOverrideConverter.class,
+    allowMultiple = true,
+    defaultValue = "",
+    documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+    effectTags = {OptionEffectTag.UNKNOWN},
+    metadataTags = {OptionMetadataTag.HIDDEN},
+    help =
+        "Override toolchain resolution for a toolchain type with a specific toolchain. "
+            + "Example: --toolchain_resolution_override=@io_bazel_rules_go//:toolchain="
+            + "@io_bazel_rules_go//:linux-arm64-toolchain"
+  )
+  public List<ToolchainResolutionOverride> toolchainResolutionOverrides;
+
   @Override
   public PlatformOptions getHost(boolean fallback) {
     PlatformOptions host = (PlatformOptions) getDefault();
     host.platforms = ImmutableList.of(this.hostPlatform);
     return host;
   }
+
+  /** Data about which toolchain instance should be used for a given toolchain type. */
+  @AutoValue
+  public abstract static class ToolchainResolutionOverride {
+    public abstract Label toolchainType();
+
+    public abstract Label toolchainLabel();
+
+    private static ToolchainResolutionOverride create(Label toolchainType, Label toolchainLabel) {
+      return new AutoValue_PlatformOptions_ToolchainResolutionOverride(
+          toolchainType, toolchainLabel);
+    }
+  }
+
+  /**
+   * {@link Converter} implementation to create {@link ToolchainResolutionOverride} instances from
+   * the value set in the flag.
+   */
+  public static class ToolchainResolutionOverrideConverter
+      implements Converter<ToolchainResolutionOverride> {
+
+    @Override
+    public ToolchainResolutionOverride convert(String input) throws OptionsParsingException {
+      int index = input.indexOf('=');
+      if (index == -1) {
+        throw new OptionsParsingException(
+            "Toolchain resolution override not in the type=toolchain format");
+      }
+      Label toolchainType = convertOptionsLabel(input.substring(0, index));
+      Label toolchain = convertOptionsLabel(input.substring(index + 1));
+
+      return ToolchainResolutionOverride.create(toolchainType, toolchain);
+    }
+
+    @Override
+    public String getTypeDescription() {
+      return "a hard-coded override for toolchain resolution, "
+          + "in the format toolchainTypeLabel=toolchainLabel";
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
index d809225acf..37c85f1d97 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
@@ -239,7 +239,7 @@ public final class BuildConfiguration implements BuildEvent {
     }
   }
 
-  private static final Label convertLabel(String input) throws OptionsParsingException {
+  public static final Label convertOptionsLabel(String input) throws OptionsParsingException {
     try {
       // Check if the input starts with '/'. We don't check for "//" so that
       // we get a better error message if the user accidentally tries to use
@@ -259,7 +259,7 @@ public final class BuildConfiguration implements BuildEvent {
   public static class LabelConverter implements Converter<Label> {
     @Override
     public Label convert(String input) throws OptionsParsingException {
-      return convertLabel(input);
+      return convertOptionsLabel(input);
     }
 
     @Override
@@ -274,7 +274,7 @@ public final class BuildConfiguration implements BuildEvent {
     public List<Label> convert(String input) throws OptionsParsingException {
       ImmutableList.Builder result = ImmutableList.builder();
       for (String label : Splitter.on(",").omitEmptyStrings().split(input)) {
-        result.add(convertLabel(label));
+        result.add(convertOptionsLabel(label));
       }
       return result.build();
     }
@@ -292,7 +292,7 @@ public final class BuildConfiguration implements BuildEvent {
   public static class EmptyToNullLabelConverter implements Converter<Label> {
     @Override
     public Label convert(String input) throws OptionsParsingException {
-      return input.isEmpty() ? null : convertLabel(input);
+      return input.isEmpty() ? null : convertOptionsLabel(input);
     }
 
     @Override
@@ -315,7 +315,7 @@ public final class BuildConfiguration implements BuildEvent {
 
     @Override
     public Label convert(String input) throws OptionsParsingException {
-      return input.isEmpty() ? defaultValue : convertLabel(input);
+      return input.isEmpty() ? defaultValue : convertOptionsLabel(input);
     }
 
     @Override
@@ -340,7 +340,7 @@ public final class BuildConfiguration implements BuildEvent {
         } else {
           key = entry.substring(0, sepIndex);
           String value = entry.substring(sepIndex + 1);
-          label = value.isEmpty() ? null : convertLabel(value);
+          label = value.isEmpty() ? null : convertOptionsLabel(value);
         }
         if (result.containsKey(key)) {
           throw new OptionsParsingException("Key '" + key + "' appears twice");
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 dbe0e9f02b..28a6ebe0ca 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
@@ -16,6 +16,8 @@ package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
+import com.google.devtools.build.lib.analysis.PlatformConfiguration;
+import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.platform.ConstraintValueInfo;
 import com.google.devtools.build.lib.analysis.platform.DeclaredToolchainInfo;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
@@ -40,6 +42,16 @@ public class ToolchainResolutionFunction implements SkyFunction {
       throws SkyFunctionException, InterruptedException {
     ToolchainResolutionKey key = (ToolchainResolutionKey) skyKey.argument();
 
+    BuildConfiguration configuration = key.configuration();
+    PlatformConfiguration platformConfiguration =
+        configuration.getFragment(PlatformConfiguration.class);
+
+    if (platformConfiguration.hasToolchainOverride(key.toolchainType())) {
+      // Short circuit everything and just return the override.
+      return ToolchainResolutionValue.create(
+          platformConfiguration.getToolchainOverride(key.toolchainType()));
+    }
+
     // Get all toolchains.
     RegisteredToolchainsValue toolchains;
     try {