src/main/java/com/google/devtools/build/lib/analysis/AspectContext.java - bazel - Git at Google

 // 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 static com.google.common.collect.ImmutableSet.toImmutableSet;

 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.google.devtools.build.lib.cmdline.Label;
 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.skyframe.ConfiguredTargetAndData;
 import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import javax.annotation.Nullable;

 /** 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;

   private final PrerequisitesCollection mainAspectPrerequisites;

   /**
    * The toolchain contexts for the base target.
    *
    * <p>It only contains the providers created by the aspects that propagate to the toolchains.
    */
   @Nullable
   private final ToolchainCollection<AspectBaseTargetResolvedToolchainContext>
       baseTargetToolchainContexts;

   /** Whether the target uses auto exec groups. */
   private final boolean targetUsesAutoExecGroups;

   AspectContext(
       RuleContext.Builder builder,
       AspectAwareAttributeMapper aspectAwareAttributeMapper,
       PrerequisitesCollection ruleAndBaseAspectsPrerequisites,
       PrerequisitesCollection mainAspectPrerequisites,
       ExecGroupCollection execGroupCollection,
       @Nullable
           ToolchainCollection<AspectBaseTargetResolvedToolchainContext> baseTargetToolchainContexts,
       boolean targetUsesAutoExecGroups) {
     super(
         builder, aspectAwareAttributeMapper, ruleAndBaseAspectsPrerequisites, execGroupCollection);

     this.aspects = builder.getAspects();
     this.aspectDescriptors = aspects.stream().map(Aspect::getDescriptor).collect(toImmutableList());
     this.mainAspectPrerequisites = mainAspectPrerequisites;
     this.baseTargetToolchainContexts = baseTargetToolchainContexts;
     this.targetUsesAutoExecGroups = targetUsesAutoExecGroups;
   }

   /**
    * Returns the toolchain contexts for the base target. Can be null if no aspect in the {@code
    * aspects} path propagate to the toolchains.
    */
   @Nullable
   public ToolchainCollection<AspectBaseTargetResolvedToolchainContext>
       getBaseTargetToolchainContexts() {
     return baseTargetToolchainContexts;
   }

   /** Returns the labels of default the toolchain types that aspects have propagated. */
   public ImmutableSet<Label> getRequestedToolchainTypesLabels() {
     if (targetUsesAutoExecGroups) {
       return baseTargetToolchainContexts.contextMap().entrySet().stream()
           .filter(e -> isAutomaticExecGroup(e.getKey()))
           .flatMap(e -> e.getValue().requestedToolchainTypeLabels().keySet().stream())
           .collect(toImmutableSet());
     } else {
       return baseTargetToolchainContexts
           .getDefaultToolchainContext()
           .requestedToolchainTypeLabels()
           .keySet();
     }
   }

   /**
    * Returns the toolchain data for the given type, or {@code null} if the toolchain type was not
    * required in this context.
    */
   @Nullable
   public AspectBaseTargetResolvedToolchainContext.ToolchainAspectsProviders getToolchainTarget(
       Label toolchainType) {
     var execGroupContext = baseTargetToolchainContexts.getDefaultToolchainContext();
     if (targetUsesAutoExecGroups) {
       execGroupContext =
           baseTargetToolchainContexts.contextMap().entrySet().stream()
               .filter(
                   e ->
                       isAutomaticExecGroup(e.getKey())
                           && e.getValue().requestedToolchainTypeLabels().containsKey(toolchainType))
               .findFirst()
               .map(e -> e.getValue())
               .orElse(null);
       if (execGroupContext == null) {
         return null;
       }
     }
     return execGroupContext
         .getToolchains()
         .get(execGroupContext.requestedToolchainTypeLabels().get(toolchainType));
   }

   /**
    * 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,
       ToolchainCollection<AspectBaseTargetResolvedToolchainContext> baseTargetToolchainContexts) {
     return createAspectContextWithSeparatedPrerequisites(
         builder, ruleAttributes, targetsMap, execGroupCollection, baseTargetToolchainContexts);
   }

   /**
    * 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 ruleContextBuilder,
       AttributeMap ruleAttributes,
       ImmutableListMultimap<DependencyKind, ConfiguredTargetAndData> prerequisitesMap,
       ExecGroupCollection execGroupCollection,
       ToolchainCollection<AspectBaseTargetResolvedToolchainContext> baseTargetToolchainContexts) {
     ImmutableSortedKeyListMultimap.Builder<String, ConfiguredTargetAndData>
         mainAspectPrerequisites = ImmutableSortedKeyListMultimap.builder();
     ImmutableSortedKeyListMultimap.Builder<String, ConfiguredTargetAndData>
         ruleAndBaseAspectsPrerequisites = ImmutableSortedKeyListMultimap.builder();

     Aspect mainAspect = Iterables.getLast(ruleContextBuilder.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());
       }
     }

     boolean targetUsesAutoExecGroups =
         ruleContextBuilder
             .getRule()
             .getRuleClassObject()
             .getAutoExecGroupsMode()
             .isEnabled(ruleAttributes, ruleContextBuilder.getConfiguration().useAutoExecGroups());

     return new AspectContext(
         ruleContextBuilder,
         new AspectAwareAttributeMapper(
             ruleAttributes, mergeAspectsAttributes(ruleContextBuilder.getAspects())),
         new PrerequisitesCollection(
             ruleAndBaseAspectsPrerequisites.build(),
             mergeRuleAndBaseAspectsAttributes(ruleAttributes, ruleContextBuilder.getAspects()),
             ruleContextBuilder.getErrorConsumer(),
             ruleContextBuilder.getRule(),
             ruleContextBuilder.getRuleClassNameForLogging()),
         new PrerequisitesCollection(
             mainAspectPrerequisites.build(),
             mainAspect.getDefinition().getAttributes(),
             ruleContextBuilder.getErrorConsumer(),
             ruleContextBuilder.getRule(),
             ruleContextBuilder.getRuleClassNameForLogging()),
         execGroupCollection,
         baseTargetToolchainContexts,
         targetUsesAutoExecGroups);
   }

   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() {
     // TODO: b/370558813 - Use AutoExecGroupsMode for aspects, as well.
     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() {
     return Streams.concat(
             mainAspectPrerequisites.getAllPrerequisites().stream(),
             getRulePrerequisitesCollection().getAllPrerequisites().stream())
         .collect(toImmutableList());
   }
 }
	// 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 static com.google.common.collect.ImmutableSet.toImmutableSet;

	import com.google.common.collect.ImmutableList;
	import com.google.common.collect.ImmutableListMultimap;
	import com.google.common.collect.ImmutableMap;
	import com.google.common.collect.ImmutableSet;
	import com.google.common.collect.Iterables;
	import com.google.common.collect.Streams;
	import com.google.devtools.build.lib.cmdline.Label;
	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.skyframe.ConfiguredTargetAndData;
	import java.util.Collection;
	import java.util.LinkedHashMap;
	import java.util.Map;
	import javax.annotation.Nullable;

	/** 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;

	private final PrerequisitesCollection mainAspectPrerequisites;

	/**
	* The toolchain contexts for the base target.
	*
	* <p>It only contains the providers created by the aspects that propagate to the toolchains.
	*/
	@Nullable
	private final ToolchainCollection<AspectBaseTargetResolvedToolchainContext>
	baseTargetToolchainContexts;

	/** Whether the target uses auto exec groups. */
	private final boolean targetUsesAutoExecGroups;

	AspectContext(
	RuleContext.Builder builder,
	AspectAwareAttributeMapper aspectAwareAttributeMapper,
	PrerequisitesCollection ruleAndBaseAspectsPrerequisites,
	PrerequisitesCollection mainAspectPrerequisites,
	ExecGroupCollection execGroupCollection,
	@Nullable
	ToolchainCollection<AspectBaseTargetResolvedToolchainContext> baseTargetToolchainContexts,
	boolean targetUsesAutoExecGroups) {
	super(
	builder, aspectAwareAttributeMapper, ruleAndBaseAspectsPrerequisites, execGroupCollection);

	this.aspects = builder.getAspects();
	this.aspectDescriptors = aspects.stream().map(Aspect::getDescriptor).collect(toImmutableList());
	this.mainAspectPrerequisites = mainAspectPrerequisites;
	this.baseTargetToolchainContexts = baseTargetToolchainContexts;
	this.targetUsesAutoExecGroups = targetUsesAutoExecGroups;
	}

	/**
	* Returns the toolchain contexts for the base target. Can be null if no aspect in the {@code
	* aspects} path propagate to the toolchains.
	*/
	@Nullable
	public ToolchainCollection<AspectBaseTargetResolvedToolchainContext>
	getBaseTargetToolchainContexts() {
	return baseTargetToolchainContexts;
	}

	/** Returns the labels of default the toolchain types that aspects have propagated. */
	public ImmutableSet<Label> getRequestedToolchainTypesLabels() {
	if (targetUsesAutoExecGroups) {
	return baseTargetToolchainContexts.contextMap().entrySet().stream()
	.filter(e -> isAutomaticExecGroup(e.getKey()))
	.flatMap(e -> e.getValue().requestedToolchainTypeLabels().keySet().stream())
	.collect(toImmutableSet());
	} else {
	return baseTargetToolchainContexts
	.getDefaultToolchainContext()
	.requestedToolchainTypeLabels()
	.keySet();
	}
	}

	/**
	* Returns the toolchain data for the given type, or {@code null} if the toolchain type was not
	* required in this context.
	*/
	@Nullable
	public AspectBaseTargetResolvedToolchainContext.ToolchainAspectsProviders getToolchainTarget(
	Label toolchainType) {
	var execGroupContext = baseTargetToolchainContexts.getDefaultToolchainContext();
	if (targetUsesAutoExecGroups) {
	execGroupContext =
	baseTargetToolchainContexts.contextMap().entrySet().stream()
	.filter(
	e ->
	isAutomaticExecGroup(e.getKey())
	&& e.getValue().requestedToolchainTypeLabels().containsKey(toolchainType))
	.findFirst()
	.map(e -> e.getValue())
	.orElse(null);
	if (execGroupContext == null) {
	return null;
	}
	}
	return execGroupContext
	.getToolchains()
	.get(execGroupContext.requestedToolchainTypeLabels().get(toolchainType));
	}

	/**
	* 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,
	ToolchainCollection<AspectBaseTargetResolvedToolchainContext> baseTargetToolchainContexts) {
	return createAspectContextWithSeparatedPrerequisites(
	builder, ruleAttributes, targetsMap, execGroupCollection, baseTargetToolchainContexts);
	}

	/**
	* 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 ruleContextBuilder,
	AttributeMap ruleAttributes,
	ImmutableListMultimap<DependencyKind, ConfiguredTargetAndData> prerequisitesMap,
	ExecGroupCollection execGroupCollection,
	ToolchainCollection<AspectBaseTargetResolvedToolchainContext> baseTargetToolchainContexts) {
	ImmutableSortedKeyListMultimap.Builder<String, ConfiguredTargetAndData>
	mainAspectPrerequisites = ImmutableSortedKeyListMultimap.builder();
	ImmutableSortedKeyListMultimap.Builder<String, ConfiguredTargetAndData>
	ruleAndBaseAspectsPrerequisites = ImmutableSortedKeyListMultimap.builder();

	Aspect mainAspect = Iterables.getLast(ruleContextBuilder.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());
	}
	}

	boolean targetUsesAutoExecGroups =
	ruleContextBuilder
	.getRule()
	.getRuleClassObject()
	.getAutoExecGroupsMode()
	.isEnabled(ruleAttributes, ruleContextBuilder.getConfiguration().useAutoExecGroups());

	return new AspectContext(
	ruleContextBuilder,
	new AspectAwareAttributeMapper(
	ruleAttributes, mergeAspectsAttributes(ruleContextBuilder.getAspects())),
	new PrerequisitesCollection(
	ruleAndBaseAspectsPrerequisites.build(),
	mergeRuleAndBaseAspectsAttributes(ruleAttributes, ruleContextBuilder.getAspects()),
	ruleContextBuilder.getErrorConsumer(),
	ruleContextBuilder.getRule(),
	ruleContextBuilder.getRuleClassNameForLogging()),
	new PrerequisitesCollection(
	mainAspectPrerequisites.build(),
	mainAspect.getDefinition().getAttributes(),
	ruleContextBuilder.getErrorConsumer(),
	ruleContextBuilder.getRule(),
	ruleContextBuilder.getRuleClassNameForLogging()),
	execGroupCollection,
	baseTargetToolchainContexts,
	targetUsesAutoExecGroups);
	}

	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() {
	// TODO: b/370558813 - Use AutoExecGroupsMode for aspects, as well.
	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() {
	return Streams.concat(
	mainAspectPrerequisites.getAllPrerequisites().stream(),
	getRulePrerequisitesCollection().getAllPrerequisites().stream())
	.collect(toImmutableList());
	}
	}