Add InterimModule to represent a Module before resolution finishes

The class Module currently holds data that's no longer needed after resolution finishes (such as the compatibility level, bazel compatibility, the registry where it comes from, etc). Conversely, the repo spec field is not computed until the end of resolution.

To remove runtime checks and reduce cognitive overhead, this CL splits the Module class into two; one only used after resolution finishes (Module), and one only used before (InterimModule). This allows us to introduce max_compatibility_level in a follow CL.

Work towards bazelbuild#17378

Co-authored-by: Brentley Jones <github@brentleyjones.com>
PiperOrigin-RevId: 525780111
Change-Id: I2df8d78d324b3c8744ba0a4eda405d162e9bbb8c

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

@@ -94,8 +94,10 @@ java_library(
         "BazelModuleResolutionValue.java",
         "BzlmodFlagsAndEnvVars.java",
         "GitOverride.java",
+        "InterimModule.java",
         "LocalPathOverride.java",
         "Module.java",
+        "ModuleBase.java",
         "ModuleExtensionEvalStarlarkThreadContext.java",
         "ModuleFileValue.java",
         "ModuleOverride.java",

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

@@ -55,7 +55,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       return null;
     }
     ImmutableMap<String, ModuleOverride> overrides = root.getOverrides();
-    ImmutableMap<ModuleKey, Module> unprunedDepGraph = resolutionValue.getUnprunedDepGraph();
+    ImmutableMap<ModuleKey, InterimModule> unprunedDepGraph = resolutionValue.getUnprunedDepGraph();
     ImmutableMap<ModuleKey, Module> resolvedDepGraph = resolutionValue.getResolvedDepGraph();
 
     ImmutableMap<ModuleKey, AugmentedModule> depGraph =
@@ -74,7 +74,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
   }
 
   public static ImmutableMap<ModuleKey, AugmentedModule> computeAugmentedGraph(
-      ImmutableMap<ModuleKey, Module> unprunedDepGraph,
+      ImmutableMap<ModuleKey, InterimModule> unprunedDepGraph,
       ImmutableSet<ModuleKey> usedModules,
       ImmutableMap<String, ModuleOverride> overrides) {
     Map<ModuleKey, AugmentedModule.Builder> depGraphAugmentBuilder = new HashMap<>();
@@ -83,9 +83,9 @@ public static ImmutableMap<ModuleKey, AugmentedModule> computeAugmentedGraph(
     // to their children AugmentedModule as dependant. Also fill in their own AugmentedModule
     // with a map from their dependencies to the resolution reason that was applied to each.
     // The newly created graph will also contain ModuleAugments for non-loaded modules.
-    for (Entry<ModuleKey, Module> e : unprunedDepGraph.entrySet()) {
+    for (Entry<ModuleKey, InterimModule> e : unprunedDepGraph.entrySet()) {
       ModuleKey parentKey = e.getKey();
-      Module parentModule = e.getValue();
+      InterimModule parentModule = e.getValue();
 
       AugmentedModule.Builder parentBuilder =
           depGraphAugmentBuilder
@@ -95,10 +95,10 @@ public static ImmutableMap<ModuleKey, AugmentedModule> computeAugmentedGraph(
               .setLoaded(true);
 
       for (String childDep : parentModule.getDeps().keySet()) {
-        ModuleKey originalKey = parentModule.getOriginalDeps().get(childDep);
-        Module originalModule = unprunedDepGraph.get(originalKey);
-        ModuleKey key = parentModule.getDeps().get(childDep);
-        Module module = unprunedDepGraph.get(key);
+        ModuleKey originalKey = parentModule.getOriginalDeps().get(childDep).toModuleKey();
+        InterimModule originalModule = unprunedDepGraph.get(originalKey);
+        ModuleKey key = parentModule.getDeps().get(childDep).toModuleKey();
+        InterimModule module = unprunedDepGraph.get(key);
 
         AugmentedModule.Builder originalChildBuilder =
             depGraphAugmentBuilder.computeIfAbsent(originalKey, AugmentedModule::builder);

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

@@ -21,8 +21,10 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.analysis.BlazeVersionInfo;
 import com.google.devtools.build.lib.bazel.BazelVersion;
+import com.google.devtools.build.lib.bazel.bzlmod.InterimModule.DepSpec;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.BazelCompatibilityMode;
@@ -78,18 +80,18 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     if (root == null) {
       return null;
     }
-    ImmutableMap<ModuleKey, Module> initialDepGraph = Discovery.run(env, root);
+    ImmutableMap<ModuleKey, InterimModule> initialDepGraph = Discovery.run(env, root);
     if (initialDepGraph == null) {
       return null;
     }
 
-    BazelModuleResolutionValue selectionResultValue;
+    Selection.Result selectionResult;
     try {
-      selectionResultValue = Selection.run(initialDepGraph, root.getOverrides());
+      selectionResult = Selection.run(initialDepGraph, root.getOverrides());
     } catch (ExternalDepsException e) {
       throw new BazelModuleResolutionFunctionException(e, Transience.PERSISTENT);
     }
-    ImmutableMap<ModuleKey, Module> resolvedDepGraph = selectionResultValue.getResolvedDepGraph();
+    ImmutableMap<ModuleKey, InterimModule> resolvedDepGraph = selectionResult.getResolvedDepGraph();
 
     verifyRootModuleDirectDepsAreAccurate(
         initialDepGraph.get(ModuleKey.ROOT),
@@ -109,40 +111,15 @@ public SkyValue compute(SkyKey skyKey, Environment env)
             Objects.requireNonNull(ALLOWED_YANKED_VERSIONS.get(env))),
         env.getListener());
 
-    // Add repo spec to each module and remove registry
-    try {
-      ImmutableMap.Builder<ModuleKey, Module> mapBuilder = ImmutableMap.builder();
-      for (Map.Entry<ModuleKey, Module> entry : resolvedDepGraph.entrySet()) {
-        Module module = entry.getValue();
-        // Only change modules with registry (not overridden)
-        if (module.getRegistry() != null) {
-          RepoSpec moduleRepoSpec =
-              module
-                  .getRegistry()
-                  .getRepoSpec(module.getKey(), module.getCanonicalRepoName(), env.getListener());
-          ModuleOverride override = root.getOverrides().get(entry.getKey().getName());
-          moduleRepoSpec = maybeAppendAdditionalPatches(moduleRepoSpec, override);
-          module = module.toBuilder().setRepoSpec(moduleRepoSpec).setRegistry(null).build();
-        }
-        mapBuilder.put(entry.getKey(), module);
-      }
-      resolvedDepGraph = mapBuilder.buildOrThrow();
-    } catch (IOException e) {
-      throw new BazelModuleResolutionFunctionException(
-          ExternalDepsException.withMessage(
-              Code.ERROR_ACCESSING_REGISTRY,
-              "Unable to get module repo spec from registry: %s",
-              e.getMessage()),
-          Transience.PERSISTENT);
-    }
+    ImmutableMap<ModuleKey, Module> finalDepGraph =
+        computeFinalDepGraph(resolvedDepGraph, root.getOverrides(), env.getListener());
 
-    return BazelModuleResolutionValue.create(
-        resolvedDepGraph, selectionResultValue.getUnprunedDepGraph());
+    return BazelModuleResolutionValue.create(finalDepGraph, selectionResult.getUnprunedDepGraph());
   }
 
   private static void verifyRootModuleDirectDepsAreAccurate(
-      Module discoveredRootModule,
-      Module resolvedRootModule,
+      InterimModule discoveredRootModule,
+      InterimModule resolvedRootModule,
       CheckDirectDepsMode mode,
       EventHandler eventHandler)
       throws BazelModuleResolutionFunctionException {
@@ -151,14 +128,14 @@ private static void verifyRootModuleDirectDepsAreAccurate(
     }
 
     boolean failure = false;
-    for (Map.Entry<String, ModuleKey> dep : discoveredRootModule.getDeps().entrySet()) {
-      ModuleKey resolved = resolvedRootModule.getDeps().get(dep.getKey());
-      if (!dep.getValue().equals(resolved)) {
+    for (Map.Entry<String, DepSpec> dep : discoveredRootModule.getDeps().entrySet()) {
+      ModuleKey resolved = resolvedRootModule.getDeps().get(dep.getKey()).toModuleKey();
+      if (!dep.getValue().toModuleKey().equals(resolved)) {
         String message =
             String.format(
                 "For repository '%s', the root module requires module version %s, but got %s in the"
                     + " resolved dependency graph.",
-                dep.getKey(), dep.getValue(), resolved);
+                dep.getKey(), dep.getValue().toModuleKey(), resolved);
         if (mode == CheckDirectDepsMode.WARNING) {
           eventHandler.handle(Event.warn(message));
         } else {
@@ -177,7 +154,9 @@ private static void verifyRootModuleDirectDepsAreAccurate(
   }
 
   public static void checkBazelCompatibility(
-      ImmutableCollection<Module> modules, BazelCompatibilityMode mode, EventHandler eventHandler)
+      ImmutableCollection<InterimModule> modules,
+      BazelCompatibilityMode mode,
+      EventHandler eventHandler)
       throws BazelModuleResolutionFunctionException {
     if (mode == BazelCompatibilityMode.OFF) {
       return;
@@ -189,7 +168,7 @@ public static void checkBazelCompatibility(
     }
 
     BazelVersion curVersion = BazelVersion.parse(currentBazelVersion);
-    for (Module module : modules) {
+    for (InterimModule module : modules) {
       for (String compatVersion : module.getBazelCompatibility()) {
         if (!curVersion.satisfiesCompatibility(compatVersion)) {
           String message =
@@ -306,14 +285,14 @@ private boolean parseModuleKeysFromString(
     return false;
   }
 
-  private void verifyYankedVersions(
-      ImmutableMap<ModuleKey, Module> depGraph,
+  private static void verifyYankedVersions(
+      ImmutableMap<ModuleKey, InterimModule> depGraph,
       Optional<ImmutableSet<ModuleKey>> allowedYankedVersions,
       ExtendedEventHandler eventHandler)
       throws BazelModuleResolutionFunctionException, InterruptedException {
     // Check whether all resolved modules are either not yanked or allowed. Modules with a
     // NonRegistryOverride are ignored as their metadata is not available whatsoever.
-    for (Module m : depGraph.values()) {
+    for (InterimModule m : depGraph.values()) {
       if (m.getKey().equals(ModuleKey.ROOT) || m.getRegistry() == null) {
         continue;
       }
@@ -349,7 +328,7 @@ private void verifyYankedVersions(
     }
   }
 
-  private RepoSpec maybeAppendAdditionalPatches(RepoSpec repoSpec, ModuleOverride override) {
+  private static RepoSpec maybeAppendAdditionalPatches(RepoSpec repoSpec, ModuleOverride override) {
     if (!(override instanceof SingleVersionOverride)) {
       return repoSpec;
     }
@@ -369,6 +348,66 @@ private RepoSpec maybeAppendAdditionalPatches(RepoSpec repoSpec, ModuleOverride
         .build();
   }
 
+  @Nullable
+  private static RepoSpec computeRepoSpec(
+      InterimModule interimModule, ModuleOverride override, ExtendedEventHandler eventHandler)
+      throws BazelModuleResolutionFunctionException, InterruptedException {
+    if (interimModule.getRegistry() == null) {
+      // This module has a non-registry override. We don't need to store the repo spec in this case.
+      return null;
+    }
+    try {
+      RepoSpec moduleRepoSpec =
+          interimModule
+              .getRegistry()
+              .getRepoSpec(
+                  interimModule.getKey(), interimModule.getCanonicalRepoName(), eventHandler);
+      return maybeAppendAdditionalPatches(moduleRepoSpec, override);
+    } catch (IOException e) {
+      throw new BazelModuleResolutionFunctionException(
+          ExternalDepsException.withMessage(
+              Code.ERROR_ACCESSING_REGISTRY,
+              "Unable to get module repo spec from registry: %s",
+              e.getMessage()),
+          Transience.PERSISTENT);
+    }
+  }
+
+  /**
+   * Builds a {@link Module} from an {@link InterimModule}, discarding unnecessary fields and adding
+   * extra necessary ones (such as the repo spec).
+   */
+  static Module moduleFromInterimModule(
+      InterimModule interim, ModuleOverride override, ExtendedEventHandler eventHandler)
+      throws BazelModuleResolutionFunctionException, InterruptedException {
+    return Module.builder()
+        .setName(interim.getName())
+        .setVersion(interim.getVersion())
+        .setKey(interim.getKey())
+        .setRepoName(interim.getRepoName())
+        .setExecutionPlatformsToRegister(interim.getExecutionPlatformsToRegister())
+        .setToolchainsToRegister(interim.getToolchainsToRegister())
+        .setDeps(ImmutableMap.copyOf(Maps.transformValues(interim.getDeps(), DepSpec::toModuleKey)))
+        .setRepoSpec(computeRepoSpec(interim, override, eventHandler))
+        .setExtensionUsages(interim.getExtensionUsages())
+        .build();
+  }
+
+  private static ImmutableMap<ModuleKey, Module> computeFinalDepGraph(
+      ImmutableMap<ModuleKey, InterimModule> resolvedDepGraph,
+      ImmutableMap<String, ModuleOverride> overrides,
+      ExtendedEventHandler eventHandler)
+      throws BazelModuleResolutionFunctionException, InterruptedException {
+    ImmutableMap.Builder<ModuleKey, Module> finalDepGraph = ImmutableMap.builder();
+    for (Map.Entry<ModuleKey, InterimModule> entry : resolvedDepGraph.entrySet()) {
+      finalDepGraph.put(
+          entry.getKey(),
+          moduleFromInterimModule(
+              entry.getValue(), overrides.get(entry.getKey().getName()), eventHandler));
+    }
+    return finalDepGraph.buildOrThrow();
+  }
+
   static class BazelModuleResolutionFunctionException extends SkyFunctionException {
     BazelModuleResolutionFunctionException(ExternalDepsException e, Transience transience) {
       super(e, transience);

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

@@ -44,11 +44,11 @@ abstract class BazelModuleResolutionValue implements SkyValue {
    * overridden by {@code single_version_override} or {@link NonRegistryOverride}, only by {@code
    * multiple_version_override}.
    */
-  abstract ImmutableMap<ModuleKey, Module> getUnprunedDepGraph();
+  abstract ImmutableMap<ModuleKey, InterimModule> getUnprunedDepGraph();
 
   static BazelModuleResolutionValue create(
       ImmutableMap<ModuleKey, Module> resolvedDepGraph,
-      ImmutableMap<ModuleKey, Module> unprunedDepGraph) {
+      ImmutableMap<ModuleKey, InterimModule> unprunedDepGraph) {
     return new AutoValue_BazelModuleResolutionValue(resolvedDepGraph, unprunedDepGraph);
   }
 }

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

@@ -16,9 +16,9 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.bazel.bzlmod.InterimModule.DepSpec;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
-import com.google.devtools.build.skyframe.SkyFunctionException;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyframeLookupResult;
 import java.util.ArrayDeque;
@@ -42,23 +42,24 @@ private Discovery() {}
    * dependency is missing and this function needs a restart).
    */
   @Nullable
-  public static ImmutableMap<ModuleKey, Module> run(Environment env, RootModuleFileValue root)
-      throws SkyFunctionException, InterruptedException {
+  public static ImmutableMap<ModuleKey, InterimModule> run(
+      Environment env, RootModuleFileValue root) throws InterruptedException {
     String rootModuleName = root.getModule().getName();
     ImmutableMap<String, ModuleOverride> overrides = root.getOverrides();
-    Map<ModuleKey, Module> depGraph = new HashMap<>();
-    depGraph.put(ModuleKey.ROOT, rewriteDepKeys(root.getModule(), overrides, rootModuleName));
+    Map<ModuleKey, InterimModule> depGraph = new HashMap<>();
+    depGraph.put(ModuleKey.ROOT, rewriteDepSpecs(root.getModule(), overrides, rootModuleName));
     Queue<ModuleKey> unexpanded = new ArrayDeque<>();
     unexpanded.add(ModuleKey.ROOT);
     while (!unexpanded.isEmpty()) {
       Set<SkyKey> unexpandedSkyKeys = new HashSet<>();
       while (!unexpanded.isEmpty()) {
-        Module module = depGraph.get(unexpanded.remove());
-        for (ModuleKey depKey : module.getDeps().values()) {
-          if (depGraph.containsKey(depKey)) {
+        InterimModule module = depGraph.get(unexpanded.remove());
+        for (DepSpec depSpec : module.getDeps().values()) {
+          if (depGraph.containsKey(depSpec.toModuleKey())) {
             continue;
           }
-          unexpandedSkyKeys.add(ModuleFileValue.key(depKey, overrides.get(depKey.getName())));
+          unexpandedSkyKeys.add(
+              ModuleFileValue.key(depSpec.toModuleKey(), overrides.get(depSpec.getName())));
         }
       }
       SkyframeLookupResult result = env.getValuesAndExceptions(unexpandedSkyKeys);
@@ -70,7 +71,7 @@ public static ImmutableMap<ModuleKey, Module> run(Environment env, RootModuleFil
           depGraph.put(depKey, null);
         } else {
           depGraph.put(
-              depKey, rewriteDepKeys(moduleFileValue.getModule(), overrides, rootModuleName));
+              depKey, rewriteDepSpecs(moduleFileValue.getModule(), overrides, rootModuleName));
           unexpanded.add(depKey);
         }
       }
@@ -81,16 +82,16 @@ public static ImmutableMap<ModuleKey, Module> run(Environment env, RootModuleFil
     return ImmutableMap.copyOf(depGraph);
   }
 
-  private static Module rewriteDepKeys(
-      Module module, ImmutableMap<String, ModuleOverride> overrides, String rootModuleName) {
-    return module.withDepKeysTransformed(
-        depKey -> {
-          if (rootModuleName.equals(depKey.getName())) {
-            return ModuleKey.ROOT;
+  private static InterimModule rewriteDepSpecs(
+      InterimModule module, ImmutableMap<String, ModuleOverride> overrides, String rootModuleName) {
+    return module.withDepSpecsTransformed(
+        depSpec -> {
+          if (rootModuleName.equals(depSpec.getName())) {
+            return DepSpec.fromModuleKey(ModuleKey.ROOT);
           }
 
-          Version newVersion = depKey.getVersion();
-          @Nullable ModuleOverride override = overrides.get(depKey.getName());
+          Version newVersion = depSpec.getVersion();
+          @Nullable ModuleOverride override = overrides.get(depSpec.getName());
           if (override instanceof NonRegistryOverride) {
             newVersion = Version.EMPTY;
           } else if (override instanceof SingleVersionOverride) {
@@ -100,7 +101,7 @@ private static Module rewriteDepKeys(
             }
           }
 
-          return ModuleKey.create(depKey.getName(), newVersion);
+          return DepSpec.create(depSpec.getName(), newVersion, depSpec.getMaxCompatibilityLevel());
         });
   }
 }