Add FDS preprocessor support

Cargo.lock

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

nrpc-build/Cargo.toml

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

nrpc-build/src/builder.rs

@@ -6,14 +6,17 @@ use prost_build::Config;
 use prost_build::{Service, ServiceGenerator};
 use prost_types::FileDescriptorSet;
 
+use super::Preprocessor;
+
 /// Proto -> Rust transpiler configurator
-pub struct Transpiler {
+pub struct Transpiler<'a> {
     prost_config: Config,
     files: FileDescriptorSet,
     service_generator: MergedServiceGenerator,
+    preprocessors: Vec<Box<dyn Preprocessor + 'a>>,
 }
 
-impl Transpiler {
+impl<'a> Transpiler<'a> {
     pub fn new(
         files: impl IntoIterator<Item = impl AsRef<Path>>,
         includes: impl IntoIterator<Item = impl AsRef<Path>>
@@ -21,7 +24,8 @@ impl Transpiler {
         Ok::<_, protox::Error>(Self {
             prost_config: Config::new(),
             files: protox::compile(files, includes)?,
-            service_generator: MergedServiceGenerator::empty()
+            service_generator: MergedServiceGenerator::empty(),
+            preprocessors: Vec::new(),
         })
     }
 
@@ -49,11 +53,34 @@ impl Transpiler {
         self
     }
 
+    /// Add a proto file descriptor preprocessor
+    pub fn with_preprocessor<P: Preprocessor + 'a>(mut self, pp: P) -> Self {
+        self.preprocessors.push(Box::new(pp));
+        self
+    }
+
     /// Actually generate code
     pub fn transpile(mut self) -> std::io::Result<()> {
+        let mut files = self.files;
+        let mut generated = String::new();
+        for mut pp in self.preprocessors {
+            pp.process(&mut files, &mut generated);
+        }
+        self.service_generator.add_service(PreprocessedCodeGenInjector { generated_str: generated });
+
         self.prost_config
             .service_generator(Box::new(self.service_generator))
-            .compile_fds(self.files)
+            .compile_fds(files)
+    }
+}
+
+struct PreprocessedCodeGenInjector {
+    generated_str: String,
+}
+
+impl ServiceGenerator for PreprocessedCodeGenInjector {
+    fn generate(&mut self, _: Service, buf: &mut String) {
+        buf.insert_str(0, &self.generated_str);
     }
 }
 

nrpc-build/src/lib.rs

@@ -1,5 +1,7 @@
 mod builder;
+mod preprocessor;
 mod service_gen;
 
 pub use builder::{compile, compile_servers, compile_clients, Transpiler};
+pub use preprocessor::Preprocessor;
 pub(crate) use service_gen::ProtobufServiceGenerator;

nrpc-build/src/preprocessor.rs

@@ -0,0 +1,5 @@
+use prost_types::FileDescriptorSet;
+
+pub trait Preprocessor {
+    fn process(&mut self, fds: &mut FileDescriptorSet, buf: &mut String);
+}