Fix two bazel mod tidy crashes

* Fixes a crash when using a non-registry override with a specified version:

```
Caused by: java.lang.NullPointerException: null value in entry: foo=null
	at com.google.common.collect.CollectPreconditions.checkEntryNotNull(CollectPreconditions.java:33)
	at com.google.common.collect.ImmutableMapEntry.<init>(ImmutableMapEntry.java:54)
	at com.google.common.collect.ImmutableMap.entryOf(ImmutableMap.java:345)
	at com.google.common.collect.ImmutableMap$Builder.put(ImmutableMap.java:454)
	at com.google.devtools.build.lib.bazel.bzlmod.Module.getRepoMappingWithBazelDepsOnly(Module.java:67)
	at com.google.devtools.build.lib.bazel.bzlmod.BazelDepGraphFunction.getExtensionUsagesById(BazelDepGraphFunction.java:233)
	at com.google.devtools.build.lib.bazel.bzlmod.BazelDepGraphFunction.compute(BazelDepGraphFunction.java:126)
	at com.google.devtools.build.skyframe.AbstractParallelEvaluator$Evaluate.run(AbstractParallelEvaluator.java:464)
```

* Fixes a crash for root modules with no extension usages:

```
Caused by: java.lang.NullPointerException: Cannot invoke "java.util.Map.keySet()" because the return value of "com.google.common.collect.ImmutableMap.get(Object)" is null
	at com.google.devtools.build.lib.bazel.bzlmod.BazelModTidyFunction.compute(BazelModTidyFunction.java:85)
	at com.google.devtools.build.skyframe.AbstractParallelEvaluator$Evaluate.run(AbstractParallelEvaluator.java:464)
```

Fixes #21651

Closes #21686.

PiperOrigin-RevId: 615826860
Change-Id: I22be3fd53d0dc97aec92afe3dc51a9d6b7e60c98

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

@@ -132,7 +132,14 @@ public BazelLockFileValue withShallowlyReplacedRootModule(
  newDepGraph.putAll(getModuleDepGraph());
  newDepGraph.put(
  ModuleKey.ROOT,
- toModule(value.getModule(), /* override= */ null, /* remoteRepoSpec= */ null));
+ toModule(
+ value
+ .getModule()
+ .withDepSpecsTransformed(
+ InterimModule.applyOverrides(
+ value.getOverrides(), value.getModule().getName())),
+ /* override= */ null,
+ /* remoteRepoSpec= */ null));
  return toBuilder()
  .setModuleFileHash(value.getModuleFileHash())
  .setModuleDepGraph(newDepGraph.buildKeepingLast())

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

@@ -21,6 +21,7 @@
 import static com.google.devtools.build.lib.bazel.bzlmod.ModuleFileFunction.MODULE_OVERRIDES;
 import static com.google.devtools.build.lib.skyframe.PrecomputedValue.STARLARK_SEMANTICS;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
@@ -81,7 +82,12 @@ public SkyValue compute(SkyKey skyKey, Environment env)
  }
 
  ImmutableSet<SkyKey> extensionsUsedByRootModule =
- depGraphValue.getExtensionUsagesTable().columnMap().get(ModuleKey.ROOT).keySet().stream()
+ depGraphValue
+ .getExtensionUsagesTable()
+ .columnMap()
+ .getOrDefault(ModuleKey.ROOT, ImmutableMap.of())
+ .keySet()
+ .stream()
  .map(SingleExtensionEvalValue::key)
  .collect(toImmutableSet());
  SkyframeLookupResult result = env.getValuesAndExceptions(extensionsUsedByRootModule);

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

@@ -54,7 +54,10 @@ public static ImmutableMap<ModuleKey, InterimModule> run(
  String rootModuleName = root.getModule().getName();
  ImmutableMap<String, ModuleOverride> overrides = root.getOverrides();
  Map<ModuleKey, InterimModule> depGraph = new HashMap<>();
- depGraph.put(ModuleKey.ROOT, rewriteDepSpecs(root.getModule(), overrides, rootModuleName));
+ depGraph.put(
+ ModuleKey.ROOT,
+ root.getModule()
+ .withDepSpecsTransformed(InterimModule.applyOverrides(overrides, rootModuleName)));
  Queue<ModuleKey> unexpanded = new ArrayDeque<>();
  Map<ModuleKey, ModuleKey> predecessors = new HashMap<>();
  unexpanded.add(ModuleKey.ROOT);
@@ -101,7 +104,11 @@ public static ImmutableMap<ModuleKey, InterimModule> run(
  depGraph.put(depKey, null);
  } else {
  depGraph.put(
- depKey, rewriteDepSpecs(moduleFileValue.getModule(), overrides, rootModuleName));
+ depKey,
+ moduleFileValue
+ .getModule()
+ .withDepSpecsTransformed(
+ InterimModule.applyOverrides(overrides, rootModuleName)));
  unexpanded.add(depKey);
  }
  }
@@ -111,27 +118,4 @@ public static ImmutableMap<ModuleKey, InterimModule> run(
  }
  return ImmutableMap.copyOf(depGraph);
  }
-
- private static InterimModule rewriteDepSpecs(
- InterimModule module, ImmutableMap<String, ModuleOverride> overrides, String rootModuleName) {
- return module.withDepSpecsTransformed(
- depSpec -> {
- if (rootModuleName.equals(depSpec.getName())) {
- return DepSpec.fromModuleKey(ModuleKey.ROOT);
- }
-
- Version newVersion = depSpec.getVersion();
- @Nullable ModuleOverride override = overrides.get(depSpec.getName());
- if (override instanceof NonRegistryOverride) {
- newVersion = Version.EMPTY;
- } else if (override instanceof SingleVersionOverride) {
- Version overrideVersion = ((SingleVersionOverride) override).getVersion();
- if (!overrideVersion.isEmpty()) {
- newVersion = overrideVersion;
- }
- }
-
- return DepSpec.create(depSpec.getName(), newVersion, depSpec.getMaxCompatibilityLevel());
- });
- }
 }

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

@@ -260,4 +260,26 @@ private static RepoSpec maybeAppendAdditionalPatches(
  .setAttributes(AttributeValues.create(attrBuilder.buildOrThrow()))
  .build();
  }
+
+ static UnaryOperator<DepSpec> applyOverrides(
+ ImmutableMap<String, ModuleOverride> overrides, String rootModuleName) {
+ return depSpec -> {
+ if (rootModuleName.equals(depSpec.getName())) {
+ return DepSpec.fromModuleKey(ModuleKey.ROOT);
+ }
+
+ Version newVersion = depSpec.getVersion();
+ @Nullable ModuleOverride override = overrides.get(depSpec.getName());
+ if (override instanceof NonRegistryOverride) {
+ newVersion = Version.EMPTY;
+ } else if (override instanceof SingleVersionOverride) {
+ Version overrideVersion = ((SingleVersionOverride) override).getVersion();
+ if (!overrideVersion.isEmpty()) {
+ newVersion = overrideVersion;
+ }
+ }
+
+ return DepSpec.create(depSpec.getName(), newVersion, depSpec.getMaxCompatibilityLevel());
+ };
+ }
 }

src/test/py/bazel/bzlmod/mod_command_test.py

@@ -800,6 +800,76 @@ def testModTidyNoop(self):
  module_file.read().split('\n'),
  )
 
+ def testModTidyWithNonRegistryOverride(self):
+ self.ScratchFile(
+ 'MODULE.bazel',
+ [
+ 'bazel_dep(name = "foo", version = "1.2.3")',
+ 'local_path_override(module_name = "foo", path = "foo")',
+ 'ext = use_extension("//:extension.bzl", "ext")',
+ 'use_repo(ext, "dep")',
+ ],
+ )
+ self.ScratchFile('BUILD.bazel')
+ self.ScratchFile(
+ 'extension.bzl',
+ [
+ 'def _ext_impl(ctx):',
+ ' pass',
+ '',
+ 'ext = module_extension(implementation=_ext_impl)',
+ ],
+ )
+ self.ScratchFile(
+ 'foo/MODULE.bazel', ['module(name = "foo", version = "1.2.3")']
+ )
+
+ # Verify that bazel mod tidy doesn't fail without the lockfile.
+ self.RunBazel(['mod', 'tidy'])
+
+ with open('MODULE.bazel', 'r') as module_file:
+ self.assertEqual(
+ [
+ 'bazel_dep(name = "foo", version = "1.2.3")',
+ 'local_path_override(',
+ ' module_name = "foo",',
+ ' path = "foo",',
+ ')',
+ '',
+ 'ext = use_extension("//:extension.bzl", "ext")',
+ 'use_repo(ext, "dep")',
+ # This newline is from ScratchFile.
+ '',
+ ],
+ module_file.read().split('\n'),
+ )
+
+ # Verify that bazel mod tidy doesn't fail with the lockfile.
+ self.RunBazel(['mod', 'tidy'])
+
+ def testModTidyWithoutUsages(self):
+ self.ScratchFile(
+ 'MODULE.bazel',
+ [
+ 'module( name = "foo", version = "1.2.3")',
+ ],
+ )
+
+ self.RunBazel(['mod', 'tidy'])
+
+ with open('MODULE.bazel', 'r') as module_file:
+ self.assertEqual(
+ [
+ 'module(',
+ ' name = "foo",',
+ ' version = "1.2.3",',
+ ')',
+ # This newline is from ScratchFile.
+ '',
+ ],
+ module_file.read().split('\n'),
+ )
+
  def testModTidyFailsOnExtensionFailure(self):
  self.ScratchFile(
  'MODULE.bazel',