Bzlmod: Allow multiple use_extensions on the same extension (bazelbuild#14945)

Fixes bazelbuild#14635.

PiperOrigin-RevId: 432155351

Author: Wyverald

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

@@ -31,7 +31,11 @@ public abstract class ModuleExtensionUsage {
   /** The name of the extension. */
   public abstract String getExtensionName();
 
-  /** The location where this proxy object was created (by the {@code use_extension} call). */
+  /**
+   * The location where this proxy object was created (by the {@code use_extension} call). Note that
+   * if there were multiple {@code use_extension} calls on same extension, then this only stores the
+   * location of the first one.
+   */
   public abstract Location getLocation();
 
   /**

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

@@ -273,20 +273,25 @@ public void bazelDep(
   public ModuleExtensionProxy useExtension(
       String extensionBzlFile, String extensionName, boolean devDependency, StarlarkThread thread)
       throws EvalException {
+    ModuleExtensionProxy newProxy =
+        new ModuleExtensionProxy(extensionBzlFile, extensionName, thread.getCallerLocation());
+
+    if (ignoreDevDeps && devDependency) {
+      // This is a no-op proxy.
+      return newProxy;
+    }
+
+    // Find an existing proxy object corresponding to this extension.
     for (ModuleExtensionProxy proxy : extensionProxies) {
       if (proxy.extensionBzlFile.equals(extensionBzlFile)
           && proxy.extensionName.equals(extensionName)) {
-        throw Starlark.errorf("this extension is already being used at %s", proxy.location);
+        return proxy;
       }
     }
-    ModuleExtensionProxy proxy =
-        new ModuleExtensionProxy(extensionBzlFile, extensionName, thread.getCallerLocation());
 
-    if (!(ignoreDevDeps && devDependency)) {
-      extensionProxies.add(proxy);
-    }
-
-    return proxy;
+    // If no such proxy exists, we can just use a new one.
+    extensionProxies.add(newProxy);
+    return newProxy;
   }
 
   @StarlarkBuiltin(name = "module_extension_proxy", documented = false)
@@ -368,9 +373,7 @@ public String getErrorMessageForUnknownField(String field) {
       parameters = {
         @Param(
             name = "extension_proxy",
-            doc =
-                "A module extension proxy object returned by a <code>get_module_extension</code>"
-                    + " call."),
+            doc = "A module extension proxy object returned by a <code>use_extension</code> call."),
       },
       extraPositionals = @Param(name = "args", doc = "The names of the repos to import."),
       extraKeywords =

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunctionTest.java

@@ -499,21 +499,83 @@ public void testModuleExtensions_good() throws Exception {
   }
 
   @Test
-  public void testModuleExtensions_duplicateProxy() throws Exception {
+  public void testModuleExtensions_duplicateProxy_asRoot() throws Exception {
+    scratch.file(
+        rootDirectory.getRelative("MODULE.bazel").getPathString(),
+        "myext1 = use_extension('//:defs.bzl','myext',dev_dependency=True)",
+        "use_repo(myext1, 'alpha')",
+        "myext2 = use_extension('//:defs.bzl','myext')",
+        "use_repo(myext2, 'beta')",
+        "myext3 = use_extension('//:defs.bzl','myext',dev_dependency=True)",
+        "use_repo(myext3, 'gamma')",
+        "myext4 = use_extension('//:defs.bzl','myext')",
+        "use_repo(myext4, 'delta')");
+    ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of());
+
+    SkyKey skyKey = ModuleFileValue.KEY_FOR_ROOT_MODULE;
+    EvaluationResult<ModuleFileValue> result =
+        driver.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    if (result.hasError()) {
+      throw result.getError().getException();
+    }
+    ModuleFileValue moduleFileValue = result.get(skyKey);
+    assertThat(moduleFileValue.getModule())
+        .isEqualTo(
+            Module.builder()
+                .setKey(ModuleKey.ROOT)
+                .addExtensionUsage(
+                    ModuleExtensionUsage.builder()
+                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionName("myext")
+                        .setLocation(Location.fromFileLineColumn("<root>/MODULE.bazel", 1, 23))
+                        .setImports(
+                            ImmutableBiMap.of(
+                                "alpha", "alpha", "beta", "beta", "gamma", "gamma", "delta",
+                                "delta"))
+                        .build())
+                .build());
+  }
+
+  @Test
+  public void testModuleExtensions_duplicateProxy_asDep() throws Exception {
     FakeRegistry registry =
         registryFactory
             .newFakeRegistry("/foo")
             .addModule(
                 createModuleKey("mymod", "1.0"),
                 "module(name='mymod',version='1.0')",
-                "myext1 = use_extension('//:defs.bzl','myext')",
-                "myext2 = use_extension('//:defs.bzl','myext')");
+                "myext1 = use_extension('//:defs.bzl','myext',dev_dependency=True)",
+                "use_repo(myext1, 'alpha')",
+                "myext2 = use_extension('//:defs.bzl','myext')",
+                "use_repo(myext2, 'beta')",
+                "myext3 = use_extension('//:defs.bzl','myext',dev_dependency=True)",
+                "use_repo(myext3, 'gamma')",
+                "myext4 = use_extension('//:defs.bzl','myext')",
+                "use_repo(myext4, 'delta')");
     ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of(registry.getUrl()));
 
     SkyKey skyKey = ModuleFileValue.key(createModuleKey("mymod", "1.0"), null);
-    reporter.removeHandler(failFastHandler); // expect failures
-    driver.evaluate(ImmutableList.of(skyKey), evaluationContext);
-    assertContainsEvent("this extension is already being used at");
+    EvaluationResult<ModuleFileValue> result =
+        driver.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    if (result.hasError()) {
+      throw result.getError().getException();
+    }
+    ModuleFileValue moduleFileValue = result.get(skyKey);
+    assertThat(moduleFileValue.getModule())
+        .isEqualTo(
+            Module.builder()
+                .setName("mymod")
+                .setVersion(Version.parse("1.0"))
+                .setKey(createModuleKey("mymod", "1.0"))
+                .setRegistry(registry)
+                .addExtensionUsage(
+                    ModuleExtensionUsage.builder()
+                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionName("myext")
+                        .setLocation(Location.fromFileLineColumn("[email protected]/MODULE.bazel", 4, 23))
+                        .setImports(ImmutableBiMap.of("beta", "beta", "delta", "delta"))
+                        .build())
+                .build());
   }
 
   @Test