From 87325e5ac476736d83420888d53ced3631b350c9 Mon Sep 17 00:00:00 2001
From: pcloudy <pcloudy@google.com>
Date: Wed, 7 Jul 2021 04:03:14 -0700
Subject: [PATCH] bzlmod: Support creating Starlark repo rule

Related: https://github.com/bazelbuild/bazel/issues/13316

RELNOTES: None.
PiperOrigin-RevId: 383380283
---
 .../google/devtools/build/lib/skyframe/BUILD  |   1 +
 .../lib/skyframe/BzlmodRepoRuleFunction.java  |  90 +++++++++++++-
 .../bzlmod/BzlmodRepoRuleFunctionTest.java    | 110 ++++++++++++++++--
 3 files changed, 190 insertions(+), 11 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BUILD b/src/main/java/com/google/devtools/build/lib/skyframe/BUILD
index 7ec0aeda47b93f..2fd26d1221d53f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/BUILD
@@ -274,6 +274,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis/platform",
         "//src/main/java/com/google/devtools/build/lib/analysis/platform:utils",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:repo_rule_creator",
         "//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/bugreport",
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/BzlmodRepoRuleFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/BzlmodRepoRuleFunction.java
index 6e5ab7bd1fc3ad..f01c3e48669001 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/BzlmodRepoRuleFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/BzlmodRepoRuleFunction.java
@@ -14,13 +14,18 @@
 
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
+import com.google.devtools.build.lib.bazel.bzlmod.BzlmodRepoRuleCreator;
 import com.google.devtools.build.lib.bazel.bzlmod.BzlmodRepoRuleHelper;
 import com.google.devtools.build.lib.bazel.bzlmod.BzlmodRepoRuleValue;
 import com.google.devtools.build.lib.bazel.bzlmod.RepoSpec;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
+import com.google.devtools.build.lib.cmdline.PackageIdentifier;
+import com.google.devtools.build.lib.events.StoredEventHandler;
 import com.google.devtools.build.lib.packages.NoSuchPackageException;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.PackageFactory;
@@ -30,6 +35,8 @@
 import com.google.devtools.build.lib.packages.RuleFactory;
 import com.google.devtools.build.lib.packages.RuleFactory.BuildLangTypedAttributeValuesMap;
 import com.google.devtools.build.lib.packages.RuleFactory.InvalidRuleException;
+import com.google.devtools.build.lib.server.FailureDetails.PackageLoading;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -39,6 +46,7 @@
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.Optional;
 import javax.annotation.Nullable;
+import net.starlark.java.eval.Module;
 import net.starlark.java.eval.StarlarkSemantics;
 import net.starlark.java.eval.StarlarkThread.CallStackEntry;
 import net.starlark.java.syntax.Location;
@@ -54,6 +62,7 @@ public final class BzlmodRepoRuleFunction implements SkyFunction {
   private final RuleClassProvider ruleClassProvider;
   private final BlazeDirectories directories;
   private final BzlmodRepoRuleHelper bzlmodRepoRuleHelper;
+  private static final PackageIdentifier ROOT_PACKAGE = PackageIdentifier.createInMainRepo("");
 
   public BzlmodRepoRuleFunction(
       PackageFactory packageFactory,
@@ -112,8 +121,7 @@ private BzlmodRepoRuleValue createRuleFromSpec(
     if (repoSpec.isNativeRepoRule()) {
       return createNativeRepoRule(repoSpec, starlarkSemantics, env);
     }
-    // TODO(pcloudy): Implement creating starlark repository rule
-    return BzlmodRepoRuleValue.REPO_RULE_NOT_FOUND_VALUE;
+    return createStarlarkRepoRule(repoSpec, starlarkSemantics, env);
   }
 
   private BzlmodRepoRuleValue createNativeRepoRule(
@@ -145,6 +153,84 @@ private BzlmodRepoRuleValue createNativeRepoRule(
     return new BzlmodRepoRuleValue(rule);
   }
 
+  /** Loads modules from the given bzl file. */
+  private ImmutableMap<String, Module> loadBzlModules(Environment env, String bzlFile)
+      throws InterruptedException, BzlmodRepoRuleFunctionException {
+    ImmutableList<Pair<String, Location>> programLoads =
+        ImmutableList.of(Pair.of(bzlFile, Location.BUILTIN));
+
+    ImmutableList<Label> loadLabels =
+        BzlLoadFunction.getLoadLabels(
+            env.getListener(), programLoads, ROOT_PACKAGE, /*repoMapping=*/ ImmutableMap.of());
+    if (loadLabels == null) {
+      NoSuchPackageException e =
+          PackageFunction.PackageFunctionException.builder()
+              .setType(PackageFunction.PackageFunctionException.Type.BUILD_FILE_CONTAINS_ERRORS)
+              .setPackageIdentifier(ROOT_PACKAGE)
+              .setMessage("malformed load statements")
+              .setPackageLoadingCode(PackageLoading.Code.IMPORT_STARLARK_FILE_ERROR)
+              .buildCause();
+      throw new BzlmodRepoRuleFunctionException(e, Transience.PERSISTENT);
+    }
+
+    Preconditions.checkArgument(loadLabels.size() == 1);
+    ImmutableList<BzlLoadValue.Key> keys =
+        ImmutableList.of(BzlLoadValue.keyForBzlmod(loadLabels.get(0)));
+
+    // Load the .bzl module.
+    try {
+      return PackageFunction.loadBzlModules(env, ROOT_PACKAGE, programLoads, keys, null);
+    } catch (NoSuchPackageException e) {
+      throw new BzlmodRepoRuleFunctionException(e, Transience.PERSISTENT);
+    }
+  }
+
+  private BzlmodRepoRuleValue createStarlarkRepoRule(
+      RepoSpec repoSpec, StarlarkSemantics semantics, Environment env)
+      throws InterruptedException, BzlmodRepoRuleFunctionException {
+
+    ImmutableMap<String, Module> loadedModules = loadBzlModules(env, repoSpec.bzlFile().get());
+
+    if (env.valuesMissing()) {
+      return null;
+    }
+
+    BzlmodRepoRuleCreator repoRuleCreator = getRepoRuleCreator(repoSpec, loadedModules);
+
+    // TODO(bazel-team): Don't use the {@link Rule} class for repository rule.
+    // Currently, the repository rule is represented with the {@link Rule} class that's designed
+    // for build rules. Therefore, we have to create a package instance for it, which doesn't make
+    // sense. We should migrate away from this implementation so that we don't refer to any build
+    // rule specific things in repository rule.
+    Rule rule;
+    Package.Builder pkg = createExternalPackageBuilder(semantics);
+    StoredEventHandler eventHandler = new StoredEventHandler();
+    try {
+      rule = repoRuleCreator.createRule(pkg, semantics, repoSpec.attributes(), eventHandler);
+      // We need to actually build the package so that the rule has the correct package reference.
+      pkg.build();
+    } catch (InvalidRuleException e) {
+      throw new BzlmodRepoRuleFunctionException(e, Transience.PERSISTENT);
+    } catch (NoSuchPackageException e) {
+      throw new BzlmodRepoRuleFunctionException(e, Transience.PERSISTENT);
+    }
+
+    return new BzlmodRepoRuleValue(rule);
+  }
+
+  private BzlmodRepoRuleCreator getRepoRuleCreator(
+      RepoSpec repoSpec, ImmutableMap<String, Module> loadedModules)
+      throws BzlmodRepoRuleFunctionException {
+    Object object = loadedModules.get(repoSpec.bzlFile().get()).getGlobal(repoSpec.ruleClassName());
+    if (object instanceof BzlmodRepoRuleCreator) {
+      return (BzlmodRepoRuleCreator) object;
+    } else {
+      InvalidRuleException e =
+          new InvalidRuleException("Invalid repository rule: " + repoSpec.getRuleClass());
+      throw new BzlmodRepoRuleFunctionException(e, Transience.PERSISTENT);
+    }
+  }
+
   /**
    * Create the external package builder, which is only for the convenience of creating repository
    * rules.
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodRepoRuleFunctionTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodRepoRuleFunctionTest.java
index 0a383beee5b862..2fe5b069d6fb1e 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodRepoRuleFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BzlmodRepoRuleFunctionTest.java
@@ -175,6 +175,33 @@ pkgFactory, ruleClassProvider, directories, getFakeBzlmodRepoRuleHelper()))
 
     PrecomputedValue.STARLARK_SEMANTICS.set(differencer, StarlarkSemantics.DEFAULT);
     PrecomputedValue.PATH_PACKAGE_LOCATOR.set(differencer, packageLocator.get());
+
+    setupRepoRules();
+  }
+
+  private void setupRepoRules() throws Exception {
+    scratch.file(rootDirectory.getRelative("tools/build_defs/repo/BUILD").getPathString());
+    scratch.file(
+        rootDirectory.getRelative("tools/build_defs/repo/http.bzl").getPathString(),
+        "def _http_archive_impl(ctx): pass",
+        "",
+        "http_archive = repository_rule(",
+        "    implementation = _http_archive_impl,",
+        "    attrs = {",
+        "      \"url\": attr.string(),",
+        "      \"sha256\": attr.string(),",
+        "    })");
+    scratch.file(rootDirectory.getRelative("maven/BUILD").getPathString());
+    scratch.file(
+        rootDirectory.getRelative("maven/repo.bzl").getPathString(),
+        "def _maven_repo_impl(ctx): pass",
+        "",
+        "maven_repo = repository_rule(",
+        "    implementation = _maven_repo_impl,",
+        "    attrs = {",
+        "      \"artifacts\": attr.string_list(),",
+        "      \"repositories\": attr.string_list(),",
+        "    })");
   }
 
   private FakeBzlmodRepoRuleHelper getFakeBzlmodRepoRuleHelper() {
@@ -182,14 +209,43 @@ private FakeBzlmodRepoRuleHelper getFakeBzlmodRepoRuleHelper() {
     repoSpecs
         // repos from non-registry overrides
         .put(
-        "A",
-        RepoSpec.builder()
-            .setRuleClassName("local_repository")
-            .setAttributes(
-                ImmutableMap.of(
-                    "name", "A",
-                    "path", "/foo/bar/A"))
-            .build());
+            "A",
+            RepoSpec.builder()
+                .setRuleClassName("local_repository")
+                .setAttributes(
+                    ImmutableMap.of(
+                        "name", "A",
+                        "path", "/foo/bar/A"))
+                .build())
+        // repos from Bazel modules
+        .put(
+            "B",
+            RepoSpec.builder()
+                .setBzlFile(
+                    // In real world, this will be @bazel_tools//tools/build_defs/repo:http.bzl,
+                    "//tools/build_defs/repo:http.bzl")
+                .setRuleClassName("http_archive")
+                .setAttributes(
+                    ImmutableMap.of(
+                        "name", "B",
+                        "url", "https://foo/bar/B.zip",
+                        "sha256", "1234abcd"))
+                .build())
+        // repos from module rules
+        .put(
+            "C",
+            RepoSpec.builder()
+                .setBzlFile("//maven:repo.bzl")
+                .setRuleClassName("maven_repo")
+                .setAttributes(
+                    ImmutableMap.of(
+                        "name", "C",
+                        "artifacts",
+                            ImmutableList.of("junit:junit:4.12", "com.google.guava:guava:19.0"),
+                        "repositories",
+                            ImmutableList.of(
+                                "https://maven.google.com", "https://repo1.maven.org/maven2")))
+                .build());
     return new FakeBzlmodRepoRuleHelper(repoSpecs.build());
   }
 
@@ -227,7 +283,43 @@ public void createRepoRule_overrides() throws Exception {
   }
 
   @Test
-  public void repoRule_notFound() throws Exception {
+  public void createRepoRule_bazelModules() throws Exception {
+    EvaluationResult<BzlmodRepoRuleValue> result =
+        driver.evaluate(ImmutableList.of(BzlmodRepoRuleValue.key("B")), evaluationContext);
+    if (result.hasError()) {
+      fail(result.getError().toString());
+    }
+    BzlmodRepoRuleValue bzlmodRepoRuleValue = result.get(BzlmodRepoRuleValue.key("B"));
+    Rule repoRule = bzlmodRepoRuleValue.getRule();
+
+    assertThat(repoRule.getRuleClassObject().isStarlark()).isTrue();
+    assertThat(repoRule.getRuleClass()).isEqualTo("http_archive");
+    assertThat(repoRule.getName()).isEqualTo("B");
+    assertThat(repoRule.getAttr("url", Type.STRING)).isEqualTo("https://foo/bar/B.zip");
+    assertThat(repoRule.getAttr("sha256", Type.STRING)).isEqualTo("1234abcd");
+  }
+
+  @Test
+  public void createRepoRule_moduleRules() throws Exception {
+    EvaluationResult<BzlmodRepoRuleValue> result =
+        driver.evaluate(ImmutableList.of(BzlmodRepoRuleValue.key("C")), evaluationContext);
+    if (result.hasError()) {
+      fail(result.getError().toString());
+    }
+    BzlmodRepoRuleValue bzlmodRepoRuleValue = result.get(BzlmodRepoRuleValue.key("C"));
+    Rule repoRule = bzlmodRepoRuleValue.getRule();
+
+    assertThat(repoRule.getRuleClassObject().isStarlark()).isTrue();
+    assertThat(repoRule.getRuleClass()).isEqualTo("maven_repo");
+    assertThat(repoRule.getName()).isEqualTo("C");
+    assertThat(repoRule.getAttr("artifacts", Type.STRING_LIST))
+        .isEqualTo(ImmutableList.of("junit:junit:4.12", "com.google.guava:guava:19.0"));
+    assertThat(repoRule.getAttr("repositories", Type.STRING_LIST))
+        .isEqualTo(ImmutableList.of("https://maven.google.com", "https://repo1.maven.org/maven2"));
+  }
+
+  @Test
+  public void createRepoRule_notFound() throws Exception {
     EvaluationResult<BzlmodRepoRuleValue> result =
         driver.evaluate(ImmutableList.of(BzlmodRepoRuleValue.key("unknown")), evaluationContext);
     if (result.hasError()) {