nrpc_build/
builder.rs

1use std::convert::AsRef;
2use std::iter::IntoIterator;
3use std::path::Path;
4
5use prost_build::Config;
6use prost_build::{Service, ServiceGenerator};
7use prost_types::FileDescriptorSet;
8
9use super::Preprocessor;
10
11/// Proto -> Rust transpiler configurator
12pub struct Transpiler<'a> {
13    prost_config: Config,
14    files: FileDescriptorSet,
15    service_generator: MergedServiceGenerator,
16    preprocessors: Vec<Box<dyn Preprocessor + 'a>>,
17}
18
19impl<'a> Transpiler<'a> {
20    pub fn new(
21        files: impl IntoIterator<Item = impl AsRef<Path>>,
22        includes: impl IntoIterator<Item = impl AsRef<Path>>,
23    ) -> Result<Self, impl std::error::Error> {
24        let files: Vec<_> = files.into_iter().collect();
25        for f in &files {
26            println!("cargo:rerun-if-changed={}", f.as_ref().display());
27        }
28        Ok::<_, protox::Error>(Self {
29            prost_config: Config::new(),
30            files: protox::compile(files, includes)?,
31            service_generator: MergedServiceGenerator::empty(),
32            preprocessors: Vec::new(),
33        })
34    }
35
36    /// Generate client and server service implementations
37    pub fn generate_all(mut self) -> Self {
38        self.service_generator
39            .add_service(super::ProtobufServiceGenerator::all(
40                std::env::var("OUT_DIR").unwrap().into(),
41            ));
42        self
43    }
44
45    /// Generate server services implementations
46    pub fn generate_server(mut self) -> Self {
47        self.service_generator
48            .add_service(super::ProtobufServiceGenerator::server(
49                std::env::var("OUT_DIR").unwrap().into(),
50            ));
51        self
52    }
53
54    /// Generate client services implementations
55    pub fn generate_client(mut self) -> Self {
56        self.service_generator
57            .add_service(super::ProtobufServiceGenerator::client(
58                std::env::var("OUT_DIR").unwrap().into(),
59            ));
60        self
61    }
62
63    /// Add additional custom service generator
64    pub fn with_service_generator<S: ServiceGenerator + 'static>(mut self, gen: S) -> Self {
65        self.service_generator.add_service(gen);
66        self
67    }
68
69    /// Add a proto file descriptor preprocessor
70    pub fn with_preprocessor<P: Preprocessor + 'a>(mut self, pp: P) -> Self {
71        self.preprocessors.push(Box::new(pp));
72        self
73    }
74
75    /// Actually generate code
76    pub fn transpile(mut self) -> std::io::Result<()> {
77        let mut files = self.files;
78        let mut generated = String::new();
79        for mut pp in self.preprocessors {
80            pp.process(&mut files, &mut generated);
81        }
82        self.service_generator
83            .add_service(PreprocessedCodeGenInjector {
84                generated_str: generated,
85            });
86
87        self.prost_config
88            .service_generator(Box::new(self.service_generator))
89            .compile_fds(files)
90    }
91}
92
93struct PreprocessedCodeGenInjector {
94    generated_str: String,
95}
96
97impl ServiceGenerator for PreprocessedCodeGenInjector {
98    fn generate(&mut self, _: Service, buf: &mut String) {
99        buf.insert_str(0, &self.generated_str);
100    }
101}
102
103struct MergedServiceGenerator {
104    generators: Vec<Box<dyn ServiceGenerator + 'static>>,
105}
106
107impl MergedServiceGenerator {
108    fn empty() -> Self {
109        Self {
110            generators: Vec::new(),
111        }
112    }
113
114    fn add_service<S: ServiceGenerator + 'static>(&mut self, service: S) -> &mut Self {
115        self.generators.push(Box::new(service));
116        self
117    }
118}
119
120impl ServiceGenerator for MergedServiceGenerator {
121    fn generate(&mut self, service: Service, buf: &mut String) {
122        for gen in &mut self.generators {
123            gen.generate(service.clone(), buf);
124        }
125    }
126}
127
128/// Compile proto files into Rust with server and client implementations
129pub fn compile(
130    files: impl IntoIterator<Item = impl AsRef<Path>>,
131    includes: impl IntoIterator<Item = impl AsRef<Path>>,
132) {
133    Transpiler::new(files, includes)
134        .unwrap()
135        .generate_all()
136        .transpile()
137        .unwrap();
138}
139
140/// Compile proto files into Rust with only client implementations
141pub fn compile_clients(
142    files: impl IntoIterator<Item = impl AsRef<Path>>,
143    includes: impl IntoIterator<Item = impl AsRef<Path>>,
144) {
145    Transpiler::new(files, includes)
146        .unwrap()
147        .generate_client()
148        .transpile()
149        .unwrap();
150}
151
152/// Compile proto files into Rust with only server implementations
153pub fn compile_servers(
154    files: impl IntoIterator<Item = impl AsRef<Path>>,
155    includes: impl IntoIterator<Item = impl AsRef<Path>>,
156) {
157    Transpiler::new(files, includes)
158        .unwrap()
159        .generate_server()
160        .transpile()
161        .unwrap();
162}
nrpc_build/builder.rs

nrpc_build/
builder.rs
