Split ModuleThreadContext off from ModuleFileGlobals

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

@@ -379,6 +379,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:runfiles_tree",
         "//src/main/java/com/google/devtools/build/lib/analysis/platform",
         "//src/main/java/com/google/devtools/build/lib/analysis/stringtemplate",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution",
         "//src/main/java/com/google/devtools/build/lib/bugreport",
         "//src/main/java/com/google/devtools/build/lib/buildeventstream",
         "//src/main/java/com/google/devtools/build/lib/buildeventstream/proto:build_event_stream_java_proto",

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

@@ -19,6 +19,7 @@
 import com.google.devtools.build.lib.analysis.DefaultInfo;
 import com.google.devtools.build.lib.analysis.OutputGroupInfo;
 import com.google.devtools.build.lib.analysis.RunEnvironmentInfo;
+import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileGlobals;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.packages.BuildGlobals;
 import com.google.devtools.build.lib.packages.Proto;
@@ -121,6 +122,13 @@ public ImmutableMap<String, Object> getSclToplevels() {
     return env.buildOrThrow();
   }
 
+  @Override
+  public ImmutableMap<String, Object> getModuleToplevels() {
+    var env = ImmutableMap.<String, Object>builder();
+    Starlark.addMethods(env, new ModuleFileGlobals());
+    return env.buildOrThrow();
+  }
+
   @Override
   public ImmutableMap<String, Object> getRepoToplevels() {
     ImmutableMap.Builder<String, Object> env = ImmutableMap.builder();

src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java

@@ -254,7 +254,11 @@ public void workspaceInit(
         .addSkyFunction(SkyFunctions.REPOSITORY_DIRECTORY, repositoryDelegatorFunction)
         .addSkyFunction(
             SkyFunctions.MODULE_FILE,
-            new ModuleFileFunction(registryFactory, directories.getWorkspace(), builtinModules))
+            new ModuleFileFunction(
+                runtime.getRuleClassProvider().getBazelStarlarkEnvironment(),
+                registryFactory,
+                directories.getWorkspace(),
+                builtinModules))
         .addSkyFunction(SkyFunctions.BAZEL_DEP_GRAPH, new BazelDepGraphFunction())
         .addSkyFunction(
             SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace()))

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -134,8 +134,10 @@ java_library(
         "ModuleExtensionEvalStarlarkThreadContext.java",
         "ModuleExtensionRepoMappingEntriesValue.java",
         "ModuleExtensionResolutionEvent.java",
+        "ModuleFileGlobals.java",
         "ModuleFileValue.java",
         "ModuleOverride.java",
+        "ModuleThreadContext.java",
         "MultipleVersionOverride.java",
         "NonRegistryOverride.java",
         "RegistryOverride.java",
@@ -151,6 +153,7 @@ java_library(
         ":module_extension_metadata",
         ":registry",
         ":repo_rule_creator",
+        "//src/main/java/com/google/devtools/build/docgen/annot",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
@@ -160,7 +163,9 @@ java_library(
         "//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/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
+        "//src/main/java/net/starlark/java/annot",
         "//src/main/java/net/starlark/java/eval",
         "//src/main/java/net/starlark/java/syntax",
         "//third_party:auto_value",
@@ -186,7 +191,6 @@ java_library(
         "ModuleExtensionEvaluationProgress.java",
         "ModuleExtensionRepoMappingEntriesFunction.java",
         "ModuleFileFunction.java",
-        "ModuleFileGlobals.java",
         "RepoSpecFunction.java",
         "Selection.java",
         "SingleExtensionEvalFunction.java",
@@ -207,6 +211,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/docgen/annot",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
+        "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_version_info",
         "//src/main/java/com/google/devtools/build/lib/bazel:bazel_version",

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java

@@ -31,6 +31,7 @@
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.lib.packages.BazelStarlarkEnvironment;
 import com.google.devtools.build.lib.packages.DotBazelFileSyntaxChecker;
 import com.google.devtools.build.lib.packages.StarlarkExportable;
 import com.google.devtools.build.lib.profiler.Profiler;
@@ -84,6 +85,7 @@ public class ModuleFileFunction implements SkyFunction {
   public static final Precomputed<Map<String, ModuleOverride>> MODULE_OVERRIDES =
       new Precomputed<>("module_overrides");
 
+  private final BazelStarlarkEnvironment starlarkEnv;
   private final RegistryFactory registryFactory;
   private final Path workspaceRoot;
   private final ImmutableMap<String, NonRegistryOverride> builtinModules;
@@ -103,9 +105,11 @@ public class ModuleFileFunction implements SkyFunction {
    *     non-registry overrides.
    */
   public ModuleFileFunction(
+      BazelStarlarkEnvironment starlarkEnv,
       RegistryFactory registryFactory,
       Path workspaceRoot,
       ImmutableMap<String, NonRegistryOverride> builtinModules) {
+    this.starlarkEnv = starlarkEnv;
     this.registryFactory = registryFactory;
     this.workspaceRoot = workspaceRoot;
     this.builtinModules = builtinModules;
@@ -157,7 +161,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     String moduleFileHash =
         new Fingerprint().addBytes(getModuleFileResult.moduleFile.getContent()).hexDigestAndReset();
 
-    ModuleFileGlobals moduleFileGlobals =
+    ModuleThreadContext moduleThreadContext =
         execModuleFile(
             getModuleFileResult.moduleFile,
             getModuleFileResult.registry,
@@ -168,12 +172,13 @@ public SkyValue compute(SkyKey skyKey, Environment env)
             // We try to prevent most side effects of yanked modules, in particular print().
             /* printIsNoop= */ getModuleFileResult.yankedInfo != null,
             starlarkSemantics,
+            starlarkEnv,
             env.getListener());
 
     // Perform some sanity checks.
     InterimModule module;
     try {
-      module = moduleFileGlobals.buildModule();
+      module = moduleThreadContext.buildModule();
     } catch (EvalException e) {
       env.getListener().handle(Event.error(e.getMessageWithStack()));
       throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for %s", moduleKey);
@@ -242,6 +247,7 @@ private SkyValue computeForRootModule(StarlarkSemantics starlarkSemantics, Envir
         MODULE_OVERRIDES.get(env),
         IGNORE_DEV_DEPS.get(env),
         starlarkSemantics,
+        starlarkEnv,
         env.getListener());
   }
 
@@ -257,22 +263,24 @@ public static RootModuleFileValue evaluateRootModuleFile(
       Map<String, ModuleOverride> commandOverrides,
       boolean ignoreDevDeps,
       StarlarkSemantics starlarkSemantics,
+      BazelStarlarkEnvironment starlarkEnv,
       ExtendedEventHandler eventHandler)
       throws ModuleFileFunctionException, InterruptedException {
     String moduleFileHash = new Fingerprint().addBytes(moduleFileContents).hexDigestAndReset();
-    ModuleFileGlobals moduleFileGlobals =
+    ModuleThreadContext moduleThreadContext =
         execModuleFile(
             ModuleFile.create(moduleFileContents, moduleFilePath.asPath().toString()),
             /* registry= */ null,
             ModuleKey.ROOT,
             ignoreDevDeps,
-            /* printIsNoop= */ builtinModules,
-            false,
+            builtinModules,
+            /* printIsNoop= */ false,
             starlarkSemantics,
+            starlarkEnv,
             eventHandler);
     InterimModule module;
     try {
-      module = moduleFileGlobals.buildModule();
+      module = moduleThreadContext.buildModule();
     } catch (EvalException e) {
       eventHandler.handle(Event.error(e.getMessageWithStack()));
       throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for the root module");
@@ -290,7 +298,7 @@ public static RootModuleFileValue evaluateRootModuleFile(
       }
     }
 
-    ImmutableMap<String, ModuleOverride> moduleOverrides = moduleFileGlobals.buildOverrides();
+    ImmutableMap<String, ModuleOverride> moduleOverrides = moduleThreadContext.buildOverrides();
     ImmutableMap<String, ModuleOverride> overrides =
         ImmutableMap.<String, ModuleOverride>builder()
             .putAll(moduleOverrides)
@@ -317,14 +325,15 @@ public static RootModuleFileValue evaluateRootModuleFile(
         module, moduleFileHash, overrides, nonRegistryOverrideCanonicalRepoNameLookup);
   }
 
-  private static ModuleFileGlobals execModuleFile(
+  private static ModuleThreadContext execModuleFile(
       ModuleFile moduleFile,
       @Nullable Registry registry,
       ModuleKey moduleKey,
       boolean ignoreDevDeps,
       ImmutableMap<String, NonRegistryOverride> builtinModules,
       boolean printIsNoop,
       StarlarkSemantics starlarkSemantics,
+      BazelStarlarkEnvironment starlarkEnv,
       ExtendedEventHandler eventHandler)
       throws ModuleFileFunctionException, InterruptedException {
     StarlarkFile starlarkFile =
@@ -334,18 +343,20 @@ private static ModuleFileGlobals execModuleFile(
       throw errorf(Code.BAD_MODULE, "error parsing MODULE.bazel file for %s", moduleKey);
     }
 
-    ModuleFileGlobals moduleFileGlobals =
-        new ModuleFileGlobals(builtinModules, moduleKey, registry, ignoreDevDeps);
+    ModuleThreadContext context =
+        new ModuleThreadContext(builtinModules, moduleKey, registry, ignoreDevDeps);
     try (SilentCloseable c =
             Profiler.instance()
                 .profile(ProfilerTask.BZLMOD, () -> "evaluate module file: " + moduleKey);
         Mutability mu = Mutability.create("module file", moduleKey)) {
       new DotBazelFileSyntaxChecker("MODULE.bazel files", /* canLoadBzl= */ false)
           .check(starlarkFile);
       net.starlark.java.eval.Module predeclaredEnv =
-          getPredeclaredEnv(moduleFileGlobals, starlarkSemantics);
+          net.starlark.java.eval.Module.withPredeclared(
+              starlarkSemantics, starlarkEnv.getStarlarkGlobals().getModuleToplevels());
       Program program = Program.compileFile(starlarkFile, predeclaredEnv);
       StarlarkThread thread = new StarlarkThread(mu, starlarkSemantics);
+      context.storeInThread(thread);
       if (printIsNoop) {
         thread.setPrintHandler((t, msg) -> {});
       } else {
@@ -368,7 +379,7 @@ private static ModuleFileGlobals execModuleFile(
       eventHandler.handle(Event.error(e.getMessageWithStack()));
       throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for %s", moduleKey);
     }
-    return moduleFileGlobals;
+    return context;
   }
 
   private static class GetModuleFileResult {
@@ -500,13 +511,6 @@ private static void createModuleFile(Path path, byte[] bytes) throws ModuleFileF
     }
   }
 
-  private static net.starlark.java.eval.Module getPredeclaredEnv(
-      ModuleFileGlobals moduleFileGlobals, StarlarkSemantics starlarkSemantics) {
-    ImmutableMap.Builder<String, Object> env = ImmutableMap.builder();
-    Starlark.addMethods(env, moduleFileGlobals, starlarkSemantics);
-    return net.starlark.java.eval.Module.withPredeclared(starlarkSemantics, env.buildOrThrow());
-  }
-
   @FormatMethod
   private static ModuleFileFunctionException errorf(Code code, String format, Object... args) {
     return new ModuleFileFunctionException(ExternalDepsException.withMessage(code, format, args));