Allow module extension usages to be isolated (#18727)

If `isolate = True` is specified on `use_extension`, that particular usage will be isolated from all other usages, both in the same and other modules.

Module extensions can check whether they are isolated (e.g. in case they can only be used in this way) via `module_ctx.is_isolated`.

Closes #18529.

PiperOrigin-RevId: 541823020
Change-Id: I68a7b49914bbc1fd50df2fda7a0af1e47421bb92

Co-authored-by: Fabian Meumertzheim <[email protected]>

Author: iancha1992

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);