Add support for custom proto paths

Cargo.lock

@@ -290,7 +290,7 @@ dependencies = [
 
 [[package]]
 name = "nrpc-build"
-version = "0.6.0"
+version = "0.7.0"
 dependencies = [
  "nrpc",
  "prettyplease 0.2.4",

nrpc-build/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nrpc-build"
-version = "0.6.0"
+version = "0.7.0"
 edition = "2021"
 license = "Apache-2.0"
 repository = "https://github.com/NGnius/nRPC"

nrpc-build/src/builder.rs

@@ -21,6 +21,10 @@ impl<'a> Transpiler<'a> {
         files: impl IntoIterator<Item = impl AsRef<Path>>,
         includes: impl IntoIterator<Item = impl AsRef<Path>>
     ) -> Result<Self, impl std::error::Error> {
+        let files: Vec<_> = files.into_iter().collect();
+        for f in &files {
+            println!("cargo:rerun-if-changed={}", f.as_ref().display());
+        }
         Ok::<_, protox::Error>(Self {
             prost_config: Config::new(),
             files: protox::compile(files, includes)?,
@@ -31,19 +35,19 @@ impl<'a> Transpiler<'a> {
 
     /// Generate client and server service implementations
     pub fn generate_all(mut self) -> Self {
-        self.service_generator.add_service(super::ProtobufServiceGenerator::all());
+        self.service_generator.add_service(super::ProtobufServiceGenerator::all(std::env::var("OUT_DIR").unwrap().into()));
         self
     }
 
     /// Generate server services implementations
     pub fn generate_server(mut self) -> Self {
-        self.service_generator.add_service(super::ProtobufServiceGenerator::server());
+        self.service_generator.add_service(super::ProtobufServiceGenerator::server(std::env::var("OUT_DIR").unwrap().into()));
         self
     }
 
     /// Generate client services implementations
     pub fn generate_client(mut self) -> Self {
-        self.service_generator.add_service(super::ProtobufServiceGenerator::client());
+        self.service_generator.add_service(super::ProtobufServiceGenerator::client(std::env::var("OUT_DIR").unwrap().into()));
         self
     }
 

nrpc-build/src/service_gen.rs

@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use prost_build::{Service, ServiceGenerator};
 use quote::quote;
 
@@ -6,33 +8,41 @@ pub(crate) struct ProtobufServiceGenerator {
     generate_client: bool,
     client_reexports: Vec<proc_macro2::TokenStream>,
     server_reexports: Vec<proc_macro2::TokenStream>,
+    modules: Vec<String>,
+    out_dir: PathBuf,
 }
 
 impl ProtobufServiceGenerator {
-    pub fn all() -> Self {
+    pub fn all(out_dir: PathBuf) -> Self {
         Self {
             generate_server: true,
             generate_client: true,
             client_reexports: Vec::new(),
             server_reexports: Vec::new(),
+            modules: Vec::new(),
+            out_dir: out_dir,
         }
     }
 
-    pub fn client() -> Self {
+    pub fn client(out_dir: PathBuf) -> Self {
         Self {
             generate_server: false,
             generate_client: true,
             client_reexports: Vec::new(),
             server_reexports: Vec::new(),
+            modules: Vec::new(),
+            out_dir: out_dir,
         }
     }
 
-    pub fn server() -> Self {
+    pub fn server(out_dir: PathBuf) -> Self {
         Self {
             generate_server: true,
             generate_client: false,
             client_reexports: Vec::new(),
             server_reexports: Vec::new(),
+            modules: Vec::new(),
+            out_dir: out_dir,
         }
     }
 }
@@ -125,6 +135,24 @@ fn struct_methods_client(package_name: &str, service_name: &str, descriptors: &V
     }
 }
 
+fn generate_mod_rs(module_names: &Vec<String>, out_dir: &PathBuf) {
+    // generate mod.rs
+    let modules = module_names.iter().map(|m| {
+        let mod_ident = quote::format_ident!("{}", m);
+        quote! { pub mod #mod_ident; }
+    });
+    let gen_mods: syn::File = syn::parse2(quote! {
+        #(#modules)*
+    }).expect("invalid tokenstream");
+    let mod_str = prettyplease::unparse(&gen_mods);
+    std::fs::write(
+        out_dir
+            .join("mod.rs"),
+        &mod_str
+    ).expect("Failed to write to $OUT_DIR/mod.rs");
+    //std::fs::write("/home/ngnius/potato.rs", &mod_str).unwrap();
+}
+
 impl ServiceGenerator for ProtobufServiceGenerator {
     fn generate(&mut self, service: Service, buf: &mut String) {
         if self.generate_server {
@@ -223,10 +251,10 @@ impl ServiceGenerator for ProtobufServiceGenerator {
             let code_str = prettyplease::unparse(&gen_code);
             buf.push_str(&code_str);
         }
-    }
-
-    fn finalize_package(&mut self, _package: &str, buf: &mut String) {
-        self.finalize(buf);
+        if !self.modules.contains(&service.package) {
+            self.modules.push(service.package.clone());
+            generate_mod_rs(&self.modules, &self.out_dir);
+        }
     }
 
     fn finalize(&mut self, buf: &mut String) {
@@ -258,5 +286,9 @@ impl ServiceGenerator for ProtobufServiceGenerator {
         let gen_code: syn::File = syn::parse2(gen_code).expect("invalid tokenstream");
         let code_str = prettyplease::unparse(&gen_code);
         buf.push_str(&code_str);
+
+        self.modules.clear();
+        self.client_reexports.clear();
+        self.server_reexports.clear();
     }
 }

nrpc-codegen-test/src/main.rs

@@ -3,10 +3,12 @@ use std::error::Error;
 use prost::Message;
 use nrpc::ServerService;
 
-pub mod helloworld {
-    include!(concat!(env!("OUT_DIR"), "/helloworld.rs"));
+pub mod generated {
+    include!(concat!(env!("OUT_DIR"), "/mod.rs"));
 }
 
+pub use generated::*;
+
 #[tokio::main]
 async fn main() {
     let req = helloworld::HelloRequest {