Store information about each proxy in ModuleExtensionUsage

The current ModuleExtensionUsage class was designed in the days before we allowed multiple `use_extension` calls for the same extension in the same module. This PR brings the class up to speed to better reflect the actual API.

This will also lay the groundwork for `bazel mod tidy` working with `include()`d segments.

Work towards bazelbuild#22063

Closes bazelbuild#22307.

PiperOrigin-RevId: 632377764
Change-Id: I282a68bc7962088ae4583418f73b2e60a0ec88f0

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

@@ -237,13 +237,13 @@ static BzlmodFlagsAndEnvVars getFlagsAndEnvVars(Environment env) throws Interrup
               Code.BAD_MODULE,
               e,
               "invalid label for module extension found at %s",
-              usage.getLocation());
+              usage.getProxies().getFirst().getLocation());
         }
         if (!moduleExtensionId.getBzlFileLabel().getRepository().isVisible()) {
           throw ExternalDepsException.withMessage(
               Code.BAD_MODULE,
               "invalid label for module extension found at %s: no repo visible as '@%s' here",
-              usage.getLocation(),
+              usage.getProxies().getFirst().getLocation(),
               moduleExtensionId.getBzlFileLabel().getRepository().getName());
         }
         extensionUsagesTableBuilder.put(moduleExtensionId, module.getKey(), usage);

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

@@ -118,14 +118,16 @@ public static BazelDepGraphValue createEmptyDepGraph() {
    */
   public final RepositoryMapping getFullRepoMapping(ModuleKey key) {
     ImmutableMap.Builder<String, RepositoryName> mapping = ImmutableMap.builder();
-    for (Map.Entry<ModuleExtensionId, ModuleExtensionUsage> e :
+    for (Map.Entry<ModuleExtensionId, ModuleExtensionUsage> extIdAndUsage :
         getExtensionUsagesTable().column(key).entrySet()) {
-      ModuleExtensionId extensionId = e.getKey();
-      ModuleExtensionUsage usage = e.getValue();
-      for (Map.Entry<String, String> entry : usage.getImports().entrySet()) {
-        String canonicalRepoName =
-            getExtensionUniqueNames().get(extensionId) + "~" + entry.getValue();
-        mapping.put(entry.getKey(), RepositoryName.createUnvalidated(canonicalRepoName));
+      ModuleExtensionId extensionId = extIdAndUsage.getKey();
+      ModuleExtensionUsage usage = extIdAndUsage.getValue();
+      String repoNamePrefix = getExtensionUniqueNames().get(extensionId) + "~";
+      for (ModuleExtensionUsage.Proxy proxy : usage.getProxies()) {
+        for (Map.Entry<String, String> entry : proxy.getImports().entrySet()) {
+          String canonicalRepoName = repoNamePrefix + entry.getValue();
+          mapping.put(entry.getKey(), RepositoryName.createUnvalidated(canonicalRepoName));
+        }
       }
     }
     return getDepGraph()

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

@@ -39,7 +39,7 @@
 @GenerateTypeAdapter
 public abstract class BazelLockFileValue implements SkyValue, Postable {
 
-  public static final int LOCK_FILE_VERSION = 8;
+  public static final int LOCK_FILE_VERSION = 9;
 
   @SerializationConstant public static final SkyKey KEY = () -> SkyFunctions.BAZEL_LOCK_FILE;
 

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

@@ -41,7 +41,6 @@
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkList;
 import net.starlark.java.eval.StarlarkValue;
-import net.starlark.java.syntax.Location;
 
 /** The Starlark object passed to the implementation function of module extension metadata. */
 @StarlarkBuiltin(
@@ -219,15 +218,19 @@ private static Optional<RootModuleFileFixup> generateFixup(
       Set<String> expectedImports,
       Set<String> expectedDevImports,
       EventHandler eventHandler) {
-    var actualDevImports = ImmutableSet.copyOf(rootUsage.getDevImports());
+    var actualDevImports =
+        rootUsage.getProxies().stream()
+            .filter(p -> p.isDevDependency())
+            .flatMap(p -> p.getImports().values().stream())
+            .collect(toImmutableSet());
     var actualImports =
-        rootUsage.getImports().values().stream()
-            .filter(repo -> !actualDevImports.contains(repo))
+        rootUsage.getProxies().stream()
+            .filter(p -> !p.isDevDependency())
+            .flatMap(p -> p.getImports().values().stream())
             .collect(toImmutableSet());
 
     String extensionBzlFile = rootUsage.getExtensionBzlFile();
     String extensionName = rootUsage.getExtensionName();
-    Location location = rootUsage.getLocation();
 
     var importsToAdd = ImmutableSortedSet.copyOf(Sets.difference(expectedImports, actualImports));
     var importsToRemove =
@@ -341,7 +344,7 @@ private static Optional<RootModuleFileFixup> generateFixup(
             .flatMap(Optional::stream)
             .collect(toImmutableList());
 
-    eventHandler.handle(Event.warn(location, message));
+    eventHandler.handle(Event.warn(rootUsage.getProxies().getFirst().getLocation(), message));
     return Optional.of(new RootModuleFileFixup(buildozerCommands, rootUsage));
   }
 

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

@@ -17,17 +17,17 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
 import com.google.auto.value.AutoValue;
-import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableMap;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import java.util.Optional;
 import net.starlark.java.syntax.Location;
 
 /**
- * Represents one usage of a module extension in one MODULE.bazel file. This class records all the
- * information pertinent to the proxy object returned from the {@code use_extension} call.
+ * Represents the usage of a module extension in one module. This class records all the information
+ * pertinent to all the proxy objects returned from any {@code use_extension} calls in this module
+ * that refer to the same extension (or isolate, when applicable).
  *
  * <p>When adding new fields, make sure to update {@link #trimForEvaluation()} as well.
  */
@@ -49,40 +49,80 @@ public abstract class ModuleExtensionUsage {
   /** 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
-   * location of the first one.
-   */
-  public abstract Location getLocation();
+  /** Represents one "proxy object" returned from one {@code use_extension} call. */
+  @AutoValue
+  @GenerateTypeAdapter
+  public abstract static class Proxy {
+    /** The location of the {@code use_extension} call. */
+    public abstract Location getLocation();
+
+    /**
+     * The name of the proxy object; as in, the name that the return value of {@code use_extension}
+     * is bound to. Is the empty string if the return value is not bound to any name (e.g. {@code
+     * use_repo(use_extension(...))}).
+     */
+    public abstract String getProxyName();
+
+    /** Whether {@code dev_dependency} is set to true. */
+    public abstract boolean isDevDependency();
+
+    /**
+     * All the repos imported, through this proxy, from this module extension into the scope of the
+     * current module. The key is the local repo name (in the scope of the current module), and the
+     * value is the name exported by the module extension.
+     */
+    public abstract ImmutableMap<String, String> getImports();
+
+    public static Builder builder() {
+      return new AutoValue_ModuleExtensionUsage_Proxy.Builder().setProxyName("");
+    }
 
-  /**
-   * All the repos imported from this module extension into the scope of the current module. The key
-   * is the local repo name (in the scope of the current module), and the value is the name exported
-   * by the module extension.
-   */
-  public abstract ImmutableBiMap<String, String> getImports();
+    /** Builder for {@link ModuleExtensionUsage.Proxy}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+      public abstract Builder setLocation(Location value);
 
-  /**
-   * The repo names as exported by the module extension that were imported using a proxy marked as a
-   * dev dependency.
-   */
-  public abstract ImmutableSet<String> getDevImports();
+      public abstract String getProxyName();
+
+      public abstract Builder setProxyName(String value);
+
+      public abstract boolean isDevDependency();
+
+      public abstract Builder setDevDependency(boolean value);
+
+      abstract ImmutableMap.Builder<String, String> importsBuilder();
+
+      @CanIgnoreReturnValue
+      public final Builder addImport(String key, String value) {
+        importsBuilder().put(key, value);
+        return this;
+      }
+
+      public abstract Builder setImports(ImmutableMap<String, String> value);
+
+      public abstract Proxy build();
+    }
+  }
+
+  /** The list of proxy objects that constitute */
+  public abstract ImmutableList<Proxy> getProxies();
 
   /** All the tags specified by this module for this extension. */
   public abstract ImmutableList<Tag> getTags();
 
   /**
-   * Whether any <code>use_extension</code> calls for this usage had <code>dev_dependency = True
-   * </code> set.*
+   * Whether any {@code use_extension} calls for this usage had {@code dev_dependency = True} set.
    */
-  public abstract boolean getHasDevUseExtension();
+  public final boolean getHasDevUseExtension() {
+    return getProxies().stream().anyMatch(p -> p.isDevDependency());
+  }
 
   /**
-   * Whether any <code>use_extension</code> calls for this usage had <code>dev_dependency = False
-   * </code> set.*
+   * Whether any {@code use_extension} calls for this usage had {@code dev_dependency = False} set.
    */
-  public abstract boolean getHasNonDevUseExtension();
+  public final boolean getHasNonDevUseExtension() {
+    return getProxies().stream().anyMatch(p -> !p.isDevDependency());
+  }
 
   public abstract Builder toBuilder();
 
@@ -100,13 +140,12 @@ ModuleExtensionUsage trimForEvaluation() {
     // preserves correctness in case new fields are added without updating this code.
     return toBuilder()
         .setTags(getTags().stream().map(Tag::trimForEvaluation).collect(toImmutableList()))
+        // Clear out all proxies as information contained therein isn't useful for evaluation.
         // Locations are only used for error reporting and thus don't influence whether the
         // evaluation of the extension is successful and what its result is in case of success.
-        .setLocation(Location.BUILTIN)
         // Extension implementation functions do not see the imports, they are only validated
         // against the set of generated repos in a validation step that comes afterward.
-        .setImports(ImmutableBiMap.of())
-        .setDevImports(ImmutableSet.of())
+        .setProxies(ImmutableList.of())
         .build();
   }
 
@@ -122,26 +161,26 @@ public abstract static class Builder {
 
     public abstract Builder setUsingModule(ModuleKey value);
 
-    public abstract Builder setLocation(Location value);
+    public abstract Builder setProxies(ImmutableList<Proxy> value);
 
-    public abstract Builder setImports(ImmutableBiMap<String, String> value);
+    abstract ImmutableList.Builder<Proxy> proxiesBuilder();
 
-    public abstract Builder setDevImports(ImmutableSet<String> value);
+    @CanIgnoreReturnValue
+    public Builder addProxy(Proxy value) {
+      proxiesBuilder().add(value);
+      return this;
+    }
 
     public abstract Builder setTags(ImmutableList<Tag> value);
 
     abstract ImmutableList.Builder<Tag> tagsBuilder();
 
     @CanIgnoreReturnValue
-    public ModuleExtensionUsage.Builder addTag(Tag value) {
+    public Builder addTag(Tag value) {
       tagsBuilder().add(value);
       return this;
     }
 
-    public abstract Builder setHasDevUseExtension(boolean hasDevUseExtension);
-
-    public abstract Builder setHasNonDevUseExtension(boolean hasNonDevUseExtension);
-
     public abstract ModuleExtensionUsage build();
   }
 }

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

@@ -445,13 +445,14 @@ public static RootModuleFileValue evaluateRootModuleFile(
       throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for the root module");
     }
     for (ModuleExtensionUsage usage : module.getExtensionUsages()) {
-      if (usage.getIsolationKey().isPresent() && usage.getImports().isEmpty()) {
+      ModuleExtensionUsage.Proxy firstProxy = usage.getProxies().getFirst();
+      if (usage.getIsolationKey().isPresent() && firstProxy.getImports().isEmpty()) {
         throw errorf(
             Code.BAD_MODULE,
             "the isolated usage at %s of extension %s defined in %s has no effect as no "
                 + "repositories are imported from it. Either import one or more repositories "
                 + "generated by the extension with use_repo or remove the usage.",
-            usage.getLocation(),
+            firstProxy.getLocation(),
             usage.getExtensionName(),
             usage.getExtensionBzlFile());
       }