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#[derive(Default)]
13pub struct Codegen {
14 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}