rlink_buffer_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};
7
8/// code gen
9/// struct Demo {
10///   timestamp: u64
11/// }
12#[derive(Default)]
13pub struct Codegen {
14    /// --lang_out= param
15    out_dir: PathBuf,
16    schema: String,
17    schema_snake: String,
18    fields: Vec<Filed>,
19}
20
21impl Codegen {
22    pub fn new(out_dir: impl AsRef<Path>, schema: &str) -> Self {
23        let schema_snake = to_snake(schema);
24        Codegen {
25            out_dir: out_dir.as_ref().to_owned(),
26            schema: schema.to_string(),
27            schema_snake,
28            fields: Vec::new(),
29        }
30    }
31
32    pub fn field(&mut self, name: &str, data_type: DataType) -> &mut Self {
33        self.fields.push(Filed {
34            name: name.to_string(),
35            data_type,
36        });
37
38        self
39    }
40
41    pub fn gen(&self) -> std::io::Result<()> {
42        let script = self.build_script();
43
44        let file_name = format!("{}.rs", self.schema_snake);
45        let file_path = self.out_dir.join(file_name.as_str());
46        let mut file_writer = File::create(&file_path)?;
47        file_writer.write_all(script.as_bytes())?;
48        file_writer.flush()?;
49
50        Ok(())
51    }
52
53    pub(crate) fn build_script(&self) -> String {
54        let use_script = self.build_use();
55        let data_type = self.build_data_type();
56        let field_reader = self.build_field_reader();
57        let field_writer = self.build_field_writer();
58        let entity = self.build_entity();
59
60        format!(
61            r#"#![allow(unknown_lints)]
62#![allow(clippy::all)]
63
64#![allow(unused_attributes)]
65#![rustfmt::skip]
66
67#![allow(box_pointers)]
68#![allow(dead_code)]
69#![allow(missing_docs)]
70#![allow(non_camel_case_types)]
71#![allow(non_snake_case)]
72#![allow(non_upper_case_globals)]
73#![allow(trivial_casts)]
74#![allow(unused_imports)]
75#![allow(unused_results)]
76//! Generated file by schema {}
77
78{}
79{}
80{}
81{}
82{}
83"#,
84            self.schema,
85            use_script.trim_end(),
86            data_type.trim_end(),
87            field_reader.trim_end(),
88            field_writer.trim_end(),
89            entity.trim_end(),
90        )
91    }
92
93    fn build_use(&self) -> String {
94        "use rlink_buffer::{types, BufferReader, BufferWriter, Buffer};\n".to_string()
95    }
96
97    fn build_data_type(&self) -> String {
98        let mut field_script = "".to_string();
99        for index in 0..self.fields.len() {
100            let field = self.fields.get(index).unwrap();
101            let dt = format!("{}", field.data_type);
102            let data_type = format!(
103                r#"    // {}: {}
104    types::{},
105"#,
106                index,
107                field.name,
108                dt.to_uppercase()
109            );
110            field_script = format!("{}{}", field_script, data_type);
111        }
112
113        let script = format!(
114            r#"
115pub const DATA_TYPE: [u8; {}] = [
116{}
117];"#,
118            self.fields.len(),
119            field_script.trim_end(),
120        );
121
122        script
123    }
124
125    fn build_field_reader(&self) -> String {
126        let mut field_read_method = "".to_string();
127        for index in 0..self.fields.len() {
128            let field = self.fields.get(index).unwrap();
129            let method_script = match field.data_type {
130                DataType::BYTES => format!(
131                    r#"
132    pub fn get_{}(&mut self) -> Result<&[u8], std::io::Error> {{
133        self.reader.get_bytes({})
134    }}
135"#,
136                    field.name, index
137                ),
138                DataType::STRING => format!(
139                    r#"
140    pub fn get_{}(&mut self) -> Result<String, std::io::Error> {{
141        self.reader.get_str({})
142    }}
143"#,
144                    field.name, index
145                ),
146                _ => format!(
147                    r#"
148    pub fn get_{}(&mut self) -> Result<{}, std::io::Error> {{
149        self.reader.get_{}({})
150    }}
151"#,
152                    field.name, field.data_type, field.data_type, index
153                ),
154            };
155
156            field_read_method = format!("{}{}", field_read_method, method_script);
157        }
158
159        format!(
160            r#"
161pub struct FieldReader<'a> {{
162    reader: BufferReader<'a, 'static>,
163}}
164
165impl<'a> FieldReader<'a> {{
166    pub fn new(b: &'a mut Buffer) -> Self {{
167        let reader = b.as_reader(&DATA_TYPE);
168        FieldReader {{ reader }}
169    }}
170{}
171}}"#,
172            field_read_method.trim_end()
173        )
174    }
175
176    fn build_field_writer(&self) -> String {
177        let mut field_read_method = "".to_string();
178        for index in 0..self.fields.len() {
179            let field = self.fields.get(index).unwrap();
180            let method_script = match field.data_type {
181                DataType::BYTES => format!(
182                    r#"
183    pub fn set_{}(&mut self, {}: &[u8]) -> Result<(), std::io::Error> {{
184        if self.writer_pos == {} {{
185            self.writer_pos += 1;
186            self.writer.set_bytes({})
187        }} else {{
188            Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "`{}` must be set sequentially"))
189        }}
190    }}
191"#,
192                    field.name, field.name, index, field.name, field.name,
193                ),
194                DataType::STRING => format!(
195                    r#"
196    pub fn set_{}(&mut self, {}: &str) -> Result<(), std::io::Error> {{
197        if self.writer_pos == {} {{
198            self.writer_pos += 1;
199            self.writer.set_str({})
200        }} else {{
201            Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "`{}` must be set sequentially"))
202        }}
203    }}
204"#,
205                    field.name, field.name, index, field.name, field.name,
206                ),
207                _ => format!(
208                    r#"
209    pub fn set_{}(&mut self, {}: {}) -> Result<(), std::io::Error> {{
210        if self.writer_pos == {} {{
211            self.writer_pos += 1;
212            self.writer.set_{}({})
213        }} else {{
214            Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "`{}` must be set sequentially"))
215        }}
216    }}
217"#,
218                    field.name,
219                    field.name,
220                    field.data_type,
221                    index,
222                    field.data_type,
223                    field.name,
224                    field.name,
225                ),
226            };
227
228            field_read_method = format!("{}{}", field_read_method, method_script);
229        }
230
231        format!(
232            r#"
233pub struct FieldWriter<'a> {{
234    writer: BufferWriter<'a, 'static>,
235    writer_pos: usize,
236}}
237
238impl<'a> FieldWriter<'a> {{
239    pub fn new(b: &'a mut Buffer) -> Self {{
240        let writer = b.as_writer(&DATA_TYPE);
241        FieldWriter {{
242            writer,
243            writer_pos: 0,
244        }}
245    }}
246{}
247}}"#,
248            field_read_method.trim_end()
249        )
250    }
251
252    fn build_entity(&self) -> String {
253        let mut fields = "".to_string();
254        let mut writers = "".to_string();
255        let mut readers = "".to_string();
256
257        for index in 0..self.fields.len() {
258            let field = self.fields.get(index).unwrap();
259
260            match field.data_type {
261                DataType::BYTES => {
262                    fields = format!("{}\n    pub {}: Vec<u8>,", fields, field.name);
263                    writers = format!(
264                        "{}\n        writer.set_bytes(self.{}.as_slice())?;",
265                        writers, field.name
266                    );
267                    readers = format!(
268                        "{}\n            {}: reader.get_bytes({})?.to_vec(),",
269                        readers, field.name, index
270                    );
271                }
272                DataType::STRING => {
273                    fields = format!("{}\n    pub {}: String,", fields, field.name);
274                    writers = format!(
275                        "{}\n        writer.set_str(self.{}.as_str())?;",
276                        writers, field.name
277                    );
278                    readers = format!(
279                        "{}\n            {}: reader.get_str({})?,",
280                        readers, field.name, index
281                    );
282                }
283                _ => {
284                    fields = format!("{}\n    pub {}: {},", fields, field.name, field.data_type);
285                    writers = format!(
286                        "{}\n        writer.set_{}(self.{})?;",
287                        writers, field.data_type, field.name
288                    );
289                    readers = format!(
290                        "{}\n            {}: reader.get_{}({})?,",
291                        readers, field.name, field.data_type, index
292                    );
293                }
294            };
295        }
296
297        format!(
298            r#"
299#[derive(Clone, Debug)]
300pub struct Entity {{
301    {}
302}}
303
304impl Entity {{
305    pub fn to_buffer(&self, b: &mut Buffer) -> Result<(), std::io::Error> {{
306        let mut writer = b.as_writer(&DATA_TYPE);
307        
308        {}
309
310        Ok(())
311    }}
312    
313    pub fn parse(b: &mut Buffer) -> Result<Self, std::io::Error> {{
314        let mut reader = b.as_reader(&DATA_TYPE);
315
316        let entity = Entity {{
317            {}
318        }};
319
320        Ok(entity)
321    }}
322}}
323            "#,
324            fields.trim_start(),
325            writers.trim_start(),
326            readers.trim_start()
327        )
328    }
329}
330
331fn to_snake(s: &str) -> String {
332    let mut v = Vec::new();
333    for c in s.chars() {
334        if c.is_uppercase() {
335            v.push('_');
336            v.push(c.to_ascii_lowercase());
337        } else {
338            v.push(c);
339        }
340    }
341
342    let snake = String::from_iter(v.iter());
343    if snake.starts_with("_") {
344        let ss = snake.as_str();
345        ss[1..snake.len()].to_string()
346    } else {
347        snake
348    }
349}
350
351pub enum DataType {
352    BOOL,
353    U8,
354    I8,
355    U16,
356    I16,
357    U32,
358    I32,
359    U64,
360    I64,
361    F32,
362    F64,
363    BYTES,
364    STRING,
365}
366
367impl Display for DataType {
368    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
369        match self {
370            DataType::BOOL => write!(f, "{}", "BOOL".to_lowercase()),
371            DataType::U8 => write!(f, "{}", "U8".to_lowercase()),
372            DataType::I8 => write!(f, "{}", "I8".to_lowercase()),
373            DataType::U16 => write!(f, "{}", "U16".to_lowercase()),
374            DataType::I16 => write!(f, "{}", "I16".to_lowercase()),
375            DataType::U32 => write!(f, "{}", "U32".to_lowercase()),
376            DataType::I32 => write!(f, "{}", "I32".to_lowercase()),
377            DataType::U64 => write!(f, "{}", "U64".to_lowercase()),
378            DataType::I64 => write!(f, "{}", "I64".to_lowercase()),
379            DataType::F32 => write!(f, "{}", "F32".to_lowercase()),
380            DataType::F64 => write!(f, "{}", "F64".to_lowercase()),
381            DataType::BYTES => write!(f, "{}", "BYTES".to_lowercase()),
382            DataType::STRING => write!(f, "{}", "STRING".to_lowercase()),
383        }
384    }
385}
386
387impl TryFrom<&str> for DataType {
388    type Error = &'static str;
389
390    fn try_from(value: &str) -> Result<Self, Self::Error> {
391        match value.to_uppercase().as_str() {
392            "BOOL" => Ok(DataType::BOOL),
393            "U8" => Ok(DataType::U8),
394            "I8" => Ok(DataType::I8),
395            "U16" => Ok(DataType::U16),
396            "I16" => Ok(DataType::I16),
397            "U32" => Ok(DataType::U32),
398            "I32" => Ok(DataType::I32),
399            "U64" => Ok(DataType::U64),
400            "I64" => Ok(DataType::I64),
401            "F32" => Ok(DataType::F32),
402            "F64" => Ok(DataType::F64),
403            "BYTES" => Ok(DataType::BYTES),
404            "STRING" => Ok(DataType::STRING),
405            _ => Err("unknown"),
406        }
407    }
408}
409
410struct Filed {
411    name: String,
412    data_type: DataType,
413}
414
415#[cfg(test)]
416mod tests {
417    use crate::{Codegen, DataType};
418
419    #[test]
420    pub fn code_gen_test() {
421        let script = Codegen::new("", "ApmDemo")
422            .field("timestamp", DataType::U64)
423            .field("application_name", DataType::STRING)
424            .field("agent_id", DataType::STRING)
425            .field("a", DataType::BOOL)
426            .field("a1", DataType::U8)
427            .field("a2", DataType::BYTES)
428            .build_script();
429
430        println!("-- Start ---");
431        println!("{}", script);
432        println!("-- End ---");
433    }
434}