Add identityToken to StarlarkDefinedAspect via its StarlarkAspectClass.

PiperOrigin-RevId: 638277768
Change-Id: Ib1e2a570fb1ea86e998ef50beb6096fac87d14de

Author: aoeui

src/main/java/com/google/devtools/build/lib/analysis/BUILD

@@ -656,6 +656,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/skyframe:aspect_key_creator",
         "//src/main/java/com/google/devtools/build/lib/skyframe:build_result_listener",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:bzl_load_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:configured_target_key",
         "//src/main/java/com/google/devtools/build/lib/skyframe:coverage_report_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:repository_mapping_value",

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

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.analysis;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.devtools.build.lib.skyframe.BzlLoadValue.keyForBuild;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -515,7 +516,8 @@ private ImmutableList<TopLevelAspectsKey> createTopLevelAspectKeys(
               e);
         }
         String starlarkFunctionName = aspect.substring(delimiterPosition + 1);
-        aspectClassesBuilder.add(new StarlarkAspectClass(starlarkFileLabel, starlarkFunctionName));
+        aspectClassesBuilder.add(
+            new StarlarkAspectClass(keyForBuild(starlarkFileLabel), starlarkFunctionName));
       } else {
         final NativeAspectClass aspectFactoryClass =
             ruleClassProvider.getNativeAspectClassMap().get(aspect);

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java

@@ -356,18 +356,22 @@ public StarlarkCallable macro(
       builder.addAttribute(attr);
     }
 
+    return new MacroFunction(
+        builder,
+        Starlark.toJavaOptional(doc, String.class).map(Starlark::trimDocString),
+        getBzlKeyToken(thread, "Macros"));
+  }
+
+  private static Symbol<BzlLoadValue.Key> getBzlKeyToken(StarlarkThread thread, String onBehalfOf) {
     Symbol<?> untypedToken = thread.getNextIdentityToken();
     checkState(
         untypedToken.getOwner() instanceof BzlLoadValue.Key,
-        "Macros may only be owned by .bzl files (owner=%s)",
+        "%s may only be owned by .bzl files (owner=%s)",
+        onBehalfOf,
         untypedToken);
     @SuppressWarnings("unchecked")
     var typedToken = (Symbol<BzlLoadValue.Key>) untypedToken;
-
-    return new MacroFunction(
-        builder,
-        Starlark.toJavaOptional(doc, String.class).map(Starlark::trimDocString),
-        typedToken);
+    return typedToken;
   }
 
   // TODO(bazel-team): implement attribute copy and other rule properties
@@ -1103,7 +1107,8 @@ public StarlarkAspect aspect(
         applyToGeneratingRules,
         execCompatibleWith,
         execGroups,
-        ImmutableSet.copyOf(subrules));
+        ImmutableSet.copyOf(subrules),
+        getBzlKeyToken(thread, "Aspects"));
   }
 
   private static ImmutableSet<String> getLegacyAnyTypeAttrs(RuleClass ruleClass) {

src/main/java/com/google/devtools/build/lib/packages/StarlarkAspectClass.java

@@ -14,24 +14,31 @@
 
 package com.google.devtools.build.lib.packages;
 
+import static com.google.devtools.build.lib.util.HashCodes.hashObjects;
+
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.skyframe.BzlLoadValue;
 
 /** {@link AspectClass} for aspects defined in Starlark. */
 @Immutable
 public final class StarlarkAspectClass implements AspectClass {
-  private final Label extensionLabel;
+  private final BzlLoadValue.Key extensionKey;
   private final String exportedName;
   private final String name;
 
-  public StarlarkAspectClass(Label extensionLabel, String exportedName) {
-    this.extensionLabel = extensionLabel;
+  public StarlarkAspectClass(BzlLoadValue.Key extensionKey, String exportedName) {
+    this.extensionKey = extensionKey;
     this.exportedName = exportedName;
-    this.name = extensionLabel + "%" + exportedName;
+    this.name = extensionKey.getLabel() + "%" + exportedName;
+  }
+
+  BzlLoadValue.Key getExtensionKey() {
+    return extensionKey;
   }
 
   public Label getExtensionLabel() {
-    return extensionLabel;
+    return extensionKey.getLabel();
   }
 
   public String getExportedName() {
@@ -54,15 +61,12 @@ public boolean equals(Object o) {
     }
 
     StarlarkAspectClass that = (StarlarkAspectClass) o;
-
-    return extensionLabel.equals(that.extensionLabel)
-        && exportedName.equals(that.exportedName);
+    return extensionKey.equals(that.extensionKey) && exportedName.equals(that.exportedName);
   }
 
   @Override
   public int hashCode() {
-    // Inlines the implementation of Objects.hashCode to avoid generating garbage.
-    return 31 * extensionLabel.hashCode() + exportedName.hashCode();
+    return hashObjects(extensionKey, exportedName);
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java

@@ -14,6 +14,8 @@
 
 package com.google.devtools.build.lib.packages;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 import com.github.benmanes.caffeine.cache.Caffeine;
 import com.github.benmanes.caffeine.cache.LoadingCache;
 import com.google.common.base.Ascii;
@@ -25,6 +27,7 @@
 import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.EventHandler;
+import com.google.devtools.build.lib.skyframe.BzlLoadValue;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
 import java.io.Serializable;
 import java.util.Objects;
@@ -35,6 +38,7 @@
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkCallable;
 import net.starlark.java.eval.StarlarkInt;
+import net.starlark.java.eval.SymbolGenerator.Symbol;
 
 /** A Starlark value that is a result of an 'aspect(..)' function call. */
 public final class StarlarkDefinedAspect implements StarlarkExportable, StarlarkAspect {
@@ -58,7 +62,8 @@ public final class StarlarkDefinedAspect implements StarlarkExportable, Starlark
   private final ImmutableMap<String, ExecGroup> execGroups;
   private final ImmutableSet<? extends StarlarkSubruleApi> subrules;
 
-  private StarlarkAspectClass aspectClass;
+  /** {@link Symbol} before {@link #export} and a {@link StarlarkAspectClass} after. */
+  private Object aspectClassOrIdentityToken;
 
   private static final ImmutableSet<String> TRUE_REPS =
       ImmutableSet.of("true", "1", "yes", "t", "y");
@@ -81,7 +86,8 @@ public StarlarkDefinedAspect(
       boolean applyToGeneratingRules,
       ImmutableSet<Label> execCompatibleWith,
       ImmutableMap<String, ExecGroup> execGroups,
-      ImmutableSet<? extends StarlarkSubruleApi> subrules) {
+      ImmutableSet<? extends StarlarkSubruleApi> subrules,
+      Symbol<BzlLoadValue.Key> identityToken) {
     this.implementation = implementation;
     this.documentation = documentation.orElse(null);
     this.attributeAspects = attributeAspects;
@@ -97,6 +103,7 @@ public StarlarkDefinedAspect(
     this.execCompatibleWith = execCompatibleWith;
     this.execGroups = execGroups;
     this.subrules = subrules;
+    this.aspectClassOrIdentityToken = identityToken;
   }
 
   public StarlarkCallable getImplementation() {
@@ -138,7 +145,7 @@ public String getName() {
   @Override
   public StarlarkAspectClass getAspectClass() {
     Preconditions.checkState(isExported());
-    return aspectClass;
+    return (StarlarkAspectClass) aspectClassOrIdentityToken;
   }
 
   @Override
@@ -149,7 +156,16 @@ public ImmutableSet<String> getParamAttributes() {
   @Override
   public void export(EventHandler handler, Label extensionLabel, String name) {
     Preconditions.checkArgument(!isExported());
-    this.aspectClass = new StarlarkAspectClass(extensionLabel, name);
+    @SuppressWarnings("unchecked")
+    var identityToken = (Symbol<BzlLoadValue.Key>) aspectClassOrIdentityToken;
+    BzlLoadValue.Key owner = identityToken.getOwner();
+    checkArgument(
+        owner.getLabel().equals(extensionLabel),
+        "Exporting aspect as (%s, %s) but label did not match owner=%s",
+        extensionLabel,
+        name,
+        owner);
+    this.aspectClassOrIdentityToken = new StarlarkAspectClass(owner, name);
   }
 
   private static final ImmutableList<String> ALL_ATTR_ASPECTS = ImmutableList.of("*");
@@ -173,7 +189,8 @@ public AspectDefinition getDefinition(AspectParameters aspectParams) {
   }
 
   private AspectDefinition buildDefinition(AspectParameters aspectParams) {
-    AspectDefinition.Builder builder = new AspectDefinition.Builder(aspectClass);
+    AspectDefinition.Builder builder =
+        new AspectDefinition.Builder((StarlarkAspectClass) aspectClassOrIdentityToken);
     if (ALL_ATTR_ASPECTS.equals(attributeAspects)) {
       builder.propagateAlongAllAttributes();
     } else {
@@ -254,7 +271,7 @@ private static Attribute addAttrValue(Attribute attr, String attrValue) {
 
   @Override
   public boolean isExported() {
-    return aspectClass != null;
+    return aspectClassOrIdentityToken instanceof StarlarkAspectClass;
   }
 
   @Override
@@ -406,7 +423,7 @@ public boolean equals(Object o) {
         && Objects.equals(requiredAspects, that.requiredAspects)
         && Objects.equals(fragments, that.fragments)
         && Objects.equals(toolchainTypes, that.toolchainTypes)
-        && Objects.equals(aspectClass, that.aspectClass);
+        && Objects.equals(aspectClassOrIdentityToken, that.aspectClassOrIdentityToken);
   }
 
   @Override
@@ -422,6 +439,6 @@ public int hashCode() {
         requiredAspects,
         fragments,
         toolchainTypes,
-        aspectClass);
+        aspectClassOrIdentityToken);
   }
 }

src/test/java/com/google/devtools/build/lib/skyframe/TargetCycleReporterTest.java

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.skyframe;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.skyframe.BzlLoadValue.keyForBuild;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -93,7 +94,7 @@ public void loadingPhaseCycleWithDifferentTopLevelKeyTypes() throws Exception {
         AspectKeyCreator.createTopLevelAspectsKey(
             ImmutableList.of(
                 new StarlarkAspectClass(
-                    Label.parseCanonicalUnchecked("//foo:b"), "my Starlark key")),
+                    keyForBuild(Label.parseCanonicalUnchecked("//foo:b")), "my Starlark key")),
             Label.parseCanonicalUnchecked("//foo:a"),
             targetConfig,
             /* topLevelAspectsParameters= */ ImmutableMap.of());

src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java

@@ -607,7 +607,6 @@ def _impl(name, xyz): pass
   @Test
   public void testSymbolicMacro_macroFunctionApi() throws Exception {
     ev.setSemantics("--experimental_enable_first_class_macros");
-    ev.setThreadOwner(keyForBuild(FAKE_LABEL));
 
     evalAndExport(
         ev,
@@ -1018,6 +1017,7 @@ public void testLabelWithAspects() throws Exception {
 
   @Test
   public void testLabelListWithAspectsError() throws Exception {
+    ev.setThreadOwner(keyForBuild(FAKE_LABEL));
     ev.checkEvalErrorContains(
         "at index 1 of aspects, got element of type int, want Aspect",
         "def _impl(target, ctx):",
@@ -1618,6 +1618,7 @@ public void testRuleAddAttribute() throws Exception {
   }
 
   private static void evalAndExport(BazelEvaluationTestCase ev, String... lines) throws Exception {
+    ev.setThreadOwner(keyForBuild(FAKE_LABEL));
     ev.execAndExport(FAKE_LABEL, lines);
   }
 
@@ -2906,7 +2907,7 @@ public void fancyExports() throws Exception {
     assertThat(p.getKey()).isEqualTo(new StarlarkProvider.Key(keyForBuild(FAKE_LABEL), "p"));
     assertThat(p1.getPrintableName()).isEqualTo("p1");
     assertThat(p1.getKey()).isEqualTo(new StarlarkProvider.Key(keyForBuild(FAKE_LABEL), "p1"));
-    assertThat(a.getAspectClass()).isEqualTo(new StarlarkAspectClass(FAKE_LABEL, "a"));
+    assertThat(a.getAspectClass()).isEqualTo(new StarlarkAspectClass(keyForBuild(FAKE_LABEL), "a"));
   }
 
   @Test