Allow module extension tags to be compared

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

@@ -19,6 +19,7 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
@@ -352,15 +353,18 @@ private ModuleExtensionContext createContext(
             .getRelative(LabelConstants.MODULE_EXTENSION_WORKING_DIRECTORY_LOCATION)
             .getRelative(usagesValue.getExtensionUniqueName());
     ArrayList<StarlarkBazelModule> modules = new ArrayList<>();
-    for (AbridgedModule abridgedModule : usagesValue.getAbridgedModules()) {
-      ModuleKey moduleKey = abridgedModule.getKey();
+    ImmutableList<AbridgedModule> abridgedModules = usagesValue.getAbridgedModules();
+    for (int i = 0; i < abridgedModules.size(); i++) {
+      var abridgedModule = abridgedModules.get(i);
+      var moduleKey = abridgedModule.getKey();
       modules.add(
           StarlarkBazelModule.create(
               abridgedModule,
               extension,
               usagesValue.getRepoMappings().get(moduleKey),
               usagesValue.getExtensionUsages().get(moduleKey),
-              repoMappingRecorder));
+              repoMappingRecorder,
+              i));
     }
     ModuleExtensionUsage rootUsage = usagesValue.getExtensionUsages().get(ModuleKey.ROOT);
     boolean rootModuleHasNonDevDependency =

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

@@ -26,8 +26,9 @@
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps.Code;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.SequencedMap;
 import javax.annotation.Nullable;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
@@ -52,15 +53,20 @@ public class StarlarkBazelModule implements StarlarkValue {
       name = "bazel_module_tags",
       category = DocCategory.BUILTIN,
       doc =
-          "Contains the tags in a module for the module extension currently being processed. This"
-              + " object has a field for each tag class of the extension, and the value of the"
-              + " field is a list containing an object for each tag instance. This \"tag instance\""
-              + " object in turn has a field for each attribute of the tag class.\n\n"
-              + "When passed as positional arguments to <code>print()</code> or <code>fail()"
-              + "</code>, tag instance objects turn into a meaningful string representation of the"
-              + " form \"'install' tag at /home/user/workspace/MODULE.bazel:3:4\". This can be used"
-              + " to construct error messages that point to the location of the tag in the module"
-              + " file, e.g. <code>fail(\"Conflict between\", tag1, \"and\", tag2)</code>.")
+          """
+          Contains the tags in a module for the module extension currently being processed. This \
+          object has a field for each tag class of the extension, and the value of the field is a \
+          list containing an object for each tag instance. This "tag instance" object in turn has \
+          a field for each attribute of the tag class.
+
+          Tag instance objects can be compared even across tag classes and sort in the order the \
+          tags are defined in the module file.
+
+          When passed as positional arguments to <code>print()</code> or <code>fail()</code>, tag \
+          instance objects turn into a meaningful string representation of the form "'install' tag \
+          at /home/user/workspace/MODULE.bazel:3:4". This can be used to construct error messages \
+          that point to the location of the tag in the module file, e.g. \
+          <code>fail("Conflict between", tag1, "and", tag2)</code>.""")
   static class Tags implements Structure {
     private final ImmutableMap<String, StarlarkList<TypeCheckedTag>> typeCheckedTags;
 
@@ -109,19 +115,21 @@ public static StarlarkBazelModule create(
       ModuleExtension extension,
       RepositoryMapping repoMapping,
       @Nullable ModuleExtensionUsage usage,
-      Label.RepoMappingRecorder repoMappingRecorder)
+      Label.RepoMappingRecorder repoMappingRecorder,
+      int moduleIndex)
       throws ExternalDepsException {
     LabelConverter labelConverter =
         new LabelConverter(
             PackageIdentifier.create(repoMapping.contextRepo(), PathFragment.EMPTY_FRAGMENT),
             repoMapping,
             repoMappingRecorder);
     ImmutableList<Tag> tags = usage == null ? ImmutableList.of() : usage.getTags();
-    HashMap<String, ArrayList<TypeCheckedTag>> typeCheckedTags = new HashMap<>();
+    SequencedMap<String, ArrayList<TypeCheckedTag>> typeCheckedTags = new LinkedHashMap<>();
     for (String tagClassName : extension.tagClasses().keySet()) {
       typeCheckedTags.put(tagClassName, new ArrayList<>());
     }
-    for (Tag tag : tags) {
+    for (int tagIndex = 0; tagIndex < tags.size(); tagIndex++) {
+      Tag tag = tags.get(tagIndex);
       TagClass tagClass = extension.tagClasses().get(tag.getTagName());
       if (tagClass == null) {
         throw ExternalDepsException.withMessage(
@@ -140,7 +148,12 @@ public static StarlarkBazelModule create(
           .get(tag.getTagName())
           .add(
               TypeCheckedTag.create(
-                  tagClass, tag, labelConverter, module.getKey().toDisplayString()));
+                  tagClass,
+                  tag,
+                  labelConverter,
+                  module.getKey().toDisplayString(),
+                  moduleIndex,
+                  tagIndex));
     }
     return new StarlarkBazelModule(
         module.getName(),

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

@@ -33,10 +33,11 @@
  * in the {@link TagClass}.
  */
 @StarlarkBuiltin(name = "bazel_module_tag", documented = false)
-public class TypeCheckedTag implements Structure {
+public class TypeCheckedTag implements Structure, Comparable<TypeCheckedTag> {
   private final TagClass tagClass;
   private final ImmutableList<Object> attrValues;
   private final boolean devDependency;
+  private final long compareKey;
 
   // The properties below are only used for error reporting.
   private final Location location;
@@ -46,18 +47,28 @@ private TypeCheckedTag(
       TagClass tagClass,
       ImmutableList<Object> attrValues,
       boolean devDependency,
+      int moduleIndex,
+      int tagIndex,
       Location location,
       String tagClassName) {
     this.tagClass = tagClass;
     this.attrValues = attrValues;
     this.devDependency = devDependency;
+    // Sort by module first, then in the order tags were defined within the module.
+    this.compareKey =
+        (Integer.toUnsignedLong(moduleIndex) << 32) + Integer.toUnsignedLong(tagIndex);
     this.location = location;
     this.tagClassName = tagClassName;
   }
 
   /** Creates a {@link TypeCheckedTag}. */
   public static TypeCheckedTag create(
-      TagClass tagClass, Tag tag, LabelConverter labelConverter, String moduleDisplayString)
+      TagClass tagClass,
+      Tag tag,
+      LabelConverter labelConverter,
+      String moduleDisplayString,
+      int moduleIndex,
+      int tagIndex)
       throws ExternalDepsException {
     ImmutableList<Object> attrValues =
         AttributeUtils.typeCheckAttrValues(
@@ -70,7 +81,13 @@ public static TypeCheckedTag create(
             "'%s' tag".formatted(tag.getTagName()),
             "to the %s".formatted(moduleDisplayString));
     return new TypeCheckedTag(
-        tagClass, attrValues, tag.isDevDependency(), tag.getLocation(), tag.getTagName());
+        tagClass,
+        attrValues,
+        tag.isDevDependency(),
+        moduleIndex,
+        tagIndex,
+        tag.getLocation(),
+        tag.getTagName());
   }
 
   /**
@@ -111,4 +128,19 @@ public String getErrorMessageForUnknownField(String field) {
   public void debugPrint(Printer printer, StarlarkThread thread) {
     printer.append(String.format("'%s' tag at %s", tagClassName, location));
   }
+
+  @Override
+  public int compareTo(TypeCheckedTag other) {
+    return Long.compareUnsigned(this.compareKey, other.compareKey);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    return this == other || (other instanceof TypeCheckedTag o && this.compareKey == o.compareKey);
+  }
+
+  @Override
+  public int hashCode() {
+    return Long.hashCode(compareKey);
+  }
 }

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

@@ -2798,6 +2798,59 @@ public void printAndFailOnTag() throws Exception {
         ImmutableSet.of(EventKind.DEBUG));
   }
 
+  @Test
+  public void tagSortOrder() throws Exception {
+    scratch.overwriteFile(
+        "MODULE.bazel",
+        "module(name='root',version='1.0')",
+        "bazel_dep(name='foo',version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('//:defs.bzl', 'ext')",
+        "ext.baz(id = 7)",
+        "ext.foo(id = 2)",
+        "ext.baz(id = 9)",
+        "ext.bar(id = 42)",
+        "ext.foo(id = -1)",
+        "ext.foo(id = 5)",
+        "use_repo(ext, 'ext_data')"
+        );
+    scratch.file(
+        "defs.bzl",
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "def _ext_impl(ctx):",
+        "  tags = []",
+        "  for module in ctx.modules:",
+        "    tags += module.tags.foo",
+        "    tags += module.tags.bar",
+        "    tags += module.tags.baz",
+        "  ids = [tag.id for tag in sorted(tags)]",
+        "  data_repo(name='ext_data',data=str(ids))",
+        "foo = tag_class(attrs = {'id': attr.int()})",
+        "bar = tag_class(attrs = {'id': attr.int()})",
+        "baz = tag_class(attrs = {'id': attr.int()})",
+        "ext = module_extension(implementation=_ext_impl,tag_classes={'foo':foo, 'bar':bar, 'baz':baz})");
+    scratch.overwriteFile("BUILD");
+    scratch.file("data.bzl", "load('@ext_data//:data.bzl', ext_data='data')", "data=ext_data");
+    registry.addModule(
+        createModuleKey("foo", "1.0"),
+        "module(name='foo',version='1.0')",
+        "bazel_dep(name='root',version='1.0')",
+        "ext = use_extension('@root//:defs.bzl','ext')",
+        "ext.bar(id = 3)",
+        "ext.foo(id = 4)",
+        "ext.baz(id = 1)");
+    invalidatePackages(false);
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    EvaluationResult<BzlLoadValue> result =
+        SkyframeExecutorTestUtils.evaluate(skyframeExecutor, skyKey, false, reporter);
+    if (result.hasError()) {
+      throw result.getError().getException();
+    }
+    assertThat(result.get(skyKey).getModule().getGlobal("data"))
+        .isEqualTo("[7, 2, 9, 42, -1, 5, 3, 4, 1]"); // sorted order
+  }
+
   @Test
   public void innate() throws Exception {
     scratch.overwriteFile(

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

@@ -106,7 +106,8 @@ public void basic() throws Exception {
                     fooKey, fooKey.getCanonicalRepoNameWithoutVersion(),
                     barKey, barKey.getCanonicalRepoNameWithoutVersion())),
             usage,
-            repoMappingRecorder);
+            repoMappingRecorder,
+            /* moduleIndex= */ 0);
 
     assertThat(moduleProxy.getName()).isEqualTo("foo");
     assertThat(moduleProxy.getVersion()).isEqualTo("1.0");
@@ -158,7 +159,8 @@ public void unknownTagClass() throws Exception {
                     module.getRepoMappingWithBazelDepsOnly(
                         ImmutableMap.of(fooKey, fooKey.getCanonicalRepoNameWithoutVersion())),
                     usage,
-                    new Label.RepoMappingRecorder()));
+                    new Label.RepoMappingRecorder(),
+                    /* moduleIndex= */ 0));
     assertThat(e).hasMessageThat().contains("does not have a tag class named blep");
   }
 }

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

@@ -64,7 +64,9 @@ public void basic() throws Exception {
             createTagClass(attr("foo", Type.INTEGER).build()),
             buildTag("tag_name").addAttr("foo", StarlarkInt.of(3)).setDevDependency().build(),
             /* labelConverter= */ null,
-            "root module");
+            "root module",
+            /* moduleIndex= */ 0,
+            /* tagIndex= */ 0);
     assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo");
     assertThat(getattr(typeCheckedTag, "foo")).isEqualTo(StarlarkInt.of(3));
     assertThat(typeCheckedTag.isDevDependency()).isTrue();
@@ -83,7 +85,9 @@ public void label() throws Exception {
             new LabelConverter(
                 PackageIdentifier.parse("@myrepo//mypkg"),
                 createRepositoryMapping(createModuleKey("test", "1.0"), "repo", "other_repo")),
-            "root module");
+            "root module",
+            /* moduleIndex= */ 0,
+            /* tagIndex= */ 0);
     assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo");
     assertThat(getattr(typeCheckedTag, "foo"))
         .isEqualTo(
@@ -104,7 +108,9 @@ public void label_withoutDefaultValue() throws Exception {
             new LabelConverter(
                 PackageIdentifier.parse("@myrepo//mypkg"),
                 createRepositoryMapping(createModuleKey("test", "1.0"), "repo", "other_repo")),
-            "root module");
+            "root module",
+            /* moduleIndex= */ 0,
+            /* tagIndex= */ 0);
     assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo");
     assertThat(getattr(typeCheckedTag, "foo")).isEqualTo(Starlark.NONE);
     assertThat(typeCheckedTag.isDevDependency()).isTrue();
@@ -120,7 +126,9 @@ public void stringListDict_default() throws Exception {
                     .build()),
             buildTag("tag_name").build(),
             null,
-            "root module");
+            "root module",
+            /* moduleIndex= */ 0,
+            /* tagIndex= */ 0);
     assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo");
     assertThat(getattr(typeCheckedTag, "foo"))
         .isEqualTo(
@@ -143,7 +151,9 @@ public void multipleAttributesAndDefaults() throws Exception {
                 .addAttr("quux", StarlarkList.immutableOf("quuxValue1", "quuxValue2"))
                 .build(),
             /* labelConverter= */ null,
-            "root module");
+            "root module",
+            /* moduleIndex= */ 0,
+            /* tagIndex= */ 0);
     assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo", "bar", "quux");
     assertThat(getattr(typeCheckedTag, "foo")).isEqualTo("fooValue");
     assertThat(getattr(typeCheckedTag, "bar")).isEqualTo(StarlarkInt.of(3));
@@ -162,7 +172,9 @@ public void mandatory() throws Exception {
                     createTagClass(attr("foo", Type.STRING).mandatory().build()),
                     buildTag("tag_name").build(),
                     /* labelConverter= */ null,
-                    "root module"));
+                    "root module",
+                    /* moduleIndex= */ 0,
+                    /* tagIndex= */ 0));
     assertThat(e).hasMessageThat().contains("mandatory attribute 'foo' isn't being specified");
   }
 
@@ -179,7 +191,9 @@ public void allowedValues() throws Exception {
                             .build()),
                     buildTag("tag_name").addAttr("foo", "maybe").build(),
                     /* labelConverter= */ null,
-                    "root module"));
+                    "root module",
+                    /* moduleIndex= */ 0,
+                    /* tagIndex= */ 0));
     assertThat(e)
         .hasMessageThat()
         .contains(
@@ -196,7 +210,52 @@ public void unknownAttr() throws Exception {
                     createTagClass(attr("foo", Type.STRING).build()),
                     buildTag("tag_name").addAttr("bar", "maybe").build(),
                     /* labelConverter= */ null,
-                    "root module"));
+                    "root module",
+                    /* moduleIndex= */ 0,
+                    /* tagIndex= */ 0));
     assertThat(e).hasMessageThat().contains("unknown attribute 'bar' provided");
   }
+
+  @Test
+  public void compareTo() throws Exception {
+    TypeCheckedTag module1Tag1 =
+        TypeCheckedTag.create(
+            createTagClass(attr("foo", Type.STRING).build()),
+            buildTag("tag_name").addAttr("foo", "value1").build(),
+            /* labelConverter= */ null,
+            "root module",
+            /* moduleIndex= */ 1,
+            /* tagIndex= */ 1);
+    TypeCheckedTag module2Tag1 =
+        TypeCheckedTag.create(
+            createTagClass(attr("foo", Type.STRING).build()),
+            buildTag("tag_name").addAttr("foo", "value2").build(),
+            /* labelConverter= */ null,
+            "root module",
+            /* moduleIndex= */ 2,
+            /* tagIndex= */ 1);
+    TypeCheckedTag module1Tag2 =
+        TypeCheckedTag.create(
+            createTagClass(attr("foo", Type.STRING).build()),
+            buildTag("tag_name").addAttr("foo", "value3").build(),
+            /* labelConverter= */ null,
+            "root module",
+            /* moduleIndex= */ 1,
+            /* tagIndex= */ 2);
+    TypeCheckedTag module2Tag2 =
+        TypeCheckedTag.create(
+            createTagClass(attr("foo", Type.STRING).build()),
+            buildTag("tag_name").addAttr("foo", "value4").build(),
+            /* labelConverter= */ null,
+            "root module",
+            /* moduleIndex= */ 2,
+            /* tagIndex= */ 2);
+
+    assertThat(module1Tag1).isLessThan(module2Tag1);
+    assertThat(module2Tag1).isGreaterThan(module1Tag1);
+    assertThat(module1Tag1).isLessThan(module1Tag2);
+    assertThat(module1Tag2).isGreaterThan(module1Tag1);
+    assertThat(module2Tag1).isLessThan(module2Tag2);
+    assertThat(module2Tag2).isGreaterThan(module2Tag1);
+  }
 }