Speed up phase_zinc_depscheck

rules/jvm/private/label.bzl

@@ -9,6 +9,11 @@ def get_labeled_jars(label, java_info, deps):
     deps_labeled_jars = [dep[_LabeledJars] for dep in deps if _LabeledJars in dep]
     return _LabeledJars(
         label = label,
+        transitive_label_objects = depset(
+            [label] if type(label) == "Label" else [],
+            order = "preorder",
+            transitive = [labeled_jars.transitive_label_objects for labeled_jars in deps_labeled_jars],
+        ),
         values = depset(
             [
                 _LabeledJarsData(

rules/private/phases/phase_zinc_depscheck.bzl

@@ -1,13 +1,10 @@
 load(
     "@rules_scala_annex//rules:providers.bzl",
-    _DepsConfiguration = "DepsConfiguration",
     _LabeledJars = "LabeledJars",
-    _ZincCompilationInfo = "ZincCompilationInfo",
 )
 load(
     "@rules_scala_annex//rules/common:private/utils.bzl",
     _make_jvm_flag_args = "make_jvm_flag_args",
-    _short_path = "short_path",
 )
 
 #
@@ -26,45 +23,55 @@ def phase_zinc_depscheck(ctx, g):
     toolchain = ctx.toolchains["//rules/scala:toolchain_type"]
     deps_configuration = toolchain.deps_configuration
     labeled_jar_groups = depset(transitive = [dep[_LabeledJars].values for dep in ctx.attr.deps])
+    transitive_label_objects = depset(transitive = [dep[_LabeledJars].transitive_label_objects for dep in ctx.attr.deps])
     outputs = []
 
     jvm_flag_args = _make_jvm_flag_args(ctx, toolchain.scala_configuration.jvm_flags)
 
-    for name in ("direct", "used"):
-        deps_check = ctx.actions.declare_file("{}/depscheck_{}.success".format(ctx.label.name, name))
-        deps_args = ctx.actions.args()
-        deps_args.add(name, format = "--check_%s=true")
-        deps_args.add_all(
-            "--direct",
-            [_label_for_dependency_checker(dependency) for dependency in ctx.attr.deps],
-            format_each = "_%s",
-        )
+    common_args = ctx.actions.args()
 
-        # Check the comment on the function we're calling here to understand why
-        # we're not using map_each
-        for labeled_jar_group in labeled_jar_groups.to_list():
-            _add_args_for_depscheck_labeled_group(labeled_jar_group, deps_args)
+    # When a `Label` is passed to `Args#add`, Bazel formats it using the apparent repository name. But when `Label`s are
+    # serialized in a `map_each` call, Bazel uses the canonical repository name. This matters because we always want to
+    # use the apparent repository name when correcting dependencies.
+    #
+    # We *could* use `Args#add` instead of `Args#add_all` with `map_each`, but there are strong efficiency gains to be had
+    # from using `Args#add_all`, so we instead serialize the `Label`s with their canonical repository names but provide a
+    # *canonical-to-apparent* mapping so the worker can use the apparent repository names when suggesting corrections.
+    common_args.add_all("--label_keys", transitive_label_objects, map_each = _canonical_label_key, format_each = "_%s")
+    common_args.add_all("--label_names", transitive_label_objects, format_each = "_%s")
+    common_args.add_all(labeled_jar_groups, map_each = _depscheck_labeled_group)
+    common_args.add_all(
+        "--direct",
+        [_label_for_dependency_checker(dependency) for dependency in ctx.attr.deps],
+        format_each = "_%s",
+    )
 
-        deps_args.add("--label", ctx.label, format = "_%s")
-        deps_args.add_all(
-            "--used_whitelist",
-            [_label_for_dependency_checker(dep) for dep in ctx.attr.deps_used_whitelist],
-            format_each = "_%s",
-        )
+    common_args.add("--label", ctx.label, format = "_%s")
+    common_args.add_all(
+        "--used_whitelist",
+        [_label_for_dependency_checker(dependency) for dependency in ctx.attr.deps_used_whitelist],
+        format_each = "_%s",
+    )
 
-        deps_args.add_all(
-            "--unused_whitelist",
-            [_label_for_dependency_checker(dep) for dep in ctx.attr.deps_unused_whitelist],
-            format_each = "_%s",
-        )
+    common_args.add_all(
+        "--unused_whitelist",
+        [_label_for_dependency_checker(dependency) for dependency in ctx.attr.deps_unused_whitelist],
+        format_each = "_%s",
+    )
+    common_args.set_param_file_format("multiline")
+    common_args.use_param_file("@%s", use_always = True)
 
+    for name in ("direct", "used"):
+        deps_check = ctx.actions.declare_file("{}/depscheck_{}.success".format(ctx.label.name, name))
+        deps_args = ctx.actions.args()
+        deps_args.add(name, format = "--check_%s=true")
         deps_args.add("--")
         deps_args.add(g.compile.used)
         deps_args.add(deps_check)
         deps_args.set_param_file_format("multiline")
         deps_args.use_param_file("@%s", use_always = True)
         ctx.actions.run(
-            arguments = [jvm_flag_args, deps_args],
+            arguments = [jvm_flag_args, common_args, deps_args],
             executable = deps_configuration.worker.files_to_run,
             execution_requirements = {
                 "supports-multiplex-workers": "1",
@@ -90,15 +97,9 @@ def phase_zinc_depscheck(ctx, g):
 
     g.out.output_groups["_validation"] = depset(outputs, transitive = validation_transitive)
 
-# If you use avoid using map_each, then labels are converted to their apparent repo name rather than
-# their canonical repo name. The apparent repo name is the human readable one that we want for use
-# with buildozer. See https://bazel.build/rules/lib/builtins/Args.html for more info
-#
-# Avoiding map_each is why we've got this odd section of add and add_all to create a --group
-def _add_args_for_depscheck_labeled_group(labeled_jar_group, deps_args):
-    deps_args.add("--group")
-    deps_args.add(labeled_jar_group.label, format = "_%s")
+def _canonical_label_key(label):
+    return str(label)
 
-    # We do want to use map_each on the jar paths as we don't want the configuration specific
-    # fragments of those paths.
-    deps_args.add_all(labeled_jar_group.jars, map_each = _short_path)
+def _depscheck_labeled_group(group):
+    # The leading underscore prevents @-prefixed labels from being interpreted as argument files
+    return ["--group", "_" + str(group.label)] + [jar.short_path for jar in group.jars.to_list()]

rules/providers.bzl

@@ -68,6 +68,7 @@ LabeledJars = provider(
     doc = "Exported jars and their labels.",
     fields = {
         "label": "The label of the target providing this provider.",
+        "transitive_label_objects": "The transitive depset of `Label` objects, excluding `deps_checker_label` overrides.",
         "values": "The preorder depset of label and jars.",
     },
 )

src/main/scala/higherkindness/rules_scala/workers/deps/DepsRunner.scala

@@ -33,13 +33,21 @@ object DepsRunner extends WorkerMain[Unit] {
 
   private object DepsRunnerRequest {
     def apply(pathResolver: PathResolver, namespace: Namespace): DepsRunnerRequest = {
+      val labelKeys = namespace.getList[String]("label_keys").asScala.map(_.tail)
+      val labelNames = namespace.getList[String]("label_names").asScala.map(_.tail)
+
+      if (labelKeys.size != labelNames.size) {
+        throw new AnnexWorkerError(1, "Dependency label mapping has different numbers of keys and names")
+      }
+
+      val labelMapping = labelKeys.view.zip(labelNames).toMap
       val groups = Option(namespace.getList[java.util.List[String]]("group"))
         .map(_.asScala)
         .getOrElse(List.empty)
         .view
         .map { group =>
           group.asScala match {
-            case Buffer(label, jars @ _*) => Group.apply(label, jars)
+            case Buffer(label, jars @ _*) => Group.apply(label, jars, labelMapping)
             case _                        => throw new Exception(s"Unexpected case in DepsRunner")
           }
         }
@@ -65,9 +73,10 @@ object DepsRunner extends WorkerMain[Unit] {
   )
 
   private object Group {
-    def apply(prependedLabel: String, jars: Seq[String]): Group = {
+    def apply(prependedLabel: String, jars: Seq[String], labelMapping: Map[String, String]): Group = {
+      val label = prependedLabel.tail
       new Group(
-        prependedLabel.tail,
+        labelMapping.getOrElse(label, label),
         jars.toSet,
       )
     }
@@ -77,6 +86,16 @@ object DepsRunner extends WorkerMain[Unit] {
     val parser = ArgumentParsers.newFor("deps").addHelp(true).fromFilePrefix("@").build
     parser.addArgument("--check_direct").`type`(Arguments.booleanType)
     parser.addArgument("--check_used").`type`(Arguments.booleanType)
+    parser
+      .addArgument("--label_keys")
+      .help("The canonical labels in every group, in the same order as `--label_names`")
+      .nargs("*")
+      .setDefault_(Collections.emptyList())
+    parser
+      .addArgument("--label_names")
+      .help("The apparent labels in every group, in the same order as `--label_keys`")
+      .nargs("*")
+      .setDefault_(Collections.emptyList())
     parser
       .addArgument("--direct")
       .help("Labels of direct deps")

tests/dependencies/deps_checker_label/test

@@ -6,3 +6,5 @@ bazel build :depends-on-library-alias |& grep "buildozer 'remove deps //dependen
 
 ! bazel build :depends-on-import-alias || false
 bazel build :depends-on-import-alias |& grep "buildozer 'remove deps //dependencies/deps_checker_label:import-alias' //dependencies/deps_checker_label:depends-on-import-alias"
+
+bazel build :depends-on-whitelisted-library-alias