Remove unused fields from ModuleExtension

We store the name and "rule definition environment" but those fields are never actually used. Removing them also simplifies the code a bit (obviating the "InStarlark" thing).

PiperOrigin-RevId: 513213080
Change-Id: I912e388e0d009b018da7d5a920d3a609e063ce6a
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java
index b9d2ffd..9ba37a4 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java
@@ -15,33 +15,28 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.auto.value.AutoValue;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.docgen.annot.DocCategory;
-import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.events.EventHandler;
-import com.google.devtools.build.lib.packages.StarlarkExportable;
-import javax.annotation.Nullable;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.eval.StarlarkCallable;
+import net.starlark.java.eval.StarlarkValue;
 import net.starlark.java.syntax.Location;
 
 /**
- * A module extension object, which can be used to perform arbitrary logic in order to create repos
- * or register toolchains and execution platforms.
+ * A module extension object, which can be used to perform arbitrary logic in order to create repos.
  */
 @AutoValue
-public abstract class ModuleExtension {
-  public abstract String getName();
-
+@StarlarkBuiltin(
+    name = "module_extension",
+    category = DocCategory.BUILTIN,
+    doc = "A module extension declared using the <code>module_extension</code> function.")
+public abstract class ModuleExtension implements StarlarkValue {
   public abstract StarlarkCallable getImplementation();
 
   public abstract ImmutableMap<String, TagClass> getTagClasses();
 
   public abstract String getDoc();
 
-  public abstract Label getDefinitionEnvironmentLabel();
-
   public abstract Location getLocation();
 
   public static Builder builder() {
@@ -51,59 +46,14 @@
   /** Builder for {@link ModuleExtension}. */
   @AutoValue.Builder
   public abstract static class Builder {
-
     public abstract Builder setDoc(String value);
 
-    public abstract Builder setDefinitionEnvironmentLabel(Label value);
-
     public abstract Builder setLocation(Location value);
 
-    public abstract Builder setName(String value);
-
     public abstract Builder setImplementation(StarlarkCallable value);
 
     public abstract Builder setTagClasses(ImmutableMap<String, TagClass> value);
 
     public abstract ModuleExtension build();
   }
-
-  /**
-   * A {@link ModuleExtension} exposed to Starlark. We can't use {@link ModuleExtension} directly
-   * because the name isn't known until the object is exported, so this class holds a builder until
-   * it's exported, at which point it sets the name and builds the underlying {@link
-   * ModuleExtension}.
-   */
-  @StarlarkBuiltin(
-      name = "module_extension",
-      category = DocCategory.BUILTIN,
-      doc = "A module extension declared using the <code>module_extension</code> function.")
-  public static class InStarlark implements StarlarkExportable {
-    private final Builder builder;
-    @Nullable private ModuleExtension built;
-
-    public InStarlark() {
-      builder = builder();
-      built = null;
-    }
-
-    public Builder getBuilder() {
-      return builder;
-    }
-
-    @Override
-    public boolean isExported() {
-      return built != null;
-    }
-
-    @Override
-    public void export(EventHandler handler, Label extensionLabel, String exportedName) {
-      built = builder.setName(exportedName).build();
-    }
-
-    /** Throws {@link IllegalStateException} if this is not exported yet. */
-    public ModuleExtension get() {
-      Preconditions.checkState(isExported(), "the module extension was never exported");
-      return built;
-    }
-  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
index 0418870..0aec2b5 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
@@ -149,10 +149,10 @@
 
     // Check that the .bzl file actually exports a module extension by our name.
     Object exported = bzlLoadValue.getModule().getGlobal(extensionId.getExtensionName());
-    if (!(exported instanceof ModuleExtension.InStarlark)) {
+    if (!(exported instanceof ModuleExtension)) {
       ImmutableSet<String> exportedExtensions =
           bzlLoadValue.getModule().getGlobals().entrySet().stream()
-              .filter(e -> e.getValue() instanceof ModuleExtension.InStarlark)
+              .filter(e -> e.getValue() instanceof ModuleExtension)
               .map(Entry::getKey)
               .collect(toImmutableSet());
       throw new SingleExtensionEvalFunctionException(
@@ -167,7 +167,7 @@
     }
 
     // Run that extension!
-    ModuleExtension extension = ((ModuleExtension.InStarlark) exported).get();
+    ModuleExtension extension = (ModuleExtension) exported;
     ModuleExtensionEvalStarlarkThreadContext threadContext =
         new ModuleExtensionEvalStarlarkThreadContext(
             usagesValue.getExtensionUniqueName() + "~",
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
index e3e58c7..01eecc2 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
@@ -231,7 +231,7 @@
             WorkspaceFactoryHelper.createAndAddRepositoryRule(
                 context.getBuilder(),
                 ruleClass,
-                /*bindRuleClass=*/ null,
+                /* bindRuleClass= */ null,
                 WorkspaceFactoryHelper.getFinalKwargs(kwargs),
                 thread.getSemantics(),
                 thread.getCallStack());
@@ -296,23 +296,19 @@
             positional = false)
       },
       useStarlarkThread = true)
-  public Object moduleExtension(
+  public ModuleExtension moduleExtension(
       StarlarkCallable implementation,
       Dict<?, ?> tagClasses, // Dict<String, TagClass>
       String doc,
       StarlarkThread thread)
       throws EvalException {
-    ModuleExtension.InStarlark inStarlark = new ModuleExtension.InStarlark();
-    inStarlark
-        .getBuilder()
+    return ModuleExtension.builder()
         .setImplementation(implementation)
         .setTagClasses(
             ImmutableMap.copyOf(Dict.cast(tagClasses, String.class, TagClass.class, "tag_classes")))
         .setDoc(doc)
-        .setDefinitionEnvironmentLabel(
-            BazelModuleContext.of(Module.ofInnermostEnclosingStarlarkFunction(thread)).label())
-        .setLocation(thread.getCallerLocation());
-    return inStarlark;
+        .setLocation(thread.getCallerLocation())
+        .build();
   }
 
   @StarlarkMethod(
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
index 25a3971..a6950ee 100644
--- 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
@@ -50,10 +50,8 @@
   /** 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.parseCanonicalUnchecked("//:rje.bzl"))
         .setImplementation(() -> "maven");
   }