Store repospec instead of package for module extension

PiperOrigin-RevId: 529679264
Change-Id: Ib9d410f5de1bac171c3c6f642b8584986181e44b

Author: SalmaSamy

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

@@ -24,7 +24,6 @@
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.packages.NoSuchPackageException;
-import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.RuleFactory.InvalidRuleException;
@@ -53,14 +52,14 @@ public static ModuleExtensionEvalStarlarkThreadContext from(StarlarkThread threa
   }
 
   @AutoValue
-  abstract static class PackageAndLocation {
-    abstract Package getPackage();
+  abstract static class RepoSpecAndLocation {
+    abstract RepoSpec getRepoSpec();
 
     abstract Location getLocation();
 
-    static PackageAndLocation create(Package pkg, Location location) {
-      return new AutoValue_ModuleExtensionEvalStarlarkThreadContext_PackageAndLocation(
-          pkg, location);
+    static RepoSpecAndLocation create(RepoSpec repoSpec, Location location) {
+      return new AutoValue_ModuleExtensionEvalStarlarkThreadContext_RepoSpecAndLocation(
+          repoSpec, location);
     }
   }
 
@@ -69,7 +68,7 @@ static PackageAndLocation create(Package pkg, Location location) {
   private final RepositoryMapping repoMapping;
   private final BlazeDirectories directories;
   private final ExtendedEventHandler eventHandler;
-  private final Map<String, PackageAndLocation> generatedRepos = new HashMap<>();
+  private final Map<String, RepoSpecAndLocation> generatedRepos = new HashMap<>();
 
   public ModuleExtensionEvalStarlarkThreadContext(
       String repoPrefix,
@@ -84,11 +83,6 @@ public ModuleExtensionEvalStarlarkThreadContext(
     this.eventHandler = eventHandler;
   }
 
-  /** The string that should be prefixed to the names of all repos generated by this extension. */
-  public String getRepoPrefix() {
-    return repoPrefix;
-  }
-
   public void createRepo(StarlarkThread thread, Dict<String, Object> kwargs, RuleClass ruleClass)
       throws InterruptedException, EvalException {
     Object nameValue = kwargs.getOrDefault("name", Starlark.NONE);
@@ -98,7 +92,7 @@ public void createRepo(StarlarkThread thread, Dict<String, Object> kwargs, RuleC
     }
     String name = (String) nameValue;
     RepositoryName.validateUserProvidedRepoName(name);
-    PackageAndLocation conflict = generatedRepos.get(name);
+    RepoSpecAndLocation conflict = generatedRepos.get(name);
     if (conflict != null) {
       throw Starlark.errorf(
           "A repo named %s is already generated by this module extension at %s",
@@ -116,8 +110,17 @@ public void createRepo(StarlarkThread thread, Dict<String, Object> kwargs, RuleC
               "RepositoryRuleFunction.createRule",
               ruleClass,
               Maps.transformEntries(kwargs, (k, v) -> k.equals("name") ? prefixedName : v));
-      generatedRepos.put(
-          name, PackageAndLocation.create(rule.getPackage(), thread.getCallerLocation()));
+
+      Map<String, Object> attributes = Maps.transformEntries(kwargs, (k, v) -> rule.getAttr(k));
+      String bzlFile = ruleClass.getRuleDefinitionEnvironmentLabel().getUnambiguousCanonicalForm();
+      RepoSpec repoSpec =
+          RepoSpec.builder()
+              .setBzlFile(bzlFile)
+              .setRuleClassName(ruleClass.getName())
+              .setAttributes(AttributeValues.create(attributes))
+              .build();
+
+      generatedRepos.put(name, RepoSpecAndLocation.create(repoSpec, thread.getCallerLocation()));
     } catch (InvalidRuleException | NoSuchPackageException e) {
       throw Starlark.errorf("%s", e.getMessage());
     }
@@ -128,8 +131,8 @@ public void createRepo(StarlarkThread thread, Dict<String, Object> kwargs, RuleC
    * specified by the extension) of the repo, and the value is the package containing (only) the
    * repo rule.
    */
-  public ImmutableMap<String, Package> getGeneratedRepos() {
+  public ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs() {
     return ImmutableMap.copyOf(
-        Maps.transformValues(generatedRepos, PackageAndLocation::getPackage));
+        Maps.transformValues(generatedRepos, RepoSpecAndLocation::getRepoSpec));
   }
 }

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

@@ -199,7 +199,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
           ModuleExtensionMetadata metadata = (ModuleExtensionMetadata) returnValue;
           metadata.evaluate(
               usagesValue.getExtensionUsages().values(),
-              threadContext.getGeneratedRepos().keySet(),
+              threadContext.getGeneratedRepoSpecs().keySet(),
               env.getListener());
         }
       } catch (NeedsSkyframeRestartException e) {
@@ -227,7 +227,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     // Check that all imported repos have been actually generated
     for (ModuleExtensionUsage usage : usagesValue.getExtensionUsages().values()) {
       for (Entry<String, String> repoImport : usage.getImports().entrySet()) {
-        if (!threadContext.getGeneratedRepos().containsKey(repoImport.getValue())) {
+        if (!threadContext.getGeneratedRepoSpecs().containsKey(repoImport.getValue())) {
           throw new SingleExtensionEvalFunctionException(
               ExternalDepsException.withMessage(
                   Code.BAD_MODULE,
@@ -239,15 +239,15 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                   repoImport.getKey(),
                   usage.getLocation(),
                   SpellChecker.didYouMean(
-                      repoImport.getValue(), threadContext.getGeneratedRepos().keySet())),
+                      repoImport.getValue(), threadContext.getGeneratedRepoSpecs().keySet())),
               Transience.PERSISTENT);
         }
       }
     }
 
     return SingleExtensionEvalValue.create(
-        threadContext.getGeneratedRepos(),
-        threadContext.getGeneratedRepos().keySet().stream()
+        threadContext.getGeneratedRepoSpecs(),
+        threadContext.getGeneratedRepoSpecs().keySet().stream()
             .collect(
                 toImmutableBiMap(
                     e ->

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

@@ -33,9 +33,9 @@
 public abstract class SingleExtensionEvalValue implements SkyValue {
   /**
    * Returns the repos generated by the extension. The key is the "internal name" (as specified by
-   * the extension) of the repo, and the value is the package containing (only) the repo rule.
+   * the extension) of the repo, and the value is the the repo specs that defins the repository .
    */
-  public abstract ImmutableMap<String, Package> getGeneratedRepos();
+  public abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
 
   /**
    * Returns the mapping from the canonical repo names of the repos generated by this extension to
@@ -45,9 +45,10 @@ public abstract class SingleExtensionEvalValue implements SkyValue {
 
   @AutoCodec.Instantiator
   public static SingleExtensionEvalValue create(
-      ImmutableMap<String, Package> generatedRepos,
+      ImmutableMap<String, RepoSpec> generatedRepoSpecs,
       ImmutableBiMap<RepositoryName, String> canonicalRepoNameToInternalNames) {
-    return new AutoValue_SingleExtensionEvalValue(generatedRepos, canonicalRepoNameToInternalNames);
+    return new AutoValue_SingleExtensionEvalValue(
+        generatedRepoSpecs, canonicalRepoNameToInternalNames);
   }
 
   public static Key key(ModuleExtensionId id) {

src/main/java/com/google/devtools/build/lib/skyframe/BzlmodRepoRuleFunction.java

@@ -140,7 +140,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     if (internalRepo == null) {
       return BzlmodRepoRuleValue.REPO_RULE_NOT_FOUND_VALUE;
     }
-    Package pkg = extensionEval.getGeneratedRepos().get(internalRepo);
+    RepoSpec extRepoSpec = extensionEval.getGeneratedRepoSpecs().get(internalRepo);
+    Package pkg = createRuleFromSpec(extRepoSpec, starlarkSemantics, env).getRule().getPackage();
     Preconditions.checkNotNull(pkg);
 
     return new BzlmodRepoRuleValue(pkg, repositoryName.getName());