Omit module version from repo name if it is unique

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

@@ -30,10 +30,6 @@ public abstract class AbridgedModule {
 
   public abstract ModuleKey getKey();
 
-  public final RepositoryName getCanonicalRepoName() {
-    return getKey().getCanonicalRepoName();
-  }
-
   public static AbridgedModule from(Module module) {
     return new AutoValue_AbridgedModule(module.getName(), module.getVersion(), module.getKey());
   }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -281,6 +281,7 @@ java_library(
     deps = [
         ":common",
         ":module_extension",
+        "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",

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

@@ -16,8 +16,12 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import static com.google.common.base.Strings.nullToEmpty;
+import static com.google.common.collect.ImmutableBiMap.toImmutableBiMap;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static java.util.stream.Collectors.counting;
+import static java.util.stream.Collectors.groupingBy;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -26,11 +30,13 @@
 import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableTable;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.packages.LabelConverter;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps.Code;
@@ -44,6 +50,7 @@
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /**
@@ -113,13 +120,11 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       depGraph = selectionResult.getResolvedDepGraph();
     }
 
-    ImmutableMap<RepositoryName, ModuleKey> canonicalRepoNameLookup =
-        depGraph.keySet().stream()
-            .collect(toImmutableMap(ModuleKey::getCanonicalRepoName, key -> key));
-
+    ImmutableBiMap<RepositoryName, ModuleKey> canonicalRepoNameLookup =
+        computeCanonicalRepoNameLookup(depGraph);
     ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> extensionUsagesById;
     try {
-      extensionUsagesById = getExtensionUsagesById(depGraph);
+      extensionUsagesById = getExtensionUsagesById(depGraph, canonicalRepoNameLookup.inverse());
     } catch (ExternalDepsException e) {
       throw new BazelDepGraphFunctionException(e, Transience.PERSISTENT);
     }
@@ -214,13 +219,23 @@ static BzlmodFlagsAndEnvVars getFlagsAndEnvVars(Environment env) throws Interrup
   public static ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
       getExtensionUsagesById(ImmutableMap<ModuleKey, Module> depGraph)
           throws ExternalDepsException {
+    return getExtensionUsagesById(depGraph, computeCanonicalRepoNameLookup(depGraph).inverse());
+  }
+
+  private static ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
+      getExtensionUsagesById(
+          ImmutableMap<ModuleKey, Module> depGraph,
+          ImmutableMap<ModuleKey, RepositoryName> moduleKeyToRepositoryNames)
+          throws ExternalDepsException {
     ImmutableTable.Builder<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
         extensionUsagesTableBuilder = ImmutableTable.builder();
     for (Module module : depGraph.values()) {
+      RepositoryMapping repoMapping =
+          module.getRepoMappingWithBazelDepsOnly(moduleKeyToRepositoryNames);
       LabelConverter labelConverter =
           new LabelConverter(
-              PackageIdentifier.create(module.getCanonicalRepoName(), PathFragment.EMPTY_FRAGMENT),
-              module.getRepoMappingWithBazelDepsOnly());
+              PackageIdentifier.create(repoMapping.ownerRepo(), PathFragment.EMPTY_FRAGMENT),
+              module.getRepoMappingWithBazelDepsOnly(moduleKeyToRepositoryNames));
       for (ModuleExtensionUsage usage : module.getExtensionUsages()) {
         ModuleExtensionId moduleExtensionId;
         try {
@@ -249,6 +264,27 @@ static BzlmodFlagsAndEnvVars getFlagsAndEnvVars(Environment env) throws Interrup
     return extensionUsagesTableBuilder.buildOrThrow();
   }
 
+  private static ImmutableBiMap<RepositoryName, ModuleKey> computeCanonicalRepoNameLookup(
+      ImmutableMap<ModuleKey, Module> depGraph) {
+    // Find those modules of which there is only a single version in the dep graph. Currently, the
+    // only way to have multiple versions of a module in the dep graph is via
+    // multiple_version_override.
+    ImmutableSet<String> uniqueVersionModules =
+        depGraph.keySet().stream()
+            .collect(groupingBy(ModuleKey::getName, counting()))
+            .entrySet()
+            .stream()
+            .filter(entry -> entry.getValue() == 1)
+            .map(Entry::getKey)
+            .collect(toImmutableSet());
+
+    return depGraph.keySet().stream()
+        .collect(
+            toImmutableBiMap(
+                key -> key.getCanonicalRepoName(uniqueVersionModules.contains(key.getName())),
+                key -> key));
+  }
+
   private ImmutableBiMap<String, ModuleExtensionId> calculateUniqueNameForUsedExtensionId(
       ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> extensionUsagesById) {
     // Calculate a unique name for each used extension id with the following property that is

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

@@ -18,6 +18,7 @@
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 
 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.common.collect.ImmutableTable;
@@ -45,7 +46,7 @@ public static BazelDepGraphValue create(
       ImmutableMap<ModuleExtensionId, String> extensionUniqueNames) {
     return new AutoValue_BazelDepGraphValue(
         depGraph,
-        canonicalRepoNameLookup,
+        ImmutableBiMap.copyOf(canonicalRepoNameLookup),
         abridgedModules,
         extensionUsagesTable,
         extensionUniqueNames);
@@ -67,7 +68,9 @@ public static BazelDepGraphValue createEmptyDepGraph() {
 
     ImmutableMap<RepositoryName, ModuleKey> canonicalRepoNameLookup =
         emptyDepGraph.keySet().stream()
-            .collect(toImmutableMap(ModuleKey::getCanonicalRepoName, key -> key));
+            .collect(
+                toImmutableMap(
+                    key -> key.getCanonicalRepoName(/* hasUniqueVersion */ true), key -> key));
 
     return BazelDepGraphValue.create(
         emptyDepGraph,
@@ -83,8 +86,8 @@ public static BazelDepGraphValue createEmptyDepGraph() {
    */
   public abstract ImmutableMap<ModuleKey, Module> getDepGraph();
 
-  /** A mapping from a canonical repo name to the key of the module backing it. */
-  public abstract ImmutableMap<RepositoryName, ModuleKey> getCanonicalRepoNameLookup();
+  /** A mapping from a canonical repo name to the key of the module backing it and back. */
+  public abstract ImmutableBiMap<RepositoryName, ModuleKey> getCanonicalRepoNameLookup();
 
   /** All modules in the same order as {@link #getDepGraph}, but with limited information. */
   public abstract ImmutableList<AbridgedModule> getAbridgedModules();
@@ -124,7 +127,7 @@ public final RepositoryMapping getFullRepoMapping(ModuleKey key) {
     }
     return getDepGraph()
         .get(key)
-        .getRepoMappingWithBazelDepsOnly()
+        .getRepoMappingWithBazelDepsOnly(getCanonicalRepoNameLookup().inverse())
         .withAdditionalMappings(mapping.buildOrThrow());
   }
 }

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

@@ -82,7 +82,11 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
                         AugmentedModule::getName,
                         Collectors.mapping(AugmentedModule::getKey, toImmutableSet()))));
 
-    return BazelModuleInspectorValue.create(depGraph, modulesIndex, extensionToRepoInternalNames);
+    return BazelModuleInspectorValue.create(
+        depGraph,
+        modulesIndex,
+        extensionToRepoInternalNames,
+        depGraphValue.getCanonicalRepoNameLookup().inverse());
   }
 
   public static ImmutableMap<ModuleKey, AugmentedModule> computeAugmentedGraph(

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

@@ -21,6 +21,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -42,9 +43,10 @@ public abstract class BazelModuleInspectorValue implements SkyValue {
   public static BazelModuleInspectorValue create(
       ImmutableMap<ModuleKey, AugmentedModule> depGraph,
       ImmutableMap<String, ImmutableSet<ModuleKey>> modulesIndex,
-      ImmutableSetMultimap<ModuleExtensionId, String> extensionToRepoInternalNames) {
+      ImmutableSetMultimap<ModuleExtensionId, String> extensionToRepoInternalNames,
+      ImmutableMap<ModuleKey, RepositoryName> moduleKeyToCanonicalNames) {
     return new AutoValue_BazelModuleInspectorValue(
-        depGraph, modulesIndex, extensionToRepoInternalNames);
+        depGraph, modulesIndex, extensionToRepoInternalNames, moduleKeyToCanonicalNames);
   }
 
   /**
@@ -69,6 +71,9 @@ public static BazelModuleInspectorValue create(
    */
   public abstract ImmutableSetMultimap<ModuleExtensionId, String> getExtensionToRepoInternalNames();
 
+  /** A mapping from a module key to the canonical repository name of the module repository. */
+  public abstract ImmutableMap<ModuleKey, RepositoryName> getModuleKeyToCanonicalNames();
+
   /**
    * A wrapper for {@link Module}, augmented with references to dependants (and also those who are
    * not used in the final dep graph).

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

@@ -23,6 +23,7 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import java.util.Map;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /**
@@ -48,23 +49,25 @@ public abstract class Module extends ModuleBase {
    * Returns a {@link RepositoryMapping} with only Bazel module repos and no repos from module
    * extensions. For the full mapping, see {@link BazelDepGraphValue#getFullRepoMapping}.
    */
-  public final RepositoryMapping getRepoMappingWithBazelDepsOnly() {
+  public final RepositoryMapping getRepoMappingWithBazelDepsOnly(
+      ImmutableMap<ModuleKey, RepositoryName> moduleKeyToRepositoryNames) {
     ImmutableMap.Builder<String, RepositoryName> mapping = ImmutableMap.builder();
     // If this is the root module, then the main repository should be visible as `@`.
     if (getKey().equals(ModuleKey.ROOT)) {
       mapping.put("", RepositoryName.MAIN);
     }
     // Every module should be able to reference itself as @<module repo name>.
     // If this is the root module, this perfectly falls into @<module repo name> => @
+    RepositoryName owner = moduleKeyToRepositoryNames.get(getKey());
     if (!getRepoName().isEmpty()) {
-      mapping.put(getRepoName(), getCanonicalRepoName());
+      mapping.put(getRepoName(), owner);
     }
     for (Map.Entry<String, ModuleKey> dep : getDeps().entrySet()) {
       // Special note: if `dep` is actually the root module, its ModuleKey would be ROOT whose
       // canonicalRepoName is the empty string. This perfectly maps to the main repo ("@").
-      mapping.put(dep.getKey(), dep.getValue().getCanonicalRepoName());
+      mapping.put(dep.getKey(), moduleKeyToRepositoryNames.get(dep.getValue()));
     }
-    return RepositoryMapping.create(mapping.buildOrThrow(), getCanonicalRepoName());
+    return RepositoryMapping.create(mapping.buildOrThrow(), owner);
   }
 
   /**

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

@@ -49,10 +49,6 @@ abstract class ModuleBase {
    */
   public abstract ModuleKey getKey();
 
-  public final RepositoryName getCanonicalRepoName() {
-    return getKey().getCanonicalRepoName();
-  }
-
   /**
    * The name of the repository representing this module, as seen by the module itself. By default,
    * the name of the repo is the name of the module. This can be specified to ease migration for

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

@@ -25,7 +25,9 @@
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.NonRootModuleFileValue;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
+import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.DotBazelFileSyntaxChecker;
@@ -42,6 +44,7 @@
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -275,13 +278,17 @@ private SkyValue computeForRootModule(StarlarkSemantics starlarkSemantics, Envir
     if (rootOverride != null) {
       throw errorf(Code.BAD_MODULE, "invalid override for the root module found: %s", rootOverride);
     }
+    // A module with a non-registry override always has a unique version across the entire dep
+    // graph.
     ImmutableMap<RepositoryName, String> nonRegistryOverrideCanonicalRepoNameLookup =
         Maps.filterValues(overrides, override -> override instanceof NonRegistryOverride)
             .keySet()
             .stream()
             .collect(
                 toImmutableMap(
-                    name -> ModuleKey.create(name, Version.EMPTY).getCanonicalRepoName(),
+                    name ->
+                        ModuleKey.create(name, Version.EMPTY)
+                            .getCanonicalRepoName(/* hasUniqueVersion= */ true),
                     name -> name));
     return RootModuleFileValue.create(
         module, moduleFileHash, overrides, nonRegistryOverrideCanonicalRepoNameLookup);
@@ -359,7 +366,9 @@ private GetModuleFileResult getModuleFile(
     // If there is a non-registry override for this module, we need to fetch the corresponding repo
     // first and read the module file from there.
     if (override instanceof NonRegistryOverride) {
-      RepositoryName canonicalRepoName = key.getCanonicalRepoName();
+      // A module with a non-registry override always has a unique version across the entire dep
+      // graph.
+      RepositoryName canonicalRepoName = key.getCanonicalRepoName(/* hasUniqueVersion */ true);
       RepositoryDirectoryValue repoDir =
           (RepositoryDirectoryValue) env.getValue(RepositoryDirectoryValue.key(canonicalRepoName));
       if (repoDir == null) {
@@ -372,10 +381,14 @@ private GetModuleFileResult getModuleFile(
         return null;
       }
       GetModuleFileResult result = new GetModuleFileResult();
+      Label moduleFileLabel =
+          Label.createUnvalidated(
+              PackageIdentifier.create(canonicalRepoName, PathFragment.EMPTY_FRAGMENT),
+              LabelConstants.MODULE_DOT_BAZEL_FILE_NAME.getBaseName());
       result.moduleFile =
           ModuleFile.create(
               readModuleFile(moduleFilePath.asPath()),
-              key.moduleFileLabel().getUnambiguousCanonicalForm());
+              moduleFileLabel.getUnambiguousCanonicalForm());
       return result;
     }