Allow WORKSPACE and WORKSPACE-loaded .bzl files to see Bzlmod root module's mappings (bazelbuild#17818)

Currently, we evaluate WORKSPACE (`WorkspaceFileFunction`) and WORKSPACE-loaded .bzl files (`BzlLoadFunction` with `BzlLoadValue.KeyForWorkspace`) with the repo mapping purely computed from previous WORKSPACE chunks. This is unlike BUILD-loaded .bzl files from WORKSPACE-defined repos (`BzlLoadFunction` with `BzlLoadValue.KeyForBuild`, which is the same as `RepositoryMappingFunction`), which take the mappings of the Bzlmod root module into account. This CL fixes that discrepancy by doing the same "repo mapping composition" everywhere.

Fixes bazelbuild#17655

RELNOTES: Fixed an issue where WORKSPACE and WORKSPACE-loaded .bzl files couldn't see the Bzlmod root module's mappings when Bzlmod is enabled.
PiperOrigin-RevId: 515318590
Change-Id: I4babc922f6cdb932d17ce18d9a9d9d427dbed2eb

Co-authored-by: Googler <[email protected]>

Author: keertk

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

@@ -49,16 +49,18 @@ public abstract class RepositoryMapping {
   abstract RepositoryName ownerRepo();
 
   public static RepositoryMapping create(
-      Map<String, RepositoryName> repositoryMapping, RepositoryName ownerRepo) {
-    return new AutoValue_RepositoryMapping(
-        ImmutableMap.copyOf(Preconditions.checkNotNull(repositoryMapping)),
-        Preconditions.checkNotNull(ownerRepo));
+      Map<String, RepositoryName> entries, RepositoryName ownerRepo) {
+    return createInternal(
+        Preconditions.checkNotNull(entries), Preconditions.checkNotNull(ownerRepo));
   }
 
-  public static RepositoryMapping createAllowingFallback(
-      Map<String, RepositoryName> repositoryMapping) {
-    return new AutoValue_RepositoryMapping(
-        ImmutableMap.copyOf(Preconditions.checkNotNull(repositoryMapping)), null);
+  public static RepositoryMapping createAllowingFallback(Map<String, RepositoryName> entries) {
+    return createInternal(Preconditions.checkNotNull(entries), null);
+  }
+
+  private static RepositoryMapping createInternal(
+      Map<String, RepositoryName> entries, RepositoryName ownerRepo) {
+    return new AutoValue_RepositoryMapping(ImmutableMap.copyOf(entries), ownerRepo);
   }
 
   /**
@@ -68,7 +70,7 @@ public static RepositoryMapping createAllowingFallback(
   public RepositoryMapping withAdditionalMappings(Map<String, RepositoryName> additionalMappings) {
     HashMap<String, RepositoryName> allMappings = new HashMap<>(additionalMappings);
     allMappings.putAll(entries());
-    return new AutoValue_RepositoryMapping(ImmutableMap.copyOf(allMappings), ownerRepo());
+    return createInternal(allMappings, ownerRepo());
   }
 
   /**
@@ -114,4 +116,24 @@ public Optional<String> getInverse(RepositoryName postMappingName) {
         .map(Entry::getKey)
         .findFirst();
   }
+
+  /**
+   * Creates a new {@link RepositoryMapping} instance that is the equivalent of composing this
+   * {@link RepositoryMapping} with another one. That is, {@code a.composeWith(b).get(name) ===
+   * b.get(a.get(name))} (treating {@code b} as allowing fallback).
+   *
+   * <p>Since we're treating the result of {@code a.get(name)} as an apparent repo name instead of a
+   * canonical repo name, this only really makes sense when {@code a} does not use strict deps (i.e.
+   * allows fallback).
+   */
+  public RepositoryMapping composeWith(RepositoryMapping other) {
+    Preconditions.checkArgument(
+        !usesStrictDeps(), "only an allow-fallback mapping can be composed with other mappings");
+    HashMap<String, RepositoryName> entries = new HashMap<>(other.entries());
+    for (Map.Entry<String, RepositoryName> entry : entries().entrySet()) {
+      RepositoryName mappedName = other.get(entry.getValue().getName());
+      entries.put(entry.getKey(), mappedName.isVisible() ? mappedName : entry.getValue());
+    }
+    return createInternal(entries, null);
+  }
 }

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

@@ -165,7 +165,7 @@ public static BzlLoadFunction create(
         // (b) The memory overhead of the extra Skyframe node and edge per bzl file is pure
         // waste.
         new InliningAndCachingGetter(packageFactory, hashFunction, bzlCompileCache),
-        /*cachedBzlLoadDataManager=*/ null);
+        /* cachedBzlLoadDataManager= */ null);
   }
 
   public static BzlLoadFunction createForInlining(
@@ -188,7 +188,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       throws SkyFunctionException, InterruptedException {
     BzlLoadValue.Key key = (BzlLoadValue.Key) skyKey.argument();
     try {
-      return computeInternal(key, env, /*inliningState=*/ null);
+      return computeInternal(key, env, /* inliningState= */ null);
     } catch (BzlLoadFailedException e) {
       throw new BzlLoadFunctionException(e);
     }
@@ -455,12 +455,12 @@ private InliningState(
     static InliningState create(Environment env) {
       return new InliningState(
           new RecordingSkyFunctionEnvironment(env, x -> {}, x -> {}, x -> {}),
-          /*cachedDataBuilder=*/ null,
-          /*loadStack=*/ new LinkedHashSet<>(),
-          /*successfulLoads=*/ new HashMap<>(),
-          /*unsuccessfulLoads=*/ new HashSet<>(),
+          /* cachedDataBuilder= */ null,
+          /* loadStack= */ new LinkedHashSet<>(),
+          /* successfulLoads= */ new HashMap<>(),
+          /* unsuccessfulLoads= */ new HashSet<>(),
           // No parent value to mutate
-          /*childCachedDataHandler=*/ x -> {});
+          /* childCachedDataHandler= */ x -> {});
     }
 
     /**
@@ -598,7 +598,10 @@ private StarlarkBuiltinsValue getBuiltins(
             env.getValueOrThrow(StarlarkBuiltinsValue.key(), BuiltinsFailedException.class);
       } else {
         return StarlarkBuiltinsFunction.computeInline(
-            StarlarkBuiltinsValue.key(), inliningState, packageFactory, /*bzlLoadFunction=*/ this);
+            StarlarkBuiltinsValue.key(),
+            inliningState,
+            packageFactory,
+            /* bzlLoadFunction= */ this);
       }
     } catch (BuiltinsFailedException e) {
       throw BzlLoadFailedException.builtinsFailed(key.getLabel(), e);
@@ -776,7 +779,7 @@ private BzlLoadValue computeInternalWithCompiledBzl(
         loadValues,
         loadKeys,
         programLoads,
-        /*demoteErrorsToWarnings=*/ !builtins.starlarkSemantics.getBool(
+        /* demoteErrorsToWarnings= */ !builtins.starlarkSemantics.getBool(
             BuildLanguageOptions.CHECK_BZL_VISIBILITY),
         env.getListener());
 
@@ -845,7 +848,8 @@ private BzlLoadValue computeInternalWithCompiledBzl(
   private static RepositoryMapping getRepositoryMapping(
       BzlLoadValue.Key key, StarlarkSemantics semantics, Environment env)
       throws InterruptedException {
-    if (key.isBuiltins() && !semantics.getBool(BuildLanguageOptions.ENABLE_BZLMOD)) {
+    boolean bzlmod = semantics.getBool(BuildLanguageOptions.ENABLE_BZLMOD);
+    if (key.isBuiltins() && !bzlmod) {
       // Without Bzlmod, builtins .bzls never have a repo mapping defined for them, so return
       // without requesting a RepositoryMappingValue. (NB: In addition to being a slight
       // optimization, this avoids adding a reverse dependency on the special //external package,
@@ -862,21 +866,37 @@ private static RepositoryMapping getRepositoryMapping(
     if (key instanceof BzlLoadValue.KeyForWorkspace) {
       // Still during workspace file evaluation
       BzlLoadValue.KeyForWorkspace keyForWorkspace = (BzlLoadValue.KeyForWorkspace) key;
+      RepositoryMapping pureWorkspaceMapping;
       if (keyForWorkspace.getWorkspaceChunk() == 0) {
         // There is no previous workspace chunk
-        return RepositoryMapping.ALWAYS_FALLBACK;
+        pureWorkspaceMapping = RepositoryMapping.ALWAYS_FALLBACK;
       } else {
         SkyKey workspaceFileKey =
             WorkspaceFileValue.key(
                 keyForWorkspace.getWorkspacePath(), keyForWorkspace.getWorkspaceChunk() - 1);
         WorkspaceFileValue workspaceFileValue = (WorkspaceFileValue) env.getValue(workspaceFileKey);
         // Note: we know for sure that the requested WorkspaceFileValue is fully computed so we do
         // not need to check if it is null
-        return RepositoryMapping.createAllowingFallback(
-            workspaceFileValue.getRepositoryMapping().getOrDefault(repoName, ImmutableMap.of()));
-        // NOTE(wyv): this means that, in the WORKSPACE file, we can't load from a repo generated by
-        // bzlmod. If that's a problem, we should "fall back" to the bzlmod case below.
+        pureWorkspaceMapping =
+            RepositoryMapping.createAllowingFallback(
+                workspaceFileValue
+                    .getRepositoryMapping()
+                    .getOrDefault(repoName, ImmutableMap.of()));
+      }
+      if (!bzlmod) {
+        // Without Bzlmod, we just return the mapping purely computed from WORKSPACE stuff.
+        return pureWorkspaceMapping;
+      }
+      // If Bzlmod is in play, we need to make sure that pure WORKSPACE mapping is composed with the
+      // root module's mapping (just like how all WORKSPACE repos can see what the root module sees
+      // _after_ WORKSPACE evaluation).
+      RepositoryMappingValue rootModuleMappingValue =
+          (RepositoryMappingValue)
+              env.getValue(RepositoryMappingValue.KEY_FOR_ROOT_MODULE_WITHOUT_WORKSPACE_REPOS);
+      if (rootModuleMappingValue == null) {
+        return null;
       }
+      return pureWorkspaceMapping.composeWith(rootModuleMappingValue.getRepositoryMapping());
     }
 
     if (key instanceof BzlLoadValue.KeyForBzlmod) {
@@ -903,9 +923,8 @@ private static RepositoryMapping getRepositoryMapping(
       }
     }
 
-    // This is either a .bzl loaded from BUILD files, or a .bzl loaded for bzlmod (in which case the
-    // .bzl file *has* to be from a Bazel module anyway). So we can just use the full repo mapping
-    // from RepositoryMappingFunction.
+    // This is either a .bzl loaded from BUILD files, or a .bzl loaded for bzlmod, so we can just
+    // use the full repo mapping from RepositoryMappingFunction.
     RepositoryMappingValue repositoryMappingValue =
         (RepositoryMappingValue) env.getValue(RepositoryMappingValue.key(repoName));
     if (repositoryMappingValue == null) {

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

@@ -15,7 +15,6 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleResolutionValue;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionId;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleKey;
@@ -213,26 +212,17 @@ private SkyValue computeFromWorkspace(
     if (externalPackage.containsErrors()) {
       throw new RepositoryMappingFunctionException();
     }
+    RepositoryMapping workspaceMapping =
+        RepositoryMapping.createAllowingFallback(
+            externalPackage.getRepositoryMapping(repositoryName));
     if (rootModuleRepoMapping == null) {
       // This means Bzlmod is disabled.
-      return RepositoryMappingValue.withMapping(
-          RepositoryMapping.createAllowingFallback(
-              externalPackage.getRepositoryMapping(repositoryName)));
+      return RepositoryMappingValue.withMapping(workspaceMapping);
     }
     // If Bzlmod is in play, we need to make sure that WORKSPACE repos see all repos that root
     // module can see, taking care to compose the existing WORKSPACE mapping with the main repo
     // mapping from Bzlmod.
-    return RepositoryMappingValue.withMapping(
-        RepositoryMapping.createAllowingFallback(
-                Maps.transformValues(
-                    externalPackage.getRepositoryMapping(repositoryName),
-                    toRepo -> {
-                      RepositoryName mappedName = rootModuleRepoMapping.get(toRepo.getName());
-                      // If the Bzlmod-only main repo mapping doesn't contain this repo, don't touch
-                      // it.
-                      return mappedName.isVisible() ? mappedName : toRepo;
-                    }))
-            .withAdditionalMappings(rootModuleRepoMapping));
+    return RepositoryMappingValue.withMapping(workspaceMapping.composeWith(rootModuleRepoMapping));
   }
 
   private static Optional<ModuleExtensionId> maybeGetModuleExtensionForRepo(

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

@@ -120,14 +120,14 @@ public SkyValue compute(SkyKey skyKey, Environment env)
             RepositoryDelegatorFunction.RESOLVED_FILE_INSTEAD_OF_WORKSPACE.get(env));
     boolean useWorkspaceResolvedFile = resolvedFile.isPresent();
 
+    final boolean bzlmod = starlarkSemantics.getBool(BuildLanguageOptions.ENABLE_BZLMOD);
     boolean useWorkspaceBzlmodFile = false;
     RootedPath workspaceBzlmodFile =
         RootedPath.toRootedPath(
             workspaceFile.getRoot(),
             workspaceFile.getRootRelativePath().replaceName("WORKSPACE.bzlmod"));
     // We only need to check WORKSPACE.bzlmod when the resolved file isn't used.
-    if (!useWorkspaceResolvedFile
-        && starlarkSemantics.getBool(BuildLanguageOptions.ENABLE_BZLMOD)) {
+    if (!useWorkspaceResolvedFile && bzlmod) {
       FileValue workspaceBzlmodFileValue =
           (FileValue) env.getValue(FileValue.key(workspaceBzlmodFile));
       if (workspaceBzlmodFileValue == null) {
@@ -266,6 +266,15 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                   .getRepositoryMapping()
                   .getOrDefault(RepositoryName.MAIN, ImmutableMap.of()));
     }
+    if (bzlmod) {
+      RepositoryMappingValue rootModuleMapping =
+          (RepositoryMappingValue)
+              env.getValue(RepositoryMappingValue.KEY_FOR_ROOT_MODULE_WITHOUT_WORKSPACE_REPOS);
+      if (rootModuleMapping == null) {
+        return null;
+      }
+      repoMapping = repoMapping.composeWith(rootModuleMapping.getRepositoryMapping());
+    }
 
     Package.Builder builder =
         packageFactory.newExternalPackageBuilder(
@@ -274,12 +283,12 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     if (chunks.isEmpty()) {
       return new WorkspaceFileValue(
           buildAndReportEvents(builder, env),
-          /* loadedModules = */ ImmutableMap.<String, Module>of(),
-          /* loadToChunkMap = */ ImmutableMap.<String, Integer>of(),
-          /* bindings = */ ImmutableMap.<String, Object>of(),
+          /* loadedModules= */ ImmutableMap.<String, Module>of(),
+          /* loadToChunkMap= */ ImmutableMap.<String, Integer>of(),
+          /* bindings= */ ImmutableMap.<String, Object>of(),
           workspaceFile,
-          /* idx = */ 0, // first fragment
-          /* hasNext = */ false);
+          /* idx= */ 0, // first fragment
+          /* hasNext= */ false);
     }
 
     List<StarlarkFile> chunk = chunks.get(key.getIndex());

src/test/java/com/google/devtools/build/lib/cmdline/RepositoryMappingTest.java

@@ -47,7 +47,7 @@ public void neverFallback() throws Exception {
   }
 
   @Test
-  public void additionalMappings() throws Exception {
+  public void additionalMappings_basic() throws Exception {
     RepositoryMapping mapping =
         RepositoryMapping.create(
                 ImmutableMap.of("A", RepositoryName.create("com_foo_bar_a")),
@@ -59,4 +59,34 @@ public void additionalMappings() throws Exception {
         .isEqualTo(
             RepositoryName.create("C").toNonVisible(RepositoryName.create("fake_owner_repo")));
   }
+
+  @Test
+  public void additionalMappings_precedence() throws Exception {
+    RepositoryMapping mapping =
+        RepositoryMapping.createAllowingFallback(ImmutableMap.of("A", RepositoryName.create("A1")))
+            .withAdditionalMappings(ImmutableMap.of("A", RepositoryName.create("A2")));
+    assertThat(mapping.get("A")).isEqualTo(RepositoryName.create("A1"));
+  }
+
+  @Test
+  public void composeWith() throws Exception {
+    RepositoryMapping mapping =
+        RepositoryMapping.createAllowingFallback(
+                ImmutableMap.of(
+                    "A", RepositoryName.create("A_mapped"), "B", RepositoryName.create("B_mapped")))
+            .composeWith(
+                RepositoryMapping.create(
+                    ImmutableMap.of(
+                        "A",
+                        RepositoryName.create("A_mapped_differently"),
+                        "A_mapped",
+                        RepositoryName.create("A_mapped_twice"),
+                        "C",
+                        RepositoryName.create("C_mapped")),
+                    RepositoryName.create("blah")));
+    assertThat(mapping.get("A")).isEqualTo(RepositoryName.create("A_mapped_twice"));
+    assertThat(mapping.get("B")).isEqualTo(RepositoryName.create("B_mapped"));
+    assertThat(mapping.get("C")).isEqualTo(RepositoryName.create("C_mapped"));
+    assertThat(mapping.get("D")).isEqualTo(RepositoryName.create("D"));
+  }
 }

src/test/java/com/google/devtools/build/lib/rules/repository/BUILD

@@ -44,6 +44,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe:package_lookup_function",
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_function",
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_value",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:repository_mapping_function",
         "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
         "//src/main/java/com/google/devtools/build/lib/skyframe:skyframe_cluster",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/repository",

src/test/java/com/google/devtools/build/lib/rules/repository/RepositoryDelegatorTest.java

@@ -68,6 +68,7 @@
 import com.google.devtools.build.lib.skyframe.PackageLookupFunction.CrossRepositoryLabelViolationStrategy;
 import com.google.devtools.build.lib.skyframe.PrecomputedFunction;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
+import com.google.devtools.build.lib.skyframe.RepositoryMappingFunction;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.WorkspaceFileFunction;
 import com.google.devtools.build.lib.starlarkbuildapi.repository.RepositoryBootstrap;
@@ -220,6 +221,7 @@ public void setupDelegator() throws Exception {
                     new IgnoredPackagePrefixesFunction(
                         /*ignoredPackagePrefixesFile=*/ PathFragment.EMPTY_FRAGMENT))
                 .put(SkyFunctions.RESOLVED_HASH_VALUES, new ResolvedHashesFunction())
+                .put(SkyFunctions.REPOSITORY_MAPPING, new RepositoryMappingFunction())
                 .put(
                     SkyFunctions.MODULE_FILE,
                     new ModuleFileFunction(registryFactory, rootPath, ImmutableMap.of()))

src/test/java/com/google/devtools/build/lib/skyframe/RepositoryMappingFunctionTest.java

@@ -62,25 +62,29 @@ private EvaluationResult<RepositoryMappingValue> eval(SkyKey key)
             ModifiedFileSet.builder().modify(PathFragment.create("WORKSPACE")).build(),
             Root.fromPath(rootDirectory));
     return SkyframeExecutorTestUtils.evaluate(
-        getSkyframeExecutor(), key, /*keepGoing=*/ false, reporter);
+        getSkyframeExecutor(), key, /* keepGoing= */ false, reporter);
   }
 
   @Before
   public void setUpForBzlmod() throws Exception {
     setBuildLanguageOptions("--enable_bzlmod");
     scratch.file("MODULE.bazel");
+  }
+
+  @Override
+  protected ImmutableList<PrecomputedValue.Injected> extraPrecomputedValues() throws Exception {
     registry = FakeRegistry.DEFAULT_FACTORY.newFakeRegistry(scratch.dir("modules").getPathString());
-    ModuleFileFunction.REGISTRIES.set(
-        getSkyframeExecutor().getDifferencerForTesting(), ImmutableList.of(registry.getUrl()));
-    ModuleFileFunction.IGNORE_DEV_DEPS.set(getSkyframeExecutor().getDifferencerForTesting(), false);
-    ModuleFileFunction.MODULE_OVERRIDES.set(
-        getSkyframeExecutor().getDifferencerForTesting(), ImmutableMap.of());
-    BazelModuleResolutionFunction.ALLOWED_YANKED_VERSIONS.set(
-        getSkyframeExecutor().getDifferencerForTesting(), ImmutableList.of());
-    BazelModuleResolutionFunction.CHECK_DIRECT_DEPENDENCIES.set(
-        getSkyframeExecutor().getDifferencerForTesting(), CheckDirectDepsMode.WARNING);
-    BazelModuleResolutionFunction.BAZEL_COMPATIBILITY_MODE.set(
-        getSkyframeExecutor().getDifferencerForTesting(), BazelCompatibilityMode.ERROR);
+    return ImmutableList.of(
+        PrecomputedValue.injected(
+            ModuleFileFunction.REGISTRIES, ImmutableList.of(registry.getUrl())),
+        PrecomputedValue.injected(ModuleFileFunction.IGNORE_DEV_DEPS, false),
+        PrecomputedValue.injected(ModuleFileFunction.MODULE_OVERRIDES, ImmutableMap.of()),
+        PrecomputedValue.injected(
+            BazelModuleResolutionFunction.ALLOWED_YANKED_VERSIONS, ImmutableList.of()),
+        PrecomputedValue.injected(
+            BazelModuleResolutionFunction.CHECK_DIRECT_DEPENDENCIES, CheckDirectDepsMode.WARNING),
+        PrecomputedValue.injected(
+            BazelModuleResolutionFunction.BAZEL_COMPATIBILITY_MODE, BazelCompatibilityMode.ERROR));
   }
 
   @Override