[6.3.0] Track dev/non-dev use_extension calls

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

@@ -46,6 +46,7 @@
 public class ModuleExtensionContext extends StarlarkBaseExternalContext {
   private final ModuleExtensionId extensionId;
   private final StarlarkList<StarlarkBazelModule> modules;
+  private final boolean rootModuleHasNonDevDependency;
 
   protected ModuleExtensionContext(
       Path workingDirectory,
@@ -57,7 +58,8 @@ protected ModuleExtensionContext(
       StarlarkSemantics starlarkSemantics,
       @Nullable RepositoryRemoteExecutor remoteExecutor,
       ModuleExtensionId extensionId,
-      StarlarkList<StarlarkBazelModule> modules) {
+      StarlarkList<StarlarkBazelModule> modules,
+      boolean rootModuleHasNonDevDependency) {
     super(
         workingDirectory,
         env,
@@ -69,6 +71,7 @@ protected ModuleExtensionContext(
         remoteExecutor);
     this.extensionId = extensionId;
     this.modules = modules;
+    this.rootModuleHasNonDevDependency = rootModuleHasNonDevDependency;
   }
 
   public Path getWorkingDirectory() {
@@ -96,10 +99,11 @@ protected ImmutableMap<String, String> getRemoteExecProperties() throws EvalExce
       name = "modules",
       structField = true,
       doc =
-          "A list of all the Bazel modules in the external dependency graph, each of which is a <a"
-              + " href=\"bazel_module.html\">bazel_module</a> object that exposes all the tags it"
-              + " specified for this module extension. The iteration order of this dictionary is"
-              + " guaranteed to be the same as breadth-first search starting from the root module.")
+          "A list of all the Bazel modules in the external dependency graph that use this module "
+              + "extension, each of which is a <a href=\"../builtins/bazel_module.html\">"
+              + "bazel_module</a> object that exposes all the tags it specified for this extension."
+              + " The iteration order of this dictionary is guaranteed to be the same as"
+              + " breadth-first search starting from the root module.")
   public StarlarkList<StarlarkBazelModule> getModules() {
     return modules;
   }
@@ -130,6 +134,14 @@ public boolean isIsolated() {
     return extensionId.getIsolationKey().isPresent();
   }
 
+  @StarlarkMethod(
+      name = "root_module_has_non_dev_dependency",
+      doc = "Whether the root module uses this extension as a non-dev dependency.",
+      structField = true)
+  public boolean rootModuleHasNonDevDependency() {
+    return rootModuleHasNonDevDependency;
+  }
+
   @StarlarkMethod(
       name = "extension_metadata",
       doc =
@@ -151,7 +163,7 @@ public boolean isIsolated() {
                     + " disjoint.<p>Exactly one of <code>root_module_direct_deps</code> and"
                     + " <code>root_module_direct_dev_deps</code> can be set to the special value"
                     + " <code>\"all\"</code>, which is treated as if a list with the names of"
-                    + " allrepositories generated by the extension was specified as the value.",
+                    + " all repositories generated by the extension was specified as the value.",
             positional = false,
             named = true,
             defaultValue = "None",
@@ -177,7 +189,7 @@ public boolean isIsolated() {
                     + " disjoint.<p>Exactly one of <code>root_module_direct_deps</code> and"
                     + " <code>root_module_direct_dev_deps</code> can be set to the special value"
                     + " <code>\"all\"</code>, which is treated as if a list with the names of"
-                    + " allrepositories generated by the extension was specified as the value.",
+                    + " all repositories generated by the extension was specified as the value.",
             positional = false,
             named = true,
             defaultValue = "None",

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

@@ -20,14 +20,14 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 import com.google.devtools.build.docgen.annot.DocCategory;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import java.util.Collection;
 import java.util.LinkedHashSet;
-import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -82,23 +82,11 @@ static ModuleExtensionMetadata create(
     // root_module_direct_dev_deps = [], but not root_module_direct_dev_deps = ["some_repo"].
     if (rootModuleDirectDepsUnchecked.equals("all")
         && rootModuleDirectDevDepsUnchecked.equals(StarlarkList.immutableOf())) {
-      if (extensionId.getIsolationKey().isPresent()
-          && extensionId.getIsolationKey().get().isDevUsage()) {
-        throw Starlark.errorf(
-            "root_module_direct_deps must be empty for an isolated extension usage with "
-                + "dev_dependency = True");
-      }
       return new ModuleExtensionMetadata(null, null, UseAllRepos.REGULAR);
     }
 
     if (rootModuleDirectDevDepsUnchecked.equals("all")
         && rootModuleDirectDepsUnchecked.equals(StarlarkList.immutableOf())) {
-      if (extensionId.getIsolationKey().isPresent()
-          && !extensionId.getIsolationKey().get().isDevUsage()) {
-        throw Starlark.errorf(
-            "root_module_direct_dev_deps must be empty for an isolated extension usage with "
-                + "dev_dependency = False");
-      }
       return new ModuleExtensionMetadata(null, null, UseAllRepos.DEV);
     }
 
@@ -128,20 +116,6 @@ static ModuleExtensionMetadata create(
         Sequence.cast(
             rootModuleDirectDevDepsUnchecked, String.class, "root_module_direct_dev_deps");
 
-    if (extensionId.getIsolationKey().isPresent()) {
-      ModuleExtensionId.IsolationKey isolationKey = extensionId.getIsolationKey().get();
-      if (isolationKey.isDevUsage() && !rootModuleDirectDeps.isEmpty()) {
-        throw Starlark.errorf(
-            "root_module_direct_deps must be empty for an isolated extension usage with "
-                + "dev_dependency = True");
-      }
-      if (!isolationKey.isDevUsage() && !rootModuleDirectDevDeps.isEmpty()) {
-        throw Starlark.errorf(
-            "root_module_direct_dev_deps must be empty for an isolated extension usage with "
-                + "dev_dependency = False");
-      }
-    }
-
     Set<String> explicitRootModuleDirectDeps = new LinkedHashSet<>();
     for (String dep : rootModuleDirectDeps) {
       try {
@@ -192,40 +166,46 @@ Optional<Event> generateFixupMessage(
       // expected to modify any other module's MODULE.bazel file.
       return Optional.empty();
     }
+    // Every module only has at most a single usage of a given extension.
+    ModuleExtensionUsage rootUsage = Iterables.getOnlyElement(rootUsages);
 
     var rootModuleDirectDevDeps = getRootModuleDirectDevDeps(allRepos);
     var rootModuleDirectDeps = getRootModuleDirectDeps(allRepos);
     if (rootModuleDirectDevDeps.isEmpty() && rootModuleDirectDeps.isEmpty()) {
       return Optional.empty();
     }
-
     Preconditions.checkState(
         rootModuleDirectDevDeps.isPresent() && rootModuleDirectDeps.isPresent());
+
+    if (!rootUsage.getHasNonDevUseExtension() && !rootModuleDirectDeps.get().isEmpty()) {
+      throw Starlark.errorf(
+          "root_module_direct_deps must be empty if the root module contains no "
+              + "usages with dev_dependency = False");
+    }
+    if (!rootUsage.getHasDevUseExtension() && !rootModuleDirectDevDeps.get().isEmpty()) {
+      throw Starlark.errorf(
+          "root_module_direct_dev_deps must be empty if the root module contains no "
+              + "usages with dev_dependency = True");
+    }
+
     return generateFixupMessage(
-        rootUsages, allRepos, rootModuleDirectDeps.get(), rootModuleDirectDevDeps.get());
+        rootUsage, allRepos, rootModuleDirectDeps.get(), rootModuleDirectDevDeps.get());
   }
 
   private static Optional<Event> generateFixupMessage(
-      List<ModuleExtensionUsage> rootUsages,
+      ModuleExtensionUsage rootUsage,
       Set<String> allRepos,
       Set<String> expectedImports,
       Set<String> expectedDevImports) {
-    var actualDevImports =
-        rootUsages.stream()
-            .flatMap(usage -> usage.getDevImports().stream())
-            .collect(toImmutableSet());
+    var actualDevImports = ImmutableSet.copyOf(rootUsage.getDevImports());
     var actualImports =
-        rootUsages.stream()
-            .flatMap(usage -> usage.getImports().values().stream())
+        rootUsage.getImports().values().stream()
             .filter(repo -> !actualDevImports.contains(repo))
             .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.
-    ModuleExtensionUsage firstUsage = rootUsages.get(0);
-    String extensionBzlFile = firstUsage.getExtensionBzlFile();
-    String extensionName = firstUsage.getExtensionName();
-    Location location = firstUsage.getLocation();
+    String extensionBzlFile = rootUsage.getExtensionBzlFile();
+    String extensionName = rootUsage.getExtensionName();
+    Location location = rootUsage.getLocation();
 
     var importsToAdd = ImmutableSortedSet.copyOf(Sets.difference(expectedImports, actualImports));
     var importsToRemove =
@@ -290,28 +270,28 @@ private static Optional<Event> generateFixupMessage(
                     importsToAdd,
                     extensionBzlFile,
                     extensionName,
-                    firstUsage.getIsolationKey()),
+                    rootUsage.getIsolationKey()),
                 makeUseRepoCommand(
                     "use_repo_remove",
                     false,
                     importsToRemove,
                     extensionBzlFile,
                     extensionName,
-                    firstUsage.getIsolationKey()),
+                    rootUsage.getIsolationKey()),
                 makeUseRepoCommand(
                     "use_repo_add",
                     true,
                     devImportsToAdd,
                     extensionBzlFile,
                     extensionName,
-                    firstUsage.getIsolationKey()),
+                    rootUsage.getIsolationKey()),
                 makeUseRepoCommand(
                     "use_repo_remove",
                     true,
                     devImportsToRemove,
                     extensionBzlFile,
                     extensionName,
-                    firstUsage.getIsolationKey()))
+                    rootUsage.getIsolationKey()))
             .flatMap(Optional::stream);
 
     return Optional.of(

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

@@ -68,6 +68,18 @@ public abstract class ModuleExtensionUsage {
   /** 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.*
+   */
+  public abstract boolean getHasDevUseExtension();
+
+  /**
+   * Whether any <code>use_extension</code> calls for this usage had <code>dev_dependency = False
+   * </code> set.*
+   */
+  public abstract boolean getHasNonDevUseExtension();
+
   public static Builder builder() {
     return new AutoValue_ModuleExtensionUsage.Builder();
   }
@@ -100,6 +112,10 @@ public ModuleExtensionUsage.Builder addTag(Tag 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/ModuleFileGlobals.java

@@ -515,6 +515,9 @@ class ModuleExtensionUsageBuilder {
     private final ImmutableSet.Builder<String> devImports;
     private final ImmutableList.Builder<Tag> tags;
 
+    private boolean hasNonDevUseExtension;
+    private boolean hasDevUseExtension;
+
     ModuleExtensionUsageBuilder(
         String extensionBzlFile,
         String extensionName,
@@ -539,6 +542,8 @@ ModuleExtensionUsage buildUsage() {
               .setLocation(location)
               .setImports(ImmutableBiMap.copyOf(imports))
               .setDevImports(devImports.build())
+              .setHasDevUseExtension(hasDevUseExtension)
+              .setHasNonDevUseExtension(hasNonDevUseExtension)
               .setTags(tags.build());
       return builder.build();
     }
@@ -548,6 +553,11 @@ ModuleExtensionUsage buildUsage() {
      * tags with all other such proxies, thus preserving their order across dev/non-dev deps.
      */
     ModuleExtensionProxy getProxy(boolean devDependency) {
+      if (devDependency) {
+        hasDevUseExtension = true;
+      } else {
+        hasNonDevUseExtension = true;
+      }
       return new ModuleExtensionProxy(devDependency);
     }
 

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

@@ -400,6 +400,9 @@ private ModuleExtensionContext createContext(
         throw new SingleExtensionEvalFunctionException(e, Transience.PERSISTENT);
       }
     }
+    ModuleExtensionUsage rootUsage = usagesValue.getExtensionUsages().get(ModuleKey.ROOT);
+    boolean rootModuleHasNonDevDependency =
+        rootUsage != null && rootUsage.getHasNonDevUseExtension();
     return new ModuleExtensionContext(
         workingDirectory,
         env,
@@ -410,7 +413,8 @@ private ModuleExtensionContext createContext(
         starlarkSemantics,
         repositoryRemoteExecutor,
         extensionId,
-        StarlarkList.immutableCopyOf(modules));
+        StarlarkList.immutableCopyOf(modules),
+        rootModuleHasNonDevDependency);
   }
 
   static final class SingleExtensionEvalFunctionException extends SkyFunctionException {

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

@@ -219,6 +219,8 @@ private static ModuleExtensionUsage createModuleExtensionUsage(
         .setDevImports(ImmutableSet.of())
         .setUsingModule(ModuleKey.ROOT)
         .setLocation(Location.BUILTIN)
+        .setHasDevUseExtension(false)
+        .setHasNonDevUseExtension(true)
         .build();
   }