codegen: Introduce a new codegen using prost

This commit refactors the ttrpc-codegen and the compiler, and merges the two
crates into a single crate, named "codegen". The codegen uses prost crate, a
protobuf compiler for Rust.

Signed-off-by: Xuewei Niu <niuxuewei.nxw@antgroup.com>

Author: justxuewei

.gitignore

@@ -7,4 +7,6 @@ Cargo.lock
 *.o
 example/protocols/**/*.rs
 !example/protocols/**/mod.rs
+example2/protocols/**/*.rs
+!example2/protocols/**/mod.rs
 src/ttrpc.rs

codegen/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "ttrpc-codegen"
+version = "1.0.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+prost = "0.11"
+prost-types = "0.11"
+prost-build = "0.11"
+proc-macro2 = "1.0"
+quote = "1.0"
+anyhow = "^1.0"
+lazy_static = "1.4"
+regex = "1.7"

codegen/src/codegen.rs

@@ -0,0 +1,169 @@
+use anyhow::{anyhow, Context, Result};
+use prost_build::Config;
+use std::{
+    fs::{self, File},
+    io::{Read, Write},
+    path::Path,
+};
+
+use crate::svcgen::{AsyncMode, TtrpcServiceGenerator};
+
+pub struct Codegen<'a, P: AsRef<Path>> {
+    out_dir: &'a P,
+    protos: &'a [P],
+    includes: &'a [P],
+    /// Whether to enable serde
+    serde: bool,
+    async_mode: AsyncMode,
+}
+
+impl<'a, P> Codegen<'a, P>
+where
+    P: AsRef<Path>,
+{
+    pub fn generate(&self) -> Result<()> {
+        self.compile_protos().context("Compile protos")?;
+        self.write_header().context("Write header")?;
+
+        Ok(())
+    }
+
+    fn write_header(&self) -> Result<()> {
+        let dir = fs::read_dir(self.out_dir.as_ref()).context("Read out_dir")?;
+
+        for child in dir {
+            let entry = child.context("Unable to get entry")?;
+            let path = entry.path();
+            let ext = match path.extension() {
+                Some(ext) => match ext.to_str() {
+                    Some(ext) => ext,
+                    None => "",
+                },
+                None => "",
+            };
+            if entry.file_name() == "mod.rs" || ext != "rs" {
+                continue;
+            }
+
+            let header = format!(
+                r#"// This file is generated by ttrpc-codegen {}. Do not edit
+// @generated
+
+"#,
+                env!("CARGO_PKG_VERSION")
+            );
+
+            let mut buf = Vec::<u8>::new();
+            buf.write(header.as_bytes()).context("Write header")?;
+            let mut f = File::open(&path).context(format!("Open rust file {:?}", path))?;
+            f.read_to_end(&mut buf)
+                .context(format!("Read from rust file {:?}", path))?;
+            let mut f = File::create(&path).context(format!("Open rust file {:?}", path))?;
+            f.write_all(buf.as_slice())
+                .context(format!("Write to rust file {:?}", path))?;
+        }
+
+        Ok(())
+    }
+
+    fn compile_protos(&self) -> Result<()> {
+        let mut config = Config::new();
+        config.out_dir(self.out_dir.as_ref());
+        config.service_generator(Box::new(TtrpcServiceGenerator::new(self.async_mode)));
+        config.protoc_arg("--experimental_allow_proto3_optional");
+        config.compile_well_known_types();
+        config.include_file("_include.rs");
+        if self.serde {
+            config.message_attribute(".", "#[derive(::serde::Serialize, ::serde::Deserialize)]");
+        }
+        config
+            .compile_protos(self.protos, self.includes)
+            .context("Compile protos by prost")?;
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+pub struct CodegenBuilder<'a, P: AsRef<Path>> {
+    out_dir: Option<&'a P>,
+    protos: Option<&'a [P]>,
+    includes: Option<&'a [P]>,
+    /// Whether to enable serde
+    serde: Option<bool>,
+    async_mode: Option<AsyncMode>,
+}
+
+impl<'a, P> CodegenBuilder<'a, P>
+where
+    P: AsRef<Path>,
+{
+    pub fn new() -> Self {
+        Self {
+            out_dir: None,
+            protos: None,
+            includes: None,
+            serde: None,
+            async_mode: None,
+        }
+    }
+
+    pub fn set_out_dir(mut self, out_dir: &'a P) -> Self {
+        self.out_dir = Some(out_dir);
+        self
+    }
+
+    pub fn set_protos(mut self, protos: &'a [P]) -> Self {
+        self.protos = Some(protos);
+        self
+    }
+
+    pub fn set_includes(mut self, includes: &'a [P]) -> Self {
+        self.includes = Some(includes);
+        self
+    }
+
+    pub fn set_serde(mut self, serde: bool) -> Self {
+        self.serde = Some(serde);
+        self
+    }
+
+    pub fn set_async_mode(mut self, async_mode: AsyncMode) -> Self {
+        self.async_mode = Some(async_mode);
+        self
+    }
+
+    pub fn build(&self) -> Result<Codegen<'a, P>> {
+        let out_dir = match self.out_dir {
+            Some(out_dir) => out_dir,
+            None => return Err(anyhow!("The out dir is required.")),
+        };
+
+        let protos = match self.protos {
+            Some(protos) => protos,
+            None => return Err(anyhow!("The protos are required.")),
+        };
+
+        let includes = match self.includes {
+            Some(includes) => includes,
+            None => return Err(anyhow!("The includes are required.")),
+        };
+
+        let serde = match self.serde {
+            Some(serde) => serde,
+            None => false,
+        };
+
+        let async_mode = match self.async_mode {
+            Some(mode) => mode,
+            None => AsyncMode::None,
+        };
+
+        Ok(Codegen {
+            out_dir,
+            protos,
+            includes,
+            serde,
+            async_mode,
+        })
+    }
+}

codegen/src/lib.rs

@@ -0,0 +1,6 @@
+mod codegen;
+mod svcgen;
+mod util;
+
+pub use codegen::{Codegen, CodegenBuilder};
+pub use svcgen::AsyncMode;