Add register_toolchains function to the WORKSPACE for registering toolchains to use.

Part of #2219.

Change-Id: Id6dfe6ec102f609bb19461242a098bf977be29ae
PiperOrigin-RevId: 161527986

5 files changed, 100 insertions, 7 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/Package.java b/src/main/java/com/google/devtools/build/lib/packages/Package.java
index e5a5598a38..4ec65dd3b0 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/Package.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/Package.java
@@ -193,6 +193,8 @@ public class Package {
   private ImmutableList<Event> events;
   private ImmutableList<Postable> posts;
 
+  private ImmutableList<Label> registeredToolchainLabels;
+
   /**
    * Package initialization, part 1 of 3: instantiates a new package with the
    * given name.
@@ -317,6 +319,7 @@ public class Package {
     this.features = ImmutableSortedSet.copyOf(builder.features);
     this.events = ImmutableList.copyOf(builder.events);
     this.posts = ImmutableList.copyOf(builder.posts);
+    this.registeredToolchainLabels = ImmutableList.copyOf(builder.registeredToolchainLabels);
   }
 
   /**
@@ -643,6 +646,10 @@ public class Package {
     return defaultRestrictedTo;
   }
 
+  public ImmutableList<Label> getRegisteredToolchainLabels() {
+    return registeredToolchainLabels;
+  }
+
   @Override
   public String toString() {
     return "Package(" + name + ")="
@@ -694,6 +701,7 @@ public class Package {
    * {@link com.google.devtools.build.lib.skyframe.PackageFunction}.
    */
   public static class Builder {
+
     public static interface Helper {
       /**
        * Returns a fresh {@link Package} instance that a {@link Builder} will internally mutate
@@ -756,6 +764,8 @@ public class Package {
 
     protected ExternalPackageBuilder externalPackageData = new ExternalPackageBuilder();
 
+    protected List<Label> registeredToolchainLabels = new ArrayList<>();
+
     /**
      * True iff the "package" function has already been called in this package.
      */
@@ -1270,6 +1280,10 @@ public class Package {
       addRuleUnchecked(rule);
     }
 
+    void addRegisteredToolchainLabels(List<Label> toolchains) {
+      this.registeredToolchainLabels.addAll(toolchains);
+    }
+
     private Builder beforeBuild(boolean discoverAssumedInputFiles) throws InterruptedException {
       Preconditions.checkNotNull(pkg);
       Preconditions.checkNotNull(filename);
diff --git a/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java b/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java
index 41dbdd7822..ccb81e1201 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java
@@ -46,11 +46,14 @@ import com.google.devtools.build.lib.syntax.FunctionSignature;
 import com.google.devtools.build.lib.syntax.Mutability;
 import com.google.devtools.build.lib.syntax.ParserInputSource;
 import com.google.devtools.build.lib.syntax.Runtime;
+import com.google.devtools.build.lib.syntax.Runtime.NoneType;
 import com.google.devtools.build.lib.syntax.SkylarkList;
 import com.google.devtools.build.lib.syntax.SkylarkSignatureProcessor;
 import com.google.devtools.build.lib.vfs.Path;
 import java.io.File;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -60,7 +63,6 @@ import javax.annotation.Nullable;
  * Parser for WORKSPACE files.  Fills in an ExternalPackage.Builder
  */
 public class WorkspaceFactory {
-  public static final String BIND = "bind";
   private static final Pattern LEGAL_WORKSPACE_NAME = Pattern.compile("^\\p{Alpha}\\w*$");
 
   // List of static function added by #addWorkspaceFunctions. Used to trim them out from the
@@ -260,6 +262,7 @@ public class WorkspaceFactory {
     if (aPackage.containsErrors()) {
       builder.setContainsErrors();
     }
+    builder.addRegisteredToolchainLabels(aPackage.getRegisteredToolchainLabels());
     for (Rule rule : aPackage.getTargets(Rule.class)) {
       try {
         // The old rule references another Package instance and we wan't to keep the invariant that
@@ -382,6 +385,56 @@ public class WorkspaceFactory {
     };
   }
 
+  @SkylarkSignature(
+    name = "register_toolchains",
+    objectType = Object.class,
+    returnType = NoneType.class,
+    doc =
+        "Registers a toolchain created with the toolchain() rule so that it is available for "
+            + "toolchain resolution.",
+    extraPositionals =
+        @Param(
+          name = "toolchain_labels",
+          type = SkylarkList.class,
+          generic1 = String.class,
+          doc = "The labels of the toolchains to register."
+        ),
+    useAst = true,
+    useEnvironment = true
+  )
+  private static final BuiltinFunction.Factory newRegisterToolchainsFunction =
+      new BuiltinFunction.Factory("register_toolchains") {
+        public BuiltinFunction create(final RuleFactory ruleFactory) {
+          return new BuiltinFunction(
+              "register_toolchains", FunctionSignature.POSITIONALS, BuiltinFunction.USE_AST_ENV) {
+            public Object invoke(
+                SkylarkList<String> toolchainLabels, FuncallExpression ast, Environment env)
+                throws EvalException, InterruptedException {
+
+              // Collect the toolchain labels.
+              List<Label> toolchains = new ArrayList<>();
+              for (String rawLabel :
+                  toolchainLabels.getContents(String.class, "toolchain_labels")) {
+                try {
+                  toolchains.add(Label.parseAbsolute(rawLabel));
+                } catch (LabelSyntaxException e) {
+                  throw new EvalException(
+                      ast.getLocation(),
+                      String.format("Unable to parse toolchain %s: %s", rawLabel, e.getMessage()),
+                      e);
+                }
+              }
+
+              // Add to the package definition for later.
+              Package.Builder builder = PackageFactory.getContext(env, ast).pkgBuilder;
+              builder.addRegisteredToolchainLabels(toolchains);
+
+              return NONE;
+            }
+          };
+        }
+      };
+
   /**
    * Returns a function-value implementing the build rule "ruleClass" (e.g. cc_library) in the
    * specified package context.
@@ -426,15 +479,16 @@ public class WorkspaceFactory {
 
   private static ImmutableMap<String, BaseFunction> createWorkspaceFunctions(
       boolean allowOverride, RuleFactory ruleFactory) {
-    ImmutableMap.Builder<String, BaseFunction> mapBuilder = ImmutableMap.builder();
-    mapBuilder.put(BIND, newBindFunction(ruleFactory));
+    Map<String, BaseFunction> map = new HashMap<>();
+    map.put("bind", newBindFunction(ruleFactory));
+    map.put("register_toolchains", newRegisterToolchainsFunction.apply(ruleFactory));
     for (String ruleClass : ruleFactory.getRuleClassNames()) {
-      if (!ruleClass.equals(BIND)) {
+      if (!map.containsKey(ruleClass)) {
         BaseFunction ruleFunction = newRuleFunction(ruleFactory, ruleClass, allowOverride);
-        mapBuilder.put(ruleClass, ruleFunction);
+        map.put(ruleClass, ruleFunction);
       }
     }
-    return mapBuilder.build();
+    return ImmutableMap.copyOf(map);
   }
 
   private void addWorkspaceFunctions(Environment workspaceEnv, StoredEventHandler localReporter) {
diff --git a/src/main/java/com/google/devtools/build/lib/rules/ExternalPackageUtil.java b/src/main/java/com/google/devtools/build/lib/rules/ExternalPackageUtil.java
index 4a73d05cdb..3a68b46903 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/ExternalPackageUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/ExternalPackageUtil.java
@@ -26,6 +26,7 @@ import com.google.devtools.build.lib.packages.NoSuchPackageException;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.skyframe.PackageLookupValue;
+import com.google.devtools.build.lib.skyframe.PackageValue;
 import com.google.devtools.build.lib.skyframe.WorkspaceFileValue;
 import com.google.devtools.build.lib.syntax.EvalException;
 import com.google.devtools.build.lib.util.Preconditions;
@@ -171,6 +172,24 @@ public class ExternalPackageUtil {
     return rule;
   }
 
+  /**
+   * Loads the external package and then returns the registered toolchain labels.
+   *
+   * @param env the environment to use for lookups
+   */
+  @Nullable
+  public static List<Label> getRegisteredToolchainLabels(Environment env)
+      throws ExternalPackageException, InterruptedException {
+    PackageValue externalPackageValue =
+        (PackageValue) env.getValue(PackageValue.key(Label.EXTERNAL_PACKAGE_IDENTIFIER));
+    if (externalPackageValue == null) {
+      return null;
+    }
+
+    Package externalPackage = externalPackageValue.getPackage();
+    return externalPackage.getRegisteredToolchainLabels();
+  }
+
   /** Exception thrown when something goes wrong accessing a rule. */
   public static class ExternalPackageException extends SkyFunctionException {
     public ExternalPackageException(NoSuchPackageException cause, Transience transience) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
index 15e21e3bb9..33fb4aa5cb 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
@@ -453,7 +453,9 @@ public class PackageFunction implements SkyFunction {
       env.getListener().post(post);
     }
 
-    packageFactory.afterDoneLoadingPackage(pkg);
+    if (packageFactory != null) {
+      packageFactory.afterDoneLoadingPackage(pkg);
+    }
     return new PackageValue(pkg);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/FunctionSignature.java b/src/main/java/com/google/devtools/build/lib/syntax/FunctionSignature.java
index 5dc2a6a5c6..37e6f06de8 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/FunctionSignature.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/FunctionSignature.java
@@ -589,6 +589,10 @@ public abstract class FunctionSignature implements Serializable {
     }
   }
 
+  /** A ready-made signature to allow only positional arguments and put them in a star parameter */
+  public static final FunctionSignature POSITIONALS =
+      FunctionSignature.of(0, 0, 0, true, false, "star");
+
   /** A ready-made signature to allow only keyword arguments and put them in a kwarg parameter */
   public static final FunctionSignature KWARGS =
       FunctionSignature.of(0, 0, 0, false, true, "kwargs");