[6.3.0] Identify isolated extensions by exported name

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

@@ -252,12 +252,26 @@ 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.getIsolationKey()
+              .map(
+                  namespace ->
+                      String.format(
+                          "%s~_%s~%s~%s~%s",
+                          nonEmptyRepoPart,
+                          id.getExtensionName(),
+                          namespace.getModule().getName(),
+                          namespace.getModule().getVersion(),
+                          namespace.getUsageExportedName()))
+              .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/ModuleExtensionId.java

@@ -16,7 +16,6 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import static com.google.common.collect.Comparators.emptiesFirst;
-import static com.google.common.primitives.Booleans.falseFirst;
 import static java.util.Comparator.comparing;
 
 import com.google.auto.value.AutoValue;
@@ -39,24 +38,16 @@ public abstract class ModuleExtensionId {
   abstract static class IsolationKey {
     static final Comparator<IsolationKey> LEXICOGRAPHIC_COMPARATOR =
         comparing(IsolationKey::getModule, ModuleKey.LEXICOGRAPHIC_COMPARATOR)
-            .thenComparing(IsolationKey::isDevUsage, falseFirst())
-            .thenComparing(IsolationKey::getIsolatedUsageIndex);
+            .thenComparing(IsolationKey::getUsageExportedName);
 
     /** 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 exported name of the first extension proxy for this usage. */
+    public abstract String getUsageExportedName();
 
-    /**
-     * 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 static IsolationKey create(ModuleKey module, String usageExportedName) {
+      return new AutoValue_ModuleExtensionId_IsolationKey(module, usageExportedName);
     }
   }
 

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

@@ -26,6 +26,7 @@
 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.ArrayList;
 import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.Optional;
@@ -319,22 +320,20 @@ private static Optional<String> makeUseRepoCommand(
       return Optional.empty();
     }
 
-    String extensionUsageIdentifier = extensionName;
+    var commandParts = new ArrayList<String>();
+    commandParts.add(cmd);
     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();
+      commandParts.add(isolationKey.get().getUsageExportedName());
+    } else {
+      if (devDependency) {
+        commandParts.add("dev");
+      }
+      commandParts.add(extensionBzlFile);
+      commandParts.add(extensionName);
     }
+    commandParts.addAll(repos);
     return Optional.of(
-        String.format(
-            "buildozer '%s%s %s %s %s' //MODULE.bazel:all",
-            cmd,
-            devDependency ? " dev" : "",
-            extensionBzlFile,
-            extensionUsageIdentifier,
-            String.join(" ", repos)));
+        String.format("buildozer '%s' //MODULE.bazel:all", String.join(" ", commandParts)));
   }
 
   private Optional<ImmutableSet<String>> getRootModuleDirectDeps(Set<String> allRepos)

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

@@ -27,6 +27,7 @@
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.packages.StarlarkExportable;
 import com.google.devtools.build.lib.rules.repository.RepositoryDirectoryValue;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps.Code;
 import com.google.devtools.build.lib.skyframe.ClientEnvironmentFunction;
@@ -147,7 +148,13 @@ public SkyValue compute(SkyKey skyKey, Environment env)
             env);
 
     // Perform some sanity checks.
-    InterimModule module = moduleFileGlobals.buildModule();
+    InterimModule module;
+    try {
+      module = moduleFileGlobals.buildModule();
+    } catch (EvalException e) {
+      env.getListener().handle(Event.error(e.getMessageWithStack()));
+      throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for %s", moduleKey);
+    }
     if (!module.getName().equals(moduleKey.getName())) {
       throw errorf(
           Code.BAD_MODULE,
@@ -203,7 +210,13 @@ private SkyValue computeForRootModule(StarlarkSemantics starlarkSemantics, Envir
             /* printIsNoop= */ false,
             starlarkSemantics,
             env);
-    InterimModule module = moduleFileGlobals.buildModule();
+    InterimModule module;
+    try {
+      module = moduleFileGlobals.buildModule();
+    } catch (EvalException e) {
+      env.getListener().handle(Event.error(e.getMessageWithStack()));
+      throw errorf(Code.BAD_MODULE, "error executing MODULE.bazel file for the root module");
+    }
     for (ModuleExtensionUsage usage : module.getExtensionUsages()) {
       if (usage.getIsolationKey().isPresent() && usage.getImports().isEmpty()) {
         throw errorf(
@@ -271,6 +284,15 @@ private ModuleFileGlobals execModuleFile(
       } else {
         thread.setPrintHandler(Event.makeDebugPrintHandler(env.getListener()));
       }
+      thread.setPostAssignHook(
+          (name, value) -> {
+            if (value instanceof StarlarkExportable) {
+              StarlarkExportable exportable = (StarlarkExportable) value;
+              if (!exportable.isExported()) {
+                exportable.export(env.getListener(), null, name);
+              }
+            }
+          });
       Starlark.execFileProgram(program, predeclaredEnv, thread);
     } catch (SyntaxError.Exception e) {
       Event.replayEventsOn(env.getListener(), e.errors());

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

@@ -15,8 +15,6 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
-
 import com.google.auto.value.AutoValue;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.HashBiMap;
@@ -29,7 +27,10 @@
 import com.google.devtools.build.lib.bazel.bzlmod.InterimModule.DepSpec;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileGlobals.ModuleExtensionUsageBuilder.ModuleExtensionProxy;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
+import com.google.devtools.build.lib.events.EventHandler;
+import com.google.devtools.build.lib.packages.StarlarkExportable;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -71,8 +72,6 @@ public class ModuleFileGlobals {
   private final List<ModuleExtensionUsageBuilder> extensionUsageBuilders = new ArrayList<>();
   private final Map<String, ModuleOverride> overrides = new HashMap<>();
   private final Map<String, RepoNameUsage> repoNameUsages = new HashMap<>();
-  private int nextIsolatedNonDevUsageIndex = 0;
-  private int nextIsolatedDevUsageIndex = 0;
 
   public ModuleFileGlobals(
       ImmutableMap<String, NonRegistryOverride> builtinModules,
@@ -451,22 +450,9 @@ public ModuleExtensionProxy useExtension(
     hadNonModuleCall = true;
 
     String extensionBzlFile = normalizeLabelString(rawExtensionBzlFile);
-
-    Optional<ModuleExtensionId.IsolationKey> isolationKey;
-    if (isolate) {
-      isolationKey =
-          Optional.of(
-              ModuleExtensionId.IsolationKey.create(
-                  module.getKey(),
-                  devDependency,
-                  devDependency ? nextIsolatedDevUsageIndex++ : nextIsolatedNonDevUsageIndex++));
-    } else {
-      isolationKey = Optional.empty();
-    }
-
     ModuleExtensionUsageBuilder newUsageBuilder =
         new ModuleExtensionUsageBuilder(
-            extensionBzlFile, extensionName, isolationKey, thread.getCallerLocation());
+            extensionBzlFile, extensionName, isolate, thread.getCallerLocation());
 
     if (ignoreDevDeps && devDependency) {
       // This is a no-op proxy.
@@ -479,7 +465,7 @@ public ModuleExtensionProxy useExtension(
       for (ModuleExtensionUsageBuilder usageBuilder : extensionUsageBuilders) {
         if (usageBuilder.extensionBzlFile.equals(extensionBzlFile)
             && usageBuilder.extensionName.equals(extensionName)
-            && usageBuilder.isolationKey.isEmpty()) {
+            && !usageBuilder.isolate) {
           return usageBuilder.getProxy(devDependency);
         }
       }
@@ -509,42 +495,49 @@ private String normalizeLabelString(String rawExtensionBzlFile) {
   class ModuleExtensionUsageBuilder {
     private final String extensionBzlFile;
     private final String extensionName;
-    private final Optional<ModuleExtensionId.IsolationKey> isolationKey;
+    private final boolean isolate;
     private final Location location;
     private final HashBiMap<String, String> imports;
     private final ImmutableSet.Builder<String> devImports;
     private final ImmutableList.Builder<Tag> tags;
 
     private boolean hasNonDevUseExtension;
     private boolean hasDevUseExtension;
+    private String exportedName;
 
     ModuleExtensionUsageBuilder(
-        String extensionBzlFile,
-        String extensionName,
-        Optional<ModuleExtensionId.IsolationKey> isolationKey,
-        Location location) {
+        String extensionBzlFile, String extensionName, boolean isolate, Location location) {
       this.extensionBzlFile = extensionBzlFile;
       this.extensionName = extensionName;
-      this.isolationKey = isolationKey;
+      this.isolate = isolate;
       this.location = location;
       this.imports = HashBiMap.create();
       this.devImports = ImmutableSet.builder();
       this.tags = ImmutableList.builder();
     }
 
-    ModuleExtensionUsage buildUsage() {
+    ModuleExtensionUsage buildUsage() throws EvalException {
       var builder =
           ModuleExtensionUsage.builder()
               .setExtensionBzlFile(extensionBzlFile)
               .setExtensionName(extensionName)
-              .setIsolationKey(isolationKey)
               .setUsingModule(module.getKey())
               .setLocation(location)
               .setImports(ImmutableBiMap.copyOf(imports))
               .setDevImports(devImports.build())
               .setHasDevUseExtension(hasDevUseExtension)
               .setHasNonDevUseExtension(hasNonDevUseExtension)
               .setTags(tags.build());
+      if (isolate) {
+        if (exportedName == null) {
+          throw Starlark.errorf(
+              "Isolated extension usage at %s must be assigned to a top-level variable", location);
+        }
+        builder.setIsolationKey(
+            Optional.of(ModuleExtensionId.IsolationKey.create(module.getKey(), exportedName)));
+      } else {
+        builder.setIsolationKey(Optional.empty());
+      }
       return builder.build();
     }
 
@@ -562,7 +555,7 @@ ModuleExtensionProxy getProxy(boolean devDependency) {
     }
 
     @StarlarkBuiltin(name = "module_extension_proxy", documented = false)
-    class ModuleExtensionProxy implements Structure {
+    class ModuleExtensionProxy implements Structure, StarlarkExportable {
 
       private final boolean devDependency;
 
@@ -619,6 +612,16 @@ public ImmutableCollection<String> getFieldNames() {
       public String getErrorMessageForUnknownField(String field) {
         return null;
       }
+
+      @Override
+      public boolean isExported() {
+        return exportedName != null;
+      }
+
+      @Override
+      public void export(EventHandler handler, Label bzlFileLabel, String name) {
+        exportedName = name;
+      }
     }
   }
 
@@ -977,14 +980,15 @@ public void localPathOverride(String moduleName, String path) throws EvalExcepti
     addOverride(moduleName, LocalPathOverride.create(path));
   }
 
-  public InterimModule buildModule() {
+  public InterimModule buildModule() throws EvalException {
+    var extensionUsages = ImmutableList.<ModuleExtensionUsage>builder();
+    for (var extensionUsageBuilder : extensionUsageBuilders) {
+      extensionUsages.add(extensionUsageBuilder.buildUsage());
+    }
     return module
         .setDeps(ImmutableMap.copyOf(deps))
         .setOriginalDeps(ImmutableMap.copyOf(deps))
-        .setExtensionUsages(
-            extensionUsageBuilders.stream()
-                .map(ModuleExtensionUsageBuilder::buildUsage)
-                .collect(toImmutableList()))
+        .setExtensionUsages(extensionUsages.build())
         .build();
   }
 

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

@@ -2107,9 +2107,8 @@ public void extensionMetadata_isolated() throws Exception {
             + "\033[35m\033[1m ** You can use the following buildozer command(s) to fix these"
             + " issues:\033[0m\n"
             + "\n"
-            + "buildozer 'use_repo_add @ext//:defs.bzl ext:0 direct_dep missing_direct_dep'"
-            + " //MODULE.bazel:all\n"
-            + "buildozer 'use_repo_remove @ext//:defs.bzl ext:0 indirect_dep' //MODULE.bazel:all",
+            + "buildozer 'use_repo_add ext1 direct_dep missing_direct_dep' //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_remove ext1 indirect_dep' //MODULE.bazel:all",
         ImmutableSet.of(EventKind.WARNING));
     assertContainsEvent(
         "WARNING /ws/MODULE.bazel:8:21: The module extension ext defined in @ext//:defs.bzl"
@@ -2122,8 +2121,7 @@ public void extensionMetadata_isolated() throws Exception {
             + "\033[35m\033[1m ** You can use the following buildozer command(s) to fix these"
             + " issues:\033[0m\n"
             + "\n"
-            + "buildozer 'use_repo_add @ext//:defs.bzl ext:1 missing_direct_dep'"
-            + " //MODULE.bazel:all",
+            + "buildozer 'use_repo_add ext2 missing_direct_dep' //MODULE.bazel:all",
         ImmutableSet.of(EventKind.WARNING));
   }
 
@@ -2193,10 +2191,8 @@ public void extensionMetadata_isolatedDev() throws Exception {
             + "\033[35m\033[1m ** You can use the following buildozer command(s) to fix these"
             + " issues:\033[0m\n"
             + "\n"
-            + "buildozer 'use_repo_add dev @ext//:defs.bzl ext:0 direct_dep missing_direct_dep'"
-            + " //MODULE.bazel:all\n"
-            + "buildozer 'use_repo_remove dev @ext//:defs.bzl ext:0 indirect_dep'"
-            + " //MODULE.bazel:all",
+            + "buildozer 'use_repo_add ext1 direct_dep missing_direct_dep' //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_remove ext1 indirect_dep' //MODULE.bazel:all",
         ImmutableSet.of(EventKind.WARNING));
     assertContainsEvent(
         "WARNING /ws/MODULE.bazel:8:21: The module extension ext defined in @ext//:defs.bzl"
@@ -2209,8 +2205,7 @@ public void extensionMetadata_isolatedDev() throws Exception {
             + "\033[35m\033[1m ** You can use the following buildozer command(s) to fix these"
             + " issues:\033[0m\n"
             + "\n"
-            + "buildozer 'use_repo_add dev @ext//:defs.bzl ext:1 missing_direct_dep'"
-            + " //MODULE.bazel:all",
+            + "buildozer 'use_repo_add ext2 missing_direct_dep' //MODULE.bazel:all",
         ImmutableSet.of(EventKind.WARNING));
   }