Address comments

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

@@ -19,9 +19,9 @@
 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 static java.util.stream.Collectors.toSet;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -30,7 +30,6 @@
 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;
@@ -50,7 +49,7 @@
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.Map;
 import java.util.Map.Entry;
-import java.util.function.Function;
+import java.util.Set;
 import javax.annotation.Nullable;
 
 /**
@@ -266,22 +265,32 @@ static BzlmodFlagsAndEnvVars getFlagsAndEnvVars(Environment env) throws Interrup
 
   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 =
+    // Find modules with multiple versions in the dep graph. Currently, the only source of such
+    // modules is multiple_version_override.
+    Set<String> multipleVersionsModules =
         depGraph.keySet().stream()
             .collect(groupingBy(ModuleKey::getName, counting()))
             .entrySet()
             .stream()
-            .filter(entry -> entry.getValue() == 1)
+            .filter(entry -> entry.getValue() > 1)
             .map(Entry::getKey)
-            .collect(toImmutableSet());
+            .collect(toSet());
 
+    // If there is a unique version of this module in the entire dep graph, we elide the version
+    // from the canonical repository name. This has a number of benefits:
+    // * It prevents the output base from being polluted with repository directories corresponding
+    //   to outdated versions of modules, which can be large and would otherwise only be cleaned
+    //   up by the discouraged bazel clean --expunge.
+    // * It improves cache hit rates by ensuring that a module update doesn't e.g. cause the paths
+    //   of all toolchains provided by its extensions to change, which would result in widespread
+    //   cache misses on every update.
     return depGraph.keySet().stream()
         .collect(
             toImmutableBiMap(
-                key -> key.getCanonicalRepoName(uniqueVersionModules.contains(key.getName())),
+                key ->
+                    multipleVersionsModules.contains(key.getName())
+                        ? key.getCanonicalRepoNameWithVersion()
+                        : key.getCanonicalRepoNameWithoutVersion(),
                 key -> key));
   }
 

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

@@ -70,7 +70,10 @@ public static BazelDepGraphValue createEmptyDepGraph() {
         emptyDepGraph.keySet().stream()
             .collect(
                 toImmutableMap(
-                    key -> key.getCanonicalRepoName(/* hasUniqueVersion */ true), key -> key));
+                    // All modules in the empty dep graph (just the root module) have an empty
+                    // version, so the choice of including it in the canonical repo name does not
+                    // matter.
+                    ModuleKey::getCanonicalRepoNameWithoutVersion, key -> key));
 
     return BazelDepGraphValue.create(
         emptyDepGraph,

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

@@ -278,17 +278,16 @@ 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(
+                    // A module with a non-registry override always has a unique version across the
+                    // entire dep graph.
                     name ->
-                        ModuleKey.create(name, Version.EMPTY)
-                            .getCanonicalRepoName(/* hasUniqueVersion= */ true),
+                        ModuleKey.create(name, Version.EMPTY).getCanonicalRepoNameWithoutVersion(),
                     name -> name));
     return RootModuleFileValue.create(
         module, moduleFileHash, overrides, nonRegistryOverrideCanonicalRepoNameLookup);
@@ -368,7 +367,7 @@ private GetModuleFileResult getModuleFile(
     if (override instanceof NonRegistryOverride) {
       // A module with a non-registry override always has a unique version across the entire dep
       // graph.
-      RepositoryName canonicalRepoName = key.getCanonicalRepoName(/* hasUniqueVersion */ true);
+      RepositoryName canonicalRepoName = key.getCanonicalRepoNameWithoutVersion();
       RepositoryDirectoryValue repoDir =
           (RepositoryDirectoryValue) env.getValue(RepositoryDirectoryValue.key(canonicalRepoName));
       if (repoDir == null) {

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

@@ -74,24 +74,37 @@ public final String toString() {
     return getName() + "@" + (getVersion().isEmpty() ? "_" : getVersion().toString());
   }
 
-  /** Returns the canonical name of the repo backing this module. */
-  public RepositoryName getCanonicalRepoName(boolean hasUniqueVersion) {
+  /**
+   * Returns the canonical name of the repo backing this module, including its version. This name is
+   * always guaranteed to be unique.
+   */
+  public RepositoryName getCanonicalRepoNameWithVersion() {
+    return getCanonicalRepoName(/* includeVersion= */ true);
+  }
+
+  /**
+   * Returns the canonical name of the repo backing this module, excluding its version. This name is
+   * only guaranteed to be unique when there is a single version of the module in the entire dep
+   * graph.
+   */
+  public RepositoryName getCanonicalRepoNameWithoutVersion() {
+    return getCanonicalRepoName(/* includeVersion= */ false);
+  }
+
+  private RepositoryName getCanonicalRepoName(boolean includeVersion) {
     if (WELL_KNOWN_MODULES.containsKey(getName())) {
       return WELL_KNOWN_MODULES.get(getName());
     }
     if (ROOT.equals(this)) {
       return RepositoryName.MAIN;
     }
     String suffix;
-    if (hasUniqueVersion) {
-      // If there is a unique version of this module in the entire dep graph, we elide the version
-      // from the canonical repository name. This has a number of benefits:
-      // * It prevents the output base from being polluted with repository directories corresponding
-      //   to outdated versions of modules, which can be large and would otherwise only be cleaned
-      //   up by the discouraged bazel clean --expunge.
-      // * It improves cache hit rates by ensuring that a module update doesn't e.g. cause the paths
-      //   of all toolchains provided by its extensions to change, which would result in widespread
-      //   cache misses on every update.
+    if (includeVersion) {
+      // getVersion().isEmpty() is true only for modules with non-registry overrides, which enforce
+      // that there is a single version of the module in the dep graph.
+      Preconditions.checkState(!getVersion().isEmpty());
+      suffix = getVersion().toString();
+    } else {
       // This results in canonical repository names such as `rules_foo~` for the module `rules_foo`.
       // This particular format is chosen since:
       // * The tilde ensures that canonical and apparent repository names can be distinguished even
@@ -107,11 +120,6 @@ public RepositoryName getCanonicalRepoName(boolean hasUniqueVersion) {
       //   would only be discovered when used with a `multiple_version_override`, which is very
       //   rarely used.
       suffix = "";
-    } else {
-      // getVersion().isEmpty() is true only for modules with non-registry overrides, which imply
-      // enforce there is a single version of the module in the dep graph.
-      Preconditions.checkState(!getVersion().isEmpty());
-      suffix = getVersion().toString();
     }
     return RepositoryName.createUnvalidated(String.format("%s~%s", getName(), suffix));
   }

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

@@ -307,7 +307,7 @@ public static RepositoryMapping createRepositoryMapping(ModuleKey key, String...
       mappingBuilder.put(names[i], RepositoryName.createUnvalidated(names[i + 1]));
     }
     return RepositoryMapping.create(
-        mappingBuilder.buildOrThrow(), key.getCanonicalRepoName(/* hasUniqueVersion= */ false));
+        mappingBuilder.buildOrThrow(), key.getCanonicalRepoNameWithoutVersion());
   }
 
   public static TagClass createTagClass(Attribute... attrs) {

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

@@ -81,10 +81,7 @@ public RepoSpec getRepoSpec(ModuleKey key, ExtendedEventHandler eventHandler) {
         .setAttributes(
             AttributeValues.create(
                 ImmutableMap.of(
-                    "path",
-                    rootPath
-                        + "/"
-                        + key.getCanonicalRepoName(/* hasUniqueVersion= */ false).getName())))
+                    "path", rootPath + "/" + key.getCanonicalRepoNameWithVersion().getName())))
         .build();
   }
 

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

@@ -43,9 +43,7 @@ public void getRepoMapping() throws Exception {
     assertThat(
             module.getRepoMappingWithBazelDepsOnly(
                 Stream.of(key, fooKey, barKey, ModuleKey.ROOT)
-                    .collect(
-                        toImmutableMap(
-                            k -> k, k -> k.getCanonicalRepoName(/* hasUniqueVersion= */ false)))))
+                    .collect(toImmutableMap(k -> k, ModuleKey::getCanonicalRepoNameWithVersion))))
         .isEqualTo(
             createRepositoryMapping(
                 key,
@@ -72,9 +70,7 @@ public void getRepoMapping_asMainModule() throws Exception {
     assertThat(
             module.getRepoMappingWithBazelDepsOnly(
                 Stream.of(ModuleKey.ROOT, fooKey, barKey)
-                    .collect(
-                        toImmutableMap(
-                            k -> k, k -> k.getCanonicalRepoName(/* hasUniqueVersion= */ false)))))
+                    .collect(toImmutableMap(k -> k, ModuleKey::getCanonicalRepoNameWithVersion))))
         .isEqualTo(
             createRepositoryMapping(
                 ModuleKey.ROOT,

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

@@ -102,8 +102,8 @@ public void basic() throws Exception {
             extension,
             module.getRepoMappingWithBazelDepsOnly(
                 ImmutableMap.of(
-                    fooKey, fooKey.getCanonicalRepoName(/* hasUniqueVersion= */ false),
-                    barKey, barKey.getCanonicalRepoName(/* hasUniqueVersion= */ false))),
+                    fooKey, fooKey.getCanonicalRepoNameWithVersion(),
+                    barKey, barKey.getCanonicalRepoNameWithVersion())),
             usage);
 
     assertThat(moduleProxy.getName()).isEqualTo("foo");
@@ -150,8 +150,7 @@ public void unknownTagClass() throws Exception {
                     abridgedModule,
                     extension,
                     module.getRepoMappingWithBazelDepsOnly(
-                        ImmutableMap.of(
-                            fooKey, fooKey.getCanonicalRepoName(/* hasUniqueVersion= */ false))),
+                        ImmutableMap.of(fooKey, fooKey.getCanonicalRepoNameWithVersion())),
                     usage));
     assertThat(e).hasMessageThat().contains("does not have a tag class named blep");
   }

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

@@ -89,8 +89,7 @@ public void resolve_good() throws Exception {
             .buildOrThrow();
     ImmutableMap<ModuleKey, RepositoryName> moduleKeyToCanonicalNames =
         depGraph.keySet().stream()
-            .collect(
-                toImmutableMap(k -> k, k -> k.getCanonicalRepoName(/* hasUniqueVersion= */ false)));
+            .collect(toImmutableMap(k -> k, ModuleKey::getCanonicalRepoNameWithVersion));
     ImmutableBiMap<String, ModuleKey> baseModuleDeps = ImmutableBiMap.of("fred", key);
     ImmutableBiMap<String, ModuleKey> baseModuleUnusedDeps = ImmutableBiMap.of();
 
@@ -119,8 +118,7 @@ public void resolve_badLabel() throws Exception {
             .buildOrThrow();
     ImmutableMap<ModuleKey, RepositoryName> moduleKeyToCanonicalNames =
         depGraph.keySet().stream()
-            .collect(
-                toImmutableMap(k -> k, k -> k.getCanonicalRepoName(/* hasUniqueVersion= */ false)));
+            .collect(toImmutableMap(k -> k, ModuleKey::getCanonicalRepoNameWithVersion));
     ImmutableBiMap<String, ModuleKey> baseModuleDeps = ImmutableBiMap.of("fred", key);
     ImmutableBiMap<String, ModuleKey> baseModuleUnusedDeps = ImmutableBiMap.of();
 
@@ -164,8 +162,7 @@ public void resolve_noneOrtooManyModules() throws Exception {
             .buildOrThrow();
     ImmutableMap<ModuleKey, RepositoryName> moduleKeyToCanonicalNames =
         depGraph.keySet().stream()
-            .collect(
-                toImmutableMap(k -> k, k -> k.getCanonicalRepoName(/* hasUniqueVersion= */ false)));
+            .collect(toImmutableMap(k -> k, ModuleKey::getCanonicalRepoNameWithVersion));
     ImmutableBiMap<String, ModuleKey> baseModuleDeps =
         ImmutableBiMap.of("foo1", foo1, "foo2", foo2);
     ImmutableBiMap<String, ModuleKey> baseModuleUnusedDeps = ImmutableBiMap.of();

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

@@ -80,8 +80,7 @@ public void converter() throws Exception {
 
   ImmutableMap<ModuleKey, RepositoryName> moduleKeyToCanonicalNames =
       depGraph.keySet().stream()
-          .collect(
-              toImmutableMap(k -> k, k -> k.getCanonicalRepoName(/* hasUniqueVersion= */ false)));
+          .collect(toImmutableMap(k -> k, ModuleKey::getCanonicalRepoNameWithVersion));
   ImmutableBiMap<String, ModuleKey> baseModuleDeps = ImmutableBiMap.of("fred", foo2);
   ImmutableBiMap<String, ModuleKey> baseModuleUnusedDeps = ImmutableBiMap.of("fred", foo1);
   RepositoryMapping rootMapping = createRepositoryMapping(ModuleKey.ROOT, "fred", "foo~2.0");
@@ -271,7 +270,7 @@ public void resolve_apparentRepoName_notFound() throws Exception {
 
   @Test
   public void resolve_canonicalRepoName_good() throws Exception {
-    var arg = CanonicalRepoName.create(foo2.getCanonicalRepoName(/* hasUniqueVersion= */ false));
+    var arg = CanonicalRepoName.create(foo2.getCanonicalRepoNameWithVersion());
 
     assertThat(
             arg.resolveToModuleKeys(
@@ -312,7 +311,7 @@ public void resolve_canonicalRepoName_notFound() throws Exception {
 
   @Test
   public void resolve_canonicalRepoName_unused() throws Exception {
-    var arg = CanonicalRepoName.create(foo1.getCanonicalRepoName(/* hasUniqueVersion= */ false));
+    var arg = CanonicalRepoName.create(foo1.getCanonicalRepoNameWithVersion());
 
     // Without --include_unused, this doesn't resolve, as [email protected] has been replaced by [email protected].
     assertThat(