Allow aspects to propagate to all attributes.

--
MOS_MIGRATED_REVID=134378592

2 files changed, 57 insertions, 9 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java b/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java
index 10703ba8eb..895b6be4a0 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/DependencyResolver.java
@@ -548,7 +548,7 @@ public abstract class DependencyResolver {
     aspectCandidates.addAll(attribute.getAspects(originalRule));
     if (aspect != null) {
       for (AspectClass aspectClass :
-          aspect.getDefinition().getAttributeAspects().get(attribute.getName())) {
+          aspect.getDefinition().getAttributeAspects(attribute)) {
         if (aspectClass.equals(aspect.getAspectClass())) {
           aspectCandidates.add(aspect);
         } else if (aspectClass instanceof NativeAspectClass) {
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 20c2a6ab28..85843b879a 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
@@ -14,6 +14,8 @@
 
 package com.google.devtools.build.lib.packages;
 
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
@@ -25,14 +27,12 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.ConfigurationFragmentPolicy.MissingFragmentPolicy;
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.Preconditions;
-
 import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-
 import javax.annotation.Nullable;
 
 /**
@@ -59,14 +59,18 @@ public final class AspectDefinition {
   private final ImmutableSet<Class<?>> requiredProviders;
   private final ImmutableSet<String> requiredProviderNames;
   private final ImmutableMap<String, Attribute> attributes;
-  private final ImmutableMultimap<String, AspectClass> attributeAspects;
+  private final PropagationFunction attributeAspects;
   @Nullable private final ConfigurationFragmentPolicy configurationFragmentPolicy;
 
+  private interface PropagationFunction {
+    ImmutableCollection<AspectClass> propagate(Attribute attribute);
+  }
+
   private AspectDefinition(
       String name,
       ImmutableSet<Class<?>> requiredProviders,
       ImmutableMap<String, Attribute> attributes,
-      ImmutableMultimap<String, AspectClass> attributeAspects,
+      PropagationFunction attributeAspects,
       @Nullable ConfigurationFragmentPolicy configurationFragmentPolicy) {
     this.name = name;
     this.requiredProviders = requiredProviders;
@@ -119,10 +123,10 @@ public final class AspectDefinition {
   }
 
   /**
-   * Returns the attribute -&gt; set of required aspects map.
+   * Returns the set of required aspects for a given atribute.
    */
-  public ImmutableMultimap<String, AspectClass> getAttributeAspects() {
-    return attributeAspects;
+  public ImmutableCollection<AspectClass> getAttributeAspects(Attribute attribute) {
+    return attributeAspects.propagate(attribute);
   }
 
   /**
@@ -226,6 +230,7 @@ public final class AspectDefinition {
     private final Map<String, Attribute> attributes = new LinkedHashMap<>();
     private final Set<Class<?>> requiredProviders = new LinkedHashSet<>();
     private final Multimap<String, AspectClass> attributeAspects = LinkedHashMultimap.create();
+    private ImmutableCollection<AspectClass> allAttributesAspects = null;
     private final ConfigurationFragmentPolicy.Builder configurationFragmentPolicy =
         new ConfigurationFragmentPolicy.Builder();
 
@@ -265,12 +270,23 @@ public final class AspectDefinition {
      */
     public final Builder attributeAspect(String attribute, AspectClass aspectClass) {
       Preconditions.checkNotNull(attribute);
+      Preconditions.checkState(this.allAttributesAspects == null,
+          "Specify either aspects for all attributes, or for specific attributes, not both");
 
       this.attributeAspects.put(attribute, Preconditions.checkNotNull(aspectClass));
 
       return this;
     }
 
+    public final Builder allAttributesAspect(AspectClass... aspectClasses) {
+      Preconditions.checkState(this.attributeAspects.isEmpty(),
+          "Specify either aspects for all attributes, or for specific attributes, not both");
+      Preconditions.checkState(this.allAttributesAspects == null,
+          "Aspects for all attributes must only be specified once");
+      this.allAttributesAspects = ImmutableList.copyOf(aspectClasses);
+      return this;
+    }
+
     /**
      * Adds an attribute to the aspect.
      *
@@ -361,6 +377,35 @@ public final class AspectDefinition {
       return this;
     }
 
+    @Immutable
+    private static final class AllAttributesPropagationFunction implements PropagationFunction {
+      private final ImmutableCollection<AspectClass> aspects;
+
+      private AllAttributesPropagationFunction(ImmutableCollection<AspectClass> aspects) {
+        this.aspects = aspects;
+      }
+
+      @Override
+      public ImmutableCollection<AspectClass> propagate(Attribute attribute) {
+        return aspects;
+      }
+    }
+
+    @Immutable
+    private static final class PerAttributePropagationFunction implements PropagationFunction {
+      ImmutableSetMultimap<String, AspectClass> aspects;
+
+      public PerAttributePropagationFunction(
+          ImmutableSetMultimap<String, AspectClass> aspects) {
+        this.aspects = aspects;
+      }
+
+      @Override
+      public ImmutableCollection<AspectClass> propagate(Attribute attribute) {
+        return aspects.get(attribute.getName());
+      }
+    }
+
     /**
      * Builds the aspect definition.
      *
@@ -368,7 +413,10 @@ public final class AspectDefinition {
      */
     public AspectDefinition build() {
       return new AspectDefinition(name, ImmutableSet.copyOf(requiredProviders),
-          ImmutableMap.copyOf(attributes), ImmutableSetMultimap.copyOf(attributeAspects),
+          ImmutableMap.copyOf(attributes),
+          allAttributesAspects != null
+              ? new AllAttributesPropagationFunction(allAttributesAspects)
+              : new PerAttributePropagationFunction(ImmutableSetMultimap.copyOf(attributeAspects)),
           configurationFragmentPolicy.build());
     }
   }