bzlmod: Types for module extension *evaluation*
=================================================================

(https://github.com/bazelbuild/bazel/issues/13316)

* `ModuleExtensionContext`: The `module_ctx` object to be passed to the module extension's implementation function. For now, it only has 1 extra property `modules` which allows the module extension to access the dependency graph and all relevant tags.
* `StarlarkBazelModule`: The elements in `module_ctx.modules`. Each exposes the name and version of the module, and all the tags on it too.
* `TypeCheckedTag`: The type-checked version of `Tag`, which is exposed to Starlark through `StarlarkBazelModule` above. It contains all the attribute values passed to it in tags, but also has everything converted to native types and back to Starlark (so strings in `Tag` could become Labels in `TypeCheckedTag`).

PiperOrigin-RevId: 395101902
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/util/BUILD b/src/test/java/com/google/devtools/build/lib/analysis/util/BUILD
index ea58cab..ace84db 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/util/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/analysis/util/BUILD
@@ -85,7 +85,7 @@
         "//src/main/java/com/google/devtools/build/lib/analysis:view_creation_failed_exception",
         "//src/main/java/com/google/devtools/build/lib/analysis:workspace_status_action",
         "//src/main/java/com/google/devtools/build/lib/analysis/platform",
-        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution_impl",
         "//src/main/java/com/google/devtools/build/lib/bazel/rules/android",
         "//src/main/java/com/google/devtools/build/lib/causes",
         "//src/main/java/com/google/devtools/build/lib/clock",
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
index 1e760c9..ce7bb19 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
@@ -37,9 +37,10 @@
         "//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/bazel/bzlmod:resolution",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution_impl",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/starlark",
-        "//src/main/java/com/google/devtools/build/lib/cmdline:cmdline-primitives",
+        "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/pkgcache",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/local_repository_rule",
@@ -55,6 +56,7 @@
         "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
         "//src/main/java/com/google/devtools/build/lib/skyframe:skyframe_cluster",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/repository",
+        "//src/main/java/com/google/devtools/build/lib/util:filetype",
         "//src/main/java/com/google/devtools/build/lib/util/io",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
@@ -90,8 +92,13 @@
     ],
     deps = [
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:module_extension",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:registry",
+        "//src/main/java/com/google/devtools/build/lib/cmdline:cmdline-primitives",
         "//src/main/java/com/google/devtools/build/lib/events",
+        "//src/main/java/com/google/devtools/build/lib/packages",
+        "//src/main/java/net/starlark/java/eval",
+        "//src/main/java/net/starlark/java/syntax",
         "//third_party:guava",
     ],
 )
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodTestUtil.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodTestUtil.java
index 2a40060..842636c 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodTestUtil.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodTestUtil.java
@@ -15,6 +15,14 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.cmdline.RepositoryMapping;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
+import com.google.devtools.build.lib.packages.Attribute;
+import net.starlark.java.eval.Dict;
+import net.starlark.java.syntax.Location;
+
 /** Utilities for bzlmod tests. */
 public final class BzlmodTestUtil {
   private BzlmodTestUtil() {}
@@ -27,4 +35,45 @@
       throw new IllegalArgumentException(e);
     }
   }
+
+  public static RepositoryMapping createRepositoryMapping(String... names) {
+    ImmutableMap.Builder<RepositoryName, RepositoryName> mappingBuilder = ImmutableMap.builder();
+    for (int i = 0; i < names.length; i += 2) {
+      mappingBuilder.put(
+          RepositoryName.createFromValidStrippedName(names[i]),
+          RepositoryName.createFromValidStrippedName(names[i + 1]));
+    }
+    return RepositoryMapping.createAllowingFallback(mappingBuilder.build());
+  }
+
+  public static TagClass createTagClass(Attribute... attrs) {
+    return TagClass.create(ImmutableList.copyOf(attrs), "doc", Location.BUILTIN);
+  }
+
+  /** A builder for {@link Tag} for testing purposes. */
+  public static class TestTagBuilder {
+    private final Dict.Builder<String, Object> attrValuesBuilder = Dict.builder();
+    private final String tagName;
+
+    private TestTagBuilder(String tagName) {
+      this.tagName = tagName;
+    }
+
+    public TestTagBuilder addAttr(String attrName, Object attrValue) {
+      attrValuesBuilder.put(attrName, attrValue);
+      return this;
+    }
+
+    public Tag build() {
+      return Tag.builder()
+          .setTagName(tagName)
+          .setLocation(Location.BUILTIN)
+          .setAttributeValues(attrValuesBuilder.buildImmutable())
+          .build();
+    }
+  }
+
+  public static TestTagBuilder buildTag(String tagName) throws Exception {
+    return new TestTagBuilder(tagName);
+  }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleTest.java
index 866aa7f..54c719a 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleTest.java
@@ -16,14 +16,12 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.createModuleKey;
+import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.createRepositoryMapping;
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.bazel.bzlmod.Module.WhichRepoMappings;
-import com.google.devtools.build.lib.cmdline.RepositoryMapping;
-import com.google.devtools.build.lib.cmdline.RepositoryName;
 import net.starlark.java.syntax.Location;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -86,16 +84,6 @@
                 .build());
   }
 
-  private static RepositoryMapping createRepositoryMapping(String... names) {
-    ImmutableMap.Builder<RepositoryName, RepositoryName> mappingBuilder = ImmutableMap.builder();
-    for (int i = 0; i < names.length; i += 2) {
-      mappingBuilder.put(
-          RepositoryName.createFromValidStrippedName(names[i]),
-          RepositoryName.createFromValidStrippedName(names[i + 1]));
-    }
-    return RepositoryMapping.createAllowingFallback(mappingBuilder.build());
-  }
-
   @Test
   public void getRepoMapping() throws Exception {
     Module module =
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java
new file mode 100644
index 0000000..d880cfb
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java
@@ -0,0 +1,140 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package com.google.devtools.build.lib.bazel.bzlmod;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.buildTag;
+import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.createModuleKey;
+import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.createTagClass;
+import static com.google.devtools.build.lib.packages.Attribute.attr;
+import static org.junit.Assert.assertThrows;
+
+import com.google.common.collect.ImmutableBiMap;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.packages.BuildType;
+import com.google.devtools.build.lib.packages.Type;
+import com.google.devtools.build.lib.util.FileTypeSet;
+import net.starlark.java.eval.StarlarkList;
+import net.starlark.java.syntax.Location;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link StarlarkBazelModule}. */
+@RunWith(JUnit4.class)
+public class StarlarkBazelModuleTest {
+
+  /** A builder for ModuleExtensionUsage that sets all the mandatory but irrelevant fields. */
+  private static ModuleExtensionUsage.Builder getBaseUsageBuilder() {
+    return ModuleExtensionUsage.builder()
+        .setExtensionBzlFile("//:rje.bzl")
+        .setExtensionName("maven")
+        .setLocation(Location.BUILTIN)
+        .setImports(ImmutableBiMap.of());
+  }
+
+  /** A builder for ModuleExtension that sets all the mandatory but irrelevant fields. */
+  private static ModuleExtension.Builder getBaseExtensionBuilder() {
+    return ModuleExtension.builder()
+        .setName("maven")
+        .setDoc("")
+        .setLocation(Location.BUILTIN)
+        .setDefinitionEnvironmentLabel(Label.parseAbsoluteUnchecked("//:rje.bzl"))
+        .setImplementation(() -> "maven");
+  }
+
+  @Test
+  public void basic() throws Exception {
+    ModuleExtensionUsage usage =
+        getBaseUsageBuilder()
+            .addTag(buildTag("dep").addAttr("coord", "junit").build())
+            .addTag(buildTag("dep").addAttr("coord", "guava").build())
+            .addTag(
+                buildTag("pom")
+                    .addAttr("pom_xmls", StarlarkList.immutableOf("//:pom.xml", "@bar//:pom.xml"))
+                    .build())
+            .build();
+    ModuleExtension extension =
+        getBaseExtensionBuilder()
+            .setTagClasses(
+                ImmutableMap.of(
+                    "dep", createTagClass(attr("coord", Type.STRING).build()),
+                    "repos", createTagClass(attr("repos", Type.STRING_LIST).build()),
+                    "pom",
+                        createTagClass(
+                            attr("pom_xmls", BuildType.LABEL_LIST)
+                                .allowedFileTypes(FileTypeSet.ANY_FILE)
+                                .build())))
+            .build();
+    Module module =
+        Module.builder()
+            .setName("foo")
+            .setVersion(Version.parse("1.0"))
+            .addDep("bar", createModuleKey("bar", "2.0"))
+            .addExtensionUsage(usage)
+            .build();
+
+    StarlarkBazelModule moduleProxy =
+        StarlarkBazelModule.create(createModuleKey("foo", ""), module, extension, usage);
+
+    assertThat(moduleProxy.getName()).isEqualTo("foo");
+    assertThat(moduleProxy.getVersion()).isEqualTo("1.0");
+    assertThat(moduleProxy.getTags().getFieldNames()).containsExactly("dep", "repos", "pom");
+
+    // We have 2 "dep" tags...
+    @SuppressWarnings("unchecked")
+    ImmutableList<TypeCheckedTag> depTags =
+        (ImmutableList<TypeCheckedTag>) moduleProxy.getTags().getValue("dep");
+    assertThat(depTags).hasSize(2);
+    assertThat(depTags.get(0).getValue("coord")).isEqualTo("junit");
+    assertThat(depTags.get(1).getValue("coord")).isEqualTo("guava");
+
+    // ... zero "repos" tags...
+    assertThat(moduleProxy.getTags().getValue("repos")).isEqualTo(ImmutableList.of());
+
+    // ... and 1 "pom" tag.
+    @SuppressWarnings("unchecked")
+    ImmutableList<TypeCheckedTag> pomTags =
+        (ImmutableList<TypeCheckedTag>) moduleProxy.getTags().getValue("pom");
+    assertThat(pomTags).hasSize(1);
+    assertThat(pomTags.get(0).getValue("pom_xmls"))
+        .isEqualTo(
+            StarlarkList.immutableOf(
+                Label.parseAbsoluteUnchecked("@foo.//:pom.xml"),
+                Label.parseAbsoluteUnchecked("@bar.2.0//:pom.xml")));
+  }
+
+  @Test
+  public void unknownTagClass() throws Exception {
+    ModuleExtensionUsage usage = getBaseUsageBuilder().addTag(buildTag("blep").build()).build();
+    ModuleExtension extension =
+        getBaseExtensionBuilder().setTagClasses(ImmutableMap.of("dep", createTagClass())).build();
+    Module module =
+        Module.builder()
+            .setName("foo")
+            .setVersion(Version.parse("1.0"))
+            .addExtensionUsage(usage)
+            .build();
+
+    ExternalDepsException e =
+        assertThrows(
+            ExternalDepsException.class,
+            () -> StarlarkBazelModule.create(createModuleKey("foo", ""), module, extension, usage));
+    assertThat(e).hasMessageThat().contains("does not have a tag class named blep");
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/TypeCheckedTagTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/TypeCheckedTagTest.java
new file mode 100644
index 0000000..5fb9d4b
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/TypeCheckedTagTest.java
@@ -0,0 +1,139 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package com.google.devtools.build.lib.bazel.bzlmod;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.buildTag;
+import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.createRepositoryMapping;
+import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.createTagClass;
+import static com.google.devtools.build.lib.packages.Attribute.attr;
+import static org.junit.Assert.assertThrows;
+
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.packages.Attribute.AllowedValueSet;
+import com.google.devtools.build.lib.packages.BuildType;
+import com.google.devtools.build.lib.packages.BuildType.LabelConversionContext;
+import com.google.devtools.build.lib.packages.Type;
+import com.google.devtools.build.lib.util.FileTypeSet;
+import java.util.HashMap;
+import net.starlark.java.eval.StarlarkInt;
+import net.starlark.java.eval.StarlarkList;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link TypeCheckedTag}. */
+@RunWith(JUnit4.class)
+public class TypeCheckedTagTest {
+
+  @Test
+  public void basic() throws Exception {
+    TypeCheckedTag typeCheckedTag =
+        TypeCheckedTag.create(
+            createTagClass(attr("foo", Type.INTEGER).build()),
+            buildTag("tag_name").addAttr("foo", StarlarkInt.of(3)).build(),
+            /*labelConversionContext=*/ null);
+    assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo");
+    assertThat(typeCheckedTag.getValue("foo")).isEqualTo(StarlarkInt.of(3));
+  }
+
+  @Test
+  public void label() throws Exception {
+    TypeCheckedTag typeCheckedTag =
+        TypeCheckedTag.create(
+            createTagClass(
+                attr("foo", BuildType.LABEL_LIST).allowedFileTypes(FileTypeSet.ANY_FILE).build()),
+            buildTag("tag_name")
+                .addAttr(
+                    "foo", StarlarkList.immutableOf(":thing1", "//pkg:thing2", "@repo//pkg:thing3"))
+                .build(),
+            new LabelConversionContext(
+                Label.parseAbsoluteUnchecked("@myrepo//mypkg:defs.bzl"),
+                createRepositoryMapping("repo", "other_repo"),
+                new HashMap<>()));
+    assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo");
+    assertThat(typeCheckedTag.getValue("foo"))
+        .isEqualTo(
+            StarlarkList.immutableOf(
+                Label.parseAbsoluteUnchecked("@myrepo//mypkg:thing1"),
+                Label.parseAbsoluteUnchecked("@myrepo//pkg:thing2"),
+                Label.parseAbsoluteUnchecked("@other_repo//pkg:thing3")));
+  }
+
+  @Test
+  public void multipleAttributesAndDefaults() throws Exception {
+    TypeCheckedTag typeCheckedTag =
+        TypeCheckedTag.create(
+            createTagClass(
+                attr("foo", Type.STRING).mandatory().build(),
+                attr("bar", Type.INTEGER).value(StarlarkInt.of(3)).build(),
+                attr("quux", Type.STRING_LIST).build()),
+            buildTag("tag_name")
+                .addAttr("foo", "fooValue")
+                .addAttr("quux", StarlarkList.immutableOf("quuxValue1", "quuxValue2"))
+                .build(),
+            /*labelConversionContext=*/ null);
+    assertThat(typeCheckedTag.getFieldNames()).containsExactly("foo", "bar", "quux");
+    assertThat(typeCheckedTag.getValue("foo")).isEqualTo("fooValue");
+    assertThat(typeCheckedTag.getValue("bar")).isEqualTo(StarlarkInt.of(3));
+    assertThat(typeCheckedTag.getValue("quux"))
+        .isEqualTo(StarlarkList.immutableOf("quuxValue1", "quuxValue2"));
+  }
+
+  @Test
+  public void mandatory() throws Exception {
+    ExternalDepsException e =
+        assertThrows(
+            ExternalDepsException.class,
+            () ->
+                TypeCheckedTag.create(
+                    createTagClass(attr("foo", Type.STRING).mandatory().build()),
+                    buildTag("tag_name").build(),
+                    /*labelConversionContext=*/ null));
+    assertThat(e).hasMessageThat().contains("mandatory attribute foo isn't being specified");
+  }
+
+  @Test
+  public void allowedValues() throws Exception {
+    ExternalDepsException e =
+        assertThrows(
+            ExternalDepsException.class,
+            () ->
+                TypeCheckedTag.create(
+                    createTagClass(
+                        attr("foo", Type.STRING)
+                            .allowedValues(new AllowedValueSet("yes", "no"))
+                            .build()),
+                    buildTag("tag_name").addAttr("foo", "maybe").build(),
+                    /*labelConversionContext=*/ null));
+    assertThat(e)
+        .hasMessageThat()
+        .contains("the value for attribute foo has to be one of 'yes' or 'no' instead of 'maybe'");
+  }
+
+  @Test
+  public void unknownAttr() throws Exception {
+    ExternalDepsException e =
+        assertThrows(
+            ExternalDepsException.class,
+            () ->
+                TypeCheckedTag.create(
+                    createTagClass(attr("foo", Type.STRING).build()),
+                    buildTag("tag_name").addAttr("bar", "maybe").build(),
+                    /*labelConversionContext=*/ null));
+    assertThat(e).hasMessageThat().contains("unknown attribute bar provided");
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/rules/repository/BUILD b/src/test/java/com/google/devtools/build/lib/rules/repository/BUILD
index 0fe74eb..4465ec0 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/repository/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/rules/repository/BUILD
@@ -23,7 +23,7 @@
         "//src/main/java/com/google/devtools/build/lib/analysis:server_directories",
         "//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/bazel/bzlmod:resolution",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution_impl",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/starlark",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/BUILD b/src/test/java/com/google/devtools/build/lib/skyframe/BUILD
index 4a916e6..f7f7cdf 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/BUILD
@@ -102,7 +102,7 @@
     }) + [
         ":testutil",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
-        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:resolution_impl",
         "//src/main/java/com/google/devtools/build/lib:build-request-options",
         "//src/main/java/com/google/devtools/build/lib:keep-going-option",
         "//src/main/java/com/google/devtools/build/lib:runtime",