Added rust_prost_transform rule for modifying granular proto_library

docs/src/rust_prost.md

@@ -192,3 +192,57 @@ Rust Prost toolchain rule.
 | <a id="rust_prost_toolchain-tonic_runtime"></a>tonic_runtime |  The Tonic runtime crates to use.   | <a href="https://bazel.build/concepts/labels">Label</a> | optional |  `None`  |
 
 
+<a id="rust_prost_transform"></a>
+
+## rust_prost_transform
+
+<pre>
+rust_prost_transform(<a href="#rust_prost_transform-name">name</a>, <a href="#rust_prost_transform-deps">deps</a>, <a href="#rust_prost_transform-srcs">srcs</a>, <a href="#rust_prost_transform-prost_opts">prost_opts</a>, <a href="#rust_prost_transform-tonic_opts">tonic_opts</a>)
+</pre>
+
+A rule for transforming the outputs of `ProstGenProto` actions.
+
+This rule is used by adding it to the `data` attribute of `proto_library` targets. E.g.
+```python
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@rules_rust_prost//:defs.bzl", "rust_prost_library", "rust_prost_transform")
+
+rust_prost_transform(
+    name = "a_transform",
+    srcs = [
+        "a_src.rs",
+    ],
+)
+
+proto_library(
+    name = "a_proto",
+    srcs = [
+        "a.proto",
+    ],
+    data = [
+        ":transform",
+    ],
+)
+
+rust_prost_library(
+    name = "a_rs_proto",
+    proto = ":a_proto",
+)
+```
+
+The `rust_prost_library` will spawn an action on the `a_proto` target which consumes the
+`a_transform` rule to provide a means of granularly modifying a proto library for `ProstGenProto`
+actions with minimal impact to other consumers.
+
+**ATTRIBUTES**
+
+
+| Name  | Description | Type | Mandatory | Default |
+| :------------- | :------------- | :------------- | :------------- | :------------- |
+| <a id="rust_prost_transform-name"></a>name |  A unique name for this target.   | <a href="https://bazel.build/concepts/labels#target-names">Name</a> | required |  |
+| <a id="rust_prost_transform-deps"></a>deps |  Additional dependencies to add to the compiled crate.   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional |  `[]`  |
+| <a id="rust_prost_transform-srcs"></a>srcs |  Additional source files to include in generated Prost source code.   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional |  `[]`  |
+| <a id="rust_prost_transform-prost_opts"></a>prost_opts |  Additional options to add to Prost.   | List of strings | optional |  `[]`  |
+| <a id="rust_prost_transform-tonic_opts"></a>tonic_opts |  Additional options to add to Tonic.   | List of strings | optional |  `[]`  |
+
+

extensions/prost/defs.bzl

@@ -146,6 +146,11 @@ load(
     _rust_prost_library = "rust_prost_library",
     _rust_prost_toolchain = "rust_prost_toolchain",
 )
+load(
+    "//private:prost_transform.bzl",
+    _rust_prost_transform = "rust_prost_transform",
+)
 
 rust_prost_library = _rust_prost_library
 rust_prost_toolchain = _rust_prost_toolchain
+rust_prost_transform = _rust_prost_transform

extensions/prost/private/prost.bzl

@@ -19,6 +19,7 @@ load("@rules_rust//rust/private:rustc.bzl", "rustc_compile_action")
 # buildifier: disable=bzl-visibility
 load("@rules_rust//rust/private:utils.bzl", "can_build_metadata")
 load("//:providers.bzl", "ProstProtoInfo")
+load(":prost_transform.bzl", "ProstTransformInfo")
 
 RUST_EDITION = "2021"
 
@@ -39,7 +40,15 @@ def _create_proto_lang_toolchain(ctx, prost_toolchain):
 
     return proto_lang_toolchain
 
-def _compile_proto(ctx, crate_name, proto_info, deps, prost_toolchain, rustfmt_toolchain = None):
+def _compile_proto(
+        *,
+        ctx,
+        crate_name,
+        proto_info,
+        transform_infos,
+        deps,
+        prost_toolchain,
+        rustfmt_toolchain = None):
     deps_info_file = ctx.actions.declare_file(ctx.label.name + ".prost_deps_info")
     dep_package_infos = [dep[ProstProtoInfo].package_info for dep in deps]
     ctx.actions.write(
@@ -53,6 +62,15 @@ def _compile_proto(ctx, crate_name, proto_info, deps, prost_toolchain, rustfmt_t
     proto_compiler = prost_toolchain.proto_compiler
     tools = depset([proto_compiler.executable])
 
+    tonic_opts = []
+    prost_opts = []
+    additional_srcs = []
+    for transform_info in transform_infos:
+        tonic_opts.extend(transform_info.tonic_opts)
+        prost_opts.extend(transform_info.prost_opts)
+        additional_srcs.append(transform_info.srcs)
+
+    all_additional_srcs = depset(transitive = additional_srcs)
     direct_crate_names = [dep[ProstProtoInfo].dep_variant_info.crate_info.name for dep in deps]
     additional_args = ctx.actions.args()
 
@@ -65,22 +83,27 @@ def _compile_proto(ctx, crate_name, proto_info, deps, prost_toolchain, rustfmt_t
     additional_args.add("--direct_dep_crate_names={}".format(",".join(direct_crate_names)))
     additional_args.add("--prost_opt=compile_well_known_types")
     additional_args.add("--descriptor_set={}".format(proto_info.direct_descriptor_set.path))
-    additional_args.add_all(prost_toolchain.prost_opts, format_each = "--prost_opt=%s")
+    additional_args.add("--additional_srcs={}".format(",".join([f.path for f in all_additional_srcs.to_list()])))
+    additional_args.add_all(prost_toolchain.prost_opts + prost_opts, format_each = "--prost_opt=%s")
 
     if prost_toolchain.tonic_plugin:
         tonic_plugin = prost_toolchain.tonic_plugin[DefaultInfo].files_to_run
         additional_args.add(prost_toolchain.tonic_plugin_flag % tonic_plugin.executable.path)
         additional_args.add("--tonic_opt=no_include")
         additional_args.add("--tonic_opt=compile_well_known_types")
         additional_args.add("--is_tonic")
-        additional_args.add_all(prost_toolchain.tonic_opts, format_each = "--tonic_opt=%s")
+
+        additional_args.add_all(prost_toolchain.tonic_opts + tonic_opts, format_each = "--tonic_opt=%s")
         tools = depset([tonic_plugin.executable], transitive = [tools])
 
     if rustfmt_toolchain:
         additional_args.add("--rustfmt={}".format(rustfmt_toolchain.rustfmt.path))
         tools = depset(transitive = [tools, rustfmt_toolchain.all_files])
 
-    additional_inputs = depset([deps_info_file, proto_info.direct_descriptor_set] + [dep[ProstProtoInfo].package_info for dep in deps])
+    additional_inputs = depset(
+        [deps_info_file, proto_info.direct_descriptor_set] + [dep[ProstProtoInfo].package_info for dep in deps],
+        transitive = [all_additional_srcs],
+    )
 
     proto_common.compile(
         actions = ctx.actions,
@@ -116,7 +139,14 @@ def _get_cc_info(providers):
             return provider
     fail("Couldn't find a CcInfo in the list of providers")
 
-def _compile_rust(ctx, attr, crate_name, src, deps, edition):
+def _compile_rust(
+        *,
+        ctx,
+        attr,
+        crate_name,
+        src,
+        deps,
+        edition):
     """Compiles a Rust source file.
 
     Args:
@@ -233,7 +263,14 @@ def _rust_prost_aspect_impl(target, ctx):
         if RustAnalyzerInfo in proto_dep:
             rust_analyzer_deps.append(proto_dep[RustAnalyzerInfo])
 
-    deps = runtime_deps + direct_deps
+    transform_infos = []
+    for data_target in getattr(ctx.rule.attr, "data", []):
+        if ProstTransformInfo in data_target:
+            transform_infos.append(data_target[ProstTransformInfo])
+
+    rust_deps = runtime_deps + direct_deps
+    for transform_info in transform_infos:
+        rust_deps.extend(transform_info.deps)
 
     crate_name = ctx.label.name.replace("-", "_").replace("/", "_")
 
@@ -243,6 +280,7 @@ def _rust_prost_aspect_impl(target, ctx):
         ctx = ctx,
         crate_name = crate_name,
         proto_info = proto_info,
+        transform_infos = transform_infos,
         deps = proto_deps,
         prost_toolchain = prost_toolchain,
         rustfmt_toolchain = rustfmt_toolchain,
@@ -253,7 +291,7 @@ def _rust_prost_aspect_impl(target, ctx):
         attr = ctx.rule.attr,
         crate_name = crate_name,
         src = lib_rs,
-        deps = deps,
+        deps = rust_deps,
         edition = RUST_EDITION,
     )
 
@@ -495,7 +533,7 @@ def _current_prost_runtime_impl(ctx):
     )]
 
 current_prost_runtime = rule(
-    doc = "A rule for accessing the current Prost toolchain components needed by the process wrapper",
+    doc = "A rule for accessing the current Prost toolchain components needed by the process wrapper.",
     provides = [rust_common.crate_group_info],
     implementation = _current_prost_runtime_impl,
     toolchains = [TOOLCHAIN_TYPE],

extensions/prost/private/prost_transform.bzl

@@ -0,0 +1,88 @@
+"""Prost rules."""
+
+load("@rules_rust//rust:defs.bzl", "rust_common")
+
+ProstTransformInfo = provider(
+    doc = "Info about transformations to apply to Prost generated source code.",
+    fields = {
+        "deps": "List[DepVariantInfo]: Additional dependencies to compile into the Prost target.",
+        "prost_opts": "List[str]: Additional prost flags.",
+        "srcs": "Depset[File]: Additional source files to include in generated Prost source code.",
+        "tonic_opts": "List[str]: Additional tonic flags.",
+    },
+)
+
+def _rust_prost_transform_impl(ctx):
+    deps = []
+    for target in ctx.attr.deps:
+        deps.append(rust_common.dep_variant_info(
+            crate_info = target[rust_common.crate_info] if rust_common.crate_info in target else None,
+            dep_info = target[rust_common.dep_info] if rust_common.dep_info in target else None,
+            cc_info = target[CcInfo] if CcInfo in target else None,
+            build_info = None,
+        ))
+
+    # DefaultInfo is intentionally not returned here to avoid impacting other
+    # consumers of the `proto_library` target this rule is expected to be passed
+    # to.
+    return [ProstTransformInfo(
+        deps = deps,
+        prost_opts = ctx.attr.prost_opts,
+        srcs = depset(ctx.files.srcs),
+        tonic_opts = ctx.attr.tonic_opts,
+    )]
+
+rust_prost_transform = rule(
+    doc = """\
+A rule for transforming the outputs of `ProstGenProto` actions.
+
+This rule is used by adding it to the `data` attribute of `proto_library` targets. E.g.
+```python
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@rules_rust_prost//:defs.bzl", "rust_prost_library", "rust_prost_transform")
+
+rust_prost_transform(
+    name = "a_transform",
+    srcs = [
+        "a_src.rs",
+    ],
+)
+
+proto_library(
+    name = "a_proto",
+    srcs = [
+        "a.proto",
+    ],
+    data = [
+        ":transform",
+    ],
+)
+
+rust_prost_library(
+    name = "a_rs_proto",
+    proto = ":a_proto",
+)
+```
+
+The `rust_prost_library` will spawn an action on the `a_proto` target which consumes the
+`a_transform` rule to provide a means of granularly modifying a proto library for `ProstGenProto`
+actions with minimal impact to other consumers.
+""",
+    implementation = _rust_prost_transform_impl,
+    attrs = {
+        "deps": attr.label_list(
+            doc = "Additional dependencies to add to the compiled crate.",
+            providers = [[rust_common.crate_info], [rust_common.crate_group_info]],
+        ),
+        "prost_opts": attr.string_list(
+            doc = "Additional options to add to Prost.",
+        ),
+        "srcs": attr.label_list(
+            doc = "Additional source files to include in generated Prost source code.",
+            allow_files = True,
+        ),
+        "tonic_opts": attr.string_list(
+            doc = "Additional options to add to Tonic.",
+        ),
+    },
+)