Allow module extension usages to be isolated

If `isolated = 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`.

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

@@ -45,6 +45,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Optional;
 import javax.annotation.Nullable;
 
 /**
@@ -214,7 +215,11 @@ private ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> getEx
         try {
           moduleExtensionId =
               ModuleExtensionId.create(
-                  labelConverter.convert(usage.getExtensionBzlFile()), usage.getExtensionName());
+                  labelConverter.convert(usage.getExtensionBzlFile()),
+                  usage.getExtensionName(),
+                  usage.isIsolated()
+                      ? Optional.of(new ModuleExtensionId.Namespace(usage.getUsingModule()))
+                      : Optional.empty());
         } catch (LabelSyntaxException e) {
           throw new BazelDepGraphFunctionException(
               ExternalDepsException.withCauseAndMessage(
@@ -248,7 +253,21 @@ 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 "_".
+      String bestName =
+          id.getNamespace()
+              .map(
+                  namespace ->
+                      nonEmptyRepoPart
+                          + "~_"
+                          + id.getExtensionName()
+                          + "~"
+                          + namespace.module.getName()
+                          + "~"
+                          + namespace.module.getVersion())
+              .orElse(nonEmptyRepoPart + "~" + id.getExtensionName());
       if (extensionUniqueNames.putIfAbsent(bestName, id) == null) {
         continue;
       }

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

@@ -123,6 +123,16 @@ public boolean isDevDependency(TypeCheckedTag tag) {
     return tag.isDevDependency();
   }
 
+  @StarlarkMethod(
+      name = "is_isolated",
+      doc = "Whether this particular usage of the extension is isolated from all others, in " +
+          "particular those in other modules",
+      structField = true
+  )
+  public boolean isIsolated() {
+    return extensionId.getNamespace().isPresent();
+  }
+
   @StarlarkMethod(
       name = "extension_metadata",
       doc =

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

@@ -17,15 +17,28 @@
 
 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 {
+
+  static final class Namespace {
+    final ModuleKey module;
+
+    Namespace(ModuleKey module) {
+      this.module = module;
+    }
+  }
+
   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<Namespace> getNamespace();
+
+  public static ModuleExtensionId create(
+      Label bzlFileLabel, String extensionName, Optional<Namespace> namespace) {
+    return new AutoValue_ModuleExtensionId(bzlFileLabel, extensionName, namespace);
   }
 }

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

@@ -35,6 +35,8 @@ public abstract class ModuleExtensionUsage {
   /** The name of the extension. */
   public abstract String getExtensionName();
 
+  public abstract boolean isIsolated();
+
   /** The module that contains this particular extension usage. */
   public abstract ModuleKey getUsingModule();
 
@@ -73,6 +75,8 @@ public abstract static class Builder {
 
     public abstract Builder setExtensionName(String value);
 
+    public abstract Builder setIsolated(boolean value);
+
     public abstract Builder setUsingModule(ModuleKey value);
 
     public abstract Builder setLocation(Location value);

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

@@ -432,31 +432,42 @@ public void registerToolchains(boolean devDependency, Sequence<?> toolchainLabel
             named = true,
             positional = false,
             defaultValue = "False"),
+        @Param(
+            name = "isolated",
+            doc =
+                "If true, this usage of the module extension will be isolated from all other "
+                    + "usages, in particular those in other modules.",
+            named = true,
+            positional = false,
+            defaultValue = "False"),
       },
       useStarlarkThread = true)
   public ModuleExtensionProxy useExtension(
       String rawExtensionBzlFile,
       String extensionName,
       boolean devDependency,
+      boolean isolated,
       StarlarkThread thread) {
     hadNonModuleCall = true;
 
     String extensionBzlFile = normalizeLabelString(rawExtensionBzlFile);
 
     ModuleExtensionUsageBuilder newUsageBuilder =
         new ModuleExtensionUsageBuilder(
-            extensionBzlFile, extensionName, thread.getCallerLocation());
+            extensionBzlFile, extensionName, isolated, thread.getCallerLocation());
 
     if (ignoreDevDeps && devDependency) {
       // This is a no-op proxy.
       return newUsageBuilder.getProxy(devDependency);
     }
 
     // Find an existing usage builder corresponding to this extension.
-    for (ModuleExtensionUsageBuilder usageBuilder : extensionUsageBuilders) {
-      if (usageBuilder.extensionBzlFile.equals(extensionBzlFile)
-          && usageBuilder.extensionName.equals(extensionName)) {
-        return usageBuilder.getProxy(devDependency);
+    if (!isolated) {
+      for (ModuleExtensionUsageBuilder usageBuilder : extensionUsageBuilders) {
+        if (usageBuilder.extensionBzlFile.equals(extensionBzlFile)
+            && usageBuilder.extensionName.equals(extensionName)) {
+          return usageBuilder.getProxy(devDependency);
+        }
       }
     }
 
@@ -484,14 +495,17 @@ private String normalizeLabelString(String rawExtensionBzlFile) {
   class ModuleExtensionUsageBuilder {
     private final String extensionBzlFile;
     private final String extensionName;
+    private final boolean isolated;
     private final Location location;
     private final HashBiMap<String, String> imports;
     private final ImmutableSet.Builder<String> devImports;
     private final ImmutableList.Builder<Tag> tags;
 
-    ModuleExtensionUsageBuilder(String extensionBzlFile, String extensionName, Location location) {
+    ModuleExtensionUsageBuilder(
+        String extensionBzlFile, String extensionName, boolean isolated, Location location) {
       this.extensionBzlFile = extensionBzlFile;
       this.extensionName = extensionName;
+      this.isolated = isolated;
       this.location = location;
       this.imports = HashBiMap.create();
       this.devImports = ImmutableSet.builder();
@@ -502,6 +516,7 @@ ModuleExtensionUsage buildUsage() {
       return ModuleExtensionUsage.builder()
           .setExtensionBzlFile(extensionBzlFile)
           .setExtensionName(extensionName)
+          .setIsolated(isolated)
           .setUsingModule(module.getKey())
           .setLocation(location)
           .setImports(ImmutableBiMap.copyOf(imports))

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

@@ -64,6 +64,7 @@
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.StarlarkSemantics;
@@ -214,6 +215,7 @@ private static ModuleExtensionUsage createModuleExtensionUsage(
     return ModuleExtensionUsage.builder()
         .setExtensionBzlFile(bzlFile)
         .setExtensionName(name)
+        .setIsolated(false)
         .setImports(importsBuilder.buildOrThrow())
         .setDevImports(ImmutableSet.of())
         .setUsingModule(ModuleKey.ROOT)
@@ -253,14 +255,16 @@ public void createValue_moduleExtensions() throws Exception {
 
     ModuleExtensionId maven =
         ModuleExtensionId.create(
-            Label.parseCanonical("@@rules_jvm_external~1.0//:defs.bzl"), "maven");
+            Label.parseCanonical("@@rules_jvm_external~1.0//:defs.bzl"), "maven", Optional.empty());
     ModuleExtensionId pip =
-        ModuleExtensionId.create(Label.parseCanonical("@@rules_python~2.0//:defs.bzl"), "pip");
+        ModuleExtensionId.create(
+            Label.parseCanonical("@@rules_python~2.0//:defs.bzl"), "pip", Optional.empty());
     ModuleExtensionId myext =
-        ModuleExtensionId.create(Label.parseCanonical("@@dep~2.0//:defs.bzl"), "myext");
+        ModuleExtensionId.create(
+            Label.parseCanonical("@@dep~2.0//:defs.bzl"), "myext", Optional.empty());
     ModuleExtensionId myext2 =
         ModuleExtensionId.create(
-            Label.parseCanonical("@@dep~2.0//incredible:conflict.bzl"), "myext");
+            Label.parseCanonical("@@dep~2.0//incredible:conflict.bzl"), "myext", Optional.empty());
 
     resolutionFunctionMock.setDepGraph(depGraph);
     EvaluationResult<BazelDepGraphValue> result =

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

@@ -1900,7 +1900,7 @@ private EvaluationResult<SingleExtensionEvalValue> evaluateSimpleModuleExtension
     scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
 
     ModuleExtensionId extensionId =
-        ModuleExtensionId.create(Label.parseCanonical("//:defs.bzl"), "ext");
+        ModuleExtensionId.create(Label.parseCanonical("//:defs.bzl"), "ext", Optional.empty());
     reporter.removeHandler(failFastHandler);
     return evaluator.evaluate(
         ImmutableList.of(SingleExtensionEvalValue.key(extensionId)), evaluationContext);

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

@@ -476,6 +476,7 @@ public void testModuleExtensions_good() throws Exception {
                     ModuleExtensionUsage.builder()
                         .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext1")
+                        .setIsolated(false)
                         .setUsingModule(myMod)
                         .setLocation(
                             Location.fromFileLineColumn(
@@ -500,6 +501,7 @@ public void testModuleExtensions_good() throws Exception {
                     ModuleExtensionUsage.builder()
                         .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext2")
+                        .setIsolated(false)
                         .setUsingModule(myMod)
                         .setLocation(
                             Location.fromFileLineColumn(
@@ -537,6 +539,7 @@ public void testModuleExtensions_good() throws Exception {
                     ModuleExtensionUsage.builder()
                         .setExtensionBzlFile("@rules_jvm_external//:defs.bzl")
                         .setExtensionName("maven")
+                        .setIsolated(false)
                         .setUsingModule(myMod)
                         .setLocation(
                             Location.fromFileLineColumn(
@@ -607,6 +610,7 @@ public void testModuleExtensions_duplicateProxy_asRoot() throws Exception {
                     ModuleExtensionUsage.builder()
                         .setExtensionBzlFile("@//:defs.bzl")
                         .setExtensionName("myext")
+                        .setIsolated(false)
                         .setUsingModule(ModuleKey.ROOT)
                         .setLocation(Location.fromFileLineColumn("/workspace/MODULE.bazel", 1, 23))
                         .setImports(
@@ -704,6 +708,7 @@ public void testModuleExtensions_duplicateProxy_asDep() throws Exception {
                     ModuleExtensionUsage.builder()
                         .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext")
+                        .setIsolated(false)
                         .setUsingModule(myMod)
                         .setLocation(
                             Location.fromFileLineColumn(

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

@@ -45,6 +45,7 @@ private static ModuleExtensionUsage.Builder getBaseUsageBuilder() {
     return ModuleExtensionUsage.builder()
         .setExtensionBzlFile("//:rje.bzl")
         .setExtensionName("maven")
+        .setIsolated(false)
         .setUsingModule(ModuleKey.ROOT)
         .setLocation(Location.BUILTIN)
         .setImports(ImmutableBiMap.of())