Switch to BazelStarlarkContext

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

@@ -14,8 +14,8 @@
 
 package com.google.devtools.build.lib.analysis;
 
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.packages.BazelStarlarkContext;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
@@ -33,7 +33,7 @@ public class BazelRuleAnalysisThreadContext extends BazelStarlarkContext {
    * @param ruleContext is the {@link RuleContext} of the rule for analysis of a rule or aspect
    */
   public BazelRuleAnalysisThreadContext(RuleContext ruleContext) {
-    super(Phase.ANALYSIS);
+    super(Phase.ANALYSIS, ruleContext.getAnalysisEnvironment()::getMainRepoMapping);
     this.ruleContext = ruleContext;
   }
 

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

@@ -37,8 +37,8 @@
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
-import com.google.devtools.build.lib.packages.BazelStarlarkContext;
-import com.google.devtools.build.lib.packages.BazelStarlarkContext.Phase;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext.Phase;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.RuleTransitionData;
 import com.google.devtools.build.lib.packages.StructImpl;
@@ -552,7 +552,7 @@ public ImmutableMap<String, Map<String, Object>> evaluate(
         // Create a new {@link BazelStarlarkContext} for the new thread. We need to
         // create a new context every time because {@link BazelStarlarkContext}s
         // should be confined to a single thread.
-        new BazelStarlarkContext(Phase.ANALYSIS).storeInThread(thread);
+        new BazelStarlarkContext(Phase.ANALYSIS, /* mainRepoMapping= */ null).storeInThread(thread);
 
         result =
             Starlark.fastcall(

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

@@ -64,7 +64,7 @@
 import com.google.devtools.build.lib.packages.Attribute.StarlarkComputedDefaultTemplate;
 import com.google.devtools.build.lib.packages.AttributeTransitionData;
 import com.google.devtools.build.lib.packages.AttributeValueSource;
-import com.google.devtools.build.lib.packages.BazelStarlarkContext;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
 import com.google.devtools.build.lib.packages.BuildSetting;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.BuiltinRestriction;

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

@@ -34,6 +34,8 @@
 import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
 import com.google.devtools.build.lib.bazel.repository.starlark.StarlarkRepositoryModule.RepositoryRuleFunction;
 import com.google.devtools.build.lib.cmdline.BazelModuleContext;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext.Phase;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
@@ -132,6 +134,12 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     if (starlarkSemantics == null) {
       return null;
     }
+    RepositoryMappingValue mainRepoMappingValue =
+        (RepositoryMappingValue)
+            env.getValue(RepositoryMappingValue.KEY_FOR_ROOT_MODULE_WITHOUT_WORKSPACE_REPOS);
+    if (mainRepoMappingValue == null) {
+      return null;
+    }
 
     ModuleExtensionId extensionId = (ModuleExtensionId) skyKey.argument();
     SingleExtensionUsagesValue usagesValue =
@@ -182,7 +190,12 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     // Run that extension!
     env.getListener().post(ModuleExtensionEvaluationProgress.ongoing(extensionId, "starting"));
     RunModuleExtensionResult moduleExtensionResult =
-        extension.run(env, usagesValue, starlarkSemantics, extensionId);
+        extension.run(
+            env,
+            usagesValue,
+            starlarkSemantics,
+            extensionId,
+            mainRepoMappingValue.getRepositoryMapping());
     if (moduleExtensionResult == null) {
       return null;
     }
@@ -525,7 +538,8 @@ RunModuleExtensionResult run(
         Environment env,
         SingleExtensionUsagesValue usagesValue,
         StarlarkSemantics starlarkSemantics,
-        ModuleExtensionId extensionId)
+        ModuleExtensionId extensionId,
+        RepositoryMapping repositoryMapping)
         throws InterruptedException, SingleExtensionEvalFunctionException;
   }
 
@@ -677,7 +691,8 @@ public RunModuleExtensionResult run(
         Environment env,
         SingleExtensionUsagesValue usagesValue,
         StarlarkSemantics starlarkSemantics,
-        ModuleExtensionId extensionId)
+        ModuleExtensionId extensionId,
+        RepositoryMapping mainRepositoryMapping)
         throws InterruptedException, SingleExtensionEvalFunctionException {
       var generatedRepoSpecs = ImmutableMap.<String, RepoSpec>builderWithExpectedSize(repos.size());
       // Instantiate the repos one by one.
@@ -846,7 +861,8 @@ public RunModuleExtensionResult run(
         Environment env,
         SingleExtensionUsagesValue usagesValue,
         StarlarkSemantics starlarkSemantics,
-        ModuleExtensionId extensionId)
+        ModuleExtensionId extensionId,
+        RepositoryMapping mainRepositoryMapping)
         throws InterruptedException, SingleExtensionEvalFunctionException {
       ModuleExtensionEvalStarlarkThreadContext threadContext =
           new ModuleExtensionEvalStarlarkThreadContext(
@@ -870,6 +886,8 @@ public RunModuleExtensionResult run(
         thread.setPrintHandler(Event.makeDebugPrintHandler(env.getListener()));
         moduleContext = createContext(env, usagesValue, starlarkSemantics, extensionId);
         threadContext.storeInThread(thread);
+        new BazelStarlarkContext(Phase.WORKSPACE, () -> mainRepositoryMapping)
+            .storeInThread(thread);
         // This is used by the `Label()` constructor in Starlark, to record any attempts to resolve
         // apparent repo names to canonical repo names. See #20721 for why this is necessary.
         thread.setThreadLocal(Label.RepoMappingRecorder.class, repoMappingRecorder);

src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/BUILD

@@ -46,6 +46,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe:directory_tree_digest_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:ignored_package_prefixes_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_value",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:repository_mapping_value",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/repository",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:string",

src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryFunction.java

@@ -25,11 +25,12 @@
 import com.google.devtools.build.lib.analysis.RuleDefinition;
 import com.google.devtools.build.lib.bazel.repository.RepositoryResolvedEvent;
 import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
-import com.google.devtools.build.lib.packages.BazelStarlarkContext;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
@@ -46,6 +47,7 @@
 import com.google.devtools.build.lib.runtime.RepositoryRemoteExecutor;
 import com.google.devtools.build.lib.skyframe.IgnoredPackagePrefixesValue;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
+import com.google.devtools.build.lib.skyframe.RepositoryMappingValue;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -242,6 +244,20 @@ private RepositoryDirectoryValue.Builder fetchInternal(
       return null;
     }
 
+    boolean enableBzlmod = starlarkSemantics.getBool(BuildLanguageOptions.ENABLE_BZLMOD);
+    RepositoryMapping mainRepoMapping;
+    if (enableBzlmod || !isWorkspaceRepo(rule)) {
+      var mainRepoMappingValue =
+          (RepositoryMappingValue)
+              env.getValue(RepositoryMappingValue.KEY_FOR_ROOT_MODULE_WITHOUT_WORKSPACE_REPOS);
+      if (mainRepoMappingValue == null) {
+        return null;
+      }
+      mainRepoMapping = mainRepoMappingValue.getRepositoryMapping();
+    } else {
+      mainRepoMapping = rule.getPackage().getRepositoryMapping();
+    }
+
     IgnoredPackagePrefixesValue ignoredPackagesValue =
         (IgnoredPackagePrefixesValue) env.getValue(IgnoredPackagePrefixesValue.key());
     if (env.valuesMissing()) {
@@ -264,7 +280,8 @@ private RepositoryDirectoryValue.Builder fetchInternal(
         thread.setThreadLocal(Label.RepoMappingRecorder.class, repoMappingRecorder);
       }
 
-      new BazelStarlarkContext(BazelStarlarkContext.Phase.LOADING).storeInThread(thread); // "fetch"
+      new BazelStarlarkContext(BazelStarlarkContext.Phase.LOADING, () -> mainRepoMapping)
+          .storeInThread(thread); // "fetch"
 
       StarlarkRepositoryContext starlarkRepositoryContext =
           new StarlarkRepositoryContext(

src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java

@@ -33,7 +33,7 @@
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.packages.Attribute;
 import com.google.devtools.build.lib.packages.AttributeValueSource;
-import com.google.devtools.build.lib.packages.BazelStarlarkContext;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
 import com.google.devtools.build.lib.packages.BzlInitThreadContext;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.RuleClass;

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

@@ -17,6 +17,7 @@ java_library(
     srcs = [
         "BazelCompileContext.java",
         "BazelModuleContext.java",
+        "BazelStarlarkContext.java",
         "Label.java",
         "LabelCodec.java",
         "LabelConstants.java",

src/main/java/com/google/devtools/build/lib/cmdline/BazelModuleContext.java

@@ -44,13 +44,6 @@ public abstract class BazelModuleContext {
   /** The repository mapping applicable to the repo where the .bzl file is located in. */
   public abstract RepositoryMapping repoMapping();
 
-  /**
-   * The repository mapping applicable to the main repository. This is purely meant to support
-   * {@link Label#debugPrint}.
-   */
-  @Nullable
-  public abstract RepositoryMapping mainRepoMapping();
-
   /** Returns the name of the module's .bzl file, as provided to the parser. */
   public abstract String filename();
 
@@ -167,12 +160,11 @@ public static BazelModuleContext ofInnermostBzlOrFail(StarlarkThread thread, Str
   public static BazelModuleContext create(
       Label label,
       RepositoryMapping repoMapping,
-      @Nullable RepositoryMapping mainRepoMapping,
       String filename,
       ImmutableList<Module> loads,
       byte[] bzlTransitiveDigest) {
     return new AutoValue_BazelModuleContext(
-        label, repoMapping, mainRepoMapping, filename, loads, bzlTransitiveDigest);
+        label, repoMapping, filename, loads, bzlTransitiveDigest);
   }
 
   public final Label.PackageContext packageContext() {

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

@@ -11,9 +11,11 @@
 // 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;
+package com.google.devtools.build.lib.cmdline;
 
 import com.google.common.base.Preconditions;
+import com.google.devtools.build.lib.supplier.InterruptibleSupplier;
+import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkThread;
@@ -94,19 +96,31 @@ public void storeInThread(StarlarkThread thread) {
 
   // TODO(b/236456122): Eliminate Phase, migrate analysisRuleLabel to a separate context class.
   private final Phase phase;
+  @Nullable private final InterruptibleSupplier<RepositoryMapping> mainRepoMappingSupplier;
 
   /**
    * @param phase the phase to which this Starlark thread belongs
    */
-  public BazelStarlarkContext(Phase phase) {
+  public BazelStarlarkContext(
+      Phase phase, @Nullable InterruptibleSupplier<RepositoryMapping> mainRepoMappingSupplier) {
     this.phase = Preconditions.checkNotNull(phase);
+    this.mainRepoMappingSupplier = mainRepoMappingSupplier;
   }
 
   /** Returns the phase associated with this context. */
   public Phase getPhase() {
     return phase;
   }
 
+  /**
+   * The repository mapping applicable to the main repository. This is purely meant to support
+   * {@link Label#debugPrint}.
+   */
+  @Nullable
+  public RepositoryMapping getMainRepoMapping() throws InterruptedException {
+    return mainRepoMappingSupplier == null ? null : mainRepoMappingSupplier.get();
+  }
+
   @Override
   public String getContextForUncheckedException() {
     return phase.toString();

src/main/java/com/google/devtools/build/lib/cmdline/Label.java

@@ -661,7 +661,13 @@ public void repr(Printer printer) {
 
   @Override
   public void debugPrint(Printer printer, StarlarkThread thread) {
-    printer.append(getDisplayForm(BazelModuleContext.ofInnermostBzlOrThrow(thread).mainRepoMapping()));
+    RepositoryMapping mainRepoMapping;
+    try {
+      mainRepoMapping = BazelStarlarkContext.fromOrFail(thread).getMainRepoMapping();
+    } catch (EvalException | InterruptedException e) {
+      mainRepoMapping = null;
+    }
+    printer.append(getDisplayForm(mainRepoMapping));
   }
 
   @Override

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

@@ -15,6 +15,7 @@
 
 import com.google.devtools.build.docgen.annot.GlobalMethods;
 import com.google.devtools.build.docgen.annot.GlobalMethods.Environment;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.packages.TargetDefinitionContext.NameConflictException;

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

@@ -16,7 +16,9 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.RuleDefinitionEnvironment;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.Optional;
@@ -57,15 +59,16 @@ public final class BzlInitThreadContext extends BazelStarlarkContext
    * @param networkAllowlistForTests an allowlist for rule classes created by this thread
    * @param fragmentNameToClass a map from configuration fragment name to configuration fragment
    *     class, such as "apple" to AppleConfiguration.class
-   * @param symbolGenerator symbol generator for this context
+   * @param mainRepoMapping the repository mapping of the main repository
    */
   public BzlInitThreadContext(
       Label bzlFile,
       byte[] transitiveDigest,
       RepositoryName toolsRepository,
       Optional<Label> networkAllowlistForTests,
-      ImmutableMap<String, Class<?>> fragmentNameToClass) {
-    super(BazelStarlarkContext.Phase.LOADING);
+      ImmutableMap<String, Class<?>> fragmentNameToClass,
+      RepositoryMapping mainRepoMapping) {
+    super(BazelStarlarkContext.Phase.LOADING, () -> mainRepoMapping);
     this.bzlFile = bzlFile;
     this.transitiveDigest = transitiveDigest;
     this.toolsRepository = toolsRepository;

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

@@ -29,6 +29,7 @@
 import com.google.devtools.build.lib.bugreport.BugReport;
 import com.google.devtools.build.lib.cmdline.BazelModuleContext;
 import com.google.devtools.build.lib.cmdline.BazelModuleContext.LoadGraphVisitor;
+import com.google.devtools.build.lib.cmdline.BazelStarlarkContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
@@ -758,6 +759,7 @@ public static Builder newPackageBuilder(
       Optional<String> associatedModuleVersion,
       boolean noImplicitFileExport,
       RepositoryMapping repositoryMapping,
+      RepositoryMapping mainRepositoryMapping,
       @Nullable Semaphore cpuBoundSemaphore,
       PackageOverheadEstimator packageOverheadEstimator,
       @Nullable ImmutableMap<Location, String> generatorMap,
@@ -775,6 +777,7 @@ public static Builder newPackageBuilder(
         associatedModuleVersion,
         noImplicitFileExport,
         repositoryMapping,
+        mainRepositoryMapping,
         cpuBoundSemaphore,
         packageOverheadEstimator,
         generatorMap,
@@ -802,6 +805,7 @@ public static Builder newExternalPackageBuilder(
         /* associatedModuleVersion= */ Optional.empty(),
         noImplicitFileExport,
         /* repositoryMapping= */ mainRepoMapping,
+        /* mainRepositoryMapping= */ mainRepoMapping,
         /* cpuBoundSemaphore= */ null,
         packageOverheadEstimator,
         /* generatorMap= */ null,
@@ -825,6 +829,7 @@ public static Builder newExternalPackageBuilderForBzlmod(
             /* associatedModuleVersion= */ Optional.empty(),
             noImplicitFileExport,
             repoMapping,
+            /* mainRepoMapping= */ null,
             /* cpuBoundSemaphore= */ null,
             PackageOverheadEstimator.NOOP_ESTIMATOR,
             /* generatorMap= */ null,
@@ -1062,7 +1067,7 @@ public T intern(T sample) {
     private boolean alreadyBuilt = false;
 
     private Builder(
-        BazelStarlarkContext.Phase phase,
+        Phase phase,
         SymbolGenerator<?> symbolGenerator,
         PackageSettings packageSettings,
         PackageIdentifier id,
@@ -1072,14 +1077,15 @@ private Builder(
         Optional<String> associatedModuleVersion,
         boolean noImplicitFileExport,
         RepositoryMapping repositoryMapping,
+        @Nullable RepositoryMapping mainRepositoryMapping,
         @Nullable Semaphore cpuBoundSemaphore,
         PackageOverheadEstimator packageOverheadEstimator,
         @Nullable ImmutableMap<Location, String> generatorMap,
         // TODO(bazel-team): Config policy is an enum, what is null supposed to mean?
         // Maybe convert null -> LEGACY_OFF, assuming that's the correct default.
         @Nullable ConfigSettingVisibilityPolicy configSettingVisibilityPolicy,
         @Nullable Globber globber) {
-      super(phase);
+      super(phase, mainRepositoryMapping);
       this.symbolGenerator = symbolGenerator;
 
       Metadata metadata = new Metadata();