Add toolchains attribute to aspect.

Part of #2219.

Change-Id: I39ced1f3e2605154771df9424d6ed2f971820baf
PiperOrigin-RevId: 157002268

7 files changed, 176 insertions, 95 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredTargetFactory.java b/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredTargetFactory.java
index b1ec35db8b..3e48b7a59c 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredTargetFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredTargetFactory.java
@@ -244,6 +244,9 @@ public final class ConfiguredTargetFactory {
             .setConfigConditions(configConditions)
             .setUniversalFragment(ruleClassProvider.getUniversalFragment())
             .setSkylarkProvidersRegistry(ruleClassProvider.getRegisteredSkylarkProviders())
+            // TODO(katre): Populate the actual selected toolchains.
+            .setToolchainContext(
+                new ToolchainContext(rule.getRuleClassObject().getRequiredToolchains(), null))
             .build();
     if (ruleContext.hasErrors()) {
       return null;
@@ -350,6 +353,9 @@ public final class ConfiguredTargetFactory {
             .setAspectAttributes(aspect.getDefinition().getAttributes())
             .setConfigConditions(configConditions)
             .setUniversalFragment(ruleClassProvider.getUniversalFragment())
+            // TODO(katre): Populate the actual selected toolchains.
+            .setToolchainContext(
+                new ToolchainContext(aspect.getDefinition().getRequiredToolchains(), null))
             .build();
     if (ruleContext.hasErrors()) {
       return null;
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java b/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java
index 5989b628e4..170d323631 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/RuleContext.java
@@ -167,7 +167,7 @@ public final class RuleContext extends TargetContext
   private final ErrorReporter reporter;
   private final ImmutableBiMap<String, Class<? extends TransitiveInfoProvider>>
       skylarkProviderRegistry;
-  private final ToolchainContext toolchainContext;
+  @Nullable private final ToolchainContext toolchainContext;
 
   private ActionOwner actionOwner;
 
@@ -182,7 +182,8 @@ public final class RuleContext extends TargetContext
       ImmutableMap<Label, ConfigMatchingProvider> configConditions,
       Class<? extends BuildConfiguration.Fragment> universalFragment,
       String ruleClassNameForLogging,
-      ImmutableMap<String, Attribute> aspectAttributes) {
+      ImmutableMap<String, Attribute> aspectAttributes,
+      @Nullable ToolchainContext toolchainContext) {
     super(builder.env, builder.rule, builder.configuration, builder.prerequisiteMap.get(null),
         builder.visibility);
     this.rule = builder.rule;
@@ -198,8 +199,7 @@ public final class RuleContext extends TargetContext
     this.skylarkProviderRegistry = builder.skylarkProviderRegistry;
     this.hostConfiguration = builder.hostConfiguration;
     reporter = builder.reporter;
-    // TODO(katre): Populate the actual selected toolchains.
-    this.toolchainContext = new ToolchainContext(null);
+    this.toolchainContext = toolchainContext;
   }
 
   private ImmutableSet<String> getEnabledFeatures() {
@@ -1124,6 +1124,7 @@ public final class RuleContext extends TargetContext
     }
   }
 
+  @Nullable
   public ToolchainContext getToolchainContext() {
     return toolchainContext;
   }
@@ -1497,6 +1498,7 @@ public final class RuleContext extends TargetContext
     private ImmutableMap<String, Attribute> aspectAttributes;
     private ImmutableBiMap<String, Class<? extends TransitiveInfoProvider>> skylarkProviderRegistry;
     private ImmutableList<AspectDescriptor> aspectDescriptors;
+    private ToolchainContext toolchainContext;
 
     Builder(
         AnalysisEnvironment env,
@@ -1533,7 +1535,8 @@ public final class RuleContext extends TargetContext
           configConditions,
           universalFragment,
           getRuleClassNameForLogging(),
-          aspectAttributes != null ? aspectAttributes : ImmutableMap.<String, Attribute>of());
+          aspectAttributes != null ? aspectAttributes : ImmutableMap.<String, Attribute>of(),
+          toolchainContext);
     }
 
     private void validateAttributes(AttributeMap attributes) {
@@ -1593,6 +1596,12 @@ public final class RuleContext extends TargetContext
       return this;
     }
 
+    /** Sets the {@link ToolchainContext} used to access toolchains used by this rule. */
+    Builder setToolchainContext(ToolchainContext toolchainContext) {
+      this.toolchainContext = toolchainContext;
+      return this;
+    }
+
     private boolean validateFilesetEntry(FilesetEntry filesetEntry, ConfiguredTarget src) {
       if (src.getProvider(FilesetProvider.class) != null) {
         return true;
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 c4c9a9f36b..fc903dc2d6 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,6 +14,7 @@
 
 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.platform.ToolchainInfo;
 import com.google.devtools.build.lib.packages.ClassObjectConstructor;
@@ -23,15 +24,23 @@ import javax.annotation.Nullable;
 
 /** Contains toolchain-related information needed for a {@link RuleContext}. */
 public class ToolchainContext {
+  private final ImmutableList<ClassObjectConstructor.Key> requiredToolchains;
   private final ImmutableMap<ClassObjectConstructor.Key, ToolchainInfo> toolchains;
 
-  public ToolchainContext(@Nullable Map<ClassObjectConstructor.Key, ToolchainInfo> toolchains) {
+  public ToolchainContext(
+      ImmutableList<ClassObjectConstructor.Key> requiredToolchains,
+      @Nullable Map<ClassObjectConstructor.Key, ToolchainInfo> toolchains) {
+    this.requiredToolchains = requiredToolchains;
     this.toolchains =
         toolchains == null
             ? ImmutableMap.<ClassObjectConstructor.Key, ToolchainInfo>of()
             : ImmutableMap.copyOf(toolchains);
   }
 
+  public ImmutableList<ClassObjectConstructor.Key> getRequiredToolchains() {
+    return requiredToolchains;
+  }
+
   public SkylarkDict<ClassObjectConstructor.Key, ToolchainInfo> collectToolchains() {
     return SkylarkDict.<ClassObjectConstructor.Key, ToolchainInfo>copyOf(null, toolchains);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java b/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
index 1dbb131e06..860e3dff3d 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
@@ -28,9 +28,11 @@ import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.syntax.Type.LabelClass;
 import com.google.devtools.build.lib.syntax.Type.LabelVisitor;
 import com.google.devtools.build.lib.util.Preconditions;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -58,6 +60,7 @@ public final class AspectDefinition {
   private final RequiredProviders requiredProviders;
   private final RequiredProviders requiredProvidersForAspects;
   private final ImmutableMap<String, Attribute> attributes;
+  private final ImmutableList<ClassObjectConstructor.Key> requiredToolchains;
 
   /**
    * Which attributes aspect should propagate along:
@@ -74,13 +77,13 @@ public final class AspectDefinition {
     return advertisedProviders;
   }
 
-
   private AspectDefinition(
       AspectClass aspectClass,
       AdvertisedProviderSet advertisedProviders,
       RequiredProviders requiredProviders,
       RequiredProviders requiredAspectProviders,
       ImmutableMap<String, Attribute> attributes,
+      ImmutableList<ClassObjectConstructor.Key> requiredToolchains,
       @Nullable ImmutableSet<String> restrictToAttributes,
       @Nullable ConfigurationFragmentPolicy configurationFragmentPolicy,
       boolean applyToFiles) {
@@ -90,6 +93,7 @@ public final class AspectDefinition {
     this.requiredProvidersForAspects = requiredAspectProviders;
 
     this.attributes = attributes;
+    this.requiredToolchains = requiredToolchains;
     this.restrictToAttributes = restrictToAttributes;
     this.configurationFragmentPolicy = configurationFragmentPolicy;
     this.applyToFiles = applyToFiles;
@@ -112,6 +116,11 @@ public final class AspectDefinition {
     return attributes;
   }
 
+  /** Returns the required toolchains declared by this aspect. */
+  public ImmutableList<ClassObjectConstructor.Key> getRequiredToolchains() {
+    return requiredToolchains;
+  }
+
   /**
    * Returns {@link RequiredProviders} that a configured target must have so that
    * this aspect can be applied to it.
@@ -256,6 +265,7 @@ public final class AspectDefinition {
     private final ConfigurationFragmentPolicy.Builder configurationFragmentPolicy =
         new ConfigurationFragmentPolicy.Builder();
     private boolean applyToFiles = false;
+    private final List<ClassObjectConstructor.Key> requiredToolchains = new ArrayList<>();
 
     public Builder(AspectClass aspectClass) {
       this.aspectClass = aspectClass;
@@ -456,21 +466,25 @@ public final class AspectDefinition {
       return this;
     }
 
-
+    /** Adds the given toolchains as requirements for this aspect. */
+    public Builder addRequiredToolchains(List<ClassObjectConstructor.Key> requiredToolchains) {
+      this.requiredToolchains.addAll(requiredToolchains);
+      return this;
+    }
     /**
      * Builds the aspect definition.
      *
      * <p>The builder object is reusable afterwards.
      */
     public AspectDefinition build() {
-      return new AspectDefinition(aspectClass,
+      return new AspectDefinition(
+          aspectClass,
           advertisedProviders.build(),
           requiredProviders.build(),
           requiredAspectProviders.build(),
           ImmutableMap.copyOf(attributes),
-          propagateAlongAttributes == null
-              ? null
-              : ImmutableSet.copyOf(propagateAlongAttributes),
+          ImmutableList.copyOf(requiredToolchains),
+          propagateAlongAttributes == null ? null : ImmutableSet.copyOf(propagateAlongAttributes),
           configurationFragmentPolicy.build(),
           applyToFiles);
     }
diff --git a/src/main/java/com/google/devtools/build/lib/packages/SkylarkAspect.java b/src/main/java/com/google/devtools/build/lib/packages/SkylarkAspect.java
index 58a58015c0..895afba506 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/SkylarkAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/SkylarkAspect.java
@@ -47,6 +47,8 @@ public class SkylarkAspect implements SkylarkExportable {
   private final ImmutableSet<String> paramAttributes;
   private final ImmutableSet<String> fragments;
   private final ImmutableSet<String> hostFragments;
+  private final ImmutableList<ClassObjectConstructor.Key> requiredToolchains;
+
   private final Environment funcallEnv;
   private SkylarkAspectClass aspectClass;
 
@@ -59,6 +61,7 @@ public class SkylarkAspect implements SkylarkExportable {
       ImmutableSet<String> paramAttributes,
       ImmutableSet<String> fragments,
       ImmutableSet<String> hostFragments,
+      ImmutableList<ClassObjectConstructor.Key> requiredToolchains,
       Environment funcallEnv) {
     this.implementation = implementation;
     this.attributeAspects = attributeAspects;
@@ -68,6 +71,7 @@ public class SkylarkAspect implements SkylarkExportable {
     this.paramAttributes = paramAttributes;
     this.fragments = fragments;
     this.hostFragments = hostFragments;
+    this.requiredToolchains = requiredToolchains;
     this.funcallEnv = funcallEnv;
   }
 
@@ -153,6 +157,7 @@ public class SkylarkAspect implements SkylarkExportable {
     builder.advertiseProvider(advertisedSkylarkProviders.build());
     builder.requiresConfigurationFragmentsBySkylarkModuleName(fragments);
     builder.requiresHostConfigurationFragmentsBySkylarkModuleName(hostFragments);
+    builder.addRequiredToolchains(requiredToolchains);
     return builder.build();
   }
 
@@ -195,4 +200,8 @@ public class SkylarkAspect implements SkylarkExportable {
       }
     };
   }
+
+  public ImmutableList<ClassObjectConstructor.Key> getRequiredToolchains() {
+    return requiredToolchains;
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleClassFunctions.java b/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleClassFunctions.java
index 083cb92e63..f11b74ac80 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleClassFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleClassFunctions.java
@@ -501,71 +501,97 @@ public class SkylarkRuleClassFunctions {
     }
   }
 
-
-  @SkylarkSignature(name = "aspect", doc =
-      "Creates a new aspect. The result of this function must be stored in a global value. "
-          + "Please see the <a href=\"../aspects.md\">introduction to Aspects</a> for more "
-          + "details.",
-      returnType = SkylarkAspect.class,
-      parameters = {
-          @Param(name = "implementation", type = BaseFunction.class,
-              doc = "the function implementing this aspect. Must have two parameters: "
-                  + "<a href=\"Target.html\">Target</a> (the target to which the aspect is "
-                  + "applied) and <a href=\"ctx.html\">ctx</a>. Attributes of the target are "
-                  + "available via ctx.rule field. The function is called during the analysis "
-                  + "phase for each application of an aspect to a target."
-          ),
-          @Param(name = "attr_aspects", type = SkylarkList.class, generic1 = String.class,
-              defaultValue = "[]",
-              doc = "List of attribute names.  The aspect propagates along dependencies specified "
-                  + "by attributes of a target with this name. The list can also contain a single "
-                  + "string '*': in that case aspect propagates along all dependencies of a target."
-          ),
-          @Param(name = "attrs", type = SkylarkDict.class, noneable = true, defaultValue = "None",
-              doc = "dictionary to declare all the attributes of the aspect.  "
-                  + "It maps from an attribute name to an attribute object "
-                  + "(see <a href=\"attr.html\">attr</a> module). "
-                  + "Aspect attributes are available to implementation function as fields of ctx "
-                  + "parameter. Implicit attributes starting with <code>_</code> must have default "
-                  + "values, and have type <code>label</code> or <code>label_list</code>. "
-                  + "Explicit attributes must have type <code>string</code>, and must use the "
-                  + "<code>values</code> restriction. If explicit attributes are present, the "
-                  + "aspect can only be used with rules that have attributes of the same name and "
-                  + "type, with valid values."
-          ),
-          @Param(name = "required_aspect_providers",
-              type = SkylarkList.class,
-              defaultValue = "[]",
-              // todo(dslomov): Document once it works.
-              doc = "<not available>"
-          ),
-          @Param(name = "provides",
-              type = SkylarkList.class,
-              defaultValue = "[]",
-              // todo(dslomov): Document once it works.
-              doc = "<not available>"
-          ),
-          @Param(
-              name = "fragments",
-              type = SkylarkList.class,
-              generic1 = String.class,
-              defaultValue = "[]",
-              doc =
-                  "List of names of configuration fragments that the aspect requires "
-                      + "in target configuration."
-          ),
-          @Param(
-              name = "host_fragments",
-              type = SkylarkList.class,
-              generic1 = String.class,
-              defaultValue = "[]",
-              doc =
-                  "List of names of configuration fragments that the aspect requires "
-                      + "in host configuration."
-          )
-      },
-      useEnvironment = true,
-      useAst = true
+  @SkylarkSignature(
+    name = "aspect",
+    doc =
+        "Creates a new aspect. The result of this function must be stored in a global value. "
+            + "Please see the <a href=\"../aspects.md\">introduction to Aspects</a> for more "
+            + "details.",
+    returnType = SkylarkAspect.class,
+    parameters = {
+      @Param(
+        name = "implementation",
+        type = BaseFunction.class,
+        doc =
+            "the function implementing this aspect. Must have two parameters: "
+                + "<a href=\"Target.html\">Target</a> (the target to which the aspect is "
+                + "applied) and <a href=\"ctx.html\">ctx</a>. Attributes of the target are "
+                + "available via ctx.rule field. The function is called during the analysis "
+                + "phase for each application of an aspect to a target."
+      ),
+      @Param(
+        name = "attr_aspects",
+        type = SkylarkList.class,
+        generic1 = String.class,
+        defaultValue = "[]",
+        doc =
+            "List of attribute names.  The aspect propagates along dependencies specified "
+                + "by attributes of a target with this name. The list can also contain a single "
+                + "string '*': in that case aspect propagates along all dependencies of a target."
+      ),
+      @Param(
+        name = "attrs",
+        type = SkylarkDict.class,
+        noneable = true,
+        defaultValue = "None",
+        doc =
+            "dictionary to declare all the attributes of the aspect.  "
+                + "It maps from an attribute name to an attribute object "
+                + "(see <a href=\"attr.html\">attr</a> module). "
+                + "Aspect attributes are available to implementation function as fields of ctx "
+                + "parameter. Implicit attributes starting with <code>_</code> must have default "
+                + "values, and have type <code>label</code> or <code>label_list</code>. "
+                + "Explicit attributes must have type <code>string</code>, and must use the "
+                + "<code>values</code> restriction. If explicit attributes are present, the "
+                + "aspect can only be used with rules that have attributes of the same name and "
+                + "type, with valid values."
+      ),
+      @Param(
+        name = "required_aspect_providers",
+        type = SkylarkList.class,
+        defaultValue = "[]",
+        // todo(dslomov): Document once it works.
+        doc = "<not available>"
+      ),
+      @Param(
+        name = "provides",
+        type = SkylarkList.class,
+        defaultValue = "[]",
+        // todo(dslomov): Document once it works.
+        doc = "<not available>"
+      ),
+      @Param(
+        name = "fragments",
+        type = SkylarkList.class,
+        generic1 = String.class,
+        defaultValue = "[]",
+        doc =
+            "List of names of configuration fragments that the aspect requires "
+                + "in target configuration."
+      ),
+      @Param(
+        name = "host_fragments",
+        type = SkylarkList.class,
+        generic1 = String.class,
+        defaultValue = "[]",
+        doc =
+            "List of names of configuration fragments that the aspect requires "
+                + "in host configuration."
+      ),
+      @Param(
+        name = "toolchains",
+        type = SkylarkList.class,
+        generic1 = ToolchainConstructor.class,
+        defaultValue = "[]",
+        doc =
+            "<i>(Experimental)</i><br/><br/>"
+                + "If set, the set of toolchains this rule requires. Toolchains will be "
+                + "found by checking the current platform, and provided to the rule "
+                + "implementation via <code>ctx.toolchain</code>."
+      )
+    },
+    useEnvironment = true,
+    useAst = true
   )
   private static final BuiltinFunction aspect =
       new BuiltinFunction("aspect") {
@@ -577,6 +603,7 @@ public class SkylarkRuleClassFunctions {
             SkylarkList providesArg,
             SkylarkList fragments,
             SkylarkList hostFragments,
+            SkylarkList<ToolchainConstructor> toolchains,
             FuncallExpression ast,
             Environment funcallEnv)
             throws EvalException {
@@ -586,14 +613,12 @@ public class SkylarkRuleClassFunctions {
 
             if (attrName.equals("*") && attributeAspects.size() != 1) {
               throw new EvalException(
-                  ast.getLocation(),
-                  "'*' must be the only string in 'attr_aspects' list"
-              );
+                  ast.getLocation(), "'*' must be the only string in 'attr_aspects' list");
             }
 
             if (!attrName.startsWith("_")) {
               attrAspects.add(attrName);
-            } else  {
+            } else {
               // Implicit attribute names mean either implicit or late-bound attributes
               // (``$attr`` or ``:attr``). Depend on both.
               attrAspects.add(
@@ -613,7 +638,7 @@ public class SkylarkRuleClassFunctions {
             Attribute attribute = nameDescriptorPair.second.build(nameDescriptorPair.first);
             if (attribute.getType() == Type.STRING
                 && ((String) attribute.getDefaultValue(null)).isEmpty()) {
-              hasDefault = false;  // isValueSet() is always true for attr.string.
+              hasDefault = false; // isValueSet() is always true for attr.string.
             }
             if (!Attribute.isImplicit(nativeName)) {
               if (!attribute.checkAllowedValues() || attribute.getType() != Type.STRING) {
@@ -634,11 +659,10 @@ public class SkylarkRuleClassFunctions {
                       ast.getLocation(),
                       String.format(
                           "Aspect parameter attribute '%s' has a bad default value: %s",
-                          nativeName,
-                          allowed.getErrorReason(defaultVal)));
+                          nativeName, allowed.getErrorReason(defaultVal)));
                 }
               }
-            } else if (!hasDefault) {  // Implicit attribute
+            } else if (!hasDefault) { // Implicit attribute
               String skylarkName = "_" + nativeName.substring(1);
               throw new EvalException(
                   ast.getLocation(),
@@ -649,29 +673,38 @@ public class SkylarkRuleClassFunctions {
 
           for (Object o : providesArg) {
             if (!SkylarkAttr.isProvider(o)) {
-              throw new EvalException(ast.getLocation(),
-                  String.format("Illegal argument: element in 'provides' is of unexpected type. "
-                      + "Should be list of providers, but got %s. ",
+              throw new EvalException(
+                  ast.getLocation(),
+                  String.format(
+                      "Illegal argument: element in 'provides' is of unexpected type. "
+                          + "Should be list of providers, but got %s. ",
                       EvalUtils.getDataTypeName(o, true)));
             }
           }
 
+          // Collect the required toolchain keys.
+          ImmutableList.Builder<ClassObjectConstructor.Key> requiredToolchains =
+              new ImmutableList.Builder<>();
+          for (ToolchainConstructor toolchain :
+              toolchains.getContents(ToolchainConstructor.class, "toolchains")) {
+            requiredToolchains.add(toolchain.getKey());
+          }
+
           return new SkylarkAspect(
               implementation,
               attrAspects.build(),
               attributes.build(),
-              SkylarkAttr.buildProviderPredicate(requiredAspectProvidersArg,
-                  "required_aspect_providers", ast.getLocation()
-              ),
+              SkylarkAttr.buildProviderPredicate(
+                  requiredAspectProvidersArg, "required_aspect_providers", ast.getLocation()),
               SkylarkAttr.getSkylarkProviderIdentifiers(providesArg, ast.getLocation()),
               requiredParams.build(),
               ImmutableSet.copyOf(fragments.getContents(String.class, "fragments")),
               ImmutableSet.copyOf(hostFragments.getContents(String.class, "host_fragments")),
+              requiredToolchains.build(),
               funcallEnv);
         }
       };
 
-
   /** The implementation for the magic function "rule" that creates Skylark rule classes */
   public static final class RuleFunction extends BaseFunction implements SkylarkExportable {
     private RuleClass.Builder builder;
diff --git a/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleContext.java b/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleContext.java
index 90c121bb7d..acd7e31332 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleContext.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/SkylarkRuleContext.java
@@ -262,7 +262,6 @@ public final class SkylarkRuleContext implements SkylarkValue {
               this, attributes, ruleContext, attributeValueExtractorForRule(ruleContext));
       this.splitAttributes = buildSplitAttributeInfo(attributes, ruleContext);
       this.ruleAttributesCollection = null;
-      toolchains = ruleContext.getToolchainContext().collectToolchains();
     } else { // ASPECT
       this.artifactsLabelMap = ImmutableMap.of();
       this.outputsObject = null;
@@ -279,11 +278,13 @@ public final class SkylarkRuleContext implements SkylarkValue {
               ruleContext.getRule().getAttributes(),
               ruleContext,
               attributeValueExtractorForRule(ruleContext));
-      // TODO(katre): Collect toolchains for aspects.
-      toolchains = null;
     }
 
     makeVariables = ruleContext.getConfigurationMakeVariableContext().collectMakeVariables();
+    toolchains =
+        ruleContext.getToolchainContext() == null
+            ? SkylarkDict.<ClassObjectConstructor.Key, ToolchainInfo>of(null)
+            : ruleContext.getToolchainContext().collectToolchains();
   }
 
   /**