From 79c95a001805af8f1b1f9f1228df7361daba6543 Mon Sep 17 00:00:00 2001
From: "NGnius (Graham)" <ngniusness@gmail.com>
Date: Fri, 14 Apr 2023 23:50:43 -0400
Subject: [PATCH] Add service generator extensability

---
 Cargo.lock                |   1 +
 nrpc-build/Cargo.toml     |   1 +
 nrpc-build/src/builder.rs | 101 +++++++++++++++++++++++++++++++++-----
 nrpc-build/src/lib.rs     |   2 +-
 4 files changed, 92 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 152b00d..256bf98 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -287,6 +287,7 @@ dependencies = [
  "prettyplease 0.2.4",
  "proc-macro2",
  "prost-build",
+ "prost-types",
  "protox",
  "quote",
  "syn 2.0.13",
diff --git a/nrpc-build/Cargo.toml b/nrpc-build/Cargo.toml
index 0a8de9b..ae73f2f 100644
--- a/nrpc-build/Cargo.toml
+++ b/nrpc-build/Cargo.toml
@@ -13,6 +13,7 @@ description = "Yet another remote procedure call library - codegen"
 # gRPC/protobuf
 protox = "0.3"
 prost-build = "0.11"
+prost-types = "0.11"
 
 # code gen
 prettyplease = "0.2"
diff --git a/nrpc-build/src/builder.rs b/nrpc-build/src/builder.rs
index cb386fa..fd86a03 100644
--- a/nrpc-build/src/builder.rs
+++ b/nrpc-build/src/builder.rs
@@ -3,37 +3,114 @@ use std::convert::AsRef;
 use std::iter::IntoIterator;
 
 use prost_build::Config;
+use prost_build::{Service, ServiceGenerator};
+use prost_types::FileDescriptorSet;
+
+/// Proto -> Rust transpiler configurator
+pub struct Transpiler {
+    prost_config: Config,
+    files: FileDescriptorSet,
+    service_generator: MergedServiceGenerator,
+}
+
+impl Transpiler {
+    pub fn new(
+        files: impl IntoIterator<Item = impl AsRef<Path>>,
+        includes: impl IntoIterator<Item = impl AsRef<Path>>
+    ) -> Result<Self, impl std::error::Error> {
+        Ok::<_, protox::Error>(Self {
+            prost_config: Config::new(),
+            files: protox::compile(files, includes)?,
+            service_generator: MergedServiceGenerator::empty()
+        })
+    }
+
+    /// Generate client and server service implementations
+    pub fn generate_all(mut self) -> Self {
+        self.service_generator.add_service(super::ProtobufServiceGenerator::all());
+        self
+    }
+
+    /// Generate server services implementations
+    pub fn generate_server(mut self) -> Self {
+        self.service_generator.add_service(super::ProtobufServiceGenerator::server());
+        self
+    }
+
+    /// Generate client services implementations
+    pub fn generate_client(mut self) -> Self {
+        self.service_generator.add_service(super::ProtobufServiceGenerator::client());
+        self
+    }
+
+    /// Add additional custom service generator
+    pub fn with_service_generator<S: ServiceGenerator + 'static>(mut self, gen: S) -> Self {
+        self.service_generator.add_service(gen);
+        self
+    }
+
+    /// Actually generate code
+    pub fn transpile(mut self) -> std::io::Result<()> {
+        self.prost_config
+            .service_generator(Box::new(self.service_generator))
+            .compile_fds(self.files)
+    }
+}
+
+struct MergedServiceGenerator {
+    generators: Vec<Box<dyn ServiceGenerator + 'static>>,
+}
+
+impl MergedServiceGenerator {
+    fn empty() -> Self {
+        Self {
+            generators: Vec::new(),
+        }
+    }
+
+    fn add_service<S: ServiceGenerator + 'static>(&mut self, service: S) -> &mut Self {
+        self.generators.push(Box::new(service));
+        self
+    }
+}
+
+impl ServiceGenerator for MergedServiceGenerator {
+    fn generate(&mut self, service: Service, buf: &mut String) {
+        for gen in &mut self.generators {
+            gen.generate(service.clone(), buf);
+        }
+    }
+}
 
 /// Compile proto files into Rust with server and client implementations
 pub fn compile(
     files: impl IntoIterator<Item = impl AsRef<Path>>,
     includes: impl IntoIterator<Item = impl AsRef<Path>>
 ) {
-    let file_descriptors = protox::compile(files, includes).unwrap();
-    Config::new()
-        .service_generator(Box::new(super::ProtobufServiceGenerator::all()))
-        .compile_fds(file_descriptors)
+    Transpiler::new(files, includes).unwrap()
+        .generate_all()
+        .transpile()
         .unwrap();
 }
 
+/// Compile proto files into Rust with only client implementations
 pub fn compile_clients(
     files: impl IntoIterator<Item = impl AsRef<Path>>,
     includes: impl IntoIterator<Item = impl AsRef<Path>>,
 ) {
-    let file_descriptors = protox::compile(files, includes).unwrap();
-    Config::new()
-        .service_generator(Box::new(super::ProtobufServiceGenerator::client()))
-        .compile_fds(file_descriptors)
+    Transpiler::new(files, includes).unwrap()
+        .generate_client()
+        .transpile()
         .unwrap();
 }
 
+/// Compile proto files into Rust with only server implementations
 pub fn compile_servers(
     files: impl IntoIterator<Item = impl AsRef<Path>>,
     includes: impl IntoIterator<Item = impl AsRef<Path>>,
 ) {
-    let file_descriptors = protox::compile(files, includes).unwrap();
-    Config::new()
-        .service_generator(Box::new(super::ProtobufServiceGenerator::server()))
-        .compile_fds(file_descriptors)
+    Transpiler::new(files, includes).unwrap()
+        .generate_server()
+        .transpile()
         .unwrap();
 }
diff --git a/nrpc-build/src/lib.rs b/nrpc-build/src/lib.rs
index fbd70ea..8d31bf8 100644
--- a/nrpc-build/src/lib.rs
+++ b/nrpc-build/src/lib.rs
@@ -1,5 +1,5 @@
 mod builder;
 mod service_gen;
 
-pub use builder::{compile, compile_servers, compile_clients};
+pub use builder::{compile, compile_servers, compile_clients, Transpiler};
 pub(crate) use service_gen::ProtobufServiceGenerator;