Update ToolchainUtil to properly load and use the available execution

platforms, and correctly merge together the results from TRF.

Part of #4442.

Change-Id: I31d83fa73a93d39a0e18d05a43a1c8666ac5a2d2
PiperOrigin-RevId: 187324257

1 files changed, 46 insertions, 31 deletions

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 0211c46251..630542b5af 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
@@ -28,7 +28,6 @@ import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.packages.NoSuchThingException;
 import com.google.devtools.build.lib.skyframe.RegisteredToolchainsFunction.InvalidToolchainLabelException;
-import com.google.devtools.build.lib.skyframe.ToolchainResolutionValue.ToolchainResolutionKey;
 import com.google.devtools.build.lib.skyframe.ToolchainUtil.ToolchainContextException;
 import com.google.devtools.build.lib.syntax.EvalException;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -39,7 +38,6 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /** {@link SkyFunction} which performs toolchain resolution for a class of rules. */
@@ -49,7 +47,7 @@ public class ToolchainResolutionFunction implements SkyFunction {
   @Override
   public SkyValue compute(SkyKey skyKey, Environment env)
       throws ToolchainResolutionFunctionException, InterruptedException {
-    ToolchainResolutionKey key = (ToolchainResolutionKey) skyKey.argument();
+    ToolchainResolutionValue.Key key = (ToolchainResolutionValue.Key) skyKey.argument();
 
     // This call could be combined with the call below, but this SkyFunction is evaluated so rarely
     // it's not worth optimizing.
@@ -77,26 +75,19 @@ public class ToolchainResolutionFunction implements SkyFunction {
       throw new ToolchainResolutionFunctionException(e);
     }
 
-    Map<ConfiguredTargetKey, PlatformInfo> platforms;
-    try {
-      platforms =
-          ToolchainUtil.getPlatformInfo(
-              key.targetPlatformKey(), key.availableExecutionPlatformKeys(), env);
-    } catch (ToolchainContextException e) {
-      throw new ToolchainResolutionFunctionException(e);
-    }
     // Find the right one.
     boolean debug = configuration.getOptions().get(PlatformOptions.class).toolchainResolutionDebug;
-    ImmutableMap<PlatformInfo, Label> resolvedToolchainLabels =
+    ImmutableMap<ConfiguredTargetKey, Label> resolvedToolchainLabels =
         resolveConstraints(
             key.toolchainType(),
-            key.availableExecutionPlatformKeys()
-                .stream()
-                .map(platforms::get)
-                .collect(Collectors.toList()),
-            platforms.get(key.targetPlatformKey()),
+            key.availableExecutionPlatformKeys(),
+            key.targetPlatformKey(),
             toolchains.registeredToolchains(),
+            env,
             debug ? env.getListener() : null);
+    if (resolvedToolchainLabels == null) {
+      return null;
+    }
 
     if (resolvedToolchainLabels.isEmpty()) {
       throw new ToolchainResolutionFunctionException(
@@ -111,17 +102,39 @@ public class ToolchainResolutionFunction implements SkyFunction {
    * pairs that are compatible a) with each other, and b) with the toolchain type and target
    * platform.
    */
-  private static ImmutableMap<PlatformInfo, Label> resolveConstraints(
+  @Nullable
+  private static ImmutableMap<ConfiguredTargetKey, Label> resolveConstraints(
       Label toolchainType,
-      List<PlatformInfo> availableExecutionPlatforms,
-      PlatformInfo targetPlatform,
+      List<ConfiguredTargetKey> availableExecutionPlatformKeys,
+      ConfiguredTargetKey targetPlatformKey,
       ImmutableList<DeclaredToolchainInfo> toolchains,
-      @Nullable EventHandler eventHandler) {
+      Environment env,
+      @Nullable EventHandler eventHandler)
+      throws ToolchainResolutionFunctionException, InterruptedException {
+
+    // Load the PlatformInfo needed to check constraints.
+    Map<ConfiguredTargetKey, PlatformInfo> platforms;
+    try {
+      platforms =
+          ToolchainUtil.getPlatformInfo(
+              new ImmutableList.Builder<ConfiguredTargetKey>()
+                  .add(targetPlatformKey)
+                  .addAll(availableExecutionPlatformKeys)
+                  .build(),
+              env);
+      if (platforms == null) {
+        return null;
+      }
+    } catch (ToolchainContextException e) {
+      throw new ToolchainResolutionFunctionException(e);
+    }
+
+    PlatformInfo targetPlatform = platforms.get(targetPlatformKey);
 
-    // Platforms may exist multiple times in availableExecutionPlatforms. The Set lets this code
+    // Platforms may exist multiple times in availableExecutionPlatformKeys. The Set lets this code
     // check whether a platform has already been seen during processing.
-    Set<PlatformInfo> platformsSeen = new HashSet<>();
-    ImmutableMap.Builder<PlatformInfo, Label> builder = ImmutableMap.builder();
+    Set<ConfiguredTargetKey> platformKeysSeen = new HashSet<>();
+    ImmutableMap.Builder<ConfiguredTargetKey, Label> builder = ImmutableMap.builder();
 
     debugMessage(eventHandler, "Looking for toolchain of type %s...", toolchainType);
     for (DeclaredToolchainInfo toolchain : toolchains) {
@@ -142,31 +155,33 @@ public class ToolchainResolutionFunction implements SkyFunction {
       }
 
       // Find the matching execution platforms.
-      for (PlatformInfo executionPlatform : availableExecutionPlatforms) {
+      for (ConfiguredTargetKey executionPlatformKey : availableExecutionPlatformKeys) {
+        PlatformInfo executionPlatform = platforms.get(executionPlatformKey);
         if (!checkConstraints(
             eventHandler, toolchain.execConstraints(), "execution", executionPlatform)) {
           continue;
         }
 
         // Only add the toolchains if this is a new platform.
-        if (!platformsSeen.contains(executionPlatform)) {
-          builder.put(executionPlatform, toolchain.toolchainLabel());
-          platformsSeen.add(executionPlatform);
+        if (!platformKeysSeen.contains(executionPlatformKey)) {
+          builder.put(executionPlatformKey, toolchain.toolchainLabel());
+          platformKeysSeen.add(executionPlatformKey);
         }
       }
     }
 
-    ImmutableMap<PlatformInfo, Label> resolvedToolchainLabels = builder.build();
+    ImmutableMap<ConfiguredTargetKey, Label> resolvedToolchainLabels = builder.build();
     if (resolvedToolchainLabels.isEmpty()) {
       debugMessage(eventHandler, "  No toolchains found");
     } else {
       debugMessage(
           eventHandler,
-          "  Selected execution platforms and toolchains: {%s}",
+          "  For toolchain type %s, possible execution platforms and toolchains: {%s}",
+          toolchainType,
           resolvedToolchainLabels
               .entrySet()
               .stream()
-              .map(e -> String.format("%s -> %s", e.getKey().label(), e.getValue()))
+              .map(e -> String.format("%s -> %s", e.getKey().getLabel(), e.getValue()))
               .collect(joining(", ")));
     }