Separate main aspect prerequisites from rule and base aspects prerequ…

…isites

add `separate_aspect_deps` flag to guard this fix, when the flag is enabled the dependencies of the main aspect will be separated from those of the target and the base aspects. `ctx.attr.{attr_name}` will always get the attribute value from the main aspect and `ctx.rule.attr.{attr_name}` will get the value from the rule if it has an attribute with that name or from the base aspects attributes (first one in the aspects path wins).

PiperOrigin-RevId: 580292823
Change-Id: I719f8775c0c2aa582aed600387cb6b321cf5b639

src/main/java/com/google/devtools/build/lib/analysis/AspectContext.java

@@ -0,0 +1,278 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.analysis;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Streams;
+import com.google.devtools.build.lib.collect.ImmutableSortedKeyListMultimap;
+import com.google.devtools.build.lib.packages.Aspect;
+import com.google.devtools.build.lib.packages.AspectDescriptor;
+import com.google.devtools.build.lib.packages.Attribute;
+import com.google.devtools.build.lib.packages.AttributeMap;
+import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
+import com.google.devtools.build.lib.skyframe.ConfiguredTargetAndData;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/** Extends {@link RuleContext} to provide all data available during the analysis of an aspect. */
+public final class AspectContext extends RuleContext {
+
+  /**
+   * A list of all aspects applied to the target.
+   *
+   * <p>The last aspect in the list is the main aspect that this context is for.
+   */
+  private final ImmutableList<Aspect> aspects;
+
+  private final ImmutableList<AspectDescriptor> aspectDescriptors;
+
+  // If (separate_aspect_deps) flag is disabled, holds all merged prerequisites otherwise it will
+  // only have the main aspect's prerequisites.
+  private final PrerequisitesCollection mainAspectPrerequisites;
+
+  AspectContext(
+      RuleContext.Builder builder,
+      AspectAwareAttributeMapper aspectAwareAttributeMapper,
+      PrerequisitesCollection ruleAndBaseAspectsPrerequisites,
+      PrerequisitesCollection mainAspectPrerequisites,
+      ExecGroupCollection execGroupCollection) {
+    super(
+        builder, aspectAwareAttributeMapper, ruleAndBaseAspectsPrerequisites, execGroupCollection);
+
+    this.aspects = builder.getAspects();
+    this.aspectDescriptors = aspects.stream().map(Aspect::getDescriptor).collect(toImmutableList());
+    this.mainAspectPrerequisites = mainAspectPrerequisites;
+  }
+
+  /**
+   * Merge the attributes of the aspects in the aspects path.
+   *
+   * <p>For attributes with the same name, the one that is first encountered takes precedence.
+   */
+  private static ImmutableMap<String, Attribute> mergeAspectsAttributes(
+      ImmutableList<Aspect> aspects) {
+    if (aspects.isEmpty()) {
+      return ImmutableMap.of();
+    } else if (aspects.size() == 1) {
+      return aspects.get(0).getDefinition().getAttributes();
+    } else {
+      LinkedHashMap<String, Attribute> aspectAttributes = new LinkedHashMap<>();
+      for (Aspect aspect : aspects) {
+        ImmutableMap<String, Attribute> currentAttributes = aspect.getDefinition().getAttributes();
+        for (Map.Entry<String, Attribute> kv : currentAttributes.entrySet()) {
+          aspectAttributes.putIfAbsent(kv.getKey(), kv.getValue());
+        }
+      }
+      return ImmutableMap.copyOf(aspectAttributes);
+    }
+  }
+
+  static AspectContext create(
+      Builder builder,
+      AttributeMap ruleAttributes,
+      ImmutableListMultimap<DependencyKind, ConfiguredTargetAndData> targetsMap,
+      ExecGroupCollection execGroupCollection) {
+
+    if (builder.separateAspectDeps()) {
+      return createAspectContextWithSeparatedPrerequisites(
+          builder, ruleAttributes, targetsMap, execGroupCollection);
+    } else {
+      return createAspectContextWithMergedPrerequisites(
+          builder, ruleAttributes, targetsMap, execGroupCollection);
+    }
+  }
+
+  /**
+   * Merges the rule prerequisites with the aspects prerequisites giving precedence to aspects
+   * prerequisites if any.
+   *
+   * <p>TODO(b/293304543) merging prerequisites should be not needed after completing the solution
+   * to isolate main aspect dependencies from the underlying rule and aspects dependencies.
+   */
+  private static AspectContext createAspectContextWithMergedPrerequisites(
+      RuleContext.Builder builder,
+      AttributeMap ruleAttributes,
+      ImmutableListMultimap<DependencyKind, ConfiguredTargetAndData> targetsMap,
+      ExecGroupCollection execGroupCollection) {
+
+    ImmutableSortedKeyListMultimap.Builder<String, ConfiguredTargetAndData>
+        attributeNameToTargetsMap = ImmutableSortedKeyListMultimap.builder();
+    HashSet<String> processedAttributes = new HashSet<>();
+
+    // add aspects prerequisites
+    for (Aspect aspect : builder.getAspects()) {
+      for (Attribute attribute : aspect.getDefinition().getAttributes().values()) {
+        DependencyKind key =
+            DependencyKind.AttributeDependencyKind.forAspect(attribute, aspect.getAspectClass());
+        if (targetsMap.containsKey(key)) {
+          if (processedAttributes.add(attribute.getName())) {
+            attributeNameToTargetsMap.putAll(attribute.getName(), targetsMap.get(key));
+          }
+        }
+      }
+    }
+
+    // add rule prerequisites
+    for (var entry : targetsMap.asMap().entrySet()) {
+      if (entry.getKey().getOwningAspect() == null) {
+        if (!processedAttributes.contains(entry.getKey().getAttribute().getName())) {
+          attributeNameToTargetsMap.putAll(
+              entry.getKey().getAttribute().getName(), entry.getValue());
+        }
+      }
+    }
+    AspectAwareAttributeMapper ruleAndAspectsAttributes =
+        new AspectAwareAttributeMapper(
+            ruleAttributes, mergeAspectsAttributes(builder.getAspects()));
+    PrerequisitesCollection mergedPrerequisitesCollection =
+        new PrerequisitesCollection(
+            attributeNameToTargetsMap.build(),
+            ruleAndAspectsAttributes,
+            builder.getErrorConsumer(),
+            builder.getRule(),
+            builder.getRuleClassNameForLogging());
+    return new AspectContext(
+        builder,
+        ruleAndAspectsAttributes,
+        /* ruleAndBaseAspectsPrerequisites= */ mergedPrerequisitesCollection,
+        /* mainAspectPrerequisites= */ mergedPrerequisitesCollection,
+        execGroupCollection);
+  }
+
+  /**
+   * Create prerequisites collection for aspect evaluation separating the main aspect prerequisites
+   * from the underlying rule and base aspects prerequisites.
+   */
+  private static AspectContext createAspectContextWithSeparatedPrerequisites(
+      RuleContext.Builder builder,
+      AttributeMap ruleAttributes,
+      ImmutableListMultimap<DependencyKind, ConfiguredTargetAndData> prerequisitesMap,
+      ExecGroupCollection execGroupCollection) {
+    ImmutableSortedKeyListMultimap.Builder<String, ConfiguredTargetAndData>
+        mainAspectPrerequisites = ImmutableSortedKeyListMultimap.builder();
+    ImmutableSortedKeyListMultimap.Builder<String, ConfiguredTargetAndData>
+        ruleAndBaseAspectsPrerequisites = ImmutableSortedKeyListMultimap.builder();
+
+    Aspect mainAspect = Iterables.getLast(builder.getAspects());
+
+    for (Map.Entry<DependencyKind, Collection<ConfiguredTargetAndData>> entry :
+        prerequisitesMap.asMap().entrySet()) {
+      String attributeName = entry.getKey().getAttribute().getName();
+
+      if (mainAspect.getAspectClass().equals(entry.getKey().getOwningAspect())) {
+        mainAspectPrerequisites.putAll(attributeName, entry.getValue());
+      } else {
+        ruleAndBaseAspectsPrerequisites.putAll(attributeName, entry.getValue());
+      }
+    }
+
+    return new AspectContext(
+        builder,
+        new AspectAwareAttributeMapper(
+            ruleAttributes, mergeAspectsAttributes(builder.getAspects())),
+        new PrerequisitesCollection(
+            ruleAndBaseAspectsPrerequisites.build(),
+            mergeRuleAndBaseAspectsAttributes(ruleAttributes, builder.getAspects()),
+            builder.getErrorConsumer(),
+            builder.getRule(),
+            builder.getRuleClassNameForLogging()),
+        new PrerequisitesCollection(
+            mainAspectPrerequisites.build(),
+            mainAspect.getDefinition().getAttributes(),
+            builder.getErrorConsumer(),
+            builder.getRule(),
+            builder.getRuleClassNameForLogging()),
+        execGroupCollection);
+  }
+
+  private static AspectAwareAttributeMapper mergeRuleAndBaseAspectsAttributes(
+      AttributeMap ruleAttributes, ImmutableList<Aspect> aspects) {
+    LinkedHashMap<String, Attribute> mergedBaseAspectsAttributes = new LinkedHashMap<>();
+    for (int i = 0; i < aspects.size() - 1; i++) {
+      for (Attribute attribute : aspects.get(i).getDefinition().getAttributes().values()) {
+        mergedBaseAspectsAttributes.putIfAbsent(attribute.getName(), attribute);
+      }
+    }
+    return new AspectAwareAttributeMapper(
+        ruleAttributes, ImmutableMap.copyOf(mergedBaseAspectsAttributes));
+  }
+
+  @Override
+  PrerequisitesCollection getOwningPrerequisitesCollection(String attributeName) {
+    if (mainAspectPrerequisites.has(attributeName)) {
+      return mainAspectPrerequisites;
+    }
+    return getRulePrerequisitesCollection();
+  }
+
+  public PrerequisitesCollection getMainAspectPrerequisitesCollection() {
+    return mainAspectPrerequisites;
+  }
+
+  @Override
+  public ImmutableList<Aspect> getAspects() {
+    return aspects;
+  }
+
+  /**
+   * Return the main aspect of this context.
+   *
+   * <p>It is the last aspect in the list of aspects applied to a target; all other aspects are the
+   * ones main aspect sees as specified by its "required_aspect_providers").
+   */
+  @Override
+  public Aspect getMainAspect() {
+    return Iterables.getLast(aspects);
+  }
+
+  /** All aspects applied to the rule. */
+  @Override
+  public ImmutableList<AspectDescriptor> getAspectDescriptors() {
+    return aspectDescriptors;
+  }
+
+  @Override
+  public boolean useAutoExecGroups() {
+    ImmutableMap<String, Attribute> aspectAttributes =
+        getMainAspect().getDefinition().getAttributes();
+    if (aspectAttributes.containsKey("$use_auto_exec_groups")) {
+      return (boolean) aspectAttributes.get("$use_auto_exec_groups").getDefaultValueUnchecked();
+    } else {
+      return getConfiguration().useAutoExecGroups();
+    }
+  }
+
+  @Override
+  public ImmutableList<? extends TransitiveInfoCollection> getAllPrerequisites() {
+    boolean separateAspectDeps =
+        getAnalysisEnvironment()
+            .getStarlarkSemantics()
+            .getBool(BuildLanguageOptions.SEPARATE_ASPECT_DEPS);
+    if (separateAspectDeps) {
+      return Streams.concat(
+              mainAspectPrerequisites.getAllPrerequisites().stream(),
+              getRulePrerequisitesCollection().getAllPrerequisites().stream())
+          .collect(toImmutableList());
+    }
+    return super.getAllPrerequisites();
+  }
+}