serbuffer_gen/
lib.rs

1use std::convert::TryFrom;
2use std::fmt::{Display, Formatter};
3use std::fs::File;
4use std::io::Write;
5use std::iter::FromIterator;
6use std::path::{Path, PathBuf};
7use std::{env, fs};
8
9pub const VERSION: &'static str = env!("CARGO_PKG_VERSION");
10
11pub struct SchemaBuilder {
12    schema: String,
13    schema_snake: String,
14    fields: Vec<Filed>,
15    serde_derive: bool,
16}
17
18impl SchemaBuilder {
19    pub fn new(schema: &str) -> Self {
20        let schema_snake = to_snake(schema);
21        SchemaBuilder {
22            schema: schema.to_string(),
23            schema_snake,
24            fields: vec![],
25            serde_derive: false,
26        }
27    }
28
29    pub fn field(mut self, name: &str, data_type: DataType) -> Self {
30        self.fields.push(Filed {
31            name: name.to_string(),
32            data_type,
33        });
34
35        self
36    }
37
38    pub fn set_serde_derive(mut self) -> Self {
39        self.serde_derive = true;
40        self
41    }
42
43    pub(crate) fn build_script(&self) -> String {
44        let use_script = self.build_use();
45        let field_indies = self.build_field_index();
46        let data_type = self.build_data_type();
47        let field_name = self.build_field_name();
48        let field_metadata = self.build_field_metadata();
49        let field_reader = self.build_field_reader();
50        let field_writer = self.build_field_writer();
51        let entity = self.build_entity();
52
53        format!(
54            r#"#![allow(unknown_lints)]
55#![allow(clippy::all)]
56
57#![allow(unused_attributes)]
58#![cfg_attr(rustfmt, rustfmt_skip)]
59
60#![allow(dead_code)]
61#![allow(missing_docs)]
62#![allow(non_camel_case_types)]
63#![allow(non_snake_case)]
64#![allow(non_upper_case_globals)]
65#![allow(trivial_casts)]
66#![allow(unused_imports)]
67#![allow(unused_results)]
68//! Generated file by schema {}, version {}
69
70{}
71{}
72{}
73{}
74{}
75{}
76{}
77{}
78"#,
79            self.schema,
80            VERSION,
81            use_script.trim_end(),
82            field_indies.trim_end(),
83            data_type.trim_end(),
84            field_name.trim_end(),
85            field_metadata.trim_end(),
86            field_reader.trim_end(),
87            field_writer.trim_end(),
88            entity.trim_end(),
89        )
90    }
91
92    fn build_use(&self) -> String {
93        "use serbuffer::{types, BufferReader, BufferWriter, Buffer, FieldMetadata};\n".to_string()
94    }
95
96    fn build_field_index(&self) -> String {
97        let mut indies = Vec::new();
98        for index in 0..self.fields.len() {
99            let field_index = format!(
100                "    pub const {}: usize = {};",
101                &self.fields[index].name, index
102            );
103            indies.push(field_index);
104        }
105
106        format!(
107            r#"
108pub mod index {{
109{}
110}}
111        "#,
112            indies.join("\n")
113        )
114    }
115
116    fn build_data_type(&self) -> String {
117        let mut field_script = "".to_string();
118        for index in 0..self.fields.len() {
119            let field = self.fields.get(index).unwrap();
120            let dt = format!("{}", field.data_type);
121            let data_type = format!(
122                r#"    // {}: {}
123    types::{},
124"#,
125                index,
126                field.name,
127                dt.to_uppercase()
128            );
129            field_script = format!("{}{}", field_script, data_type);
130        }
131
132        let script = format!(
133            r#"
134pub const FIELD_TYPE: [u8; {}] = [
135{}
136];"#,
137            self.fields.len(),
138            field_script.trim_end(),
139        );
140
141        script
142    }
143
144    fn build_field_name(&self) -> String {
145        let mut field_script = "".to_string();
146        for index in 0..self.fields.len() {
147            let field = self.fields.get(index).unwrap();
148            let data_type = format!(
149                r#"    // {}: {}
150    "{}",
151"#,
152                index, field.name, field.name,
153            );
154            field_script = format!("{}{}", field_script, data_type);
155        }
156
157        let script = format!(
158            r#"
159pub const FIELD_NAME: [&'static str; {}] = [
160{}
161];"#,
162            self.fields.len(),
163            field_script.trim_end(),
164        );
165
166        script
167    }
168
169    fn build_field_metadata(&self) -> String {
170        format!(
171            r#"
172pub const FIELD_METADATA: FieldMetadata<{}> = FieldMetadata::new(&FIELD_TYPE, &FIELD_NAME);
173"#,
174            self.fields.len(),
175        )
176    }
177
178    fn build_field_reader(&self) -> String {
179        let mut field_read_method = "".to_string();
180        for index in 0..self.fields.len() {
181            let field = self.fields.get(index).unwrap();
182            let method_script = match field.data_type {
183                DataType::BINARY => format!(
184                    r#"
185    pub fn get_{}(&mut self) -> Result<&[u8], std::io::Error> {{
186        self.reader.get_binary({})
187    }}
188"#,
189                    field.name, index
190                ),
191                DataType::STRING => format!(
192                    r#"
193    pub fn get_{}(&mut self) -> Result<&str, std::io::Error> {{
194        self.reader.get_str({})
195    }}
196"#,
197                    field.name, index
198                ),
199                _ => format!(
200                    r#"
201    pub fn get_{}(&mut self) -> Result<{}, std::io::Error> {{
202        self.reader.get_{}({})
203    }}
204"#,
205                    field.name, field.data_type, field.data_type, index
206                ),
207            };
208
209            field_read_method = format!("{}{}", field_read_method, method_script);
210        }
211
212        format!(
213            r#"
214pub struct FieldReader<'a> {{
215    reader: BufferReader<'a, 'static>,
216}}
217
218impl<'a> FieldReader<'a> {{
219    pub fn new(b: &'a mut Buffer) -> Self {{
220        let reader = b.as_reader(&FIELD_TYPE);
221        FieldReader {{ reader }}
222    }}
223{}
224}}"#,
225            field_read_method.trim_end()
226        )
227    }
228
229    fn build_field_writer(&self) -> String {
230        let mut field_read_method = "".to_string();
231        for index in 0..self.fields.len() {
232            let field = self.fields.get(index).unwrap();
233            let method_script = match field.data_type {
234                DataType::BINARY => format!(
235                    r#"
236    pub fn set_{}(&mut self, {}: &[u8]) -> Result<(), std::io::Error> {{
237        if self.writer_pos == {} {{
238            self.writer_pos += 1;
239            self.writer.set_binary({})
240        }} else {{
241            Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "`{}` must be set sequentially"))
242        }}
243    }}
244"#,
245                    field.name, field.name, index, field.name, field.name,
246                ),
247                DataType::STRING => format!(
248                    r#"
249    pub fn set_{}(&mut self, {}: &str) -> Result<(), std::io::Error> {{
250        if self.writer_pos == {} {{
251            self.writer_pos += 1;
252            self.writer.set_str({})
253        }} else {{
254            Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "`{}` must be set sequentially"))
255        }}
256    }}
257"#,
258                    field.name, field.name, index, field.name, field.name,
259                ),
260                _ => format!(
261                    r#"
262    pub fn set_{}(&mut self, {}: {}) -> Result<(), std::io::Error> {{
263        if self.writer_pos == {} {{
264            self.writer_pos += 1;
265            self.writer.set_{}({})
266        }} else {{
267            Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "`{}` must be set sequentially"))
268        }}
269    }}
270"#,
271                    field.name,
272                    field.name,
273                    field.data_type,
274                    index,
275                    field.data_type,
276                    field.name,
277                    field.name,
278                ),
279            };
280
281            field_read_method = format!("{}{}", field_read_method, method_script);
282        }
283
284        format!(
285            r#"
286pub struct FieldWriter<'a> {{
287    writer: BufferWriter<'a, 'static>,
288    writer_pos: usize,
289}}
290
291impl<'a> FieldWriter<'a> {{
292    pub fn new(b: &'a mut Buffer) -> Self {{
293        let writer = b.as_writer(&FIELD_TYPE);
294        FieldWriter {{
295            writer,
296            writer_pos: 0,
297        }}
298    }}
299{}
300}}"#,
301            field_read_method.trim_end()
302        )
303    }
304
305    fn build_entity(&self) -> String {
306        let mut ref_type = false;
307        let mut fields = "".to_string();
308        let mut writers = "".to_string();
309        let mut readers = "".to_string();
310
311        for index in 0..self.fields.len() {
312            let field = self.fields.get(index).unwrap();
313
314            match field.data_type {
315                DataType::BINARY => {
316                    ref_type = true;
317                    fields = format!("{}\n    pub {}: &'a [u8],", fields, field.name);
318                    writers = format!(
319                        "{}\n        writer.set_binary(self.{})?;",
320                        writers, field.name
321                    );
322                    readers = format!(
323                        "{}\n            {}: reader.get_binary({})?,",
324                        readers, field.name, index
325                    );
326                }
327                DataType::STRING => {
328                    ref_type = true;
329                    fields = format!("{}\n    pub {}: &'a str,", fields, field.name);
330                    writers = format!("{}\n        writer.set_str(self.{})?;", writers, field.name);
331                    readers = format!(
332                        "{}\n            {}: reader.get_str({})?,",
333                        readers, field.name, index
334                    );
335                }
336                _ => {
337                    fields = format!("{}\n    pub {}: {},", fields, field.name, field.data_type);
338                    writers = format!(
339                        "{}\n        writer.set_{}(self.{})?;",
340                        writers, field.data_type, field.name
341                    );
342                    readers = format!(
343                        "{}\n            {}: reader.get_{}({})?,",
344                        readers, field.name, field.data_type, index
345                    );
346                }
347            };
348        }
349
350        let ref_type = if ref_type { "<'a>" } else { "" };
351        let serde_derive = if self.serde_derive {
352            ", Serialize, Deserialize"
353        } else {
354            ""
355        };
356
357        format!(
358            r#"
359#[derive(Clone, Debug{})]
360pub struct Entity{} {{
361    {}
362}}
363
364impl{} Entity{} {{
365    pub fn to_buffer(&self, b: &mut Buffer) -> Result<(), std::io::Error> {{
366        let mut writer = b.as_writer(&FIELD_TYPE);
367        
368        {}
369
370        Ok(())
371    }}
372    
373    pub fn parse(b: &'a mut Buffer) -> Result<Self, std::io::Error> {{
374        let reader = b.as_reader(&FIELD_TYPE);
375
376        let entity = Entity {{
377            {}
378        }};
379
380        Ok(entity)
381    }}
382}}
383            "#,
384            serde_derive,
385            ref_type,
386            fields.trim_start(),
387            ref_type,
388            ref_type,
389            writers.trim_start(),
390            readers.trim_start()
391        )
392    }
393}
394
395/// code gen
396/// struct Demo {
397///   timestamp: u64
398/// }
399#[derive(Default)]
400pub struct Codegen {
401    /// --lang_out= param
402    generated_dir: PathBuf,
403    schemas: Vec<SchemaBuilder>,
404}
405
406impl Codegen {
407    pub fn new(generated_dir: &str) -> Self {
408        if Path::new(generated_dir).exists() {
409            fs::remove_dir_all(generated_dir).unwrap();
410        }
411        fs::create_dir(generated_dir).unwrap();
412
413        Codegen {
414            generated_dir: PathBuf::from(generated_dir),
415            schemas: vec![],
416        }
417    }
418
419    pub fn out_dir(path: &str) -> Self {
420        let out_dir = env::var("OUT_DIR").unwrap();
421        let generated_dir = Path::new(out_dir.as_str()).join(path);
422
423        Self::new(generated_dir.as_path().to_str().unwrap())
424    }
425
426    pub fn schema(mut self, schema_builder: SchemaBuilder) -> Self {
427        self.schemas.push(schema_builder);
428        self
429    }
430
431    pub fn gen(&self) -> std::io::Result<()> {
432        for schema_builder in &self.schemas {
433            let script = schema_builder.build_script();
434
435            let file_name = format!("{}.rs", schema_builder.schema_snake);
436            let file_path = self.generated_dir.join(file_name.as_str());
437            let mut file_writer = File::create(&file_path)?;
438            file_writer.write_all(script.as_bytes())?;
439            file_writer.flush()?;
440        }
441
442        {
443            let mods_str: Vec<String> = self
444                .schemas
445                .iter()
446                .map(|x| format!("pub mod {};", x.schema_snake))
447                .collect();
448            let script = mods_str.join("\n");
449
450            let file_name = "mod.rs";
451            let file_path = self.generated_dir.join(file_name);
452            let mut file_writer = File::create(&file_path)?;
453            file_writer.write_all(script.as_bytes())?;
454            file_writer.flush()?;
455        }
456
457        Ok(())
458    }
459}
460
461fn to_snake(s: &str) -> String {
462    let mut v = Vec::new();
463    for c in s.chars() {
464        if c.is_uppercase() {
465            v.push('_');
466            v.push(c.to_ascii_lowercase());
467        } else {
468            v.push(c);
469        }
470    }
471
472    let snake = String::from_iter(v.iter());
473    if snake.starts_with("_") {
474        let ss = snake.as_str();
475        ss[1..snake.len()].to_string()
476    } else {
477        snake
478    }
479}
480
481pub enum DataType {
482    BOOL,
483    U8,
484    I8,
485    U16,
486    I16,
487    U32,
488    I32,
489    U64,
490    I64,
491    F32,
492    F64,
493    BINARY,
494    STRING,
495}
496
497impl Display for DataType {
498    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
499        match self {
500            DataType::BOOL => write!(f, "{}", "BOOL".to_lowercase()),
501            DataType::U8 => write!(f, "{}", "U8".to_lowercase()),
502            DataType::I8 => write!(f, "{}", "I8".to_lowercase()),
503            DataType::U16 => write!(f, "{}", "U16".to_lowercase()),
504            DataType::I16 => write!(f, "{}", "I16".to_lowercase()),
505            DataType::U32 => write!(f, "{}", "U32".to_lowercase()),
506            DataType::I32 => write!(f, "{}", "I32".to_lowercase()),
507            DataType::U64 => write!(f, "{}", "U64".to_lowercase()),
508            DataType::I64 => write!(f, "{}", "I64".to_lowercase()),
509            DataType::F32 => write!(f, "{}", "F32".to_lowercase()),
510            DataType::F64 => write!(f, "{}", "F64".to_lowercase()),
511            DataType::BINARY => write!(f, "{}", "BINARY".to_lowercase()),
512            DataType::STRING => write!(f, "{}", "STRING".to_lowercase()),
513        }
514    }
515}
516
517impl TryFrom<&str> for DataType {
518    type Error = &'static str;
519
520    fn try_from(value: &str) -> Result<Self, Self::Error> {
521        match value.to_uppercase().as_str() {
522            "BOOL" => Ok(DataType::BOOL),
523            "U8" => Ok(DataType::U8),
524            "I8" => Ok(DataType::I8),
525            "U16" => Ok(DataType::U16),
526            "I16" => Ok(DataType::I16),
527            "U32" => Ok(DataType::U32),
528            "I32" => Ok(DataType::I32),
529            "U64" => Ok(DataType::U64),
530            "I64" => Ok(DataType::I64),
531            "F32" => Ok(DataType::F32),
532            "F64" => Ok(DataType::F64),
533            "BINARY" => Ok(DataType::BINARY),
534            "STRING" => Ok(DataType::STRING),
535            _ => Err("unknown"),
536        }
537    }
538}
539
540struct Filed {
541    name: String,
542    data_type: DataType,
543}
544
545#[cfg(test)]
546mod tests {
547    use crate::{DataType, SchemaBuilder};
548
549    #[test]
550    pub fn code_gen_basic_type_test() {
551        let script = SchemaBuilder::new("DemoSchema")
552            .field("timestamp", DataType::U64)
553            .field("a", DataType::BOOL)
554            .field("a1", DataType::U8)
555            .build_script();
556
557        println!("-- Start ---");
558        println!("{}", script);
559        println!("-- End ---");
560    }
561
562    #[test]
563    pub fn code_gen_ref_type_test() {
564        let script = SchemaBuilder::new("DemoSchema")
565            .field("timestamp", DataType::U64)
566            .field("application_name", DataType::STRING)
567            .field("agent_id", DataType::STRING)
568            .field("a", DataType::BOOL)
569            .field("a1", DataType::U8)
570            .field("a2", DataType::BINARY)
571            .build_script();
572
573        println!("-- Start ---");
574        println!("{}", script);
575        println!("-- End ---");
576    }
577}
serbuffer_gen/lib.rs

serbuffer_gen/
lib.rs
