[6.4.0] Remove stale extension entries from lockfile if module order changes

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

@@ -66,6 +66,8 @@ public static BazelDepGraphValue create(
    * usage occurs. For each extension identifier ID, extensionUsagesTable[ID][moduleKey] is the
    * ModuleExtensionUsage of ID in the module keyed by moduleKey.
    */
+  // Note: Equality of BazelDepGraphValue does not check for equality of the order of the rows of
+  // this table, but it is tracked implicitly via the order of the abridged modules.
   public abstract ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
       getExtensionUsagesTable();
 

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

@@ -19,7 +19,6 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.ImmutableTable;
-import com.google.common.collect.Maps;
 import com.google.common.eventbus.Subscribe;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions;
@@ -206,14 +205,9 @@ 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.
-    var currentTrimmedUsages =
-        Maps.transformValues(
-            moduleResolutionEvent.getExtensionUsagesById().row(extensionId),
-            ModuleExtensionUsage::trimForEvaluation);
-    var lockedTrimmedUsages =
-        Maps.transformValues(
-            oldExtensionUsages.row(extensionId), ModuleExtensionUsage::trimForEvaluation);
-    return currentTrimmedUsages.equals(lockedTrimmedUsages);
+    return ModuleExtensionUsage.trimForEvaluation(
+            moduleResolutionEvent.getExtensionUsagesById().row(extensionId))
+        .equals(ModuleExtensionUsage.trimForEvaluation(oldExtensionUsages.row(extensionId)));
   }
 
   /**

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

@@ -126,8 +126,8 @@ public ImmutableList<String> getModuleExtensionDiff(
       byte[] transitiveDigest,
       boolean filesChanged,
       ImmutableMap<String, String> envVariables,
-      ImmutableMap<ModuleKey, ModuleExtensionUsage> extensionUsages,
-      ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages) {
+      ImmutableList<Map.Entry<ModuleKey, ModuleExtensionUsage>> extensionUsages,
+      ImmutableList<Map.Entry<ModuleKey, ModuleExtensionUsage>> lockedExtensionUsages) {
 
     ImmutableList.Builder<String> extDiff = new ImmutableList.Builder<>();
     if (!Arrays.equals(transitiveDigest, lockedExtension.getBzlTransitiveDigest())) {

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

@@ -19,9 +19,12 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableBiMap;
 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.errorprone.annotations.CanIgnoreReturnValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
+import java.util.Map;
 import java.util.Optional;
 import net.starlark.java.syntax.Location;
 
@@ -90,11 +93,24 @@ public static Builder builder() {
     return new AutoValue_ModuleExtensionUsage.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.
+   */
+  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());
+  }
+
   /**
    * Returns a new usage with all information removed that does not influence the evaluation of the
    * extension.
    */
-  ModuleExtensionUsage trimForEvaluation() {
+  private ModuleExtensionUsage trimForEvaluation() {
     // We start with the full usage and selectively remove information that does not influence the
     // evaluation of the extension. Compared to explicitly copying over the parts that do, this
     // preserves correctness in case new fields are added without updating this code.

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

@@ -18,7 +18,6 @@
 import static com.google.common.base.StandardSystemProperty.OS_ARCH;
 import static com.google.common.collect.ImmutableBiMap.toImmutableBiMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static com.google.common.collect.Maps.transformValues;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
@@ -273,13 +272,8 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
 
     // Check extension data in lockfile is still valid, disregarding usage information that is not
     // relevant for the evaluation of the extension.
-    var trimmedLockedUsages =
-        ImmutableMap.copyOf(
-            transformValues(lockedExtensionUsages, ModuleExtensionUsage::trimForEvaluation));
-    var trimmedUsages =
-        ImmutableMap.copyOf(
-            transformValues(
-                usagesValue.getExtensionUsages(), ModuleExtensionUsage::trimForEvaluation));
+    var trimmedLockedUsages = ModuleExtensionUsage.trimForEvaluation(lockedExtensionUsages);
+    var trimmedUsages = ModuleExtensionUsage.trimForEvaluation(usagesValue.getExtensionUsages());
     if (!filesChanged
         && Arrays.equals(bzlTransitiveDigest, lockedExtension.getBzlTransitiveDigest())
         && trimmedUsages.equals(trimmedLockedUsages)

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

@@ -31,6 +31,8 @@
 @AutoValue
 public abstract class SingleExtensionUsagesValue implements SkyValue {
   /** All usages of this extension, by the key of the module where the usage occurs. */
+  // Note: Equality of SingleExtensionUsagesValue does not check for equality of the order of the
+  // entries of this map, but it is tracked implicitly via the order of the abridged modules.
   public abstract ImmutableMap<ModuleKey, ModuleExtensionUsage> getExtensionUsages();
 
   /**

src/test/py/bazel/bzlmod/bazel_lockfile_test.py

@@ -1269,6 +1269,130 @@ def testLockfileWithNoUserSpecificPath(self):
     for key in remote_patches.keys():
       self.assertIn('%workspace%', key)
 
+  def testExtensionEvaluationRerunsIfDepGraphOrderChanges(self):
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'module(name = "root", version = "1.0")',
+            'bazel_dep(name = "aaa", version = "1.0")',
+            'bazel_dep(name = "bbb", version = "1.0")',
+            'ext = use_extension("extension.bzl", "ext")',
+            'ext.tag(value = "root")',
+            'use_repo(ext, "dep")',
+        ],
+    )
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile(
+        'extension.bzl',
+        [
+            'def _repo_rule_impl(ctx):',
+            '    ctx.file("WORKSPACE")',
+            '    ctx.file("BUILD", "exports_files([\\"data.txt\\"])")',
+            '    ctx.file("data.txt", ctx.attr.value)',
+            '    print(ctx.attr.value)',
+            '',
+            'repo_rule = repository_rule(',
+            '    implementation=_repo_rule_impl,',
+            '    attrs = {"value": attr.string()},)',
+            '',
+            'def _ext_impl(ctx):',
+            '    print("Ext is being evaluated")',
+            (
+                '    values = ",".join([tag.value for mod in ctx.modules for'
+                ' tag in mod.tags.tag])'
+            ),
+            '    repo_rule(name="dep", value="Ext saw values: " + values)',
+            '',
+            'ext = module_extension(',
+            '    implementation=_ext_impl,',
+            (
+                '    tag_classes={"tag": tag_class(attrs={"value":'
+                ' attr.string()})}'
+            ),
+            ')',
+        ],
+    )
+    self.main_registry.createCcModule(
+        'aaa',
+        '1.0',
+        extra_module_file_contents=[
+            'bazel_dep(name = "root", version = "1.0")',
+            'ext = use_extension("@root//:extension.bzl", "ext")',
+            'ext.tag(value = "aaa")',
+        ],
+    )
+    self.main_registry.createCcModule(
+        'bbb',
+        '1.0',
+        extra_module_file_contents=[
+            'bazel_dep(name = "root", version = "1.0")',
+            'ext = use_extension("@root//:extension.bzl", "ext")',
+            'ext.tag(value = "bbb")',
+        ],
+    )
+
+    _, _, stderr = self.RunBazel(['build', '@dep//:all'])
+    stderr = '\n'.join(stderr)
+
+    self.assertIn('Ext is being evaluated', stderr)
+    self.assertIn('Ext saw values: root,aaa,bbb', stderr)
+    ext_key = '//:extension.bzl%ext'
+    with open('MODULE.bazel.lock', 'r') as json_file:
+      lockfile = json.load(json_file)
+    self.assertIn(ext_key, lockfile['moduleExtensions'])
+    self.assertIn(
+        'Ext saw values: root,aaa,bbb',
+        lockfile['moduleExtensions'][ext_key]['general']['generatedRepoSpecs'][
+            'dep'
+        ]['attributes']['value'],
+    )
+
+    # Shut down bazel to empty the cache, modify the MODULE.bazel
+    # file in a way that only changes the order of the bazel_deps
+    # on aaa and bbb, which results in their usages being ordered
+    # differently in module_ctx.modules, which is visible to the
+    # extension. Rerun a build that does not trigger evaluation
+    # of the extension.
+    self.RunBazel(['shutdown'])
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'module(name = "root", version = "1.0")',
+            'bazel_dep(name = "bbb", version = "1.0")',
+            'bazel_dep(name = "aaa", version = "1.0")',
+            'ext = use_extension("extension.bzl", "ext")',
+            'ext.tag(value = "root")',
+            'use_repo(ext, "dep")',
+        ],
+    )
+    _, _, stderr = self.RunBazel(['build', '//:all'])
+    stderr = '\n'.join(stderr)
+
+    self.assertNotIn('Ext is being evaluated', stderr)
+    with open('MODULE.bazel.lock', 'r') as json_file:
+      lockfile = json.load(json_file)
+    # The order of usages of ext changed, but the extension
+    # is not re-evaluated, so its previous, now stale
+    # resolution result must have been removed.
+    self.assertNotIn(ext_key, lockfile['moduleExtensions'])
+
+    # Trigger evaluation of the extension.
+    _, _, stderr = self.RunBazel(['build', '@dep//:all'])
+    stderr = '\n'.join(stderr)
+
+    self.assertIn('Ext is being evaluated', stderr)
+    self.assertIn('Ext saw values: root,bbb,aaa', stderr)
+    ext_key = '//:extension.bzl%ext'
+    with open('MODULE.bazel.lock', 'r') as json_file:
+      lockfile = json.load(json_file)
+    self.assertIn(ext_key, lockfile['moduleExtensions'])
+    self.assertIn(
+        'Ext saw values: root,bbb,aaa',
+        lockfile['moduleExtensions'][ext_key]['general']['generatedRepoSpecs'][
+            'dep'
+        ]['attributes']['value'],
+    )
+
 
 if __name__ == '__main__':
   unittest.main()