BazelModuleInspectorValue SkyValue added

- `BazelModuleInspectorValue` - `SkyValue` which stores:
  - an unpruned dep graph with `ModuleAugment` wrapper nodes which give information about dependants and resolution rules applied
  - a `modulesIndex` map from *module name* to the corresponding set of `ModuleKeys`

- `BazelModuleInspectorFunction` - `SkyFunction` computes the above based on the information inside `BazelModuleResolutionValue`

- `BazelModuleInspectorFunctionTest` - UnitTests for the `computeAugmentedGraph` core method of the inspector function

#15365

PiperOrigin-RevId: 448029252

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

@@ -215,3 +215,21 @@ java_library(
         "//third_party:guava",
     ],
 )
+
+java_library(
+    name = "inspection",
+    srcs = [
+        "BazelModuleInspectorFunction.java",
+        "BazelModuleInspectorValue.java",
+    ],
+    deps = [
+        ":common",
+        ":resolution",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant",
+        "//src/main/java/com/google/devtools/build/skyframe",
+        "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
+        "//third_party:auto_value",
+        "//third_party:guava",
+    ],
+)

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

@@ -0,0 +1,154 @@
+// Copyright 2022 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.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleInspectorValue.AugmentedModule;
+import com.google.devtools.build.lib.bazel.bzlmod.BazelModuleInspectorValue.AugmentedModule.ResolutionReason;
+import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
+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;
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
+
+/**
+ * Precomputes an augmented version of the un-pruned dep graph that is used for dep graph
+ * inspection. By this stage, the Bazel module resolution should have been completed.
+ */
+public class BazelModuleInspectorFunction implements SkyFunction {
+
+  @Override
+  public SkyValue compute(SkyKey skyKey, Environment env)
+      throws SkyFunctionException, InterruptedException {
+    RootModuleFileValue root =
+        (RootModuleFileValue) env.getValue(ModuleFileValue.KEY_FOR_ROOT_MODULE);
+    if (root == null) {
+      return null;
+    }
+    BazelModuleResolutionValue resolutionValue =
+        (BazelModuleResolutionValue) env.getValue(BazelModuleResolutionValue.KEY);
+    if (resolutionValue == null) {
+      return null;
+    }
+    ImmutableMap<String, ModuleOverride> overrides = root.getOverrides();
+    ImmutableMap<ModuleKey, Module> unprunedDepGraph = resolutionValue.getUnprunedDepGraph();
+    ImmutableMap<ModuleKey, Module> resolvedDepGraph = resolutionValue.getDepGraph();
+
+    ImmutableMap<ModuleKey, AugmentedModule> depGraph =
+        computeAugmentedGraph(unprunedDepGraph, resolvedDepGraph.keySet(), overrides);
+
+    // Group all ModuleKeys seen by their module name for easy lookup
+    ImmutableMap<String, ImmutableSet<ModuleKey>> modulesIndex =
+        ImmutableMap.copyOf(
+            depGraph.values().stream()
+                .collect(
+                    Collectors.groupingBy(
+                        AugmentedModule::getName,
+                        Collectors.mapping(AugmentedModule::getKey, toImmutableSet()))));
+
+    return BazelModuleInspectorValue.create(depGraph, modulesIndex);
+  }
+
+  public static ImmutableMap<ModuleKey, AugmentedModule> computeAugmentedGraph(
+      ImmutableMap<ModuleKey, Module> unprunedDepGraph,
+      ImmutableSet<ModuleKey> usedModules,
+      ImmutableMap<String, ModuleOverride> overrides) {
+    Map<ModuleKey, AugmentedModule.Builder> depGraphAugmentBuilder = new HashMap<>();
+
+    // For all Modules in the un-pruned dep graph, inspect their dependencies and add themselves
+    // to their children AugmentedModule as dependant. Also fill in their own AugmentedModule
+    // with a map from their dependencies to the resolution reason that was applied to each.
+    // The newly created graph will also contain ModuleAugments for non-loaded modules.
+    for (Entry<ModuleKey, Module> e : unprunedDepGraph.entrySet()) {
+      ModuleKey parentKey = e.getKey();
+      Module parentModule = e.getValue();
+
+      AugmentedModule.Builder parentBuilder =
+          depGraphAugmentBuilder
+              .computeIfAbsent(
+                  parentKey, k -> AugmentedModule.builder(k).setName(parentModule.getName()))
+              .setVersion(parentModule.getVersion())
+              .setLoaded(true);
+
+      for (String childDep : parentModule.getDeps().keySet()) {
+        ModuleKey originalKey = parentModule.getOriginalDeps().get(childDep);
+        Module originalModule = unprunedDepGraph.get(originalKey);
+        ModuleKey key = parentModule.getDeps().get(childDep);
+        Module module = unprunedDepGraph.get(key);
+
+        AugmentedModule.Builder originalChildBuilder =
+            depGraphAugmentBuilder.computeIfAbsent(originalKey, AugmentedModule::builder);
+        if (originalModule != null) {
+          originalChildBuilder
+              .setName(originalModule.getName())
+              .setVersion(originalModule.getVersion())
+              .setLoaded(true);
+        }
+
+        AugmentedModule.Builder newChildBuilder =
+            depGraphAugmentBuilder.computeIfAbsent(
+                key,
+                k ->
+                    AugmentedModule.builder(k)
+                        .setName(module.getName())
+                        .setVersion(module.getVersion())
+                        .setLoaded(true));
+
+        // originalDependants and dependants can differ because
+        // parentModule could have had originalChild in the unresolved graph, but in the resolved
+        // graph the originalChild could have become orphan due to an override or selection
+        originalChildBuilder.addOriginalDependant(parentKey);
+        // also, even if the dep has not changed, the parentModule may not be referenced
+        // anymore in the resolved graph, so parentModule will only be added above
+        if (usedModules.contains(parentKey)) {
+          newChildBuilder.addDependant(parentKey);
+        }
+
+        ResolutionReason reason = ResolutionReason.ORIGINAL;
+        if (!key.getVersion().equals(originalKey.getVersion())) {
+          ModuleOverride override = overrides.get(key.getName());
+          if (override != null) {
+            if (override instanceof SingleVersionOverride) {
+              reason = ResolutionReason.SINGLE_VERSION_OVERRIDE;
+            } else if (override instanceof MultipleVersionOverride) {
+              reason = ResolutionReason.MULTIPLE_VERSION_OVERRIDE;
+            } else {
+              // There is no other possible override
+              Preconditions.checkArgument(override instanceof NonRegistryOverride);
+              reason = ResolutionReason.NON_REGISTRY_OVERRIDE;
+            }
+          } else {
+            reason = ResolutionReason.MINIMAL_VERSION_SELECTION;
+          }
+        }
+
+        parentBuilder.addDep(key, reason);
+      }
+    }
+
+    return depGraphAugmentBuilder.entrySet().stream()
+        .collect(toImmutableMap(Entry::getKey, e -> e.getValue().build()));
+  }
+}

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

@@ -0,0 +1,168 @@
+// Copyright 2022 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.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.skyframe.SkyFunctions;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+
+/**
+ * The result of running Bazel module inspection pre-processing, containing the un-pruned and
+ * augmented wrappers of the Bazel module dependency graph (post-version-resolution).
+ */
+@AutoValue
+public abstract class BazelModuleInspectorValue implements SkyValue {
+
+  @SerializationConstant
+  public static final SkyKey KEY = () -> SkyFunctions.BAZEL_MODULE_INSPECTION;
+
+  public static BazelModuleInspectorValue create(
+      ImmutableMap<ModuleKey, AugmentedModule> depGraph,
+      ImmutableMap<String, ImmutableSet<ModuleKey>> modulesIndex) {
+    return new AutoValue_BazelModuleInspectorValue(depGraph, modulesIndex);
+  }
+
+  /**
+   * The (bidirectional) inspection dep graph, containing wrappers of the {@link Module}, augmented
+   * with references to dependants. The order is non-deterministic, inherited from the {@code
+   * completeDepGraph} of {@link BazelModuleResolutionValue}. For any KEY in the returned map, it's
+   * guaranteed that {@code depGraph[KEY].getKey() == KEY}.
+   */
+  public abstract ImmutableMap<ModuleKey, AugmentedModule> getDepGraph();
+
+  /**
+   * Index of all module keys mentioned in the un-pruned dep graph (loaded or not) for easy lookup.
+   * It is a map from <i>module name</i> to the set of {@link ModuleKey}s that point to a version of
+   * that module.
+   */
+  public abstract ImmutableMap<String, ImmutableSet<ModuleKey>> getModulesIndex();
+
+  /**
+   * A wrapper for {@link Module}, augmented with references to dependants (and also those who are
+   * not used in the final dep graph).
+   */
+  @AutoValue
+  abstract static class AugmentedModule {
+    /** Name of the module. Same as in {@link Module}. */
+    abstract String getName();
+
+    /** Version of the module. Same as in {@link Module}. */
+    abstract Version getVersion();
+
+    /** {@link ModuleKey} of this module. Same as in {@link Module} */
+    abstract ModuleKey getKey();
+
+    /**
+     * The set of modules in the resolved dep graph that depend on this module
+     * <strong>after</strong> the module resolution.
+     */
+    abstract ImmutableSet<ModuleKey> getDependants();
+
+    /**
+     * The set of modules in the complete dep graph that originally depended on this module *before*
+     * the module resolution (can contain unused nodes).
+     */
+    abstract ImmutableSet<ModuleKey> getOriginalDependants();
+
+    /**
+     * A map from the resolved dependencies of this module to the rules that were used for their
+     * resolution (can be either the original dependency, changed by the Minimal-Version Selection
+     * algorithm or by an override rule
+     */
+    abstract ImmutableMap<ModuleKey, ResolutionReason> getDeps();
+
+    /**
+     * Flag that tell whether the module was loaded and added to the dependency graph. Modules
+     * overridden by {@code single_version_override} and {@link NonRegistryOverride} are not loaded
+     * so their {@code originalDeps} are (yet) unknown.
+     */
+    abstract boolean isLoaded();
+
+    /** Flag for checking whether the module is present in the resolved dep graph. */
+    boolean isUsed() {
+      return !getDependants().isEmpty();
+    }
+
+    /** Returns a new {@link AugmentedModule.Builder} with {@code key} set. */
+    public static AugmentedModule.Builder builder(ModuleKey key) {
+      return new AutoValue_BazelModuleInspectorValue_AugmentedModule.Builder()
+          .setName(key.getName())
+          .setVersion(key.getVersion())
+          .setKey(key)
+          .setLoaded(false);
+    }
+
+    /** Builder type for {@link AugmentedModule}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+      public abstract AugmentedModule.Builder setName(String value);
+
+      public abstract AugmentedModule.Builder setVersion(Version value);
+
+      public abstract AugmentedModule.Builder setKey(ModuleKey value);
+
+      public abstract AugmentedModule.Builder setLoaded(boolean value);
+
+      public abstract AugmentedModule.Builder setOriginalDependants(ImmutableSet<ModuleKey> value);
+
+      public abstract AugmentedModule.Builder setDependants(ImmutableSet<ModuleKey> value);
+
+      public abstract AugmentedModule.Builder setDeps(
+          ImmutableMap<ModuleKey, ResolutionReason> value);
+
+      abstract ImmutableSet.Builder<ModuleKey> originalDependantsBuilder();
+
+      public AugmentedModule.Builder addOriginalDependant(ModuleKey depKey) {
+        originalDependantsBuilder().add(depKey);
+        return this;
+      }
+
+      abstract ImmutableSet.Builder<ModuleKey> dependantsBuilder();
+
+      public AugmentedModule.Builder addDependant(ModuleKey depKey) {
+        dependantsBuilder().add(depKey);
+        return this;
+      }
+
+      abstract ImmutableMap.Builder<ModuleKey, ResolutionReason> depsBuilder();
+
+      public AugmentedModule.Builder addDep(ModuleKey depKey, ResolutionReason reason) {
+        depsBuilder().put(depKey, reason);
+        return this;
+      }
+
+      abstract AugmentedModule build();
+    }
+
+    /** The reason why a final dependency of a module was resolved the way it was. */
+    enum ResolutionReason {
+      /** The dependency is the original dependency defined in the MODULE.bazel file. */
+      ORIGINAL,
+      /** The dependency was replaced by the Minimal-Version Selection algorithm. */
+      MINIMAL_VERSION_SELECTION,
+      /** The dependency was replaced by a {@code single_version_override} rule. */
+      SINGLE_VERSION_OVERRIDE,
+      /** The dependency was replaced by a {@code multiple_version_override} rule. */
+      MULTIPLE_VERSION_OVERRIDE,
+      /** The dependency was replaced by a {@link NonRegistryOverride} rule. */
+      NON_REGISTRY_OVERRIDE
+    }
+  }
+}

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

@@ -141,6 +141,8 @@ public final class SkyFunctions {
       SkyFunctionName.createNonHermetic("BUILD_DRIVER");
   public static final SkyFunctionName BAZEL_MODULE_RESOLUTION =
       SkyFunctionName.createHermetic("BAZEL_MODULE_RESOLUTION");
+  public static final SkyFunctionName BAZEL_MODULE_INSPECTION =
+      SkyFunctionName.createHermetic("BAZEL_MODULE_INSPECTION");
   public static final SkyFunctionName MODULE_EXTENSION_RESOLUTION =
       SkyFunctionName.createHermetic("MODULE_EXTENSION_RESOLUTION");
   public static final SkyFunctionName SINGLE_EXTENSION_USAGES =

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

@@ -32,6 +32,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:server_directories",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:common",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:exception",
+        "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:inspection",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:module_extension",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:registry",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:repo_rule_helper",