Use a SymbolGenerator ID instead of reference equality for StarlarkPr…

…ovider.

PiperOrigin-RevId: 628537022
Change-Id: I680aceef40af478bfd4ca2f46bd4a29cec6c71c4

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

@@ -296,16 +296,15 @@ public Object provider(Object doc, Object fields, Object init, StarlarkThread th
               Dict.cast(fields, String.class, String.class, "fields"), Starlark::trimDocString));
     }
     if (init == Starlark.NONE) {
-      return builder.build();
+      return builder.buildWithIdentityToken(thread.getNextIdentityToken());
+    }
+    if (init instanceof StarlarkCallable callable) {
+      builder.setInit(callable);
     } else {
-      if (init instanceof StarlarkCallable) {
-        builder.setInit((StarlarkCallable) init);
-      } else {
-        throw Starlark.errorf("got %s for init, want callable value", Starlark.type(init));
-      }
-      StarlarkProvider provider = builder.build();
-      return Tuple.of(provider, provider.createRawConstructor());
+      throw Starlark.errorf("got %s for init, want callable value", Starlark.type(init));
     }
+    StarlarkProvider provider = builder.buildWithIdentityToken(thread.getNextIdentityToken());
+    return Tuple.of(provider, provider.createRawConstructor());
   }
 
   @FormatMethod

src/main/java/com/google/devtools/build/lib/bazel/rules/cpp/BUILD

@@ -47,16 +47,14 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_configuration",
-        "//src/main/java/com/google/devtools/build/lib/analysis:rule_error_consumer",
-        "//src/main/java/com/google/devtools/build/lib/analysis:transitive_info_collection",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/rules/cpp",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:bzl_load_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant",
         "//src/main/java/net/starlark/java/eval",
         "//third_party:guava",
-        "//third_party:jsr305",
     ],
 )
 

src/main/java/com/google/devtools/build/lib/bazel/rules/cpp/BazelCppSemantics.java

@@ -14,6 +14,9 @@
 
 package com.google.devtools.build.lib.bazel.rules.cpp;
 
+import static com.google.devtools.build.lib.skyframe.BzlLoadValue.keyForBuild;
+import static com.google.devtools.build.lib.skyframe.BzlLoadValue.keyForBuiltins;
+
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.RuleContext;
@@ -50,18 +53,22 @@ public class BazelCppSemantics implements AspectLegalCppSemantics {
   // the repo name set.
   public static final Provider.Key CC_SHARED_INFO_PROVIDER_RULES_CC =
       new StarlarkProvider.Key(
-          Label.parseCanonicalUnchecked("@rules_cc//examples:experimental_cc_shared_library.bzl"),
+          keyForBuild(
+              Label.parseCanonicalUnchecked(
+                  "@rules_cc//examples:experimental_cc_shared_library.bzl")),
           "CcSharedLibraryInfo");
 
   public static final Provider.Key CC_SHARED_INFO_PROVIDER =
       new StarlarkProvider.Key(
-          Label.parseCanonicalUnchecked("//examples:experimental_cc_shared_library.bzl"),
+          keyForBuild(
+              Label.parseCanonicalUnchecked("//examples:experimental_cc_shared_library.bzl")),
           "CcSharedLibraryInfo");
 
   public static final Provider.Key CC_SHARED_INFO_PROVIDER_BUILT_INS =
       new StarlarkProvider.Key(
-          Label.parseCanonicalUnchecked(
-              "@_builtins//:common/cc/experimental_cc_shared_library.bzl"),
+          keyForBuiltins(
+              Label.parseCanonicalUnchecked(
+                  "@_builtins//:common/cc/experimental_cc_shared_library.bzl")),
           "CcSharedLibraryInfo");
 
   private final Language language;

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

@@ -85,14 +85,17 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/io:file_symlink_exception",
         "//src/main/java/com/google/devtools/build/lib/packages/semantics",
         "//src/main/java/com/google/devtools/build/lib/profiler",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:bzl_load_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:detailed_exceptions",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:starlark_builtins_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization:visible-for-serialization",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/core",
         "//src/main/java/com/google/devtools/build/lib/util",
+        "//src/main/java/com/google/devtools/build/lib/util:TestType",
         "//src/main/java/com/google/devtools/build/lib/util:detailed_exit_code",
         "//src/main/java/com/google/devtools/build/lib/util:filetype",
         "//src/main/java/com/google/devtools/build/lib/util:hash_codes",

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

@@ -14,13 +14,23 @@
 
 package com.google.devtools.build.lib.packages;
 
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.devtools.build.lib.bugreport.BugReport.sendNonFatalBugReport;
+import static com.google.devtools.build.lib.skyframe.BzlLoadValue.keyForBuild;
+import static com.google.devtools.build.lib.skyframe.BzlLoadValue.keyForBuiltins;
+import static com.google.devtools.build.lib.skyframe.StarlarkBuiltinsValue.isBuiltinsRepo;
+import static com.google.devtools.build.lib.util.HashCodes.hashObjects;
+import static com.google.devtools.build.lib.util.TestType.isInTest;
+
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.events.EventHandler;
+import com.google.devtools.build.lib.skyframe.BzlLoadValue;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.Collection;
@@ -35,6 +45,7 @@
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkCallable;
 import net.starlark.java.eval.StarlarkThread;
+import net.starlark.java.eval.SymbolGenerator;
 import net.starlark.java.syntax.Location;
 
 /**
@@ -78,8 +89,14 @@ public final class StarlarkProvider implements StarlarkCallable, StarlarkExporta
   // Starlark dict mapping field names (string keys) to their values.
   @Nullable private final StarlarkCallable init;
 
-  /** Null iff this provider has not yet been exported. Mutated by {@link export}. */
-  @Nullable private Key key;
+  /**
+   * An identifier for this provider.
+   *
+   * <p>This is a {@link Key} if exported and a {@link SymbolGenerator.Symbol} otherwise.
+   *
+   * <p>Mutated by {@link #export}.
+   */
+  private Object keyOrIdentityToken;
 
   /**
    * For schemaful providers, an array of metadata concerning depset optimization.
@@ -133,8 +150,6 @@ public static final class Builder {
 
     @Nullable private StarlarkCallable init;
 
-    @Nullable private Key key;
-
     private Builder(Location location) {
       this.location = location;
     }
@@ -193,16 +208,14 @@ public Builder setInit(StarlarkCallable init) {
       return this;
     }
 
-    /** Sets the provider built by this builder to be exported with the given key. */
-    @CanIgnoreReturnValue
-    public Builder setExported(Key key) {
-      this.key = key;
-      return this;
+    /** Builds an exported StarlarkProvider. */
+    public StarlarkProvider buildExported(Key key) {
+      return new StarlarkProvider(location, documentation, schema, init, key);
     }
 
-    /** Builds a StarlarkProvider. */
-    public StarlarkProvider build() {
-      return new StarlarkProvider(location, documentation, schema, init, key);
+    /** Builds a unexported StarlarkProvider. */
+    public StarlarkProvider buildWithIdentityToken(SymbolGenerator.Symbol<?> identityToken) {
+      return new StarlarkProvider(location, documentation, schema, init, identityToken);
     }
   }
 
@@ -218,13 +231,13 @@ private StarlarkProvider(
       @Nullable String documentation,
       @Nullable ImmutableMap<String, Optional<String>> schema,
       @Nullable StarlarkCallable init,
-      @Nullable Key key) {
+      Object keyOrIdentityToken) {
     this.location = location;
     this.documentation = documentation;
     this.fields = schema != null ? ImmutableList.sortedCopyOf(schema.keySet()) : null;
     this.schema = schema;
     this.init = init;
-    this.key = key;
+    this.keyOrIdentityToken = keyOrIdentityToken;
     if (schema != null) {
       depsetTypePredictor = new AtomicReferenceArray<>(schema.size());
     }
@@ -327,18 +340,21 @@ public Optional<String> getDocumentation() {
 
   @Override
   public boolean isExported() {
-    return key != null;
+    return keyOrIdentityToken instanceof Key;
   }
 
   @Override
   public Key getKey() {
     Preconditions.checkState(isExported());
-    return key;
+    return (Key) keyOrIdentityToken;
   }
 
   @Override
   public String getName() {
-    return key != null ? key.getExportedName() : "<no name>";
+    if (keyOrIdentityToken instanceof Key key) {
+      return key.getExportedName();
+    }
+    return "<no name>";
   }
 
   @Override
@@ -370,36 +386,55 @@ public ImmutableMap<String, Optional<String>> getSchema() {
   @Override
   public String getErrorMessageForUnknownField(String name) {
     return String.format(
-        "'%s' value has no field or method '%s'",
-        isExported() ? key.getExportedName() : "struct", name);
+        "'%s' value has no field or method '%s'", isExported() ? getName() : "struct", name);
   }
 
   @Override
   public void export(EventHandler handler, Label extensionLabel, String exportedName) {
     Preconditions.checkState(!isExported());
-    this.key = new Key(extensionLabel, exportedName);
+    SymbolGenerator.Symbol<?> identifier = (SymbolGenerator.Symbol<?>) keyOrIdentityToken;
+    if (identifier.getOwner() instanceof BzlLoadValue.Key bzlKey) {
+      // In production code, StarlarkProviders are created only when loading .bzl files so the owner
+      // of the Symbol should be a BzlLoadValue.Key.
+      checkArgument(
+          extensionLabel.equals(bzlKey.getLabel()),
+          "export extensionLabel=%s, but owner=%s",
+          extensionLabel,
+          bzlKey);
+      this.keyOrIdentityToken = new Key(bzlKey, exportedName);
+    } else {
+      // In tests, the symbol may be arbitrary.
+      if (!isInTest()) {
+        sendNonFatalBugReport(
+            new IllegalStateException(
+                String.format(
+                    "exporting StarlarkProvider defined at %s as %s:%s but thread owner=%s was not"
+                        + " a BzlLoadValue.Key",
+                    location, extensionLabel, exportedName, identifier.getOwner())));
+      }
+      this.keyOrIdentityToken =
+          new Key(
+              isBuiltinsRepo(extensionLabel.getRepository())
+                  ? keyForBuiltins(extensionLabel)
+                  : keyForBuild(extensionLabel),
+              exportedName);
+    }
   }
 
   @Override
   public int hashCode() {
-    if (isExported()) {
-      return getKey().hashCode();
-    }
-    return System.identityHashCode(this);
+    return keyOrIdentityToken.hashCode();
   }
 
   @Override
   public boolean equals(@Nullable Object otherObject) {
-    if (!(otherObject instanceof StarlarkProvider)) {
-      return false;
+    if (this == otherObject) {
+      return true;
     }
-    StarlarkProvider other = (StarlarkProvider) otherObject;
-
-    if (this.isExported() && other.isExported()) {
-      return this.getKey().equals(other.getKey());
-    } else {
-      return this == other;
+    if (otherObject instanceof StarlarkProvider that) {
+      return this.keyOrIdentityToken.equals(that.keyOrIdentityToken);
     }
+    return false;
   }
 
   @Override
@@ -478,43 +513,36 @@ boolean isOptimised(int index, Object value) {
     return value instanceof NestedSet<?>;
   }
 
+  Object getKeyOrIdentityToken() {
+    return keyOrIdentityToken;
+  }
+
   /**
    * A serializable representation of Starlark-defined {@link StarlarkProvider} that uniquely
    * identifies all {@link StarlarkProvider}s that are exposed to SkyFrame.
    */
-  public static final class Key extends Provider.Key {
-    private final Label extensionLabel;
+  // TODO: b/335901349 - this is identical to SymbolGenerator.GlobalSymbol<BzlLoadValue.Key> and
+  // serves essentially the same purpose. Consider unifying these types.
+  public static class Key extends Provider.Key {
+    private final BzlLoadValue.Key key;
     private final String exportedName;
 
-    public Key(Label extensionLabel, String exportedName) {
-      this.extensionLabel = Preconditions.checkNotNull(extensionLabel);
-      this.exportedName = Preconditions.checkNotNull(exportedName);
+    public Key(BzlLoadValue.Key key, String exportedName) {
+      this.key = checkNotNull(key);
+      this.exportedName = checkNotNull(exportedName);
     }
 
     public Label getExtensionLabel() {
-      return extensionLabel;
+      return key.getLabel();
     }
 
-    public String getExportedName() {
+    public final String getExportedName() {
       return exportedName;
     }
 
-    @Override
-    public String toString() {
-      return exportedName;
-    }
-
-    @Override
-    void fingerprint(Fingerprint fp) {
-      // False => Not native.
-      fp.addBoolean(false);
-      fp.addString(extensionLabel.getCanonicalForm());
-      fp.addString(exportedName);
-    }
-
     @Override
     public int hashCode() {
-      return Objects.hash(extensionLabel, exportedName);
+      return hashObjects(key, exportedName);
     }
 
     @Override
@@ -527,8 +555,21 @@ public boolean equals(Object obj) {
         return false;
       }
       Key other = (Key) obj;
-      return Objects.equals(this.extensionLabel, other.extensionLabel)
+      return Objects.equals(this.key, other.key)
           && Objects.equals(this.exportedName, other.exportedName);
     }
+
+    @Override
+    final void fingerprint(Fingerprint fp) {
+      // False => Not native.
+      fp.addBoolean(false);
+      fp.addString(getExtensionLabel().getCanonicalForm());
+      fp.addString(getExportedName());
+    }
+
+    @Override
+    public String toString() {
+      return exportedName;
+    }
   }
 }

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

@@ -13,9 +13,9 @@
 // limitations under the License.
 package com.google.devtools.build.lib.packages;
 
-import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.RuleClass.ConfiguredTargetFactory.RuleErrorException;
+import com.google.devtools.build.lib.skyframe.BzlLoadValue;
 
 /**
  * A helper for wrapping an instance of a Starlark-defined provider with a native class {@code T}.
@@ -36,11 +36,11 @@
 public abstract class StarlarkProviderWrapper<T> {
   private final StarlarkProvider.Key key;
 
-  protected StarlarkProviderWrapper(Label label, String name) {
-    this.key = new StarlarkProvider.Key(label, name);
+  protected StarlarkProviderWrapper(BzlLoadValue.Key loadKey, String name) {
+    this.key = new StarlarkProvider.Key(loadKey, name);
   }
 
-  /*
+  /**
    * Converts an instance of the Starlark-defined provider to an instance of the wrapping class
    * {@code T}.
    *