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);
  }