From 4677c02eabb6feef9d2bc0710e8f70a5febf0fb8 Mon Sep 17 00:00:00 2001
From: "NGnius (Graham)" <ngniusness@gmail.com>
Date: Sat, 15 Apr 2023 09:38:12 -0400
Subject: [PATCH] Add FDS preprocessor support

---
 Cargo.lock                     |  2 +-
 nrpc-build/Cargo.toml          |  2 +-
 nrpc-build/src/builder.rs      | 35 ++++++++++++++++++++++++++++++----
 nrpc-build/src/lib.rs          |  2 ++
 nrpc-build/src/preprocessor.rs |  5 +++++
 5 files changed, 40 insertions(+), 6 deletions(-)
 create mode 100644 nrpc-build/src/preprocessor.rs

diff --git a/Cargo.lock b/Cargo.lock
index 09aeacc..7421b97 100644
--- a/Cargo.lock
+++ b/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",
diff --git a/nrpc-build/Cargo.toml b/nrpc-build/Cargo.toml
index 58190aa..b0253c4 100644
--- a/nrpc-build/Cargo.toml
+++ b/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"
diff --git a/nrpc-build/src/builder.rs b/nrpc-build/src/builder.rs
index fd86a03..28dd837 100644
--- a/nrpc-build/src/builder.rs
+++ b/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);
     }
 }
 
diff --git a/nrpc-build/src/lib.rs b/nrpc-build/src/lib.rs
index 8d31bf8..206ab8f 100644
--- a/nrpc-build/src/lib.rs
+++ b/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;
diff --git a/nrpc-build/src/preprocessor.rs b/nrpc-build/src/preprocessor.rs
new file mode 100644
index 0000000..41673e7
--- /dev/null
+++ b/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);
+}