Make Package.Builder a BazelStarlarkContext descendant

Previously, when evaluating a BUILD file we would store three types of thread-locals in the StarlarkThread:

  1) a BazelStarlarkContext to hold a SymbolGenerator (for reference equality semantics) and a Phase enum (should ideally be replaced by helpers like fromOrFail());
  2) a Package.Builder (or, prior to 29318bb, a PackageContext) to support defining targets and things like native.glob(); and
  3) a RuleDefinitionEnvironment to support the analysis_test() feature.

This CL merges (1) and (2) by having Package.Builder also be a BazelStarlarkContext. This simplifies the setup of a StarlarkThread for evaluating BUILD files and symbolic macros, and paves the way for distinguishing symbolic macro construction with its own separate builder type in a later CL.

In Package.java:

  - Add Phase and SymbolGenerator parameters to the Builder constructor, to pass along to super(). Hide this complexity from callers by making the Builder constructor private and adding Package.newPackageBuilder(), complementing the two existing builder factory methods in Package. (This results in a lint for adding a 13-parameter method, but I'm going with the lesser of evils here. What does the lint want me to do, make a builder for the builder?)

  - The Phase is always LOADING and the SymbolGenerator is based on the package id, except for in newExternalPackageBuilder() where it's WORKSPACE and the SymbolGenerator uses the skykey to distinguish evaluations of different chunks of the same WORKSPACE file. This is all just movement of existing caller logic into Package.java. newExternalPackageBuilder() now takes in the whole skykey instead of just a path.

  - Add fromOrNull() and fromOrFail() helpers, which is more idiomatic for BazelStarlarkContext subclasses. Callers now use these in place of calling getThreadLocal() directly. Callers use storeInThread() instead of calling setThreadLocal() directly.

In WorkspaceFactory, the symbol generator construction is logically moved from execute() to when the builder is first constructed in WorkspaceFileFunction. So execute() loses the key as a parameter. Some tests that call newExternalPackageBuilder() directly now pass in a key there instead of where execute() is called.

Work toward #19922.

PiperOrigin-RevId: 602187597
Change-Id: Id650be08f82c0b0e3baf89e8b89db7bd1f401f79

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

@@ -1026,9 +1026,12 @@ public String getName() {
     public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwargs)
         throws EvalException, InterruptedException {
       BazelStarlarkContext.checkLoadingPhase(thread, getName());
-      Package.Builder pkgBuilder = thread.getThreadLocal(Package.Builder.class);
+      Package.Builder pkgBuilder = Package.Builder.fromOrNull(thread);
       if (pkgBuilder == null) {
-        throw new EvalException(
+        throw Starlark.errorf(
+            // TODO: #19922 - Clarify error. Maybe we weren't called during .bzl loading but at some
+            // other bad time. Also, it's ambiguous to the user whether, strictly speaking,
+            // evaluating a symbolic macro happens while evaluating a BUILD file.
             "Cannot instantiate a macro when loading a .bzl file. "
                 + "Macros may only be instantiated while evaluating a BUILD file.");
       }
@@ -1180,9 +1183,10 @@ public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwarg
       if (ruleClass == null) {
         throw new EvalException("Invalid rule class hasn't been exported by a bzl file");
       }
-      Package.Builder pkgBuilder = thread.getThreadLocal(Package.Builder.class);
+      Package.Builder pkgBuilder = Package.Builder.fromOrNull(thread);
       if (pkgBuilder == null) {
         throw new EvalException(
+            // TODO: #19922 - Clarify message. See analogous TODO for macros, above.
             "Cannot instantiate a rule when loading a .bzl file. "
                 + "Rules may be instantiated only in a BUILD thread.");
       }

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

@@ -88,17 +88,10 @@ public static void executeMacroImplementation(
       StarlarkThread thread = new StarlarkThread(mu, semantics);
       thread.setPrintHandler(Event.makeDebugPrintHandler(builder.getLocalEventHandler()));
 
-      new BazelStarlarkContext(
-              BazelStarlarkContext.Phase.LOADING,
-              // TODO: #19922 - Technically we should use a different key than the one in the main
-              // BUILD thread, but that'll be fixed when we change the builder type to
-              // PackagePiece.Builder.
-              new SymbolGenerator<>(builder.getPackageIdentifier()))
-          .storeInThread(thread);
-
-      // TODO: #19922 - Have Package.Builder inherit from BazelStarlarkContext and only store one
-      // thread-local object.
-      thread.setThreadLocal(Package.Builder.class, builder);
+      // TODO: #19922 - Technically the embedded SymbolGenerator field should use a different key
+      // than the one in the main BUILD thread, but that'll be fixed when we change the type to
+      // PackagePiece.Builder.
+      builder.storeInThread(thread);
 
       // TODO: #19922 - If we want to support creating analysis_test rules inside symbolic macros,
       // we'd need to call `thread.setThreadLocal(RuleDefinitionEnvironment.class,

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

@@ -43,6 +43,7 @@
 import com.google.devtools.build.lib.events.EventKind;
 import com.google.devtools.build.lib.events.StoredEventHandler;
 import com.google.devtools.build.lib.packages.Package.Builder.PackageSettings;
+import com.google.devtools.build.lib.packages.WorkspaceFileValue.WorkspaceFileKey;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
 import com.google.devtools.build.lib.server.FailureDetails.PackageLoading;
 import com.google.devtools.build.lib.server.FailureDetails.PackageLoading.Code;
@@ -76,7 +77,9 @@
 import java.util.TreeMap;
 import java.util.concurrent.Semaphore;
 import javax.annotation.Nullable;
+import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Module;
+import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkThread;
 import net.starlark.java.syntax.Location;
 
@@ -753,17 +756,58 @@ public void dump(PrintStream out) {
     }
   }
 
+  /**
+   * Returns a new {@link Builder} suitable for constructing an ordinary package (i.e. not one for
+   * WORKSPACE or bzlmod).
+   */
+  public static Builder newPackageBuilder(
+      PackageSettings packageSettings,
+      PackageIdentifier id,
+      RootedPath filename,
+      String workspaceName,
+      Optional<String> associatedModuleName,
+      Optional<String> associatedModuleVersion,
+      boolean noImplicitFileExport,
+      RepositoryMapping repositoryMapping,
+      @Nullable Semaphore cpuBoundSemaphore,
+      PackageOverheadEstimator packageOverheadEstimator,
+      @Nullable ImmutableMap<Location, String> generatorMap,
+      // TODO(bazel-team): See Builder() constructor comment about use of null for this param.
+      @Nullable ConfigSettingVisibilityPolicy configSettingVisibilityPolicy,
+      @Nullable Globber globber) {
+    return new Builder(
+        BazelStarlarkContext.Phase.LOADING,
+        new SymbolGenerator<>(id),
+        packageSettings,
+        id,
+        filename,
+        workspaceName,
+        associatedModuleName,
+        associatedModuleVersion,
+        noImplicitFileExport,
+        repositoryMapping,
+        cpuBoundSemaphore,
+        packageOverheadEstimator,
+        generatorMap,
+        configSettingVisibilityPolicy,
+        globber);
+  }
+
   public static Builder newExternalPackageBuilder(
-      PackageSettings helper,
-      RootedPath workspacePath,
+      PackageSettings packageSettings,
+      WorkspaceFileKey workspaceFileKey,
       String workspaceName,
       RepositoryMapping mainRepoMapping,
       boolean noImplicitFileExport,
       PackageOverheadEstimator packageOverheadEstimator) {
     return new Builder(
-        helper,
+        BazelStarlarkContext.Phase.WORKSPACE,
+        // The SymbolGenerator is based on workspaceFileKey rather than a package id or path,
+        // in order to distinguish different chunks of the same WORKSPACE file.
+        new SymbolGenerator<>(workspaceFileKey),
+        packageSettings,
         LabelConstants.EXTERNAL_PACKAGE_IDENTIFIER,
-        /* filename= */ workspacePath,
+        /* filename= */ workspaceFileKey.getPath(),
         workspaceName,
         /* associatedModuleName= */ Optional.empty(),
         /* associatedModuleVersion= */ Optional.empty(),
@@ -782,6 +826,8 @@ public static Builder newExternalPackageBuilderForBzlmod(
       PackageIdentifier basePackageId,
       RepositoryMapping repoMapping) {
     return new Builder(
+            BazelStarlarkContext.Phase.LOADING,
+            new SymbolGenerator<>(basePackageId),
             PackageSettings.DEFAULTS,
             basePackageId,
             /* filename= */ moduleFilePath,
@@ -818,7 +864,7 @@ private static DetailedExitCode createDetailedCode(String errorMessage, Code cod
    * A builder for {@link Package} objects. Only intended to be used by {@link PackageFactory} and
    * {@link com.google.devtools.build.lib.skyframe.PackageFunction}.
    */
-  public static class Builder {
+  public static class Builder extends BazelStarlarkContext {
 
     /**
      * A bundle of options affecting package construction, that is not specific to any particular
@@ -1016,7 +1062,9 @@ public T intern(T sample) {
 
     private boolean alreadyBuilt = false;
 
-    Builder(
+    private Builder(
+        BazelStarlarkContext.Phase phase,
+        SymbolGenerator<?> symbolGenerator,
         PackageSettings packageSettings,
         PackageIdentifier id,
         RootedPath filename,
@@ -1032,6 +1080,8 @@ public T intern(T sample) {
         // Maybe convert null -> LEGACY_OFF, assuming that's the correct default.
         @Nullable ConfigSettingVisibilityPolicy configSettingVisibilityPolicy,
         @Nullable Globber globber) {
+      super(phase, symbolGenerator);
+
       Metadata metadata = new Metadata();
       metadata.packageIdentifier = Preconditions.checkNotNull(id);
 
@@ -1071,6 +1121,29 @@ public T intern(T sample) {
               pkg, metadata.buildFileLabel, Location.fromFile(filename.asPath().toString())));
     }
 
+    /** Retrieves this object from a Starlark thread. Returns null if not present. */
+    @Nullable
+    public static Builder fromOrNull(StarlarkThread thread) {
+      BazelStarlarkContext ctx = thread.getThreadLocal(BazelStarlarkContext.class);
+      return (ctx instanceof Builder) ? (Builder) ctx : null;
+    }
+
+    /**
+     * Retrieves this object from a Starlark thread. If not present, throws {@code EvalException}
+     * with an error message indicating that {@code what} can't be used in this Starlark
+     * environment.
+     */
+    @CanIgnoreReturnValue
+    public static Builder fromOrFail(StarlarkThread thread, String what) throws EvalException {
+      @Nullable BazelStarlarkContext ctx = thread.getThreadLocal(BazelStarlarkContext.class);
+      if (!(ctx instanceof Builder)) {
+        // TODO: #19922 - Clarify in the message that we can't be in a symbolic ("first-class")
+        // macro.
+        throw Starlark.errorf("%s can only be used while evaluating a BUILD file", what);
+      }
+      return (Builder) ctx;
+    }
+
     PackageIdentifier getPackageIdentifier() {
       return pkg.getPackageIdentifier();
     }

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

@@ -31,6 +31,7 @@
 import com.google.devtools.build.lib.packages.Package.Builder.PackageSettings;
 import com.google.devtools.build.lib.packages.Package.ConfigSettingVisibilityPolicy;
 import com.google.devtools.build.lib.packages.PackageValidator.InvalidPackageException;
+import com.google.devtools.build.lib.packages.WorkspaceFileValue.WorkspaceFileKey;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
@@ -215,7 +216,7 @@ public RuleClassProvider getRuleClassProvider() {
   // TODO(#19922): The name is a holdover from when we had PackageContext. Migrate this to a static
   // fromOrFail method on Package.Builder or a new parent interface of it.
   public static Package.Builder getContext(StarlarkThread thread) throws EvalException {
-    Package.Builder value = thread.getThreadLocal(Package.Builder.class);
+    Package.Builder value = Package.Builder.fromOrNull(thread);
     if (value == null) {
       // if PackageBuilder is missing, we're not called from a BUILD file. This happens if someone
       // uses native.some_func() in the wrong place.
@@ -227,13 +228,13 @@ public static Package.Builder getContext(StarlarkThread thread) throws EvalExcep
   }
 
   public Package.Builder newExternalPackageBuilder(
-      RootedPath workspacePath,
+      WorkspaceFileKey workspaceFileKey,
       String workspaceName,
       RepositoryMapping mainRepoMapping,
       StarlarkSemantics starlarkSemantics) {
     return Package.newExternalPackageBuilder(
         packageSettings,
-        workspacePath,
+        workspaceFileKey,
         workspaceName,
         mainRepoMapping,
         starlarkSemantics.getBool(BuildLanguageOptions.INCOMPATIBLE_NO_IMPLICIT_FILE_EXPORT),
@@ -256,7 +257,7 @@ public Package.Builder newPackageBuilder(
       @Nullable ImmutableMap<Location, String> generatorMap,
       @Nullable ConfigSettingVisibilityPolicy configSettingVisibilityPolicy,
       @Nullable Globber globber) {
-    return new Package.Builder(
+    return Package.newPackageBuilder(
         packageSettings,
         packageId,
         filename,
@@ -412,15 +413,7 @@ private void executeBuildFileImpl(
       StarlarkThread thread = new StarlarkThread(mu, semantics);
       thread.setLoader(loadedModules::get);
       thread.setPrintHandler(Event.makeDebugPrintHandler(pkgBuilder.getLocalEventHandler()));
-
-      new BazelStarlarkContext(
-              BazelStarlarkContext.Phase.LOADING,
-              new SymbolGenerator<>(pkgBuilder.getPackageIdentifier()))
-          .storeInThread(thread);
-
-      // TODO(#19922): Have Package.Builder inherit from BazelStarlarkContext and only store one
-      // thread-local object.
-      thread.setThreadLocal(Package.Builder.class, pkgBuilder);
+      pkgBuilder.storeInThread(thread);
 
       // TODO(b/291752414): The rule definition environment shouldn't be needed at BUILD evaluation
       // time EXCEPT for analysis_test, which needs the tools repository for use in

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

@@ -98,8 +98,7 @@ public WorkspaceFactory(
    */
   public void execute(
       StarlarkFile file, // becomes resolved as a side effect
-      Map<String, Module> additionalLoadedModules,
-      WorkspaceFileValue.WorkspaceFileKey workspaceFileKey)
+      Map<String, Module> additionalLoadedModules)
       throws InterruptedException {
     loadedModules.putAll(additionalLoadedModules);
 
@@ -118,15 +117,7 @@ public void execute(
       StarlarkThread thread = new StarlarkThread(mutability, starlarkSemantics);
       thread.setLoader(loadedModules::get);
       thread.setPrintHandler(Event.makeDebugPrintHandler(builder.getLocalEventHandler()));
-      thread.setThreadLocal(Package.Builder.class, builder);
-
-      // The workspace environment doesn't need the tools repository or the fragment map
-      // because executing workspace rules happens before analysis and it doesn't need a
-      // repository mapping because calls to the Label constructor in the WORKSPACE file
-      // are, by definition, not in an external repository and so they don't need the mapping
-      new BazelStarlarkContext(
-              BazelStarlarkContext.Phase.WORKSPACE, new SymbolGenerator<>(workspaceFileKey))
-          .storeInThread(thread);
+      builder.storeInThread(thread);
 
       try {
         Starlark.execFileProgram(prog, module, thread);

src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java

@@ -68,7 +68,7 @@ public void analysisTest(
       Object attrValuesApi,
       StarlarkThread thread)
       throws EvalException, InterruptedException {
-    Package.Builder pkgBuilder = thread.getThreadLocal(Package.Builder.class);
+    Package.Builder pkgBuilder = Package.Builder.fromOrNull(thread);
     RuleDefinitionEnvironment ruleDefinitionEnvironment =
         thread.getThreadLocal(RuleDefinitionEnvironment.class);
     // TODO(b/236456122): Refactor this check into a standard helper / error message

src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java

@@ -293,8 +293,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     // The default 'workspace name' is "__main__". Note that this is different from the "workspace
     // name" returned by WorkspaceNameFunction, which is a fixed string when Bzlmod is enabled.
     Package.Builder builder =
-        packageFactory.newExternalPackageBuilder(
-            workspaceFile, "__main__", repoMapping, starlarkSemantics);
+        packageFactory.newExternalPackageBuilder(key, "__main__", repoMapping, starlarkSemantics);
 
     if (chunks.isEmpty()) {
       builder.setLoads(ImmutableList.of());
@@ -390,7 +389,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       }
       // Execute the partial files that comprise this chunk.
       for (StarlarkFile partialFile : chunk) {
-        parser.execute(partialFile, loadedModules, key);
+        parser.execute(partialFile, loadedModules);
       }
     }
 

src/test/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryContextTest.java

@@ -37,6 +37,7 @@
 import com.google.devtools.build.lib.packages.RuleClass.Builder.RuleClassType;
 import com.google.devtools.build.lib.packages.Type;
 import com.google.devtools.build.lib.packages.WorkspaceFactoryHelper;
+import com.google.devtools.build.lib.packages.WorkspaceFileValue;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
 import com.google.devtools.build.lib.rules.repository.RepositoryFunction.RepositoryFunctionException;
@@ -137,7 +138,7 @@ private void setUpContextForRule(
     Package.Builder packageBuilder =
         Package.newExternalPackageBuilder(
             PackageSettings.DEFAULTS,
-            RootedPath.toRootedPath(root, workspaceFile),
+            WorkspaceFileValue.key(RootedPath.toRootedPath(root, workspaceFile)),
             "runfiles",
             RepositoryMapping.ALWAYS_FALLBACK,
             starlarkSemantics.getBool(BuildLanguageOptions.INCOMPATIBLE_NO_IMPLICIT_FILE_EXPORT),

src/test/java/com/google/devtools/build/lib/packages/PackageTest.java

@@ -158,7 +158,7 @@ public void testBuildPartialPopulatesImplicitTestSuiteValueIdempotently() throws
   }
 
   private Package.Builder pkgBuilder(String name) {
-    return new Package.Builder(
+    return Package.newPackageBuilder(
         PackageSettings.DEFAULTS,
         PackageIdentifier.createInMainRepo(name),
         /* filename= */ RootedPath.toRootedPath(

src/test/java/com/google/devtools/build/lib/packages/RuleFactoryTest.java

@@ -137,7 +137,7 @@ public void testCreateWorkspaceRule() throws Exception {
     Path myPkgPath = scratch.resolve("/workspace/WORKSPACE");
     Package.Builder pkgBuilder =
         packageFactory.newExternalPackageBuilder(
-            RootedPath.toRootedPath(root, myPkgPath),
+            WorkspaceFileValue.key(RootedPath.toRootedPath(root, myPkgPath)),
             "TESTING",
             RepositoryMapping.ALWAYS_FALLBACK,
             StarlarkSemantics.DEFAULT);

src/test/java/com/google/devtools/build/lib/packages/WorkspaceFactoryTestHelper.java

@@ -67,7 +67,7 @@ void parse(String... args) {
     builder =
         Package.newExternalPackageBuilder(
                 PackageSettings.DEFAULTS,
-                RootedPath.toRootedPath(root, workspaceFilePath),
+                WorkspaceFileValue.key(RootedPath.toRootedPath(root, workspaceFilePath)),
                 "",
                 RepositoryMapping.ALWAYS_FALLBACK,
                 starlarkSemantics.getBool(
@@ -86,10 +86,7 @@ void parse(String... args) {
             /* defaultSystemJavabaseDir= */ null,
             starlarkSemantics);
     try {
-      factory.execute(
-          file,
-          /* additionalLoadedModules= */ ImmutableMap.of(),
-          WorkspaceFileValue.key(RootedPath.toRootedPath(root, workspaceFilePath)));
+      factory.execute(file, /* additionalLoadedModules= */ ImmutableMap.of());
     } catch (InterruptedException e) {
       fail("Shouldn't happen: " + e.getMessage());
     }