Begin to factor BazelStarlarkContext into subclasses

BazelStarlarkContext is a hodgepodge of Bazel-specific state needed for or manipulated by various Starlark evaluations. A minor refactoring goal is to break this class up into subclasses, each one representing a different kind of Starlark evaluation that needs different kinds of data. This avoids the proliferation of new @nullable fields that may or may not be set and used for a given evaluation.

This CL begins the process by breaking out the case of initializing a .bzl module (i.e. running its top-level code). The other cases should follow a similar pathway, but I don't intend to pursue that in the immediate future. My motivation for tackling .bzl initialization now is that I want to add a field needed for the .bzl visibility without feeling guilty.

I also took this opportunity to clarify that BazelStarlarkContext ought not to be shared or reused between multiple StarlarkThreads, and fixed a location where this was violated by StarlarkCallbackHelper.

Work towards #11261.

PiperOrigin-RevId: 457732549
Change-Id: I3ee869865c47457275f1b84819f0f154290aa22a

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

@@ -47,7 +47,7 @@
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.graph.Digraph;
 import com.google.devtools.build.lib.graph.Node;
-import com.google.devtools.build.lib.packages.BazelStarlarkContext;
+import com.google.devtools.build.lib.packages.BzlInitThreadContext;
 import com.google.devtools.build.lib.packages.NativeAspectClass;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.RuleClass.Builder.ThirdPartyLicenseExistencePolicy;
@@ -871,12 +871,10 @@ public ImmutableMap<String, Object> getEnvironment() {
 
   @Override
   public void setStarlarkThreadContext(StarlarkThread thread, Label fileLabel) {
-    new BazelStarlarkContext(
-            BazelStarlarkContext.Phase.LOADING,
+    new BzlInitThreadContext(
             toolsRepository,
             configurationFragmentMap,
             new SymbolGenerator<>(fileLabel),
-            /*analysisRuleLabel=*/ null,
             networkAllowlistForTests)
         .storeInThread(thread);
   }

src/main/java/com/google/devtools/build/lib/analysis/config/StarlarkDefinedConfigTransition.java

@@ -334,7 +334,7 @@ public ImmutableMap<String, Map<String, Object>> evaluate(
       try (Mutability mu = Mutability.create("eval_transition_function")) {
         StarlarkThread thread = new StarlarkThread(mu, semantics);
         thread.setPrintHandler(Event.makeDebugPrintHandler(handler));
-        // TODO: If the resulting values of Starlark transitions ever evolve to be
+        // TODO(brandjon): If the resulting values of Starlark transitions ever evolve to be
         //  complex Starlark objects like structs as opposed to the ints, strings,
         //  etc they are today then we need a real symbol generator which is used
         //  to calculate equality between instances of Starlark objects. A candidate

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

@@ -335,10 +335,6 @@ public StarlarkCallable rule(
 
     if (implicitOutputs != Starlark.NONE) {
       if (implicitOutputs instanceof StarlarkFunction) {
-        // TODO(brandjon): Embedding bazelContext in a callback is not thread safe! Instead
-        // construct a new BazelStarlarkContext with copies of the relevant fields that are safe to
-        // share. Maybe create a method on BazelStarlarkContext for safely constructing a child
-        // context.
         StarlarkCallbackHelper callback =
             new StarlarkCallbackHelper(
                 (StarlarkFunction) implicitOutputs, thread.getSemantics(), bazelContext);

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

@@ -27,39 +27,77 @@
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkThread;
 
-/** Contextual information associated with each Starlark thread created by Bazel. */
-// TODO(adonovan): rename BazelThreadContext, for symmetry with BazelModuleContext.
-// TODO(brandjon): Use composition rather than inheritance for RuleDefinitionEnvironment; clients
-// should retrieve the RDE (if it exists) from this class in order to access e.g. the network
-// allowlist. The toolsRepository info will be duplicated between this class and RDE but we can
-// enforce consistency with a precondition check.
-public final class BazelStarlarkContext
+/**
+ * Bazel-specific contextual information associated with a Starlark evaluation thread.
+ *
+ * <p>This is stored in the StarlarkThread object as a thread-local. A subclass of this class may be
+ * used for certain kinds of Starlark evaluations; in that case it is still keyed in the
+ * thread-locals under {@code BazelStarlarkContext.class}.
+ *
+ * <p>This object is mutable and should not be reused for more than one Starlark thread.
+ */
+// TODO(b/236456122): rename BazelThreadContext, for symmetry with BazelModuleContext.
+// TODO(b/236456122): We should break this class up into subclasses for each kind of evaluation, as
+// opposed to storing specialized fields on this class and setting them to null for inapplicable
+// environments. Subclasses should define {@code from(StarlarkThread)} and {@code
+// fromOrFailFunction(StarlarkThread, String)} methods to be used in place of the {@code
+// check*Phase} methods in this class. Kinds of evaluation include:
+//   - .bzl loading
+//   - BUILD evaluation (should absorb PackageFactory.PackageContext -- no need to store multiple
+//     thread-locals in StarlarkThread)
+//   - WORKSPACE evaluation (shares logic with BUILD)
+//   - rule and aspect analysis implementation (can/should we store the RuleContext here?)
+//   - implicit outputs
+//   - computed defaults
+//   - transition implementation
+//   - Args.map_each
+//   - probably others
+// TODO(b/236456122): The inheritance of RuleDefinitionEnvironment should be replaced by
+// composition, in an appropriate subclass of this class. Things like the tools repository, network
+// allowlist, etc. can be accessed from the RDE. (If any info needs to be duplicated between RDE and
+// here, we should assert consistency with a precondition check.)
+public class BazelStarlarkContext
     implements RuleDefinitionEnvironment, StarlarkThread.UncheckedExceptionContext {
 
   /** The phase to which this Starlark thread belongs. */
+  // TODO(b/236456122): Eliminate.
   public enum Phase {
     WORKSPACE,
     LOADING,
     ANALYSIS
   }
 
-  /** Return the Bazel information associated with the specified Starlark thread. */
+  /**
+   * Retrieves this context from a Starlark thread, or throws {@link IllegalStateException} if
+   * unavailable.
+   */
   public static BazelStarlarkContext from(StarlarkThread thread) {
-    return thread.getThreadLocal(BazelStarlarkContext.class);
+    BazelStarlarkContext ctx = thread.getThreadLocal(BazelStarlarkContext.class);
+    // ISE rather than NPE for symmetry with subclasses.
+    Preconditions.checkState(
+        ctx != null, "Expected BazelStarlarkContext to be available in this Starlark thread");
+    return ctx;
   }
 
-  /** Save this BazelStarlarkContext in the specified Starlark thread. */
+  /**
+   * Saves this {@code BazelStarlarkContext} in the specified Starlark thread. Call only once,
+   * before evaluation begins.
+   */
   public void storeInThread(StarlarkThread thread) {
+    Preconditions.checkState(thread.getThreadLocal(BazelStarlarkContext.class) == null);
     thread.setThreadLocal(BazelStarlarkContext.class, this);
     thread.setUncheckedExceptionContext(this);
   }
 
+  // A generic counter for uniquely identifying symbols created in this Starlark evaluation.
+  private final SymbolGenerator<?> symbolGenerator;
+
+  // TODO(b/236456122): Migrate the below fields to subclasses.
   private final Phase phase;
   // Only necessary for loading phase threads.
   @Nullable private final RepositoryName toolsRepository;
   // Only necessary for loading phase threads to construct configuration_field.
   @Nullable private final ImmutableMap<String, Class<?>> fragmentNameToClass;
-  private final SymbolGenerator<?> symbolGenerator;
   @Nullable private final Label analysisRuleLabel;
   // TODO(b/192694287): Remove once we migrate all tests from the allowlist
   @Nullable private final Label networkAllowlistForTests;
@@ -75,13 +113,6 @@ public void storeInThread(StarlarkThread thread) {
    *     compared using reference equality.
    * @param analysisRuleLabel is the label of the rule for an analysis phase (rule or aspect
    */
-  // TODO(adonovan): clearly demarcate which fields are defined in which kinds of threads (loading,
-  // analysis, workspace, implicit outputs, computed defaults, etc), perhaps by splitting these into
-  // separate structs, exactly one of which is populated (plus the common fields). And eliminate
-  // StarlarkUtils.Phase.
-  // TODO(adonovan): move PackageFactory.PackageContext in here, for loading-phase threads.
-  // TODO(adonovan): is there any reason not to put the entire RuleContext in this thread, for
-  // analysis threads?
   public BazelStarlarkContext(
       Phase phase,
       @Nullable RepositoryName toolsRepository,
@@ -97,6 +128,11 @@ public BazelStarlarkContext(
     this.networkAllowlistForTests = networkAllowlistForTests;
   }
 
+  /** Returns the phase associated with this context. */
+  public Phase getPhase() {
+    return phase;
+  }
+
   /** Returns the name of the tools repository, such as "@bazel_tools". */
   @Nullable
   @Override
@@ -138,6 +174,8 @@ public Optional<Label> getNetworkAllowlistForTests() {
    *
    * @param function name of a function that requires this check
    */
+  // TODO(b/236456122): The Phase enum is incomplete. Ex: `Args.map_each` evaluation happens at
+  // execution time. So this is a misnomer and possibly wrong in those contexts.
   public void checkLoadingOrWorkspacePhase(String function) throws EvalException {
     if (phase == Phase.ANALYSIS) {
       throw Starlark.errorf("'%s' cannot be called during the analysis phase", function);

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

@@ -0,0 +1,72 @@
+// Copyright 2022 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.packages;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
+import javax.annotation.Nullable;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkThread;
+
+/**
+ * Bazel application data for the Starlark thread that evaluates the top-level code in a .bzl module
+ * (i.e. when evaluating that module's global symbols).
+ */
+public final class BzlInitThreadContext extends BazelStarlarkContext {
+
+  // TODO(b/236456122): Are all these arguments needed for .bzl initialization?
+  public BzlInitThreadContext(
+      @Nullable RepositoryName toolsRepository,
+      @Nullable ImmutableMap<String, Class<?>> fragmentNameToClass,
+      SymbolGenerator<?> symbolGenerator,
+      @Nullable Label networkAllowlistForTests) {
+    super(
+        BazelStarlarkContext.Phase.LOADING,
+        toolsRepository,
+        fragmentNameToClass,
+        symbolGenerator,
+        /*analysisRuleLabel=*/ null,
+        networkAllowlistForTests);
+  }
+
+  /**
+   * Retrieves this context from a Starlark thread, or throws {@link IllegalStateException} if
+   * unavailable.
+   */
+  public static BzlInitThreadContext from(StarlarkThread thread) {
+    BazelStarlarkContext ctx = thread.getThreadLocal(BazelStarlarkContext.class);
+    Preconditions.checkState(
+        ctx instanceof BzlInitThreadContext,
+        "Expected to be in a .bzl initialization (top-level evaluation) Starlark thread");
+    return (BzlInitThreadContext) ctx;
+  }
+
+  /**
+   * Retrieves this context from a Starlark thread. If not present, throws {@code EvalException}
+   * with an error message indicating the failure was in a function named {@code function}.
+   */
+  public static BzlInitThreadContext fromOrFailFunction(StarlarkThread thread, String function)
+      throws EvalException {
+    @Nullable BazelStarlarkContext ctx = thread.getThreadLocal(BazelStarlarkContext.class);
+    if (!(ctx instanceof BzlInitThreadContext)) {
+      throw Starlark.errorf(
+          "'%s' can only be called during .bzl initialization (top-level evaluation)", function);
+    }
+    return (BzlInitThreadContext) ctx;
+  }
+}

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

@@ -29,10 +29,11 @@
  * A helper class for calling Starlark functions from Java, where the argument values are supplied
  * by the fields of a Structure, as in the case of computed attribute defaults and computed implicit
  * outputs.
- *
- * <p>TODO(adonovan): eliminate the need for this class by making the Starlark calls in the same
- * Starlark thread that instantiated the rule.
  */
+// TODO(brandjon): Consider eliminating this class by executing the callback in the same thread as
+// the caller, i.e. in the thread evaluating a BUILD file. This might not be possible for implicit
+// outputs without some refactoring to cache the result of the computation (currently RuleContext
+// seems to reinvoke the callback).
 public final class StarlarkCallbackHelper {
 
   private final StarlarkFunction callback;
@@ -48,15 +49,18 @@ public final class StarlarkCallbackHelper {
   // Alternatively (or additionally), we could put PackageContext
   // into BazelStarlarkContext so there's only a single blob of state.
   private final StarlarkSemantics starlarkSemantics;
-  private final BazelStarlarkContext context;
+  // The Bazel context that was in effect at the place where the StarlarkCallbackHelper was defined.
+  // This will be used as a template for creating a new context each time the callback is invoked
+  // (since BazelStarlarkContexts cannot be safely shared or reused).
+  private final BazelStarlarkContext contextTemplate;
 
   public StarlarkCallbackHelper(
       StarlarkFunction callback,
       StarlarkSemantics starlarkSemantics,
-      BazelStarlarkContext context) {
+      BazelStarlarkContext contextTemplate) {
     this.callback = callback;
     this.starlarkSemantics = starlarkSemantics;
-    this.context = context;
+    this.contextTemplate = contextTemplate;
   }
 
   public ImmutableList<String> getParameterNames() {
@@ -65,36 +69,49 @@ public ImmutableList<String> getParameterNames() {
 
   // TODO(adonovan): opt: all current callers are forced to construct a temporary Structure.
   // Instead, make them supply a map.
-  public Object call(EventHandler eventHandler, Structure ctx, Object... arguments)
+  public Object call(EventHandler eventHandler, Structure struct, Object... arguments)
       throws EvalException, InterruptedException {
     try (Mutability mu = Mutability.create("callback", callback)) {
       StarlarkThread thread = new StarlarkThread(mu, starlarkSemantics);
       thread.setPrintHandler(Event.makeDebugPrintHandler(eventHandler));
+      // TODO(b/236456122): If we don't eliminate StarlarkCallbackHelper entirely, we can give it
+      // its own BazelStarlarkContext subclass.
+      BazelStarlarkContext context =
+          new BazelStarlarkContext(
+              contextTemplate.getPhase(),
+              contextTemplate.getToolsRepository(),
+              contextTemplate.getFragmentNameToClass(),
+              // TODO(brandjon): In principle, if we're creating a new symbol generator here, we
+              // should have a unique owner object to associate it with for distinguishing
+              // reference-equality objects. But I don't think implicit outputs or computed defaults
+              // care about identity.
+              new SymbolGenerator<>(new Object()),
+              // Analysis rule label should be null for computed defaults and implicit outputs, but
+              // may as well copy the field for symmetry.
+              contextTemplate.getAnalysisRuleLabel(),
+              contextTemplate.getNetworkAllowlistForTests().orElse(null));
       context.storeInThread(thread);
       return Starlark.call(
-          thread,
-          callback,
-          buildArgumentList(ctx, arguments),
-          /*kwargs=*/ ImmutableMap.of());
+          thread, callback, buildArgumentList(struct, arguments), /*kwargs=*/ ImmutableMap.of());
     } catch (ClassCastException | IllegalArgumentException e) { // TODO(adonovan): investigate
       throw new EvalException(e);
     }
   }
 
   /**
    * Creates a list of actual arguments that contains the given arguments and all attribute values
-   * required from the specified context.
+   * required from the specified structure.
    */
-  private ImmutableList<Object> buildArgumentList(Structure ctx, Object... arguments)
+  private ImmutableList<Object> buildArgumentList(Structure struct, Object... arguments)
       throws EvalException {
     ImmutableList.Builder<Object> builder = ImmutableList.builder();
     ImmutableList<String> names = getParameterNames();
     int requiredParameters = names.size() - arguments.length;
     for (int pos = 0; pos < requiredParameters; ++pos) {
       String name = names.get(pos);
-      Object value = ctx.getValue(name);
+      Object value = struct.getValue(name);
       if (value == null) {
-          throw new IllegalArgumentException(ctx.getErrorMessageForUnknownField(name));
+        throw new IllegalArgumentException(struct.getErrorMessageForUnknownField(name));
       }
       builder.add(value);
     }