Add override_repo and inject_repo

Work towards bazelbuild#19301
Fixes bazelbuild#23580

RELNOTES: `override_repo` and `inject_repo` can be used to override and inject repos in module extensions.

Closes bazelbuild#23534.

PiperOrigin-RevId: 678139661
Change-Id: Iea7caca949c00e701f056c1037e273fee9740e93

Author: fmeum

site/en/external/extension.md

@@ -174,6 +174,63 @@ several repo visibility rules:
         the repo visible to the module instead of the extension-generated repo
         of the same name.
 
+### Overriding and injecting module extension repos
+
+The root module can use
+[`override_repo`](/rules/lib/globals/module#override_repo) and
+[`inject_repo`](/rules/lib/globals/module#inject_repo) to override or inject
+module extension repos.
+
+#### Example: Replacing `rules_java`'s `java_tools` with a vendored copy
+
+```python
+# MODULE.bazel
+local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "local_repository")
+local_repository(
+  name = "my_java_tools",
+  path = "vendor/java_tools",
+)
+
+bazel_dep(name = "rules_java", version = "7.11.1")
+java_toolchains = use_extension("@rules_java//java:extension.bzl", "toolchains")
+
+override_repo(java_toolchains, remote_java_tools = "my_java_tools")
+```
+
+#### Example: Patch a Go dependency to depend on `@zlib` instead of the system zlib
+
+```python
+# MODULE.bazel
+bazel_dep(name = "gazelle", version = "0.38.0")
+bazel_dep(name = "zlib", version = "1.3.1.bcr.3")
+
+go_deps = use_extension("@gazelle//:extensions.bzl", "go_deps")
+go_deps.from_file(go_mod = "//:go.mod")
+go_deps.module_override(
+  patches = [
+    "//patches:my_module_zlib.patch",
+  ],
+  path = "example.com/my_module",
+)
+use_repo(go_deps, ...)
+
+inject_repo(go_deps, "zlib")
+```
+
+```diff
+# patches/my_module_zlib.patch
+--- a/BUILD.bazel
++++ b/BUILD.bazel
+@@ -1,6 +1,6 @@
+ go_binary(
+     name = "my_module",
+     importpath = "example.com/my_module",
+     srcs = ["my_module.go"],
+-    copts = ["-lz"],
++    cdeps = ["@zlib"],
+ )
+```
+
 ## Best practices
 
 This section describes best practices when writing extensions so they are

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

@@ -80,7 +80,12 @@ public SkyValue compute(SkyKey skyKey, Environment env)
         canonicalRepoNameLookup,
         depGraph.values().stream().map(AbridgedModule::from).collect(toImmutableList()),
         extensionUsagesById,
-        extensionUniqueNames.inverse());
+        extensionUniqueNames.inverse(),
+        resolveRepoOverrides(
+            depGraph,
+            extensionUsagesById,
+            extensionUniqueNames.inverse(),
+            canonicalRepoNameLookup));
   }
 
   private static ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
@@ -198,6 +203,38 @@ private static String makeUniqueNameCandidate(ModuleExtensionId id, int attempt)
                 + extensionNameDisambiguator);
   }
 
+  private static ImmutableTable<ModuleExtensionId, String, RepositoryName> resolveRepoOverrides(
+      ImmutableMap<ModuleKey, Module> depGraph,
+      ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> extensionUsagesTable,
+      ImmutableMap<ModuleExtensionId, String> extensionUniqueNames,
+      ImmutableBiMap<RepositoryName, ModuleKey> canonicalRepoNameLookup) {
+    RepositoryMapping rootModuleMappingWithoutOverrides =
+        BazelDepGraphValue.getRepositoryMapping(
+            ModuleKey.ROOT,
+            depGraph,
+            extensionUsagesTable,
+            extensionUniqueNames,
+            canonicalRepoNameLookup,
+            // ModuleFileFunction ensures that repos that override other repos are not themselves
+            // overridden, so we can safely pass an empty table here instead of resolving chains
+            // of overrides.
+            ImmutableTable.of());
+    ImmutableTable.Builder<ModuleExtensionId, String, RepositoryName> repoOverridesBuilder =
+        ImmutableTable.builder();
+    for (var extensionId : extensionUsagesTable.rowKeySet()) {
+      var rootUsage = extensionUsagesTable.row(extensionId).get(ModuleKey.ROOT);
+      if (rootUsage != null) {
+        for (var override : rootUsage.getRepoOverrides().entrySet()) {
+          repoOverridesBuilder.put(
+              extensionId,
+              override.getKey(),
+              rootModuleMappingWithoutOverrides.get(override.getValue().overridingRepoName()));
+        }
+      }
+    }
+    return repoOverridesBuilder.buildOrThrow();
+  }
+
   static class BazelDepGraphFunctionException extends SkyFunctionException {
     BazelDepGraphFunctionException(ExternalDepsException e, Transience transience) {
       super(e, transience);

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

@@ -41,13 +41,15 @@ public static BazelDepGraphValue create(
       ImmutableMap<RepositoryName, ModuleKey> canonicalRepoNameLookup,
       ImmutableList<AbridgedModule> abridgedModules,
       ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> extensionUsagesTable,
-      ImmutableMap<ModuleExtensionId, String> extensionUniqueNames) {
+      ImmutableMap<ModuleExtensionId, String> extensionUniqueNames,
+      ImmutableTable<ModuleExtensionId, String, RepositoryName> repoOverrides) {
     return new AutoValue_BazelDepGraphValue(
         depGraph,
         ImmutableBiMap.copyOf(canonicalRepoNameLookup),
         abridgedModules,
         extensionUsagesTable,
-        extensionUniqueNames);
+        extensionUniqueNames,
+        repoOverrides);
   }
 
   public static BazelDepGraphValue createEmptyDepGraph() {
@@ -71,7 +73,8 @@ public static BazelDepGraphValue createEmptyDepGraph() {
         canonicalRepoNameLookup,
         ImmutableList.of(),
         ImmutableTable.of(),
-        ImmutableMap.of());
+        ImmutableMap.of(),
+        ImmutableTable.of());
   }
 
   /**
@@ -103,27 +106,54 @@ public static BazelDepGraphValue createEmptyDepGraph() {
    */
   public abstract ImmutableMap<ModuleExtensionId, String> getExtensionUniqueNames();
 
+  /**
+   * For each module extension, a mapping from the name of the repo exported by the extension to the
+   * canonical name of the repo that should override it (if any).
+   */
+  public abstract ImmutableTable<ModuleExtensionId, String, RepositoryName> getRepoOverrides();
+
   /**
    * Returns the full {@link RepositoryMapping} for the given module, including repos from Bazel
    * module deps and module extensions.
    */
   public final RepositoryMapping getFullRepoMapping(ModuleKey key) {
+    return getRepositoryMapping(
+        key,
+        getDepGraph(),
+        getExtensionUsagesTable(),
+        getExtensionUniqueNames(),
+        getCanonicalRepoNameLookup(),
+        getRepoOverrides());
+  }
+
+  static RepositoryMapping getRepositoryMapping(
+      ModuleKey key,
+      ImmutableMap<ModuleKey, Module> depGraph,
+      ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> extensionUsagesTable,
+      ImmutableMap<ModuleExtensionId, String> extensionUniqueNames,
+      ImmutableBiMap<RepositoryName, ModuleKey> canonicalRepoNameLookup,
+      ImmutableTable<ModuleExtensionId, String, RepositoryName> repoOverrides) {
     ImmutableMap.Builder<String, RepositoryName> mapping = ImmutableMap.builder();
     for (Map.Entry<ModuleExtensionId, ModuleExtensionUsage> extIdAndUsage :
-        getExtensionUsagesTable().column(key).entrySet()) {
+        extensionUsagesTable.column(key).entrySet()) {
       ModuleExtensionId extensionId = extIdAndUsage.getKey();
       ModuleExtensionUsage usage = extIdAndUsage.getValue();
-      String repoNamePrefix = getExtensionUniqueNames().get(extensionId) + "+";
+      String repoNamePrefix = extensionUniqueNames.get(extensionId) + "+";
       for (ModuleExtensionUsage.Proxy proxy : usage.getProxies()) {
         for (Map.Entry<String, String> entry : proxy.getImports().entrySet()) {
-          String canonicalRepoName = repoNamePrefix + entry.getValue();
-          mapping.put(entry.getKey(), RepositoryName.createUnvalidated(canonicalRepoName));
+          RepositoryName defaultCanonicalRepoName =
+              RepositoryName.createUnvalidated(repoNamePrefix + entry.getValue());
+          mapping.put(
+              entry.getKey(),
+              repoOverrides
+                  .row(extensionId)
+                  .getOrDefault(entry.getValue(), defaultCanonicalRepoName));
         }
       }
     }
-    return getDepGraph()
+    return depGraph
         .get(key)
-        .getRepoMappingWithBazelDepsOnly(getCanonicalRepoNameLookup().inverse())
+        .getRepoMappingWithBazelDepsOnly(canonicalRepoNameLookup.inverse())
         .withAdditionalMappings(mapping.buildOrThrow());
   }
 }

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

@@ -71,6 +71,7 @@ record RepoRuleCall(
   private final String repoPrefix;
   private final PackageIdentifier basePackageId;
   private final RepositoryMapping baseRepoMapping;
+  private final ImmutableMap<String, RepositoryName> repoOverrides;
   private final BlazeDirectories directories;
   private final ExtendedEventHandler eventHandler;
   private final Map<String, RepoRuleCall> deferredRepos = new LinkedHashMap<>();
@@ -80,6 +81,7 @@ public ModuleExtensionEvalStarlarkThreadContext(
       String repoPrefix,
       PackageIdentifier basePackageId,
       RepositoryMapping baseRepoMapping,
+      ImmutableMap<String, RepositoryName> repoOverrides,
       RepositoryMapping mainRepoMapping,
       BlazeDirectories directories,
       ExtendedEventHandler eventHandler) {
@@ -88,6 +90,7 @@ public ModuleExtensionEvalStarlarkThreadContext(
     this.repoPrefix = repoPrefix;
     this.basePackageId = basePackageId;
     this.baseRepoMapping = baseRepoMapping;
+    this.repoOverrides = repoOverrides;
     this.directories = directories;
     this.eventHandler = eventHandler;
   }
@@ -133,13 +136,15 @@ public ImmutableMap<String, RepoSpec> createRepos(StarlarkSemantics starlarkSema
     // Make it possible to refer to extension repos in the label attributes of another extension
     // repo. Wrapping a label in Label(...) ensures that it is evaluated with respect to the
     // containing module's repo mapping instead.
-    var extensionRepos =
+    ImmutableMap.Builder<String, RepositoryName> entries = ImmutableMap.builder();
+    entries.putAll(baseRepoMapping.entries());
+    entries.putAll(
         Maps.asMap(
             deferredRepos.keySet(),
-            apparentName -> RepositoryName.createUnvalidated(repoPrefix + apparentName));
+            apparentName -> RepositoryName.createUnvalidated(repoPrefix + apparentName)));
+    entries.putAll(repoOverrides);
     RepositoryMapping fullRepoMapping =
-        RepositoryMapping.create(extensionRepos, baseRepoMapping.ownerRepo())
-            .withAdditionalMappings(baseRepoMapping);
+        RepositoryMapping.create(entries.buildKeepingLast(), baseRepoMapping.ownerRepo());
     // LINT.ThenChange(//src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionRepoMappingEntriesFunction.java)
 
     ImmutableMap.Builder<String, RepoSpec> repoSpecs = ImmutableMap.builder();

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

@@ -74,6 +74,7 @@ private ModuleExtensionRepoMappingEntriesValue computeRepoMappingEntries(
     ImmutableMap.Builder<String, RepositoryName> entries = ImmutableMap.builder();
     entries.putAll(bazelDepGraphValue.getFullRepoMapping(moduleKey).entries());
     entries.putAll(extensionEvalValue.getCanonicalRepoNameToInternalNames().inverse());
+    entries.putAll(bazelDepGraphValue.getRepoOverrides().row(extensionId));
     return ModuleExtensionRepoMappingEntriesValue.create(entries.buildKeepingLast(), moduleKey);
     // LINT.ThenChange(//src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionEvalStarlarkThreadContext.java)
   }

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

@@ -19,6 +19,7 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
@@ -130,6 +131,24 @@ public final boolean getHasNonDevUseExtension() {
     return getProxies().stream().anyMatch(p -> !p.isDevDependency());
   }
 
+  /**
+   * Represents a repo that overrides another repo within the scope of the extension.
+   *
+   * @param overridingRepoName The apparent name of the overriding repo in the root module.
+   * @param mustExist Whether this override should apply to an existing repo.
+   * @param location The location of the {@code override_repo} or {@code inject_repo} call.
+   */
+  @GenerateTypeAdapter
+  public record RepoOverride(String overridingRepoName, boolean mustExist, Location location) {}
+
+  /**
+   * Contains information about overrides that apply to repos generated by this extension. Keyed by
+   * the extension-local repo name.
+   *
+   * <p>This is only non-empty for root module usages.
+   */
+  public abstract ImmutableMap<String, RepoOverride> getRepoOverrides();
+
   public abstract Builder toBuilder();
 
   public static Builder builder() {
@@ -152,6 +171,11 @@ ModuleExtensionUsage trimForEvaluation() {
         // Extension implementation functions do not see the imports, they are only validated
         // against the set of generated repos in a validation step that comes afterward.
         .setProxies(ImmutableList.of())
+        // Tracked in SingleExtensionUsagesValue instead, using canonical instead of apparent names.
+        // Whether this override must apply to an existing repo as well as its source location also
+        // don't influence the evaluation of the extension as they are checked in
+        // SingleExtensionFunction.
+        .setRepoOverrides(ImmutableMap.of())
         .build();
   }
 
@@ -185,6 +209,9 @@ public Builder addTag(Tag value) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public abstract Builder setRepoOverrides(ImmutableMap<String, RepoOverride> repoOverrides);
+
     public abstract ModuleExtensionUsage build();
   }
 }

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

@@ -435,7 +435,7 @@ public static RootModuleFileValue evaluateRootModuleFile(
     try {
       module = moduleThreadContext.buildModule(/* registry= */ null);
     } catch (EvalException e) {
-      eventHandler.handle(Event.error(e.getMessageWithStack()));
+      eventHandler.handle(Event.error(e.getInnermostLocation(), e.getMessageWithStack()));
       throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for the root module");
     }
     for (ModuleExtensionUsage usage : module.getExtensionUsages()) {
@@ -521,7 +521,7 @@ private static ModuleThreadContext execModuleFile(
           });
       compiledRootModuleFile.runOnThread(thread);
     } catch (EvalException e) {
-      eventHandler.handle(Event.error(e.getMessageWithStack()));
+      eventHandler.handle(Event.error(e.getInnermostLocation(), e.getMessageWithStack()));
       throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for %s", moduleKey);
     }
     return context;