Update ConfiguredTargetFunction.computeUnloadedToolchainContexts to

directly require the correct execution platform.

This fixes several issues around toolchains that depend on toolchains.

Fixes bazelbuild#13243.

Closes bazelbuild#13258.

PiperOrigin-RevId: 364597996

src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java

@@ -520,18 +520,31 @@ static ToolchainCollection<UnloadedToolchainContext> computeUnloadedToolchainCon
 
     Map<String, ToolchainContextKey> toolchainContextKeys = new HashMap<>();
     String targetUnloadedToolchainContext = "target-unloaded-toolchain-context";
-    ToolchainContextKey toolchainContextKey;
+
+    ToolchainContextKey.Builder toolchainContextKeyBuilder =
+        ToolchainContextKey.key()
+            .configurationKey(toolchainConfig)
+            .requiredToolchainTypeLabels(requiredDefaultToolchains)
+            .execConstraintLabels(defaultExecConstraintLabels)
+            .shouldSanityCheckConfiguration(configuration.trimConfigurationsRetroactively());
+
     if (parentToolchainContextKey != null) {
-      toolchainContextKey = parentToolchainContextKey;
-    } else {
-      toolchainContextKey =
-          ToolchainContextKey.key()
-              .configurationKey(toolchainConfig)
-              .requiredToolchainTypeLabels(requiredDefaultToolchains)
-              .execConstraintLabels(defaultExecConstraintLabels)
-              .shouldSanityCheckConfiguration(configuration.trimConfigurationsRetroactively())
-              .build();
+      // Find out what execution platform the parent used, and force that.
+      // This key should always be present, but check just in case.
+      ToolchainContext parentToolchainContext =
+          (ToolchainContext)
+              env.getValueOrThrow(parentToolchainContextKey, ToolchainException.class);
+      if (env.valuesMissing()) {
+        return null;
+      }
+
+      Label execPlatform = parentToolchainContext.executionPlatform().label();
+      if (execPlatform != null) {
+        toolchainContextKeyBuilder.forceExecutionPlatform(execPlatform);
+      }
     }
+
+    ToolchainContextKey toolchainContextKey = toolchainContextKeyBuilder.build();
     toolchainContextKeys.put(targetUnloadedToolchainContext, toolchainContextKey);
     for (Map.Entry<String, ExecGroup> group : execGroups.entrySet()) {
       ExecGroup execGroup = group.getValue();
@@ -558,14 +571,6 @@ static ToolchainCollection<UnloadedToolchainContext> computeUnloadedToolchainCon
           (UnloadedToolchainContext) values.get(unloadedToolchainContextKey.getValue()).get();
       if (!valuesMissing) {
         String execGroup = unloadedToolchainContextKey.getKey();
-        if (parentToolchainContextKey != null) {
-          // Since we inherited the toolchain context from the parent of the dependency, the current
-          // target may also be in the resolved toolchains list. We need to clear it out.
-          // TODO(configurability): When updating this for config_setting, only remove the current
-          // target, not everything, because config_setting might want to check the toolchain
-          // dependencies.
-          unloadedToolchainContext = unloadedToolchainContext.withoutResolvedToolchains();
-        }
         if (execGroup.equals(targetUnloadedToolchainContext)) {
           toolchainContexts.addDefaultContext(unloadedToolchainContext);
         } else {

src/main/java/com/google/devtools/build/lib/skyframe/ToolchainContextKey.java

@@ -14,10 +14,12 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
+import java.util.Optional;
 
 /**
  * {@link SkyKey} implementation used for {@link ToolchainResolutionFunction} to produce {@link
@@ -31,7 +33,8 @@ public static Builder key() {
     return new AutoValue_ToolchainContextKey.Builder()
         .requiredToolchainTypeLabels(ImmutableSet.of())
         .execConstraintLabels(ImmutableSet.of())
-        .shouldSanityCheckConfiguration(false);
+        .shouldSanityCheckConfiguration(false)
+        .forceExecutionPlatform(Optional.empty());
   }
 
   @Override
@@ -47,6 +50,8 @@ public SkyFunctionName functionName() {
 
   abstract boolean shouldSanityCheckConfiguration();
 
+  abstract Optional<Label> forceExecutionPlatform();
+
   /** Builder for {@link ToolchainContextKey}. */
   @AutoValue.Builder
   public interface Builder {
@@ -63,5 +68,12 @@ public interface Builder {
     Builder shouldSanityCheckConfiguration(boolean shouldSanityCheckConfiguration);
 
     ToolchainContextKey build();
+
+    Builder forceExecutionPlatform(Optional<Label> execPlatform);
+
+    default Builder forceExecutionPlatform(Label execPlatform) {
+      Preconditions.checkNotNull(execPlatform);
+      return forceExecutionPlatform(Optional.of(execPlatform));
+    }
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunction.java

@@ -117,6 +117,7 @@ public UnloadedToolchainContext compute(SkyKey skyKey, Environment env)
           env,
           key.configurationKey(),
           resolvedToolchainTypeLabels,
+          key.forceExecutionPlatform().map(platformKeys::find),
           builder,
           platformKeys,
           key.shouldSanityCheckConfiguration());
@@ -182,6 +183,24 @@ abstract static class PlatformKeys {
 
     abstract ImmutableList<ConfiguredTargetKey> executionPlatformKeys();
 
+    @Nullable
+    public ConfiguredTargetKey find(Label platformLabel) {
+      if (platformLabel.equals(hostPlatformKey().getLabel())) {
+        return hostPlatformKey();
+      }
+      if (platformLabel.equals(targetPlatformKey().getLabel())) {
+        return targetPlatformKey();
+      }
+
+      for (ConfiguredTargetKey configuredTargetKey : executionPlatformKeys()) {
+        if (platformLabel.equals(configuredTargetKey.getLabel())) {
+          return configuredTargetKey;
+        }
+      }
+
+      return null;
+    }
+
     static PlatformKeys create(
         ConfiguredTargetKey hostPlatformKey,
         ConfiguredTargetKey targetPlatformKey,
@@ -350,6 +369,7 @@ private void determineToolchainImplementations(
       Environment environment,
       BuildConfigurationValue.Key configurationKey,
       ImmutableSet<Label> requiredToolchainTypeLabels,
+      Optional<ConfiguredTargetKey> forcedExecutionPlatform,
       UnloadedToolchainContextImpl.Builder builder,
       PlatformKeys platformKeys,
       boolean shouldSanityCheckConfiguration)
@@ -415,19 +435,12 @@ private void determineToolchainImplementations(
     ImmutableSet<ToolchainTypeInfo> requiredToolchainTypes = requiredToolchainTypesBuilder.build();
 
     // Find and return the first execution platform which has all required toolchains.
-    Optional<ConfiguredTargetKey> selectedExecutionPlatformKey;
-    if (!requiredToolchainTypeLabels.isEmpty()) {
-      selectedExecutionPlatformKey =
-          findExecutionPlatformForToolchains(
-              requiredToolchainTypes,
-              platformKeys.executionPlatformKeys(),
-              resolvedToolchains);
-    } else if (!platformKeys.executionPlatformKeys().isEmpty()) {
-      // Just use the first execution platform.
-      selectedExecutionPlatformKey = Optional.of(platformKeys.executionPlatformKeys().get(0));
-    } else {
-      selectedExecutionPlatformKey = Optional.empty();
-    }
+    Optional<ConfiguredTargetKey> selectedExecutionPlatformKey =
+        findExecutionPlatformForToolchains(
+            requiredToolchainTypes,
+            forcedExecutionPlatform,
+            platformKeys.executionPlatformKeys(),
+            resolvedToolchains);
 
     if (!selectedExecutionPlatformKey.isPresent()) {
       throw new NoMatchingPlatformException(
@@ -477,25 +490,47 @@ private static Table<ConfiguredTargetKey, ToolchainTypeInfo, Label> findPlatform
    */
   private static Optional<ConfiguredTargetKey> findExecutionPlatformForToolchains(
       ImmutableSet<ToolchainTypeInfo> requiredToolchainTypes,
+      Optional<ConfiguredTargetKey> forcedExecutionPlatform,
       ImmutableList<ConfiguredTargetKey> availableExecutionPlatformKeys,
       Table<ConfiguredTargetKey, ToolchainTypeInfo, Label> resolvedToolchains) {
-    for (ConfiguredTargetKey executionPlatformKey : availableExecutionPlatformKeys) {
-      if (!resolvedToolchains.containsRow(executionPlatformKey)) {
-        continue;
-      }
 
-      Map<ToolchainTypeInfo, Label> toolchains = resolvedToolchains.row(executionPlatformKey);
-      if (!toolchains.keySet().containsAll(requiredToolchainTypes)) {
-        // Not all toolchains are present, ignore this execution platform.
-        continue;
+    if (forcedExecutionPlatform.isPresent()) {
+      // Is the forced platform suitable?
+      if (isPlatformSuitable(
+          forcedExecutionPlatform.get(), requiredToolchainTypes, resolvedToolchains)) {
+        return forcedExecutionPlatform;
       }
+    }
+
+    return availableExecutionPlatformKeys.stream()
+        .filter(epk -> isPlatformSuitable(epk, requiredToolchainTypes, resolvedToolchains))
+        .findFirst();
+  }
+
+  private static boolean isPlatformSuitable(
+      ConfiguredTargetKey executionPlatformKey,
+      ImmutableSet<ToolchainTypeInfo> requiredToolchainTypes,
+      Table<ConfiguredTargetKey, ToolchainTypeInfo, Label> resolvedToolchains) {
+    if (requiredToolchainTypes.isEmpty()) {
+      // Since there aren't any toolchains, we should be able to use any execution platform that
+      // has made it this far.
+      return true;
+    }
 
-      return Optional.of(executionPlatformKey);
+    if (!resolvedToolchains.containsRow(executionPlatformKey)) {
+      return false;
     }
 
-    return Optional.empty();
+    Map<ToolchainTypeInfo, Label> toolchains = resolvedToolchains.row(executionPlatformKey);
+    if (!toolchains.keySet().containsAll(requiredToolchainTypes)) {
+      // Not all toolchains are present, ignore this execution platform.
+      return false;
+    }
+
+    return true;
   }
 
+
   @Nullable
   @Override
   public String extractTag(SkyKey skyKey) {

src/main/java/com/google/devtools/build/lib/skyframe/UnloadedToolchainContext.java

@@ -41,7 +41,4 @@ public interface UnloadedToolchainContext extends ToolchainContext, SkyValue {
 
   @Override
   ImmutableSet<Label> resolvedToolchainLabels();
-
-  /** Returns a copy of this context, without the resolved toolchain data. */
-  UnloadedToolchainContext withoutResolvedToolchains();
 }

src/main/java/com/google/devtools/build/lib/skyframe/UnloadedToolchainContextImpl.java

@@ -74,9 +74,4 @@ public ImmutableSet<Label> resolvedToolchainLabels() {
   }
 
   protected abstract Builder toBuilder();
-
-  @Override
-  public UnloadedToolchainContext withoutResolvedToolchains() {
-    return this.toBuilder().setToolchainTypeToResolved(ImmutableSetMultimap.of()).build();
-  }
 }

src/test/java/com/google/devtools/build/lib/skyframe/ToolchainResolutionFunctionTest.java

@@ -431,4 +431,67 @@ public void resolve_noMatchingPlatform() throws Exception {
         .hasExceptionThat()
         .isInstanceOf(NoMatchingPlatformException.class);
   }
+
+  @Test
+  public void resolve_forceExecutionPlatform() throws Exception {
+    // This should select execution platform linux, toolchain extra_toolchain_linux, due to the
+    // forced execution platform, even though execution platform mac is registered first.
+    addToolchain(
+        /* packageName= */ "extra",
+        /* toolchainName= */ "extra_toolchain_linux",
+        /* execConstraints= */ ImmutableList.of("//constraints:linux"),
+        /* targetConstraints= */ ImmutableList.of("//constraints:linux"),
+        /* data= */ "baz");
+    addToolchain(
+        /* packageName= */ "extra",
+        /* toolchainName= */ "extra_toolchain_mac",
+        /* execConstraints= */ ImmutableList.of("//constraints:mac"),
+        /* targetConstraints= */ ImmutableList.of("//constraints:linux"),
+        /* data= */ "baz");
+    rewriteWorkspace(
+        "register_toolchains('//extra:extra_toolchain_linux', '//extra:extra_toolchain_mac')",
+        "register_execution_platforms('//platforms:mac', '//platforms:linux')");
+
+    useConfiguration("--platforms=//platforms:linux");
+    ToolchainContextKey key =
+        ToolchainContextKey.key()
+            .configurationKey(targetConfigKey)
+            .requiredToolchainTypeLabels(testToolchainTypeLabel)
+            .forceExecutionPlatform(Label.parseAbsoluteUnchecked("//platforms:linux"))
+            .build();
+
+    EvaluationResult<UnloadedToolchainContext> result = invokeToolchainResolution(key);
+
+    assertThatEvaluationResult(result).hasNoError();
+    UnloadedToolchainContext unloadedToolchainContext = result.get(key);
+    assertThat(unloadedToolchainContext).isNotNull();
+
+    assertThat(unloadedToolchainContext).hasToolchainType(testToolchainTypeLabel);
+    assertThat(unloadedToolchainContext).hasResolvedToolchain("//extra:extra_toolchain_linux_impl");
+    assertThat(unloadedToolchainContext).hasExecutionPlatform("//platforms:linux");
+    assertThat(unloadedToolchainContext).hasTargetPlatform("//platforms:linux");
+  }
+
+  @Test
+  public void resolve_forceExecutionPlatform_noRequiredToolchains() throws Exception {
+    // This should select execution platform linux, due to the forced execution platform, even
+    // though execution platform mac is registered first.
+    rewriteWorkspace("register_execution_platforms('//platforms:mac', '//platforms:linux')");
+
+    useConfiguration("--platforms=//platforms:linux");
+    ToolchainContextKey key =
+        ToolchainContextKey.key()
+            .configurationKey(targetConfigKey)
+            .forceExecutionPlatform(Label.parseAbsoluteUnchecked("//platforms:linux"))
+            .build();
+
+    EvaluationResult<UnloadedToolchainContext> result = invokeToolchainResolution(key);
+
+    assertThatEvaluationResult(result).hasNoError();
+    UnloadedToolchainContext unloadedToolchainContext = result.get(key);
+    assertThat(unloadedToolchainContext).isNotNull();
+
+    assertThat(unloadedToolchainContext).hasExecutionPlatform("//platforms:linux");
+    assertThat(unloadedToolchainContext).hasTargetPlatform("//platforms:linux");
+  }
 }

src/test/java/com/google/devtools/build/lib/skyframe/ToolchainsForTargetsTest.java

@@ -451,43 +451,47 @@ public void execGroups_defaultAndNamed() throws Exception {
 
   @Test
   public void keepParentToolchainContext() throws Exception {
+    // Add some platforms and custom constraints.
     scratch.file(
-        "extra/BUILD",
-        "load('//toolchain:toolchain_def.bzl', 'test_toolchain')",
-        "toolchain_type(name = 'extra_toolchain')",
-        "toolchain(",
-        "    name = 'toolchain',",
-        "    toolchain_type = '//extra:extra_toolchain',",
-        "    exec_compatible_with = [],",
-        "    target_compatible_with = [],",
-        "    toolchain = ':toolchain_impl')",
-        "test_toolchain(",
-        "    name='toolchain_impl',",
-        "    data = 'foo')");
+        "platforms/BUILD",
+        "constraint_setting(name = 'local_setting')",
+        "constraint_value(name = 'local_value_a', constraint_setting = ':local_setting')",
+        "constraint_value(name = 'local_value_b', constraint_setting = ':local_setting')",
+        "platform(name = 'local_platform_a',",
+        "    constraint_values = [':local_value_a'],",
+        ")",
+        "platform(name = 'local_platform_b',",
+        "    constraint_values = [':local_value_b'],",
+        ")");
+
+    // Test normal resolution, and with a per-target exec constraint.
     scratch.file("a/BUILD", "load('//toolchain:rule.bzl', 'my_rule')", "my_rule(name = 'a')");
 
-    useConfiguration("--extra_toolchains=//extra:toolchain");
+    useConfiguration(
+        "--extra_execution_platforms=//platforms:local_platform_a,//platforms:local_platform_b");
+
     ConfiguredTarget target = Iterables.getOnlyElement(update("//a").getTargetsToBuild());
+    ToolchainContextKey parentKey =
+        ToolchainContextKey.key()
+            .configurationKey(target.getConfigurationKey())
+            // Force the constraint label, to make the exec platform be local_platform_b.
+            .execConstraintLabels(Label.parseAbsoluteUnchecked("//platforms:local_value_b"))
+            .build();
     ToolchainCollection<UnloadedToolchainContext> toolchainCollection =
         getToolchainCollection(
             target,
             ConfiguredTargetKey.builder()
                 .setLabel(target.getOriginalLabel())
                 .setConfigurationKey(target.getConfigurationKey())
-                .setToolchainContextKey(
-                    ToolchainContextKey.key()
-                        .configurationKey(target.getConfigurationKey())
-                        .requiredToolchainTypeLabels(
-                            Label.parseAbsoluteUnchecked("//extra:extra_toolchain"))
-                        .build())
+                .setToolchainContextKey(parentKey)
                 .build());
 
     assertThat(toolchainCollection).isNotNull();
     assertThat(toolchainCollection).hasDefaultExecGroup();
+
+    // This should have the same exec platform as parentToolchainKey, which is local_platform_b.
     assertThat(toolchainCollection)
         .defaultToolchainContext()
-        .hasToolchainType("//extra:extra_toolchain");
-    // These should be cleared out.
-    assertThat(toolchainCollection).defaultToolchainContext().resolvedToolchainLabels().isEmpty();
+        .hasExecutionPlatform("//platforms:local_platform_b");
   }
 }