Canonicalize use_extension label

Canonicalize the label by adding the current module's repo_name if the
label doesn't specify a repository name. This is necessary as
ModuleExtensionUsages are grouped by the string value of this label, but
later mapped to their Label representation. If multiple strings map to
the same Label, this would result in a crash.

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

@@ -242,6 +242,8 @@ public Builder addExtensionUsage(ModuleExtensionUsage value) {
       return this;
     }
 
+    abstract ModuleKey getKey();
+
     abstract String getName();
 
     abstract Optional<String> getRepoName();

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

@@ -27,6 +27,8 @@
 import com.google.devtools.build.docgen.annot.DocumentMethods;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileGlobals.ModuleExtensionUsageBuilder.ModuleExtensionProxy;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -382,9 +384,25 @@ public void registerToolchains(Sequence<?> toolchainLabels) throws EvalException
             defaultValue = "False"),
       },
       useStarlarkThread = true)
-  public ModuleExtensionProxy useExtension(
-      String extensionBzlFile, String extensionName, boolean devDependency, StarlarkThread thread) {
+  public ModuleExtensionProxy useExtension(String rawExtensionBzlFile, String extensionName,
+      boolean devDependency, StarlarkThread thread) throws EvalException {
     hadNonModuleCall = true;
+
+    // Canonicalize the label by adding the current module's repo_name if the label doesn't
+    // specify a repository name. This is necessary as ModuleExtensionUsages are grouped by
+    // the string value of this label, but later mapped to their Label representation. If multiple
+    // strings map to the same Label, this would result in a crash.
+    String extensionBzlFile;
+    // This can't change anymore as calling module() after this results in an error.
+    String ownName = module.getRepoName().orElse(module.getName());
+    if (module.getKey().equals(ModuleKey.ROOT) && rawExtensionBzlFile.startsWith("@//")) {
+      extensionBzlFile = "@" + ownName + rawExtensionBzlFile.substring(1);
+    } else if (rawExtensionBzlFile.startsWith("//")) {
+      extensionBzlFile = "@" + ownName + rawExtensionBzlFile;
+    } else {
+      extensionBzlFile = rawExtensionBzlFile;
+    }
+
     ModuleExtensionUsageBuilder newUsageBuilder = new ModuleExtensionUsageBuilder(extensionBzlFile,
         extensionName, thread.getCallerLocation());
 

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

@@ -329,6 +329,88 @@ public void simpleExtension() throws Exception {
     assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("foo:fu bar:ba");
   }
 
+  @Test
+  public void simpleExtension_nonCanonicalLabel() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "module(name='my_module', version = '1.0')",
+        "bazel_dep(name='data_repo', version='1.0')",
+        "ext1 = use_extension('//:defs.bzl', 'ext')",
+        "ext1.tag(name='foo', data='fu')",
+        "use_repo(ext1, 'foo')",
+        "ext2 = use_extension('@my_module//:defs.bzl', 'ext')",
+        "ext2.tag(name='bar', data='ba')",
+        "use_repo(ext2, 'bar')",
+        "ext3 = use_extension('@//:defs.bzl', 'ext')",
+        "ext3.tag(name='quz', data='qu')",
+        "use_repo(ext3, 'quz')");
+    scratch.file(
+        workspaceRoot.getRelative("defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "tag = tag_class(attrs = {'name':attr.string(),'data':attr.string()})",
+        "def _ext_impl(ctx):",
+        "  for mod in ctx.modules:",
+        "    for tag in mod.tags.tag:",
+        "      data_repo(name=tag.name,data=tag.data)",
+        "ext = module_extension(implementation=_ext_impl, tag_classes={'tag':tag})");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+    scratch.file(
+        workspaceRoot.getRelative("data.bzl").getPathString(),
+        "load('@foo//:data.bzl', foo_data='data')",
+        "load('@bar//:data.bzl', bar_data='data')",
+        "load('@quz//:data.bzl', quz_data='data')",
+        "data = 'foo:'+foo_data+' bar:'+bar_data+' quz:'+quz_data");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    if (result.hasError()) {
+      throw result.getError().getException();
+    }
+    assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("foo:fu bar:ba quz:qu");
+  }
+
+  @Test
+  public void simpleExtension_nonCanonicalLabel_repoName() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "module(name='my_module', version = '1.0', repo_name='my_name')",
+        "bazel_dep(name='data_repo', version='1.0')",
+        "ext1 = use_extension('//:defs.bzl', 'ext')",
+        "ext1.tag(name='foo', data='fu')",
+        "use_repo(ext1, 'foo')",
+        "ext2 = use_extension('@my_name//:defs.bzl', 'ext')",
+        "ext2.tag(name='bar', data='ba')",
+        "use_repo(ext2, 'bar')",
+        "ext3 = use_extension('@//:defs.bzl', 'ext')",
+        "ext3.tag(name='quz', data='qu')",
+        "use_repo(ext3, 'quz')");
+    scratch.file(
+        workspaceRoot.getRelative("defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "tag = tag_class(attrs = {'name':attr.string(),'data':attr.string()})",
+        "def _ext_impl(ctx):",
+        "  for mod in ctx.modules:",
+        "    for tag in mod.tags.tag:",
+        "      data_repo(name=tag.name,data=tag.data)",
+        "ext = module_extension(implementation=_ext_impl, tag_classes={'tag':tag})");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+    scratch.file(
+        workspaceRoot.getRelative("data.bzl").getPathString(),
+        "load('@foo//:data.bzl', foo_data='data')",
+        "load('@bar//:data.bzl', bar_data='data')",
+        "load('@quz//:data.bzl', quz_data='data')",
+        "data = 'foo:'+foo_data+' bar:'+bar_data+' quz:'+quz_data");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    if (result.hasError()) {
+      throw result.getError().getException();
+    }
+    assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("foo:fu bar:ba quz:qu");
+  }
+
   @Test
   public void multipleModules() throws Exception {
     scratch.file(

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

@@ -473,7 +473,7 @@ public void testModuleExtensions_good() throws Exception {
                 .setRegistry(registry)
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext1")
                         .setLocation(Location.fromFileLineColumn("[email protected]/MODULE.bazel", 2, 23))
                         .setImports(ImmutableBiMap.of("repo1", "repo1"))
@@ -491,7 +491,7 @@ public void testModuleExtensions_good() throws Exception {
                         .build())
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext2")
                         .setLocation(Location.fromFileLineColumn("[email protected]/MODULE.bazel", 5, 23))
                         .setImports(ImmutableBiMap.of("other_repo1", "repo1", "repo2", "repo2"))
@@ -582,7 +582,7 @@ public void testModuleExtensions_duplicateProxy_asRoot() throws Exception {
                 .setKey(ModuleKey.ROOT)
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@//:defs.bzl")
                         .setExtensionName("myext")
                         .setLocation(Location.fromFileLineColumn("<root>/MODULE.bazel", 1, 23))
                         .setImports(
@@ -656,7 +656,7 @@ public void testModuleExtensions_duplicateProxy_asDep() throws Exception {
                 .setRegistry(registry)
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext")
                         .setLocation(Location.fromFileLineColumn("[email protected]/MODULE.bazel", 5, 23))
                         .setImports(ImmutableBiMap.of("beta", "beta", "delta", "delta"))