evaluate bzlmod module extensions independently

- evaluate each module extension independently, on request. this change is intended to allow loading repositories generated by other module extensions inside a module extension's implementation, as requested in bazelbuild#16143).

- `BazelModuleResolutionFunction` is removed until further need, as now `RepositoryDelegatorFunction` directly calls `SingleExtensionEvalFunction` to resolve the necessary module extension.

PiperOrigin-RevId: 469971208
Change-Id: Iad4ecd926024564209765386496a7491146c2f2b

src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java

@@ -28,7 +28,6 @@
 import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleInspectorFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleResolutionFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.LocalPathOverride;
-import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionResolutionFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.NonRegistryOverride;
 import com.google.devtools.build.lib.bazel.bzlmod.RegistryFactory;
@@ -234,9 +233,7 @@ public void workspaceInit(
         .addSkyFunction(SkyFunctions.BAZEL_MODULE_RESOLUTION, new BazelModuleResolutionFunction())
         .addSkyFunction(SkyFunctions.BAZEL_MODULE_INSPECTION, new BazelModuleInspectorFunction())
         .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_EVAL, singleExtensionEvalFunction)
-        .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction())
-        .addSkyFunction(
-            SkyFunctions.MODULE_EXTENSION_RESOLUTION, new ModuleExtensionResolutionFunction());
+        .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction());
     filesystem = runtime.getFileSystem();
   }
 

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

@@ -91,7 +91,6 @@ java_library(
         "LocalPathOverride.java",
         "Module.java",
         "ModuleExtensionEvalStarlarkThreadContext.java",
-        "ModuleExtensionResolutionValue.java",
         "ModuleFileValue.java",
         "ModuleOverride.java",
         "MultipleVersionOverride.java",
@@ -129,7 +128,6 @@ java_library(
         "BazelModuleResolutionFunction.java",
         "Discovery.java",
         "ModuleExtensionContext.java",
-        "ModuleExtensionResolutionFunction.java",
         "ModuleFileFunction.java",
         "ModuleFileGlobals.java",
         "Selection.java",

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

@@ -15,11 +15,14 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
+import static com.google.common.collect.ImmutableBiMap.toImmutableBiMap;
+
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
 import com.google.devtools.build.lib.cmdline.BazelModuleContext;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.rules.repository.NeedsSkyframeRestartException;
 import com.google.devtools.build.lib.runtime.ProcessWrapper;
@@ -39,6 +42,7 @@
 import java.util.ArrayList;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
@@ -213,7 +217,15 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       }
     }
 
-    return SingleExtensionEvalValue.create(threadContext.getGeneratedRepos());
+    return SingleExtensionEvalValue.create(
+        threadContext.getGeneratedRepos(),
+        threadContext.getGeneratedRepos().keySet().stream()
+            .collect(
+                toImmutableBiMap(
+                    e ->
+                        RepositoryName.createUnvalidated(
+                            usagesValue.getExtensionUniqueName() + "~" + e),
+                    Function.identity())));
   }
 
   private ModuleExtensionContext createContext(

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

@@ -16,8 +16,10 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Interner;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
@@ -36,9 +38,17 @@ public abstract class SingleExtensionEvalValue implements SkyValue {
    */
   public abstract ImmutableMap<String, Package> getGeneratedRepos();
 
+  /**
+   * Returns the mapping from the canonical repo names of the repos generated by this extension to
+   * their "internal name" (as specified by the extension).
+   */
+  public abstract ImmutableBiMap<RepositoryName, String> getCanonicalRepoNameToInternalNames();
+
   @AutoCodec.Instantiator
-  public static SingleExtensionEvalValue create(ImmutableMap<String, Package> generatedRepos) {
-    return new AutoValue_SingleExtensionEvalValue(generatedRepos);
+  public static SingleExtensionEvalValue create(
+      ImmutableMap<String, Package> generatedRepos,
+      ImmutableBiMap<RepositoryName, String> canonicalRepoNameToInternalNames) {
+    return new AutoValue_SingleExtensionEvalValue(generatedRepos, canonicalRepoNameToInternalNames);
   }
 
   public static Key key(ModuleExtensionId id) {

src/main/java/com/google/devtools/build/lib/skyframe/BUILD

@@ -282,6 +282,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis/platform:utils",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:exception",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:module_extension",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:repo_rule_creator",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:repo_rule_helper",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:repo_rule_value",

src/main/java/com/google/devtools/build/lib/skyframe/BzlmodRepoRuleFunction.java

@@ -18,11 +18,13 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
+import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleResolutionValue;
 import com.google.devtools.build.lib.bazel.bzlmod.BzlmodRepoRuleCreator;
 import com.google.devtools.build.lib.bazel.bzlmod.BzlmodRepoRuleHelper;
 import com.google.devtools.build.lib.bazel.bzlmod.BzlmodRepoRuleValue;
-import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionResolutionValue;
+import com.google.devtools.build.lib.bazel.bzlmod.ModuleExtensionId;
 import com.google.devtools.build.lib.bazel.bzlmod.RepoSpec;
+import com.google.devtools.build.lib.bazel.bzlmod.SingleExtensionEvalValue;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
@@ -42,6 +44,7 @@
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.io.IOException;
+import java.util.Map.Entry;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
@@ -105,17 +108,36 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     }
 
     // Otherwise, look for the repo from module extension evaluation results.
-    ModuleExtensionResolutionValue extensionResolution =
-        (ModuleExtensionResolutionValue) env.getValue(ModuleExtensionResolutionValue.KEY);
-    if (extensionResolution == null) {
+    BazelModuleResolutionValue moduleResolution =
+        (BazelModuleResolutionValue) env.getValue(BazelModuleResolutionValue.KEY);
+    if (moduleResolution == null) {
       return null;
     }
-    Package pkg = extensionResolution.getCanonicalRepoNameToPackage().get(repositoryName);
-    if (pkg != null) {
-      return new BzlmodRepoRuleValue(pkg, repositoryName.getName());
+
+    Optional<ModuleExtensionId> extensionId =
+        moduleResolution.getExtensionUniqueNames().entrySet().stream()
+            .filter(e -> repositoryName.getName().startsWith(e.getValue() + "~"))
+            .map(Entry::getKey)
+            .findFirst();
+
+    if (extensionId.isEmpty()) {
+      return BzlmodRepoRuleValue.REPO_RULE_NOT_FOUND_VALUE;
+    }
+
+    SingleExtensionEvalValue extensionEval =
+        (SingleExtensionEvalValue) env.getValue(SingleExtensionEvalValue.key(extensionId.get()));
+    if (extensionEval == null) {
+      return null;
+    }
+
+    String internalRepo = extensionEval.getCanonicalRepoNameToInternalNames().get(repositoryName);
+    if (internalRepo == null) {
+      return BzlmodRepoRuleValue.REPO_RULE_NOT_FOUND_VALUE;
     }
+    Package pkg = extensionEval.getGeneratedRepos().get(internalRepo);
+    Preconditions.checkNotNull(pkg);
 
-    return BzlmodRepoRuleValue.REPO_RULE_NOT_FOUND_VALUE;
+    return new BzlmodRepoRuleValue(pkg, repositoryName.getName());
   }
 
   @Nullable