Address review comments

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

@@ -1,6 +1,9 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -129,8 +132,8 @@ public void evaluate(Collection<ModuleExtensionUsage> usages, Set<String> allRep
   Optional<Event> generateFixupMessage(Collection<ModuleExtensionUsage> usages,
       Set<String> allRepos) throws EvalException {
     var rootUsages = usages.stream()
-        .filter(usage -> usage.getModuleKey().equals(ModuleKey.ROOT))
-        .collect(ImmutableList.toImmutableList());
+        .filter(usage -> usage.getUsingModule().equals(ModuleKey.ROOT))
+        .collect(toImmutableList());
     if (rootUsages.isEmpty()) {
       // The root module doesn't use the current extension. Do not suggest fixes as the user isn't
       // expected to modify any other module's MODULE.bazel file.
@@ -153,11 +156,11 @@ private static Optional<Event> generateFixupMessage(List<ModuleExtensionUsage> r
       Set<String> allRepos, Set<String> expectedImports, Set<String> expectedDevImports) {
     var actualDevImports = rootUsages.stream()
         .flatMap(usage -> usage.getDevImports().stream())
-        .collect(ImmutableSet.toImmutableSet());
+        .collect(toImmutableSet());
     var actualImports = rootUsages.stream()
         .flatMap(usage -> usage.getImports().keySet().stream())
         .filter(repo -> !actualDevImports.contains(repo))
-        .collect(ImmutableSet.toImmutableSet());
+        .collect(toImmutableSet());
 
     // All label strings that map to the same Label are equivalent for buildozer as it implements
     // the same normalization of label strings with no or empty repo name.
@@ -197,15 +200,15 @@ private static Optional<Event> generateFixupMessage(List<ModuleExtensionUsage> r
         Sets.difference(allExpectedImports, allActualImports));
     if (!missingImports.isEmpty()) {
       message += String.format(
-          "Not imported, but declared as direct dependencies (may cause the build to fail):\n    %s\n\n",
+          "Not imported, but marked as direct dependencies by the extension (may cause the build to fail):\n    %s\n\n",
           String.join(", ", missingImports));
     }
 
     var indirectDepImports = ImmutableSortedSet.copyOf(
         Sets.difference(Sets.intersection(allActualImports, allRepos), allExpectedImports));
     if (!indirectDepImports.isEmpty()) {
       message += String.format(
-          "Imported, but not declared as direct dependencies:\n    %s\n\n",
+          "Imported, but not marked as direct dependencies by the extension:\n    %s\n\n",
           String.join(", ", indirectDepImports));
     }
 

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

@@ -36,6 +36,9 @@ public abstract class ModuleExtensionUsage {
   /** The name of the extension. */
   public abstract String getExtensionName();
 
+  /** The module that contains this particular extension usage. */
+  public abstract ModuleKey getUsingModule();
+
   /**
    * 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
@@ -61,16 +64,6 @@ public abstract class ModuleExtensionUsage {
    */
   public abstract ImmutableList<Tag> getTags();
 
-  public ModuleKey getModuleKey() {
-    String file = getLocation().file();
-    Preconditions.checkState(file.endsWith("/MODULE.bazel"));
-    try {
-      return ModuleKey.fromString(file.substring(0, file.length() - "/MODULE.bazel".length()));
-    } catch (Version.ParseException e) {
-      throw new IllegalStateException("Unexpected location for module extension usage: " + file, e);
-    }
-  }
-
   public static Builder builder() {
     return new AutoValue_ModuleExtensionUsage.Builder();
   }
@@ -85,6 +78,8 @@ public abstract static class Builder {
 
     public abstract Builder setExtensionName(String value);
 
+    public abstract Builder setUsingModule(ModuleKey value);
+
     public abstract Builder setLocation(Location value);
 
     public abstract Builder setImports(ImmutableBiMap<String, String> value);

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

@@ -452,6 +452,7 @@ ModuleExtensionUsage buildUsage() {
       return ModuleExtensionUsage.builder()
           .setExtensionBzlFile(extensionBzlFile)
           .setExtensionName(extensionName)
+          .setUsingModule(module.getKey())
           .setLocation(location)
           .setImports(ImmutableBiMap.copyOf(imports))
           .setDevImports(devImports.build())

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

@@ -78,6 +78,7 @@ java_library(
         "//src/main/java/net/starlark/java/syntax",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
         "//src/test/java/com/google/devtools/build/lib/testutil",
+        "//src/test/java/com/google/devtools/build/lib/testutil:JunitUtils",
         "//third_party:auto_value",
         "//third_party:caffeine",
         "//third_party:gson",

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

@@ -236,6 +236,7 @@ private static ModuleExtensionUsage createModuleExtensionUsage(
         .setExtensionName(name)
         .setImports(importsBuilder.buildOrThrow())
         .setDevImports(ImmutableSet.of())
+        .setUsingModule(ModuleKey.ROOT)
         .setLocation(Location.BUILTIN)
         .build();
   }