Store usages hash in module extension lockfile entries

Work towards bazelbuild#20369

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

@@ -211,16 +211,6 @@ static BzlmodFlagsAndEnvVars getFlagsAndEnvVars(Environment env) throws Interrup
         compatabilityMode);
   }
 
-  /**
-   * For each extension usage, we resolve (i.e. canonicalize) its bzl file label. Then we can group
-   * all usages by the label + name (the ModuleExtensionId).
-   */
-  public static ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
-      getExtensionUsagesById(ImmutableMap<ModuleKey, Module> depGraph)
-          throws ExternalDepsException {
-    return getExtensionUsagesById(depGraph, computeCanonicalRepoNameLookup(depGraph).inverse());
-  }
-
   private static ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
       getExtensionUsagesById(
           ImmutableMap<ModuleKey, Module> depGraph,

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

@@ -19,7 +19,6 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
-import com.google.common.collect.ImmutableTable;
 import com.google.common.eventbus.Subscribe;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions;
@@ -33,6 +32,7 @@
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -107,16 +107,17 @@ public void afterCommand() throws AbruptExitException {
         updatedExtensionMap = new HashMap<>();
 
     // Keep old extensions if they are still valid.
-    ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> oldExtensionUsages =
-        BazelDepGraphFunction.getExtensionUsagesById(oldLockfile.getModuleDepGraph());
     for (var entry : oldLockfile.getModuleExtensions().entrySet()) {
       var moduleExtensionId = entry.getKey();
-      var factorToLockedExtension = entry.getValue();
-      ModuleExtensionEvalFactors firstEntryFactors =
-          factorToLockedExtension.keySet().iterator().next();
-      if (shouldKeepExtension(moduleExtensionId, firstEntryFactors, oldExtensionUsages)) {
-        updatedExtensionMap.put(moduleExtensionId, factorToLockedExtension);
+      ImmutableMap.Builder<ModuleExtensionEvalFactors, LockFileModuleExtension> factorsToKeep =
+          new ImmutableMap.Builder<>();
+      for (var factorAndExtension : entry.getValue().entrySet()) {
+        if (shouldKeepExtension(
+            moduleExtensionId, factorAndExtension.getKey(), factorAndExtension.getValue())) {
+          factorsToKeep.put(factorAndExtension.getKey(), factorAndExtension.getValue());
+        }
       }
+      updatedExtensionMap.put(moduleExtensionId, factorsToKeep.buildOrThrow());
     }
 
     // Add the new resolved extensions
@@ -164,7 +165,7 @@ public void afterCommand() throws AbruptExitException {
   private boolean shouldKeepExtension(
       ModuleExtensionId extensionId,
       ModuleExtensionEvalFactors lockedExtensionKey,
-      ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> oldExtensionUsages) {
+      LockFileModuleExtension lockedExtension) {
 
     // If there is a new event for this extension, compare it with the existing ones
     ModuleExtensionResolutionEvent extEvent = extensionResolutionEventsMap.get(extensionId);
@@ -185,9 +186,11 @@ private boolean shouldKeepExtension(
     // that irrelevant changes (e.g. locations or imports) don't cause the extension to be removed.
     // Note: Extension results can still be stale for other reasons, e.g. because their transitive
     // bzl hash changed, but such changes will be detected in SingleExtensionEvalFunction.
-    return ModuleExtensionUsage.trimForEvaluation(
-            moduleResolutionEvent.getExtensionUsagesById().row(extensionId))
-        .equals(ModuleExtensionUsage.trimForEvaluation(oldExtensionUsages.row(extensionId)));
+    return Arrays.equals(
+        ModuleExtensionUsage.hashForEvaluation(
+            GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
+            moduleResolutionEvent.getExtensionUsagesById().row(extensionId)),
+        lockedExtension.getUsagesDigest());
   }
 
   /**

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

@@ -476,8 +476,18 @@ public RepoRecordedInput.Dirents read(JsonReader jsonReader) throws IOException
           };
 
   public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
-    return new GsonBuilder()
+    return newGsonBuilder()
         .setPrettyPrinting()
+        .registerTypeAdapterFactory(new LocationTypeAdapterFactory(moduleFilePath, workspaceRoot))
+        .create();
+  }
+
+  public static Gson createModuleExtensionUsagesHashGson() {
+    return newGsonBuilder().create();
+  }
+
+  private static GsonBuilder newGsonBuilder() {
+    return new GsonBuilder()
         .disableHtmlEscaping()
         .enableComplexMapKeySerialization()
         .registerTypeAdapterFactory(GenerateTypeAdapter.FACTORY)
@@ -488,7 +498,6 @@ public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
         .registerTypeAdapterFactory(IMMUTABLE_SET)
         .registerTypeAdapterFactory(OPTIONAL)
         .registerTypeAdapterFactory(IMMUTABLE_TABLE)
-        .registerTypeAdapterFactory(new LocationTypeAdapterFactory(moduleFilePath, workspaceRoot))
         .registerTypeAdapter(Label.class, LABEL_TYPE_ADAPTER)
         .registerTypeAdapter(RepositoryName.class, REPOSITORY_NAME_TYPE_ADAPTER)
         .registerTypeAdapter(Version.class, VERSION_TYPE_ADAPTER)
@@ -501,8 +510,7 @@ public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
         .registerTypeAdapter(byte[].class, BYTE_ARRAY_TYPE_ADAPTER)
         .registerTypeAdapter(RepoRecordedInput.File.class, REPO_RECORDED_INPUT_FILE_TYPE_ADAPTER)
         .registerTypeAdapter(
-            RepoRecordedInput.Dirents.class, REPO_RECORDED_INPUT_DIRENTS_TYPE_ADAPTER)
-        .create();
+            RepoRecordedInput.Dirents.class, REPO_RECORDED_INPUT_DIRENTS_TYPE_ADAPTER);
   }
 
   private GsonTypeAdapterUtil() {}

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

@@ -41,6 +41,9 @@ public static Builder builder() {
   @SuppressWarnings("mutable")
   public abstract byte[] getBzlTransitiveDigest();
 
+  @SuppressWarnings("mutable")
+  public abstract byte[] getUsagesDigest();
+
   public abstract ImmutableMap<RepoRecordedInput.File, String> getRecordedFileInputs();
 
   public abstract ImmutableMap<RepoRecordedInput.Dirents, String> getRecordedDirentsInputs();
@@ -67,6 +70,8 @@ public abstract static class Builder {
 
     public abstract Builder setBzlTransitiveDigest(byte[] digest);
 
+    public abstract Builder setUsagesDigest(byte[] digest);
+
     public abstract Builder setRecordedFileInputs(
         ImmutableMap<RepoRecordedInput.File, String> value);
 

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

@@ -22,9 +22,10 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
+import com.google.common.hash.Hashing;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import com.google.gson.Gson;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
-import java.util.Map;
 import java.util.Optional;
 import net.starlark.java.syntax.Location;
 
@@ -94,16 +95,15 @@ public static Builder builder() {
   }
 
   /**
-   * Turns the given collection of usages for a particular extension into an object that can be
-   * compared for equality with another object obtained in this way and compares equal only if the
-   * two original collections of usages are equivalent for the purpose of evaluating the extension.
+   * Turns the given collection of usages for a particular extension into a hash that can be
+   * compared for equality with another hash obtained in this way and compares equal only if the two
+   * original collections of usages are equivalent for the purpose of evaluating the extension.
    */
-  static ImmutableList<Map.Entry<ModuleKey, ModuleExtensionUsage>> trimForEvaluation(
-      ImmutableMap<ModuleKey, ModuleExtensionUsage> usages) {
-    // ImmutableMap#equals doesn't compare the order of entries, but it matters for the evaluation
-    // of the extension.
-    return ImmutableList.copyOf(
-        Maps.transformValues(usages, ModuleExtensionUsage::trimForEvaluation).entrySet());
+  static byte[] hashForEvaluation(Gson gson, ImmutableMap<ModuleKey, ModuleExtensionUsage> usages) {
+    ImmutableMap<ModuleKey, ModuleExtensionUsage> trimmedUsages =
+        ImmutableMap.copyOf(Maps.transformValues(usages, ModuleExtensionUsage::trimForEvaluation));
+    String usagesJson = gson.toJson(trimmedUsages);
+    return Hashing.sha256().hashUnencodedChars(usagesJson).asBytes();
   }
 
   /**

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

@@ -163,7 +163,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
         try {
           SingleExtensionEvalValue singleExtensionEvalValue =
               tryGettingValueFromLockFile(
-                  env, extensionId, extension, usagesValue, lockfile, lockedExtension);
+                  env, extensionId, extension, usagesValue, lockedExtension);
           if (singleExtensionEvalValue != null) {
             return singleExtensionEvalValue;
           }
@@ -217,6 +217,10 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                   extension.getEvalFactors(),
                   LockFileModuleExtension.builder()
                       .setBzlTransitiveDigest(extension.getBzlTransitiveDigest())
+                      .setUsagesDigest(
+                          ModuleExtensionUsage.hashForEvaluation(
+                              GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
+                              usagesValue.getExtensionUsages()))
                       .setRecordedFileInputs(moduleExtensionResult.getRecordedFileInputs())
                       .setRecordedDirentsInputs(moduleExtensionResult.getRecordedDirentsInputs())
                       .setEnvVariables(extension.getEnvVars())
@@ -243,24 +247,11 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
       ModuleExtensionId extensionId,
       RunnableExtension extension,
       SingleExtensionUsagesValue usagesValue,
-      BazelLockFileValue lockfile,
       LockFileModuleExtension lockedExtension)
       throws SingleExtensionEvalFunctionException,
           InterruptedException,
           NeedsSkyframeRestartException {
     LockfileMode lockfileMode = BazelLockFileFunction.LOCKFILE_MODE.get(env);
-
-    ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages;
-    try {
-      // TODO(salmasamy) might be nicer to precompute this table when we construct
-      // BazelLockFileValue, without adding it to the json file
-      ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> extensionUsagesById =
-          BazelDepGraphFunction.getExtensionUsagesById(lockfile.getModuleDepGraph());
-      lockedExtensionUsages = extensionUsagesById.row(extensionId);
-    } catch (ExternalDepsException e) {
-      throw new SingleExtensionEvalFunctionException(e, Transience.PERSISTENT);
-    }
-
     DiffRecorder diffRecorder =
         new DiffRecorder(/* recordMessages= */ lockfileMode.equals(LockfileMode.ERROR));
     try {
@@ -280,8 +271,11 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
       }
       // Check extension data in lockfile is still valid, disregarding usage information that is not
       // relevant for the evaluation of the extension.
-      if (!ModuleExtensionUsage.trimForEvaluation(usagesValue.getExtensionUsages())
-          .equals(ModuleExtensionUsage.trimForEvaluation(lockedExtensionUsages))) {
+      if (!Arrays.equals(
+          ModuleExtensionUsage.hashForEvaluation(
+              GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
+              usagesValue.getExtensionUsages()),
+          lockedExtension.getUsagesDigest())) {
         diffRecorder.record("The usages of the extension '" + extensionId + "' have changed");
       }
       if (didRepoMappingsChange(env, lockedExtension.getRecordedRepoMappingEntries())) {