1 files changed, 51 insertions, 22 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/Attribute.java b/src/main/java/com/google/devtools/build/lib/packages/Attribute.java
index ffe0369938..bf4a0a744e 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/Attribute.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/Attribute.java
@@ -61,21 +61,41 @@ public final class Attribute implements Comparable<Attribute> {
 
   public static final Predicate<RuleClass> NO_RULE = Predicates.alwaysFalse();
 
-  private static final class RuleAspect {
-    private final AspectClass aspectFactory;
-    private final Function<Rule, AspectParameters> parametersExtractor;
+  private abstract static class RuleAspect<C extends AspectClass> {
+    protected final C aspectClass;
+    protected final Function<Rule, AspectParameters> parametersExtractor;
 
-    RuleAspect(AspectClass aspectFactory, Function<Rule, AspectParameters> parametersExtractor) {
-      this.aspectFactory = aspectFactory;
+    protected RuleAspect(C aspectClass, Function<Rule, AspectParameters> parametersExtractor) {
+      this.aspectClass = aspectClass;
       this.parametersExtractor = parametersExtractor;
     }
 
-    AspectClass getAspectFactory() {
-      return aspectFactory;
+    public abstract Aspect getAspect(Rule rule);
+  }
+
+  private static class NativeRuleAspect extends RuleAspect<NativeAspectClass<?>> {
+    public NativeRuleAspect(NativeAspectClass<?> aspectClass,
+        Function<Rule, AspectParameters> parametersExtractor) {
+      super(aspectClass, parametersExtractor);
+    }
+
+    @Override
+    public Aspect getAspect(Rule rule) {
+      return Aspect.forNative(aspectClass, parametersExtractor.apply(rule));
+    }
+  }
+
+  private static class SkylarkRuleAspect extends RuleAspect<SkylarkAspectClass> {
+    public SkylarkRuleAspect(SkylarkAspectClass aspectClass) {
+      super(aspectClass, NO_PARAMETERS);
     }
-    
-    Function<Rule, AspectParameters> getParametersExtractor() {
-      return parametersExtractor;
+
+    @Override
+    public Aspect getAspect(Rule rule) {
+      return Aspect.forSkylark(
+          aspectClass,
+          aspectClass.getDefinition(),
+          parametersExtractor.apply(rule));
     }
   }
 
@@ -281,6 +301,15 @@ public final class Attribute implements Comparable<Attribute> {
     }
   }
 
+  private static final Function<Rule, AspectParameters> NO_PARAMETERS =
+      new Function<Rule, AspectParameters>() {
+        @Override
+        public AspectParameters apply(Rule input) {
+          return AspectParameters.EMPTY;
+        }
+      };
+
+
   /**
    * Creates a new attribute builder.
    *
@@ -301,6 +330,7 @@ public final class Attribute implements Comparable<Attribute> {
    * already undocumented based on its name cannot be marked as undocumented.
    */
   public static class Builder <TYPE> {
+
     private String name;
     private final Type<TYPE> type;
     private Transition configTransition = ConfigurationTransition.NONE;
@@ -724,13 +754,7 @@ public final class Attribute implements Comparable<Attribute> {
      * through this attribute.
      */
     public <T extends NativeAspectFactory> Builder<TYPE> aspect(Class<T> aspect) {
-      Function<Rule, AspectParameters> noParameters = new Function<Rule, AspectParameters>() {
-        @Override
-        public AspectParameters apply(Rule input) {
-          return AspectParameters.EMPTY;
-        }
-      };
-      return this.aspect(aspect, noParameters);
+      return this.aspect(aspect, NO_PARAMETERS);
     }
 
     /**
@@ -750,8 +774,9 @@ public final class Attribute implements Comparable<Attribute> {
      *
      * @param evaluator function that extracts aspect parameters from rule.
      */
-    public Builder<TYPE> aspect(AspectClass aspect, Function<Rule, AspectParameters> evaluator) {
-      this.aspects.add(new RuleAspect(aspect, evaluator));
+    public Builder<TYPE> aspect(
+        NativeAspectClass<?> aspect, Function<Rule, AspectParameters> evaluator) {
+      this.aspects.add(new NativeRuleAspect(aspect, evaluator));
       return this;
     }
 
@@ -759,7 +784,7 @@ public final class Attribute implements Comparable<Attribute> {
      * Asserts that a particular parameterized aspect probably needs to be computed for all direct
      * dependencies through this attribute.
      */
-    public Builder<TYPE> aspect(AspectClass aspect) {
+    public Builder<TYPE> aspect(NativeAspectClass<?> aspect) {
       Function<Rule, AspectParameters> noParameters =
           new Function<Rule, AspectParameters>() {
             @Override
@@ -770,6 +795,11 @@ public final class Attribute implements Comparable<Attribute> {
       return this.aspect(aspect, noParameters);
     }
 
+    public Builder<TYPE> aspect(SkylarkAspectClass aspectClass) {
+      this.aspects.add(new SkylarkRuleAspect(aspectClass));
+      return this;
+    }
+
     /**
      * Sets the predicate-like edge validity checker.
      */
@@ -1400,8 +1430,7 @@ public final class Attribute implements Comparable<Attribute> {
   public ImmutableList<Aspect> getAspects(Rule rule) {
     ImmutableList.Builder<Aspect> builder = ImmutableList.builder();
     for (RuleAspect aspect : aspects) {
-      builder.add(
-          new Aspect(aspect.getAspectFactory(), aspect.getParametersExtractor().apply(rule)));
+      builder.add(aspect.getAspect(rule));
     }
     return builder.build();
   }