bzlmod: Implement BzlmodRepoRuleHelper

Now we can get RepoSpec for non-registry overrides and Bazel module derived repositories,
module rule support will be added later.

Related: #13316

RELNOTES: None
PiperOrigin-RevId: 384178751

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

@@ -45,4 +45,17 @@ public static ArchiveOverride create(
 
   /** The number of path segments to strip from the paths in the supplied patches. */
   public abstract int getPatchStrip();
+
+  /** Returns the {@link RepoSpec} that defines this repository. */
+  @Override
+  public RepoSpec getRepoSpec(String repoName) {
+    return new ArchiveRepoSpecBuilder()
+        .setRepoName(repoName)
+        .setUrls(getUrls())
+        .setIntegrity(getIntegrity())
+        .setStripPrefix(getStripPrefix())
+        .setPatches(getPatches())
+        .setPatchStrip(getPatchStrip())
+        .build();
+  }
 }

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -107,8 +107,12 @@ java_library(
         "BzlmodRepoRuleHelperImpl.java",
     ],
     deps = [
-        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
+        ":common",
+        ":registry",
+        ":resolution",
+        "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/skyframe",
+        "//third_party:guava",
     ],
 )
 

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

@@ -15,9 +15,11 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
+import java.io.IOException;
 import java.util.Optional;
 
 /** A helper to get {@link RepoSpec} for Bzlmod generated repositories. */
 public interface BzlmodRepoRuleHelper {
-  Optional<RepoSpec> getRepoSpec(Environment env, String repositoryName);
+  Optional<RepoSpec> getRepoSpec(Environment env, String repositoryName)
+      throws InterruptedException, IOException;
 }

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

@@ -14,15 +14,106 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
+import java.io.IOException;
 import java.util.Optional;
 
 /** A helper class to get {@link RepoSpec} for Bzlmod generated repositories. */
 public final class BzlmodRepoRuleHelperImpl implements BzlmodRepoRuleHelper {
 
   @Override
-  public Optional<RepoSpec> getRepoSpec(Environment env, String repositoryName) {
-    // TODO(pcloudy): Implement calculating RepoSpec.
+  public Optional<RepoSpec> getRepoSpec(Environment env, String repositoryName)
+      throws InterruptedException, IOException {
+
+    ModuleFileValue root = (ModuleFileValue) env.getValue(ModuleFileValue.keyForRootModule());
+    if (env.valuesMissing()) {
+      return null;
+    }
+    ImmutableMap<String, ModuleOverride> overrides = root.getOverrides();
+
+    // Step 1: Look for repositories defined by non-registry overrides.
+    Optional<RepoSpec> repoSpec = checkRepoFromNonRegistryOverrides(overrides, repositoryName);
+    if (repoSpec.isPresent()) {
+      return repoSpec;
+    }
+
+    // SelectionValue is affected by repos found in Step 1, therefore it should NOT be asked
+    // in Step 1 to avoid cycle dependency.
+    SelectionValue selectionValue = (SelectionValue) env.getValue(SelectionValue.KEY);
+    if (env.valuesMissing()) {
+      return null;
+    }
+
+    // Step 2: Look for repositories derived from Bazel Modules.
+    repoSpec =
+        checkRepoFromBazelModules(selectionValue, overrides, env.getListener(), repositoryName);
+    if (repoSpec.isPresent()) {
+      return repoSpec;
+    }
+
+    // Step 3: Look for repositories derived from module rules.
+    return checkRepoFromModuleRules();
+  }
+
+  private static Optional<RepoSpec> checkRepoFromNonRegistryOverrides(
+      ImmutableMap<String, ModuleOverride> overrides, String repositoryName) {
+    if (overrides.containsKey(repositoryName)) {
+      ModuleOverride override = overrides.get(repositoryName);
+      if (override instanceof NonRegistryOverride) {
+        return Optional.of(((NonRegistryOverride) override).getRepoSpec(repositoryName));
+      }
+    }
+    return Optional.empty();
+  }
+
+  private static Optional<RepoSpec> checkRepoFromBazelModules(
+      SelectionValue selectionValue,
+      ImmutableMap<String, ModuleOverride> overrides,
+      ExtendedEventHandler eventlistener,
+      String repositoryName)
+      throws InterruptedException, IOException {
+    for (ModuleKey moduleKey : selectionValue.getDepGraph().keySet()) {
+      // TODO(pcloudy): Support multiple version override.
+      // Currently we assume there is only one version for each module, therefore the module name is
+      // the repository name, but that's not the case if multiple version of the same module are
+      // allowed.
+      if (moduleKey.getName().equals(repositoryName)) {
+        Module module = selectionValue.getDepGraph().get(moduleKey);
+        Registry registry = checkNotNull(module.getRegistry());
+        RepoSpec repoSpec = registry.getRepoSpec(moduleKey, repositoryName, eventlistener);
+        repoSpec = maybeAppendAdditionalPatches(repoSpec, overrides.get(moduleKey.getName()));
+        return Optional.of(repoSpec);
+      }
+    }
+    return Optional.empty();
+  }
+
+  private static RepoSpec maybeAppendAdditionalPatches(RepoSpec repoSpec, ModuleOverride override) {
+    if (!(override instanceof SingleVersionOverride)) {
+      return repoSpec;
+    }
+    SingleVersionOverride singleVersion = (SingleVersionOverride) override;
+    if (singleVersion.getPatches().isEmpty()) {
+      return repoSpec;
+    }
+    ImmutableMap.Builder<String, Object> attrBuilder = ImmutableMap.builder();
+    attrBuilder.putAll(repoSpec.attributes());
+    attrBuilder.put("patches", singleVersion.getPatches());
+    attrBuilder.put("patch_args", ImmutableList.of("-p" + singleVersion.getPatchStrip()));
+    return RepoSpec.builder()
+        .setBzlFile(repoSpec.bzlFile())
+        .setRuleClassName(repoSpec.ruleClassName())
+        .setAttributes(attrBuilder.build())
+        .build();
+  }
+
+  private static Optional<RepoSpec> checkRepoFromModuleRules() {
+    // TODO(pcloudy): Implement calculating RepoSpec from module rules.
     return Optional.empty();
   }
 }

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

@@ -17,6 +17,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 
 /** Specifies that a module should be retrieved from a Git repository. */
 @AutoValue
@@ -37,4 +38,21 @@ public static GitOverride create(
 
   /** The number of path segments to strip from the paths in the supplied patches. */
   public abstract int getPatchStrip();
+
+  /** Returns the {@link RepoSpec} that defines this repository. */
+  @Override
+  public RepoSpec getRepoSpec(String repoName) {
+    ImmutableMap.Builder<String, Object> attrBuilder = ImmutableMap.builder();
+    attrBuilder
+        .put("name", repoName)
+        .put("remote", getRemote())
+        .put("commit", getCommit())
+        .put("patches", getPatches())
+        .put("patch_args", ImmutableList.of("-p" + getPatchStrip()));
+    return RepoSpec.builder()
+        .setBzlFile("@bazel_tools//tools/build_defs/repo:git.bzl")
+        .setRuleClassName("git_repository")
+        .setAttributes(attrBuilder.build())
+        .build();
+  }
 }

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

@@ -16,6 +16,7 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
 
 /** Specifies that a module should be retrieved from a local directory. */
 @AutoValue
@@ -26,4 +27,13 @@ public static LocalPathOverride create(String path) {
 
   /** The path to the local directory where the module contents should be found. */
   public abstract String getPath();
+
+  /** Returns the {@link RepoSpec} that defines this repository. */
+  @Override
+  public RepoSpec getRepoSpec(String repoName) {
+    return RepoSpec.builder()
+        .setRuleClassName("local_repository")
+        .setAttributes(ImmutableMap.of("name", repoName, "path", getPath()))
+        .build();
+  }
 }

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

@@ -23,5 +23,6 @@
  */
 public interface NonRegistryOverride extends ModuleOverride {
 
-  // TODO(wyv): getRepoSpec
+  /** Returns the {@link RepoSpec} that defines this repository. */
+  RepoSpec getRepoSpec(String repoName);
 }

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

@@ -43,6 +43,8 @@ public static Builder builder() {
   public abstract static class Builder {
     public abstract Builder setBzlFile(String bzlFile);
 
+    public abstract Builder setBzlFile(Optional<String> bzlFile);
+
     public abstract Builder setRuleClassName(String name);
 
     public abstract Builder setAttributes(ImmutableMap<String, Object> attributes);

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

@@ -44,6 +44,7 @@
 import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.io.IOException;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.Module;
@@ -104,12 +105,16 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       return createRuleFromSpec(repoSpec, starlarkSemantics, env);
     }
 
-    Optional<RepoSpec> result = bzlmodRepoRuleHelper.getRepoSpec(env, repositoryName);
-    if (env.valuesMissing()) {
-      return null;
-    }
-    if (result.isPresent()) {
-      return createRuleFromSpec(result.get(), starlarkSemantics, env);
+    try {
+      Optional<RepoSpec> result = bzlmodRepoRuleHelper.getRepoSpec(env, repositoryName);
+      if (env.valuesMissing()) {
+        return null;
+      }
+      if (result.isPresent()) {
+        return createRuleFromSpec(result.get(), starlarkSemantics, env);
+      }
+    } catch (IOException e) {
+      throw new BzlmodRepoRuleFunctionException(e, Transience.PERSISTENT);
     }
 
     return BzlmodRepoRuleValue.REPO_RULE_NOT_FOUND_VALUE;
@@ -260,5 +265,9 @@ private static final class BzlmodRepoRuleFunctionException extends SkyFunctionEx
     BzlmodRepoRuleFunctionException(NoSuchPackageException e, Transience transience) {
       super(e, transience);
     }
+
+    BzlmodRepoRuleFunctionException(IOException e, Transience transience) {
+      super(e, transience);
+    }
   }
 }

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -51,8 +51,10 @@ java_library(
         "//src/main/java/net/starlark/java/eval",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
         "//src/test/java/com/google/devtools/build/lib/testutil",
+        "//third_party:auto_value",
         "//third_party:caffeine",
         "//third_party:guava",
+        "//third_party:jsr305",
         "//third_party:junit4",
         "//third_party:truth",
     ],