bzlmod: SelectionFunction

(#13316)

This SkyFunction takes the result of DiscoveryFunction and runs MVS on it, resulting in a dep graph that contains only 1 version of each dependency (the highest requested one).

PiperOrigin-RevId: 381844359

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

@@ -50,7 +50,10 @@ java_library(
         "ModuleFileValue.java",
         "ModuleOverride.java",
         "NonRegistryOverride.java",
+        "ParsedVersion.java",
         "RegistryOverride.java",
+        "SelectionFunction.java",
+        "SelectionValue.java",
         "SingleVersionOverride.java",
     ],
     deps = [

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

@@ -77,24 +77,22 @@ public SkyValue compute(SkyKey skyKey, Environment env)
 
   private static Module rewriteDepKeys(
       Module module, ImmutableMap<String, ModuleOverride> overrides, String rootModuleName) {
-    ImmutableMap.Builder<String, ModuleKey> newDeps = new ImmutableMap.Builder<>();
-    for (Map.Entry<String, ModuleKey> entry : module.getDeps().entrySet()) {
-      ModuleKey depKey = entry.getValue();
-      String newVersion = depKey.getVersion();
+    return module.withDepKeysTransformed(
+        depKey -> {
+          String newVersion = depKey.getVersion();
 
-      @Nullable ModuleOverride override = overrides.get(depKey.getName());
-      if (override instanceof NonRegistryOverride || rootModuleName.equals(depKey.getName())) {
-        newVersion = "";
-      } else if (override instanceof SingleVersionOverride) {
-        String overrideVersion = ((SingleVersionOverride) override).getVersion();
-        if (!overrideVersion.isEmpty()) {
-          newVersion = overrideVersion;
-        }
-      }
+          @Nullable ModuleOverride override = overrides.get(depKey.getName());
+          if (override instanceof NonRegistryOverride || rootModuleName.equals(depKey.getName())) {
+            newVersion = "";
+          } else if (override instanceof SingleVersionOverride) {
+            String overrideVersion = ((SingleVersionOverride) override).getVersion();
+            if (!overrideVersion.isEmpty()) {
+              newVersion = overrideVersion;
+            }
+          }
 
-      newDeps.put(entry.getKey(), ModuleKey.create(depKey.getName(), newVersion));
-    }
-    return module.toBuilder().setDeps(newDeps.build()).build();
+          return ModuleKey.create(depKey.getName(), newVersion);
+        });
   }
 
   @Nullable

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

@@ -17,6 +17,8 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
+import java.util.Map;
+import java.util.function.UnaryOperator;
 import javax.annotation.Nullable;
 
 /**
@@ -52,10 +54,22 @@ public abstract class Module {
   public abstract Builder toBuilder();
 
   /** Returns a new, empty {@link Builder}. */
-  static Builder builder() {
+  public static Builder builder() {
     return new AutoValue_Module.Builder();
   }
 
+  /**
+   * Returns a new {@link Module} with all values in {@link #getDeps} transformed using the given
+   * function.
+   */
+  public Module withDepKeysTransformed(UnaryOperator<ModuleKey> transform) {
+    ImmutableMap.Builder<String, ModuleKey> newDeps = new ImmutableMap.Builder<>();
+    for (Map.Entry<String, ModuleKey> entry : getDeps().entrySet()) {
+      newDeps.put(entry.getKey(), transform.apply(entry.getValue()));
+    }
+    return toBuilder().setDeps(newDeps.build()).build();
+  }
+
   /** Builder type for {@link Module}. */
   @AutoValue.Builder
   public abstract static class Builder {

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

@@ -0,0 +1,182 @@
+// 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.collect.Comparators.lexicographical;
+import static com.google.common.primitives.Booleans.falseFirst;
+import static com.google.common.primitives.Booleans.trueFirst;
+import static java.util.Comparator.comparing;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.base.Splitter;
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
+import java.util.Comparator;
+import java.util.Objects;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nullable;
+
+/**
+ * Represents a parsed version string, useful for comparison. The version format we support is
+ * {@code RELEASE[-PRERELEASE][+BUILD]}, where:
+ *
+ * <ul>
+ *   <li>{@code RELEASE} is a sequence of decimal numbers separated by dots;
+ *   <li>{@code PRERELEASE} is a sequence of "identifiers" (defined as a non-empty sequence of
+ *       alphanumerical characters, hyphens, and underscores) separated by dots;
+ *   <li>and {@code BUILD} is also a sequence of "identifiers" (see above) separated by dots.
+ * </ul>
+ *
+ * Otherwise, this format is identical to SemVer, especially in terms of the comparison algorithm
+ * (https://semver.org/#spec-item-11). In other words, this format is intentionally looser than
+ * SemVer; in particular, the "release" part isn't limited to exactly 3 numbers (major, minor,
+ * patch), but can be fewer or more. Underscores are also allowed in prerelease and build for regex
+ * brevity.
+ *
+ * <p>The special "empty string" version can also be used, and compares higher than everything else.
+ * It signifies that there is a {@link NonRegistryOverride} for a module.
+ */
+@AutoValue
+abstract class ParsedVersion implements Comparable<ParsedVersion> {
+
+  // We don't care about the "build" part at all so don't capture it.
+  private static final Pattern PATTERN =
+      Pattern.compile("(?<release>(?:\\d+\\.)*\\d+)(?:-(?<prerelease>[\\w.-]*))?(?:\\+[\\w.-]*)?");
+
+  private static final Splitter DOT_SPLITTER = Splitter.on('.');
+
+  /**
+   * Represents a segment in the prerelease part of the version string. This is separated from other
+   * "Identifier"s by a dot. An identifier is compared differently based on whether it's digits-only
+   * or not.
+   */
+  @AutoValue
+  abstract static class Identifier {
+
+    abstract boolean isDigitsOnly();
+
+    abstract int asNumber();
+
+    abstract String asString();
+
+    static Identifier from(String string) throws ParseException {
+      if (Strings.isNullOrEmpty(string)) {
+        throw new ParseException("identifier is empty");
+      }
+      if (string.chars().allMatch(Character::isDigit)) {
+        return new AutoValue_ParsedVersion_Identifier(true, Integer.parseInt(string), string);
+      } else {
+        return new AutoValue_ParsedVersion_Identifier(false, 0, string);
+      }
+    }
+  }
+
+  /** Returns the "release" part of the version string as a list of integers. */
+  abstract ImmutableList<Integer> getRelease();
+
+  /** Returns the "prerelease" part of the version string as a list of {@link Identifier}s. */
+  abstract ImmutableList<Identifier> getPrerelease();
+
+  /** Returns the original version string. */
+  public abstract String getOriginal();
+
+  /**
+   * Whether this is just the "empty string" version, which signifies a non-registry override for
+   * the module.
+   */
+  boolean isOverride() {
+    return getRelease().isEmpty();
+  }
+
+  /**
+   * Whether this is a prerelease version (i.e. the prerelease part of the version string is
+   * non-empty). A prerelease version compares lower than the same version without the prerelease
+   * part.
+   */
+  boolean isPrerelease() {
+    return !getPrerelease().isEmpty();
+  }
+
+  /** Parses a version string into a {@link ParsedVersion} object. */
+  public static ParsedVersion parse(String version) throws ParseException {
+    if (version.isEmpty()) {
+      return new AutoValue_ParsedVersion(ImmutableList.of(), ImmutableList.of(), version);
+    }
+    Matcher matcher = PATTERN.matcher(version);
+    if (!matcher.matches()) {
+      throw new ParseException("bad version (does not match regex): " + version);
+    }
+    String release = matcher.group("release");
+    @Nullable String prerelease = matcher.group("prerelease");
+
+    ImmutableList.Builder<Integer> releaseSplit = new ImmutableList.Builder<>();
+    for (String number : DOT_SPLITTER.split(release)) {
+      try {
+        releaseSplit.add(Integer.valueOf(number));
+      } catch (NumberFormatException e) {
+        throw new ParseException("error parsing version: " + version, e);
+      }
+    }
+
+    ImmutableList.Builder<Identifier> prereleaseSplit = new ImmutableList.Builder<>();
+    if (!Strings.isNullOrEmpty(prerelease)) {
+      for (String ident : DOT_SPLITTER.split(prerelease)) {
+        try {
+          prereleaseSplit.add(Identifier.from(ident));
+        } catch (ParseException e) {
+          throw new ParseException("error parsing version: " + version, e);
+        }
+      }
+    }
+
+    return new AutoValue_ParsedVersion(releaseSplit.build(), prereleaseSplit.build(), version);
+  }
+
+  private static final Comparator<ParsedVersion> COMPARATOR =
+      Comparator.nullsFirst(
+          comparing(ParsedVersion::isOverride, falseFirst())
+              .thenComparing(
+                  ParsedVersion::getRelease, lexicographical(Comparator.<Integer>naturalOrder()))
+              .thenComparing(ParsedVersion::isPrerelease, trueFirst())
+              .thenComparing(
+                  ParsedVersion::getPrerelease,
+                  lexicographical(
+                      comparing(Identifier::isDigitsOnly, trueFirst())
+                          .thenComparingInt(Identifier::asNumber)
+                          .thenComparing(Identifier::asString))));
+
+  @Override
+  public int compareTo(ParsedVersion o) {
+    return Objects.compare(this, o, COMPARATOR);
+  }
+
+  /** Returns the higher of two versions. */
+  public static ParsedVersion max(@Nullable ParsedVersion a, @Nullable ParsedVersion b) {
+    return Objects.compare(a, b, COMPARATOR) >= 0 ? a : b;
+  }
+
+  /** An exception encountered while trying to {@link ParsedVersion#parse parse} a version. */
+  public static class ParseException extends Exception {
+    public ParseException(String message) {
+      super(message);
+    }
+
+    public ParseException(String message, Throwable cause) {
+      super(message, cause);
+    }
+  }
+}

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

@@ -0,0 +1,111 @@
+// 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 com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyFunctionException;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Runs module selection. This step of module resolution reads the output of {@link
+ * DiscoveryFunction} and applies the Minimal Version Selection algorithm to it, removing unselected
+ * modules from the dependency graph and rewriting dependencies to point to the selected versions.
+ */
+public class SelectionFunction implements SkyFunction {
+
+  @Override
+  public SkyValue compute(SkyKey skyKey, Environment env)
+      throws SkyFunctionException, InterruptedException {
+    DiscoveryValue discovery = (DiscoveryValue) env.getValue(DiscoveryValue.KEY);
+    if (discovery == null) {
+      return null;
+    }
+
+    // TODO(wyv): compatibility_level, multiple_version_override
+
+    // First figure out the version to select for every module.
+    ImmutableMap<ModuleKey, Module> depGraph = discovery.getDepGraph();
+    Map<String, ParsedVersion> selectedVersionForEachModule = new HashMap<>();
+    for (ModuleKey key : depGraph.keySet()) {
+      try {
+        ParsedVersion parsedVersion = ParsedVersion.parse(key.getVersion());
+        selectedVersionForEachModule.merge(key.getName(), parsedVersion, ParsedVersion::max);
+      } catch (ParsedVersion.ParseException e) {
+        throw new SelectionFunctionException(e);
+      }
+    }
+
+    // Now build a new dep graph where deps with unselected versions are removed.
+    ImmutableMap.Builder<ModuleKey, Module> newDepGraphBuilder = new ImmutableMap.Builder<>();
+    for (Map.Entry<ModuleKey, Module> entry : depGraph.entrySet()) {
+      ModuleKey moduleKey = entry.getKey();
+      Module module = entry.getValue();
+      // Remove any dep whose version isn't selected.
+      String selectedVersion = selectedVersionForEachModule.get(moduleKey.getName()).getOriginal();
+      if (!moduleKey.getVersion().equals(selectedVersion)) {
+        continue;
+      }
+
+      // Rewrite deps to point to the selected version.
+      newDepGraphBuilder.put(
+          moduleKey,
+          module.withDepKeysTransformed(
+              depKey ->
+                  ModuleKey.create(
+                      depKey.getName(),
+                      selectedVersionForEachModule.get(depKey.getName()).getOriginal())));
+    }
+    ImmutableMap<ModuleKey, Module> newDepGraph = newDepGraphBuilder.build();
+
+    // Further remove unreferenced modules from the graph. We can find out which modules are
+    // referenced by collecting deps transitively from the root.
+    HashMap<ModuleKey, Module> finalDepGraph = new HashMap<>();
+    collectDeps(ModuleKey.create(discovery.getRootModuleName(), ""), newDepGraph, finalDepGraph);
+    return SelectionValue.create(
+        discovery.getRootModuleName(),
+        ImmutableMap.copyOf(finalDepGraph),
+        discovery.getOverrides());
+  }
+
+  private void collectDeps(
+      ModuleKey key,
+      ImmutableMap<ModuleKey, Module> oldDepGraph,
+      HashMap<ModuleKey, Module> newDepGraph) {
+    if (newDepGraph.containsKey(key)) {
+      return;
+    }
+    Module module = oldDepGraph.get(key);
+    newDepGraph.put(key, module);
+    for (ModuleKey depKey : module.getDeps().values()) {
+      collectDeps(depKey, oldDepGraph, newDepGraph);
+    }
+  }
+
+  @Override
+  public String extractTag(SkyKey skyKey) {
+    return null;
+  }
+
+  static final class SelectionFunctionException extends SkyFunctionException {
+    SelectionFunctionException(Exception cause) {
+      super(cause, Transience.PERSISTENT);
+    }
+  }
+}