[6.3.0] Allow module extension usages to be isolated

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

@@ -214,7 +214,9 @@ private ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> getEx
         try {
           moduleExtensionId =
               ModuleExtensionId.create(
-                  labelConverter.convert(usage.getExtensionBzlFile()), usage.getExtensionName());
+                  labelConverter.convert(usage.getExtensionBzlFile()),
+                  usage.getExtensionName(),
+                  usage.getIsolationKey());
         } catch (LabelSyntaxException e) {
           throw new BazelDepGraphFunctionException(
               ExternalDepsException.withCauseAndMessage(
@@ -248,12 +250,31 @@ private ImmutableBiMap<String, ModuleExtensionId> calculateUniqueNameForUsedExte
       // not start with a tilde.
       RepositoryName repository = id.getBzlFileLabel().getRepository();
       String nonEmptyRepoPart = repository.isMain() ? "_main" : repository.getName();
-      String bestName = nonEmptyRepoPart + "~" + id.getExtensionName();
+      // When using a namespace, prefix the extension name with "_" to distinguish the prefix from
+      // those generated by non-namespaced extension usages. Extension names are identified by their
+      // Starlark identifier, which in the case of an exported symbol cannot start with "_".
+      // We also include whether the isolated usage is a dev usage as well as its index in the
+      // MODULE.bazel file to ensure that canonical repository names don't change depending on
+      // whether dev dependencies are ignored. This removes potential for confusion and also
+      // prevents unnecessary refetches when --ignore_dev_dependency is toggled.
+      String bestName =
+          id.getIsolationKey()
+              .map(
+                  namespace ->
+                      String.format(
+                          "%s~_%s~%s~%s~%s%d",
+                          nonEmptyRepoPart,
+                          id.getExtensionName(),
+                          namespace.getModule().getName(),
+                          namespace.getModule().getVersion(),
+                          namespace.isDevUsage() ? "dev" : "",
+                          namespace.getIsolatedUsageIndex()))
+              .orElse(nonEmptyRepoPart + "~" + id.getExtensionName());
       if (extensionUniqueNames.putIfAbsent(bestName, id) == null) {
         continue;
       }
       int suffix = 2;
-      while (extensionUniqueNames.putIfAbsent(bestName + suffix, id) != null) {
+      while (extensionUniqueNames.putIfAbsent(bestName + "~" + suffix, id) != null) {
         suffix++;
       }
     }

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

@@ -20,17 +20,25 @@
 import static com.google.devtools.build.lib.bazel.bzlmod.DelegateTypeAdapterFactory.IMMUTABLE_MAP;
 import static com.google.devtools.build.lib.bazel.bzlmod.DelegateTypeAdapterFactory.IMMUTABLE_SET;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Splitter;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.google.gson.JsonParseException;
 import com.google.gson.TypeAdapter;
+import com.google.gson.TypeAdapterFactory;
+import com.google.gson.reflect.TypeToken;
 import com.google.gson.stream.JsonReader;
+import com.google.gson.stream.JsonToken;
 import com.google.gson.stream.JsonWriter;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import java.io.IOException;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
 import java.util.List;
+import java.util.Optional;
+import javax.annotation.Nullable;
 
 /**
  * Utility class to hold type adapters and helper methods to get gson registered with type adapters
@@ -88,6 +96,56 @@ public ModuleKey read(JsonReader jsonReader) throws IOException {
         }
       };
 
+  public static final TypeAdapterFactory OPTIONAL =
+      new TypeAdapterFactory() {
+        @Nullable
+        @Override
+        @SuppressWarnings("unchecked")
+        public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> typeToken) {
+          if (typeToken.getRawType() != Optional.class) {
+            return null;
+          }
+          Type type = typeToken.getType();
+          if (!(type instanceof ParameterizedType)) {
+            return null;
+          }
+          Type elementType = ((ParameterizedType) typeToken.getType()).getActualTypeArguments()[0];
+          var elementTypeAdapter = gson.getAdapter(TypeToken.get(elementType));
+          if (elementTypeAdapter == null) {
+            return null;
+          }
+          return (TypeAdapter<T>) new OptionalTypeAdapter<>(elementTypeAdapter);
+        }
+      };
+
+  private static final class OptionalTypeAdapter<T> extends TypeAdapter<Optional<T>> {
+    private final TypeAdapter<T> elementTypeAdapter;
+
+    public OptionalTypeAdapter(TypeAdapter<T> elementTypeAdapter) {
+      this.elementTypeAdapter = elementTypeAdapter;
+    }
+
+    @Override
+    public void write(JsonWriter jsonWriter, Optional<T> t) throws IOException {
+      Preconditions.checkNotNull(t);
+      if (t.isEmpty()) {
+        jsonWriter.nullValue();
+      } else {
+        elementTypeAdapter.write(jsonWriter, t.get());
+      }
+    }
+
+    @Override
+    public Optional<T> read(JsonReader jsonReader) throws IOException {
+      if (jsonReader.peek() == JsonToken.NULL) {
+        jsonReader.nextNull();
+        return Optional.empty();
+      } else {
+        return Optional.of(elementTypeAdapter.read(jsonReader));
+      }
+    }
+  }
+
   public static final Gson LOCKFILE_GSON =
       new GsonBuilder()
           .setPrettyPrinting()
@@ -97,6 +155,7 @@ public ModuleKey read(JsonReader jsonReader) throws IOException {
           .registerTypeAdapterFactory(IMMUTABLE_LIST)
           .registerTypeAdapterFactory(IMMUTABLE_BIMAP)
           .registerTypeAdapterFactory(IMMUTABLE_SET)
+          .registerTypeAdapterFactory(OPTIONAL)
           .registerTypeAdapter(Version.class, VERSION_TYPE_ADAPTER)
           .registerTypeAdapter(ModuleKey.class, MODULE_KEY_TYPE_ADAPTER)
           .registerTypeAdapter(AttributeValues.class, new AttributeValuesAdapter())

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

@@ -120,6 +120,16 @@ public boolean isDevDependency(TypeCheckedTag tag) {
     return tag.isDevDependency();
   }
 
+  @StarlarkMethod(
+      name = "is_isolated",
+      doc =
+          "Whether this particular usage of the extension had <code>isolate = True</code> "
+              + "specified and is thus isolated from all other usages.",
+      structField = true)
+  public boolean isIsolated() {
+    return extensionId.getIsolationKey().isPresent();
+  }
+
   @StarlarkMethod(
       name = "extension_metadata",
       doc =
@@ -181,6 +191,6 @@ public ModuleExtensionMetadata extensionMetadata(
       Object rootModuleDirectDepsUnchecked, Object rootModuleDirectDevDepsUnchecked)
       throws EvalException {
     return ModuleExtensionMetadata.create(
-        rootModuleDirectDepsUnchecked, rootModuleDirectDevDepsUnchecked);
+        rootModuleDirectDepsUnchecked, rootModuleDirectDevDepsUnchecked, extensionId);
   }
 }

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

@@ -17,15 +17,41 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.devtools.build.lib.cmdline.Label;
+import java.util.Optional;
 
 /** A unique identifier for a {@link ModuleExtension}. */
 @AutoValue
 public abstract class ModuleExtensionId {
+
+  /** A unique identifier for a single isolated usage of a fixed module extension. */
+  @AutoValue
+  abstract static class IsolationKey {
+    /** The module which contains this isolated usage of a module extension. */
+    public abstract ModuleKey getModule();
+
+    /** Whether this isolated usage specified {@code dev_dependency = True}. */
+    public abstract boolean isDevUsage();
+
+    /**
+     * The 0-based index of this isolated usage within the module's isolated usages of the same
+     * module extension and with the same {@link #isDevUsage()} value.
+     */
+    public abstract int getIsolatedUsageIndex();
+
+    public static IsolationKey create(
+        ModuleKey module, boolean isDevUsage, int isolatedUsageIndex) {
+      return new AutoValue_ModuleExtensionId_IsolationKey(module, isDevUsage, isolatedUsageIndex);
+    }
+  }
+
   public abstract Label getBzlFileLabel();
 
   public abstract String getExtensionName();
 
-  public static ModuleExtensionId create(Label bzlFileLabel, String extensionName) {
-    return new AutoValue_ModuleExtensionId(bzlFileLabel, extensionName);
+  public abstract Optional<IsolationKey> getIsolationKey();
+
+  public static ModuleExtensionId create(
+      Label bzlFileLabel, String extensionName, Optional<IsolationKey> isolationKey) {
+    return new AutoValue_ModuleExtensionId(bzlFileLabel, extensionName, isolationKey);
   }
 }

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

@@ -69,7 +69,9 @@ private ModuleExtensionMetadata(
   }
 
   static ModuleExtensionMetadata create(
-      Object rootModuleDirectDepsUnchecked, Object rootModuleDirectDevDepsUnchecked)
+      Object rootModuleDirectDepsUnchecked,
+      Object rootModuleDirectDevDepsUnchecked,
+      ModuleExtensionId extensionId)
       throws EvalException {
     if (rootModuleDirectDepsUnchecked == Starlark.NONE
         && rootModuleDirectDevDepsUnchecked == Starlark.NONE) {
@@ -80,11 +82,23 @@ 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);
     }
 
@@ -114,6 +128,20 @@ 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 {
@@ -257,13 +285,33 @@ private static Optional<Event> generateFixupMessage(
     var fixupCommands =
         Stream.of(
                 makeUseRepoCommand(
-                    "use_repo_add", false, importsToAdd, extensionBzlFile, extensionName),
+                    "use_repo_add",
+                    false,
+                    importsToAdd,
+                    extensionBzlFile,
+                    extensionName,
+                    firstUsage.getIsolationKey()),
                 makeUseRepoCommand(
-                    "use_repo_remove", false, importsToRemove, extensionBzlFile, extensionName),
+                    "use_repo_remove",
+                    false,
+                    importsToRemove,
+                    extensionBzlFile,
+                    extensionName,
+                    firstUsage.getIsolationKey()),
                 makeUseRepoCommand(
-                    "use_repo_add", true, devImportsToAdd, extensionBzlFile, extensionName),
+                    "use_repo_add",
+                    true,
+                    devImportsToAdd,
+                    extensionBzlFile,
+                    extensionName,
+                    firstUsage.getIsolationKey()),
                 makeUseRepoCommand(
-                    "use_repo_remove", true, devImportsToRemove, extensionBzlFile, extensionName))
+                    "use_repo_remove",
+                    true,
+                    devImportsToRemove,
+                    extensionBzlFile,
+                    extensionName,
+                    firstUsage.getIsolationKey()))
             .flatMap(Optional::stream);
 
     return Optional.of(
@@ -284,17 +332,28 @@ private static Optional<String> makeUseRepoCommand(
       boolean devDependency,
       Collection<String> repos,
       String extensionBzlFile,
-      String extensionName) {
+      String extensionName,
+      Optional<ModuleExtensionId.IsolationKey> isolationKey) {
+
     if (repos.isEmpty()) {
       return Optional.empty();
     }
+
+    String extensionUsageIdentifier = extensionName;
+    if (isolationKey.isPresent()) {
+      // We verified in create() that the extension did not report root module deps of a kind that
+      // does not match the isolated (and hence only) usage. It also isn't possible for users to
+      // specify repo usages of the wrong kind, so we can't get here.
+      Preconditions.checkState(isolationKey.get().isDevUsage() == devDependency);
+      extensionUsageIdentifier += ":" + isolationKey.get().getIsolatedUsageIndex();
+    }
     return Optional.of(
         String.format(
             "buildozer '%s%s %s %s %s' //MODULE.bazel:all",
             cmd,
             devDependency ? " dev" : "",
             extensionBzlFile,
-            extensionName,
+            extensionUsageIdentifier,
             String.join(" ", repos)));
   }
 

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

@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
+import java.util.Optional;
 import net.starlark.java.syntax.Location;
 
 /**
@@ -35,6 +36,12 @@ public abstract class ModuleExtensionUsage {
   /** The name of the extension. */
   public abstract String getExtensionName();
 
+  /**
+   * The isolation key of this module extension usage. This is present if and only if the usage is
+   * created with {@code isolate = True}.
+   */
+  public abstract Optional<ModuleExtensionId.IsolationKey> getIsolationKey();
+
   /** The module that contains this particular extension usage. */
   public abstract ModuleKey getUsingModule();
 
@@ -73,6 +80,8 @@ public abstract static class Builder {
 
     public abstract Builder setExtensionName(String value);
 
+    public abstract Builder setIsolationKey(Optional<ModuleExtensionId.IsolationKey> value);
+
     public abstract Builder setUsingModule(ModuleKey value);
 
     public abstract Builder setLocation(Location value);