Make bazel mod tidy work with include()

We now generate fixup commands targeting `include()`d segments in addition to the root MODULE.bazel file itself. Fixup commands are grouped by file, and then passed to Buildozer using `-f -` (i.e. multiple commands for multiple files, passed in using stdin).

- For repos to add, we find the first usage with the correct prod/dev type and add them there.
- For repos to remove, we remove them from whichever usage had them.
- At the end, all module files and included segments are formatted.

Fixes #22063

Closes #22455.

PiperOrigin-RevId: 636680730
Change-Id: Id894a449d8d1cdf39271ea3a724a91aebc51befb

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

@@ -60,6 +60,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/repository",
+        "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/net/starlark/java/eval",
         "//src/main/java/net/starlark/java/syntax",
         "//third_party:auto_value",
@@ -300,6 +301,7 @@ java_library(
         "//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/lib/vfs",
+        "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//third_party:auto_value",
         "//third_party:guava",
@@ -374,6 +376,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/docgen/annot",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
+        "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/net/starlark/java/annot",
         "//src/main/java/net/starlark/java/eval",
         "//src/main/java/net/starlark/java/syntax",
@@ -391,6 +394,7 @@ java_library(
     ],
     deps = [
         ":module_extension",
+        "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//third_party:guava",
     ],
 )

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

@@ -18,7 +18,6 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -83,12 +82,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     }
 
     ImmutableSet<SkyKey> extensionsUsedByRootModule =
-        depGraphValue
-            .getExtensionUsagesTable()
-            .columnMap()
-            .getOrDefault(ModuleKey.ROOT, ImmutableMap.of())
-            .keySet()
-            .stream()
+        depGraphValue.getExtensionUsagesTable().column(ModuleKey.ROOT).keySet().stream()
             // Use the eval-only key to avoid errors caused by incorrect imports - we can fix them.
             .map(SingleExtensionValue::evalKey)
             .collect(toImmutableSet());
@@ -108,8 +102,6 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     }
 
     return BazelModTidyValue.create(
-        fixups.build(),
-        buildozer.asPath(),
-        rootModuleFileValue.getIncludeLabelToCompiledModuleFile());
+        fixups.build(), buildozer.asPath(), rootModuleFileValue.getModuleFilePaths());
   }
 }

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

@@ -17,10 +17,11 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.List;
@@ -37,13 +38,14 @@ public abstract class BazelModTidyValue implements SkyValue {
   /** The path of the buildozer binary provided by the "buildozer" module. */
   public abstract Path buildozer();
 
-  public abstract ImmutableMap<String, CompiledModuleFile> includeLabelToCompiledModuleFile();
+  /** The set of paths to the root MODULE.bazel file and all its includes. */
+  public abstract ImmutableSet<PathFragment> moduleFilePaths();
 
   static BazelModTidyValue create(
       List<RootModuleFileFixup> fixups,
       Path buildozer,
-      ImmutableMap<String, CompiledModuleFile> includeLabelToCompiledModuleFile) {
+      ImmutableSet<PathFragment> moduleFilePaths) {
     return new AutoValue_BazelModTidyValue(
-        ImmutableList.copyOf(fixups), buildozer, includeLabelToCompiledModuleFile);
+        ImmutableList.copyOf(fixups), buildozer, moduleFilePaths);
   }
 }

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

@@ -14,26 +14,26 @@
 
 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.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableListMultimap;
 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.bazel.bzlmod.ModuleExtensionUsage.Proxy;
 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 com.google.devtools.build.lib.vfs.PathFragment;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.eval.EvalException;
@@ -171,20 +171,8 @@ static ModuleExtensionMetadata create(
   }
 
   public Optional<RootModuleFileFixup> generateFixup(
-      Collection<ModuleExtensionUsage> usages, Set<String> allRepos, EventHandler eventHandler)
+      ModuleExtensionUsage rootUsage, Set<String> allRepos, EventHandler eventHandler)
       throws EvalException {
-    var rootUsages =
-        usages.stream()
-            .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.
-      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()) {
@@ -311,73 +299,52 @@ private static Optional<RootModuleFileFixup> generateFixup(
 
     message += "Fix the use_repo calls by running 'bazel mod tidy'.";
 
-    var buildozerCommands =
-        Stream.of(
-                makeUseRepoCommand(
-                    "use_repo_add",
-                    false,
-                    importsToAdd,
-                    extensionBzlFile,
-                    extensionName,
-                    rootUsage.getIsolationKey()),
-                makeUseRepoCommand(
-                    "use_repo_remove",
-                    false,
-                    importsToRemove,
-                    extensionBzlFile,
-                    extensionName,
-                    rootUsage.getIsolationKey()),
-                makeUseRepoCommand(
-                    "use_repo_add",
-                    true,
-                    devImportsToAdd,
-                    extensionBzlFile,
-                    extensionName,
-                    rootUsage.getIsolationKey()),
-                makeUseRepoCommand(
-                    "use_repo_remove",
-                    true,
-                    devImportsToRemove,
-                    extensionBzlFile,
-                    extensionName,
-                    rootUsage.getIsolationKey()))
-            .flatMap(Optional::stream)
-            .collect(toImmutableList());
+    var moduleFilePathToCommandsBuilder = ImmutableListMultimap.<PathFragment, String>builder();
+    // Repos to add are easy: always add them to the first proxy of the correct type.
+    if (!importsToAdd.isEmpty()) {
+      Proxy firstNonDevProxy =
+          rootUsage.getProxies().stream().filter(p -> !p.isDevDependency()).findFirst().get();
+      moduleFilePathToCommandsBuilder.put(
+          firstNonDevProxy.getContainingModuleFilePath(),
+          makeUseRepoCommand("use_repo_add", firstNonDevProxy.getProxyName(), importsToAdd));
+    }
+    if (!devImportsToAdd.isEmpty()) {
+      Proxy firstDevProxy =
+          rootUsage.getProxies().stream().filter(p -> p.isDevDependency()).findFirst().get();
+      moduleFilePathToCommandsBuilder.put(
+          firstDevProxy.getContainingModuleFilePath(),
+          makeUseRepoCommand("use_repo_add", firstDevProxy.getProxyName(), devImportsToAdd));
+    }
+    // Repos to remove are a bit trickier: remove them from the proxy that actually imported them.
+    for (Proxy proxy : rootUsage.getProxies()) {
+      var toRemove =
+          ImmutableSortedSet.copyOf(
+              Sets.intersection(
+                  proxy.getImports().values(),
+                  proxy.isDevDependency() ? devImportsToRemove : importsToRemove));
+      if (!toRemove.isEmpty()) {
+        moduleFilePathToCommandsBuilder.put(
+            proxy.getContainingModuleFilePath(),
+            makeUseRepoCommand("use_repo_remove", proxy.getProxyName(), toRemove));
+      }
+    }
 
     eventHandler.handle(Event.warn(rootUsage.getProxies().getFirst().getLocation(), message));
-    return Optional.of(new RootModuleFileFixup(buildozerCommands, rootUsage));
+    return Optional.of(new RootModuleFileFixup(moduleFilePathToCommandsBuilder.build(), rootUsage));
   }
 
-  private static Optional<String> makeUseRepoCommand(
-      String cmd,
-      boolean devDependency,
-      Collection<String> repos,
-      String extensionBzlFile,
-      String extensionName,
-      Optional<ModuleExtensionId.IsolationKey> isolationKey) {
-    if (repos.isEmpty()) {
-      return Optional.empty();
-    }
-
+  private static String makeUseRepoCommand(String cmd, String proxyName, Collection<String> repos) {
     var commandParts = new ArrayList<String>();
     commandParts.add(cmd);
-    if (isolationKey.isPresent()) {
-      commandParts.add(isolationKey.get().getUsageExportedName());
-    } else {
-      if (devDependency) {
-        commandParts.add("dev");
-      }
-      commandParts.add(extensionBzlFile);
-      commandParts.add(extensionName);
-    }
+    commandParts.add(proxyName.isEmpty() ? "_unnamed_usage" : proxyName);
     commandParts.addAll(repos);
-    return Optional.of(String.join(" ", commandParts));
+    return String.join(" ", commandParts);
   }
 
   private Optional<ImmutableSet<String>> getRootModuleDirectDeps(Set<String> allRepos)
       throws EvalException {
-    switch (getUseAllRepos()) {
-      case NO:
+    return switch (getUseAllRepos()) {
+      case NO -> {
         if (getExplicitRootModuleDirectDeps() != null) {
           Set<String> invalidRepos = Sets.difference(getExplicitRootModuleDirectDeps(), allRepos);
           if (!invalidRepos.isEmpty()) {
@@ -387,13 +354,11 @@ private Optional<ImmutableSet<String>> getRootModuleDirectDeps(Set<String> allRe
                 String.join(", ", invalidRepos));
           }
         }
-        return Optional.ofNullable(getExplicitRootModuleDirectDeps());
-      case REGULAR:
-        return Optional.of(ImmutableSet.copyOf(allRepos));
-      case DEV:
-        return Optional.of(ImmutableSet.of());
-    }
-    throw new IllegalStateException("not reached");
+        yield Optional.ofNullable(getExplicitRootModuleDirectDeps());
+      }
+      case REGULAR -> Optional.of(ImmutableSet.copyOf(allRepos));
+      case DEV -> Optional.of(ImmutableSet.of());
+    };
   }
 
   private Optional<ImmutableSet<String>> getRootModuleDirectDevDeps(Set<String> allRepos)

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

@@ -17,8 +17,9 @@
 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.ImmutableMap;
+import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import java.util.Optional;
@@ -63,6 +64,12 @@ public abstract static class Proxy {
      */
     public abstract String getProxyName();
 
+    /**
+     * The path to the MODULE.bazel file (or one of its includes) that contains this proxy object.
+     * This path should be relative to the workspace root.
+     */
+    public abstract PathFragment getContainingModuleFilePath();
+
     /** Whether {@code dev_dependency} is set to true. */
     public abstract boolean isDevDependency();
 
@@ -71,7 +78,7 @@ public abstract static class Proxy {
      * 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 abstract ImmutableBiMap<String, String> getImports();
 
     public static Builder builder() {
       return new AutoValue_ModuleExtensionUsage_Proxy.Builder().setProxyName("");
@@ -86,19 +93,21 @@ public abstract static class Builder {
 
       public abstract Builder setProxyName(String value);
 
+      public abstract Builder setContainingModuleFilePath(PathFragment value);
+
       public abstract boolean isDevDependency();
 
       public abstract Builder setDevDependency(boolean value);
 
-      abstract ImmutableMap.Builder<String, String> importsBuilder();
+      abstract ImmutableBiMap.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 Builder setImports(ImmutableBiMap<String, String> value);
 
       public abstract Proxy build();
     }

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

@@ -70,6 +70,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Mutability;
@@ -476,12 +477,18 @@ public static RootModuleFileValue evaluateRootModuleFile(
                     name ->
                         ModuleKey.create(name, Version.EMPTY).getCanonicalRepoNameWithoutVersion(),
                     name -> name));
+    ImmutableSet<PathFragment> moduleFilePaths =
+        Stream.concat(
+                Stream.of(LabelConstants.MODULE_DOT_BAZEL_FILE_NAME),
+                includeLabelToCompiledModuleFile.keySet().stream()
+                    .map(label -> Label.parseCanonicalUnchecked(label).toPathFragment()))
+            .collect(toImmutableSet());
     return RootModuleFileValue.create(
         module,
         moduleFileHash,
         overrides,
         nonRegistryOverrideCanonicalRepoNameLookup,
-        includeLabelToCompiledModuleFile);
+        moduleFilePaths);
   }
 
   private static ModuleThreadContext execModuleFile(

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

@@ -430,7 +430,8 @@ public ModuleExtensionProxy useExtension(
     var proxyBuilder =
         ModuleExtensionUsage.Proxy.builder()
             .setLocation(thread.getCallerLocation())
-            .setDevDependency(devDependency);
+            .setDevDependency(devDependency)
+            .setContainingModuleFilePath(context.getCurrentModuleFilePath());
 
     String extensionBzlFile = normalizeLabelString(context.getModuleBuilder(), rawExtensionBzlFile);
     var newUsageBuilder =
@@ -692,7 +693,9 @@ public void call(
               usageBuilder,
               ModuleExtensionUsage.Proxy.builder()
                   .setDevDependency(devDependency)
-                  .setLocation(thread.getCallerLocation()));
+                  .setLocation(thread.getCallerLocation())
+                  .setContainingModuleFilePath(
+                      usageBuilder.getContext().getCurrentModuleFilePath()));
       extensionProxy.getValue(tagName).call(kwargs, thread);
       extensionProxy.addImport(name, name, "by a repo rule", thread.getCallerLocation());
     }