bzlmod: DiscoveryFunction

(#13316)

This SkyFunction performs module discovery by starting out from the root module and iteratively adding direct dependencies into the graph.

PiperOrigin-RevId: 380563418

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

@@ -34,6 +34,8 @@ java_library(
     name = "resolution",
     srcs = [
         "ArchiveOverride.java",
+        "DiscoveryFunction.java",
+        "DiscoveryValue.java",
         "GitOverride.java",
         "LocalPathOverride.java",
         "Module.java",

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

@@ -0,0 +1,105 @@
+// 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.ArrayDeque;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Queue;
+import javax.annotation.Nullable;
+
+/**
+ * Runs module discovery. This step of module resolution reads the module file of the root module
+ * (i.e. the current workspace), adds its direct {@code bazel_dep}s to the dependency graph, and
+ * repeats the step for any added dependencies until the entire graph is discovered.
+ */
+public class DiscoveryFunction implements SkyFunction {
+
+  @Nullable
+  @Override
+  public SkyValue compute(SkyKey skyKey, Environment env)
+      throws SkyFunctionException, InterruptedException {
+    ModuleFileValue root = (ModuleFileValue) env.getValue(ModuleFileValue.keyForRootModule());
+    if (root == null) {
+      return null;
+    }
+    ModuleKey rootModuleKey = ModuleKey.create(root.getModule().getName(), "");
+    ImmutableMap<String, ModuleOverride> overrides = root.getOverrides();
+    Map<ModuleKey, Module> depGraph = new HashMap<>();
+    depGraph.put(
+        rootModuleKey, rewriteDepKeys(root.getModule(), overrides, rootModuleKey.getName()));
+    Queue<ModuleKey> unexpanded = new ArrayDeque<>();
+    unexpanded.add(rootModuleKey);
+    // TODO(wyv): currently we expand the "unexpanded" keys one by one. We should try to expand them
+    //   all at once, using `env.getValues`.
+    while (!unexpanded.isEmpty()) {
+      Module module = depGraph.get(unexpanded.remove());
+      for (ModuleKey depKey : module.getDeps().values()) {
+        if (depGraph.containsKey(depKey)) {
+          continue;
+        }
+        ModuleFileValue dep =
+            (ModuleFileValue)
+                env.getValue(ModuleFileValue.key(depKey, overrides.get(depKey.getName())));
+        if (dep == null) {
+          // Don't return yet. Try to expand any other unexpanded nodes before returning.
+          depGraph.put(depKey, null);
+        } else {
+          depGraph.put(depKey, rewriteDepKeys(dep.getModule(), overrides, rootModuleKey.getName()));
+          unexpanded.add(depKey);
+        }
+      }
+    }
+    if (env.valuesMissing()) {
+      return null;
+    }
+    return DiscoveryValue.create(
+        root.getModule().getName(), ImmutableMap.copyOf(depGraph), overrides);
+  }
+
+  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();
+
+      @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();
+  }
+
+  @Nullable
+  @Override
+  public String extractTag(SkyKey skyKey) {
+    return null;
+  }
+}

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

@@ -0,0 +1,43 @@
+// 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.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.skyframe.SkyFunctions;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+
+/** The result of discovery, containing the dependency graph and overrides. */
+@AutoValue
+public abstract class DiscoveryValue implements SkyValue {
+
+  @AutoCodec public static final SkyKey KEY = () -> SkyFunctions.DISCOVERY;
+
+  public static DiscoveryValue create(
+      String rootModuleName,
+      ImmutableMap<ModuleKey, Module> depGraph,
+      ImmutableMap<String, ModuleOverride> overrides) {
+    return new AutoValue_DiscoveryValue(rootModuleName, depGraph, overrides);
+  }
+
+  public abstract String getRootModuleName();
+
+  public abstract ImmutableMap<ModuleKey, Module> getDepGraph();
+
+  public abstract ImmutableMap<String, ModuleOverride> getOverrides();
+}

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

@@ -149,6 +149,7 @@ public final class SkyFunctions {
       SkyFunctionName.createHermetic("LOCAL_CONFIG_PLATFORM");
   public static final SkyFunctionName MODULE_FILE =
       SkyFunctionName.createNonHermetic("MODULE_FILE");
+  public static final SkyFunctionName DISCOVERY = SkyFunctionName.createHermetic("DISCOVERY");
 
   public static Predicate<SkyKey> isSkyFunction(final SkyFunctionName functionName) {
     return new Predicate<SkyKey>() {