Support protoc-gen-prost-serde in prost_toolchain

This allows users to add serde support for their
generated prost proto bindings.

proto/prost/private/prost.bzl

@@ -73,6 +73,14 @@ def _compile_proto(ctx, crate_name, proto_info, deps, prost_toolchain, rustfmt_t
         additional_args.add_all(prost_toolchain.tonic_opts, format_each = "--tonic_opt=%s")
         tools = depset([tonic_plugin.executable], transitive = [tools])
 
+    if prost_toolchain.prost_serde_plugin:
+        prost_serde_plugin = prost_toolchain.prost_serde_plugin[DefaultInfo].files_to_run
+        additional_args.add(prost_toolchain.prost_serde_plugin_flag % prost_serde_plugin.executable.path)
+        additional_args.add("--prost-serde_opt=no_include")
+        additional_args.add("--is_prost_serde")
+        additional_args.add_all(prost_toolchain.prost_serde_opts, format_each = "--prost-serde_opt=%s")
+        tools = depset([prost_serde_plugin.executable], transitive = [tools])
+
     if rustfmt_toolchain:
         additional_args.add("--rustfmt={}".format(rustfmt_toolchain.rustfmt.path))
         tools = depset(transitive = [tools, rustfmt_toolchain.all_files])
@@ -359,6 +367,9 @@ def _rust_prost_toolchain_impl(ctx):
         tonic_plugin = ctx.attr.tonic_plugin,
         tonic_plugin_flag = ctx.attr.tonic_plugin_flag,
         tonic_runtime = ctx.attr.tonic_runtime,
+        prost_serde_opts = ctx.attr.prost_serde_opts,
+        prost_serde_plugin = ctx.attr.prost_serde_plugin,
+        prost_serde_plugin_flag = ctx.attr.prost_serde_plugin_flag,
     )]
 
 rust_prost_toolchain = rule(
@@ -411,6 +422,18 @@ rust_prost_toolchain = rule(
             doc = "The Tonic runtime crates to use.",
             providers = [[rust_common.crate_info], [rust_common.crate_group_info]],
         ),
+        "prost_serde_opts": attr.string_list(
+            doc = "Additional options to add to prost-serde.",
+        ),
+        "prost_serde_plugin": attr.label(
+            doc = "Additional plugins to add to prost-serde.",
+            cfg = "exec",
+            executable = True,
+        ),
+        "prost_serde_plugin_flag": attr.string(
+            doc = "prost-serde plugin flag format. (e.g. `--plugin=protoc-gen-prost-serde=%s`))",
+            default = "--plugin=protoc-gen-prost-serde=%s",
+        ),
     },
 )
 

proto/prost/private/protoc_wrapper.rs

@@ -117,7 +117,11 @@ struct Module {
 ///     }
 /// }
 /// ```
-fn generate_lib_rs(prost_outputs: &BTreeSet<PathBuf>, is_tonic: bool) -> String {
+fn generate_lib_rs(
+    prost_outputs: &BTreeSet<PathBuf>,
+    is_tonic: bool,
+    is_prost_serde: bool,
+) -> String {
     let mut module_info = BTreeMap::new();
 
     for path in prost_outputs.iter() {
@@ -133,7 +137,13 @@ fn generate_lib_rs(prost_outputs: &BTreeSet<PathBuf>, is_tonic: bool) -> String
                 .strip_suffix(".tonic")
                 .expect("Failed to strip suffix")
                 .to_string()
-        };
+        }
+        if is_prost_serde && package.ends_with(".serde") {
+            package = package
+                .strip_suffix(".serde")
+                .expect("Failed to strip suffix")
+                .to_string()
+        }
 
         if package.is_empty() {
             continue;
@@ -155,14 +165,21 @@ fn generate_lib_rs(prost_outputs: &BTreeSet<PathBuf>, is_tonic: bool) -> String
         // Avoid a stack overflow by skipping a known bad package name
         let module_name = snake_cased_package_name(&package);
 
-        module_info.insert(
-            module_name.clone(),
-            Module {
+        module_info
+            .entry(module_name.clone())
+            // Serde module will have same name as module its bindings are for.
+            .and_modify(|module: &mut Module| {
+                module.contents = format!(
+                    "{}\n{}",
+                    module.contents,
+                    fs::read_to_string(path).expect("Failed to read file")
+                );
+            })
+            .or_insert(Module {
                 name,
                 contents: fs::read_to_string(path).expect("Failed to read file"),
                 submodules: BTreeSet::new(),
-            },
-        );
+            });
 
         let module_parts = module_name.split('.').collect::<Vec<_>>();
         for parent_module_index in 0..module_parts.len() {
@@ -443,6 +460,9 @@ struct Args {
     /// Whether to generate tonic code.
     is_tonic: bool,
 
+    /// Whether to generate prost serde code.
+    is_prost_serde: bool,
+
     /// Extra arguments to pass to protoc.
     extra_args: Vec<String>,
 }
@@ -463,6 +483,7 @@ impl Args {
         let mut label: Option<String> = None;
         let mut tonic_or_prost_opts = Vec::new();
         let mut is_tonic = false;
+        let mut is_prost_serde = false;
 
         let mut extra_args = Vec::new();
 
@@ -486,6 +507,11 @@ impl Args {
                 return;
             }
 
+            if arg == "--is_prost_serde" {
+                is_prost_serde = true;
+                return;
+            }
+
             if !arg.contains('=') {
                 extra_args.push(arg);
                 return;
@@ -563,6 +589,9 @@ impl Args {
             if is_tonic {
                 extra_args.push(format!("--tonic_opt={}", tonic_or_prost_opt));
             }
+            if is_prost_serde {
+                extra_args.push(format!("--prost-serde_opt={}", tonic_or_prost_opt));
+            }
         }
 
         if protoc.is_none() {
@@ -613,6 +642,7 @@ impl Args {
             rustfmt,
             proto_paths,
             is_tonic,
+            is_prost_serde,
             label: label.unwrap(),
             extra_args,
         })
@@ -718,6 +748,7 @@ fn main() {
         rustfmt,
         proto_paths,
         is_tonic,
+        is_prost_serde,
         extra_args,
     } = Args::parse().expect("Failed to parse args");
 
@@ -737,6 +768,9 @@ fn main() {
     if is_tonic {
         cmd.arg(format!("--tonic_out={}", out_dir.display()));
     }
+    if is_prost_serde {
+        cmd.arg(format!("--prost-serde_out={}", out_dir.display()));
+    }
     cmd.args(extra_args);
     cmd.args(
         proto_paths
@@ -830,7 +864,11 @@ fn main() {
         .expect("Failed to compute proto package info");
 
     // Write outputs
-    fs::write(&out_librs, generate_lib_rs(&rust_files, is_tonic)).expect("Failed to write file.");
+    fs::write(
+        &out_librs,
+        generate_lib_rs(&rust_files, is_tonic, is_prost_serde),
+    )
+    .expect("Failed to write file.");
     fs::write(
         package_info_file,
         extern_paths