irox_structs_derive/
lib.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2024 IROX Contributors
3//
4
5use proc_macro::{Literal, TokenStream};
6use std::fmt::Display;
7
8use quote::{quote, ToTokens};
9use syn::spanned::Spanned;
10use syn::{parse_macro_input, Data, DeriveInput, Error, Fields, FieldsNamed};
11
12use irox_types::{PrimitiveType, Primitives, VariableType};
13
14use irox_derive_helpers::DeriveMethods;
15
16const TYPES_STRICT_SIZING_INCOMPATIBLE: [Primitives; 1] = [Primitives::null];
17
18struct Config {
19    strict_sizing: bool,
20    big_endian: bool,
21}
22
23fn compile_error<T: Spanned, D: Display>(span: &T, msg: D) -> TokenStream {
24    Error::new(span.span(), msg).into_compile_error().into()
25}
26
27fn get_endian_method_for_prim(ty: Primitives, read: bool, big_endian: bool) -> String {
28    let rw = if read { "read" } else { "write" };
29    let be = if big_endian { "be" } else { "le" };
30    let base = match ty {
31        Primitives::u8 | Primitives::bool | Primitives::char => "u8".to_string(),
32        Primitives::i8 => "i8".to_string(),
33
34        _ => {
35            format!("{be}_{ty:?}")
36        }
37    };
38    format!("{rw}_{base}")
39}
40
41fn get_endian_method_for_varbl(ty: VariableType, read: bool, big_endian: bool) -> String {
42    let rw = if read { "read" } else { "write" };
43    let be = if big_endian { "be" } else { "le" };
44    let base = match ty {
45        VariableType::str => "str_u32_blob".to_string(),
46        _ => {
47            format!("{be}_{ty:?}")
48        }
49    };
50    format!("{rw}_{base}")
51}
52
53fn create_write_to_fn(n: &FieldsNamed, config: &Config, sizing: &mut StructSizing) -> TokenStream {
54    let mut ts = TokenStream::new();
55    ts.extend::<TokenStream>(
56        quote!(
57            fn write_to<T: irox_structs::MutBits>(&self, out: &mut T) -> Result<(), irox_structs::Error>
58        )
59        .into(),
60    );
61
62    let mut method = TokenStream::new();
63
64    for x in &n.named {
65        let Some(ident) = &x.ident else {
66            return compile_error(&x, "No ident");
67        };
68        match PrimitiveType::try_from(x) {
69            Ok(field) => {
70                if let Some(size) = field.bytes_length() {
71                    sizing.size += size;
72                }
73                match field {
74                    PrimitiveType::Primitive(input) => {
75                        if config.strict_sizing && TYPES_STRICT_SIZING_INCOMPATIBLE.contains(&input)
76                        {
77                            return compile_error(&x, "Type is not compatible with strict sizing");
78                        };
79                        method.add_ident("out");
80                        method.add_punc('.');
81                        method.add_ident(&get_endian_method_for_prim(
82                            input,
83                            false,
84                            config.big_endian,
85                        ));
86                        method.add_parens({
87                            let mut ts = TokenStream::new();
88                            ts.add_ident("self");
89                            ts.add_punc('.');
90                            ts.add_ident(&ident.to_string());
91                            ts
92                        });
93                        method.add_punc('?');
94                        method.add_punc(';');
95                    }
96                    PrimitiveType::Array(input, len) => {
97                        if config.strict_sizing && TYPES_STRICT_SIZING_INCOMPATIBLE.contains(&input)
98                        {
99                            return compile_error(&x, "Type is not compatible with strict sizing");
100                        };
101                        method.add_ident("for");
102                        method.add_ident("elem");
103                        method.add_ident("in");
104                        method.add_ident("self");
105                        method.add_punc('.');
106                        method.add_ident(&ident.to_string());
107                        method.wrap_braces({
108                            let mut ts = TokenStream::new();
109                            for _ in 0..len {
110                                ts.add_ident("out");
111                                ts.add_punc('.');
112                                ts.add_ident(&get_endian_method_for_prim(
113                                    input,
114                                    false,
115                                    config.big_endian,
116                                ));
117                                ts.add_parens(TokenStream::create_ident("elem"));
118                                ts.add_punc('?');
119                                ts.add_punc(';');
120                            }
121                            ts
122                        })
123                    }
124                    PrimitiveType::DynamicallySized(dy) => {
125                        if config.strict_sizing {
126                            return compile_error(&x, "Type is not compatible with strict sizing");
127                        };
128                        method.add_ident("out");
129                        method.add_punc('.');
130                        method.add_ident(&get_endian_method_for_varbl(
131                            dy,
132                            false,
133                            config.big_endian,
134                        ));
135                        method.add_parens({
136                            let mut ts = TokenStream::new();
137                            ts.add_punc('&');
138                            ts.add_ident("self");
139                            ts.add_punc('.');
140                            ts.add_ident(&ident.to_string());
141                            ts
142                        });
143                        method.add_punc('?');
144                        method.add_punc(';');
145                    }
146                }
147            }
148            Err(_e) => {
149                // <ty as irox_structs::Struct>::write_to(&self.varbl, out)?;
150                let mut ts = TokenStream::new();
151                ts.wrap_generics({
152                    let mut ts = TokenStream::new();
153                    ts.extend::<TokenStream>(x.ty.to_token_stream().into());
154                    ts.add_ident("as");
155                    ts.add_ident("irox_structs");
156                    ts.add_punc2(':', ':');
157                    ts.add_ident("Struct");
158                    ts
159                });
160                ts.add_punc2(':', ':');
161                ts.add_ident("write_to");
162                ts.add_parens({
163                    let mut ts = TokenStream::new();
164                    ts.add_punc('&');
165                    ts.add_ident("self");
166                    ts.add_punc('.');
167                    ts.add_ident(&ident.to_string());
168                    ts.add_punc(',');
169                    ts.add_ident("out");
170                    ts
171                });
172                ts.add_punc('?');
173                ts.add_punc(';');
174                method.extend(ts);
175            }
176        }
177    }
178    method.add_ident("Ok");
179    method.add_parens(TokenStream::create_empty_type());
180    ts.wrap_braces(method);
181    ts
182}
183
184fn create_parse_from_fn(n: &FieldsNamed, config: &Config) -> TokenStream {
185    let mut ts = TokenStream::new();
186    ts.extend::<TokenStream>(quote!(
187        fn parse_from<T: irox_structs::Bits>(input: &mut T) -> Result<Self::ImplType, irox_structs::Error>
188    ).into());
189
190    let mut inits = TokenStream::new();
191
192    for x in &n.named {
193        let Some(ident) = &x.ident else {
194            return compile_error(&x, "No ident");
195        };
196
197        match PrimitiveType::try_from(x) {
198            Ok(field) => match field {
199                PrimitiveType::Primitive(input) => {
200                    if config.strict_sizing && TYPES_STRICT_SIZING_INCOMPATIBLE.contains(&input) {
201                        return compile_error(&x, "Type is not compatible with strict sizing");
202                    };
203                    inits.add_ident(&ident.to_string());
204                    inits.add_punc(':');
205                    inits.add_ident("input");
206                    inits.add_punc('.');
207                    inits.add_ident(&get_endian_method_for_prim(input, true, config.big_endian));
208                    inits.add_parens(TokenStream::new());
209                    inits.add_punc('?');
210                    inits.add_punc(',');
211                }
212                PrimitiveType::Array(input, len) => {
213                    if config.strict_sizing && TYPES_STRICT_SIZING_INCOMPATIBLE.contains(&input) {
214                        return compile_error(&x, "Type is not compatible with strict sizing");
215                    };
216                    inits.add_ident(&ident.to_string());
217                    inits.add_punc(':');
218                    inits.wrap_brackets({
219                        let mut ts = TokenStream::new();
220                        for _ in 0..len {
221                            ts.add_ident("input");
222                            ts.add_punc('.');
223                            ts.add_ident(&get_endian_method_for_prim(
224                                input,
225                                true,
226                                config.big_endian,
227                            ));
228                            ts.add_parens(TokenStream::new());
229                            ts.add_punc('?');
230                            ts.add_punc(',');
231                        }
232                        ts
233                    });
234                    inits.add_punc('.');
235                    inits.add_ident("into");
236                    inits.add_parens(TokenStream::new());
237                    inits.add_punc(',');
238                }
239                PrimitiveType::DynamicallySized(ds) => {
240                    if config.strict_sizing {
241                        return compile_error(&x, "Type is not compatible with strict sizing");
242                    };
243
244                    inits.add_ident(&ident.to_string());
245                    inits.add_punc(':');
246                    inits.add_ident("input");
247                    inits.add_punc('.');
248                    inits.add_ident(&get_endian_method_for_varbl(ds, true, config.big_endian));
249                    inits.add_parens(TokenStream::new());
250                    inits.add_punc('?');
251                    inits.add_punc(',');
252                }
253            },
254            Err(_e) => {
255                // <ty as irox_structs::Struct>::parse_from(input)?;
256                let ty = x.ty.to_token_stream();
257                inits.add_ident(&ident.to_string());
258                inits.add_punc(':');
259                inits.extend::<TokenStream>(
260                    quote! {
261                        <#ty as irox_structs::Struct>::parse_from(input)?,
262                    }
263                    .into(),
264                );
265            }
266        }
267    }
268
269    let mut method = TokenStream::new();
270    method.add_ident("Ok");
271    method.add_parens({
272        let mut ts = TokenStream::new();
273        ts.add_ident("Self");
274        ts.wrap_braces(inits);
275        ts
276    });
277    ts.wrap_braces(method);
278    ts
279}
280
281#[derive(Default)]
282struct StructSizing {
283    size: usize,
284}
285
286#[proc_macro_derive(Struct, attributes(little_endian, big_endian, strict_sizing))]
287pub fn struct_derive(input: TokenStream) -> TokenStream {
288    let input = parse_macro_input!(input as DeriveInput);
289    let mut config = Config {
290        big_endian: true,
291        strict_sizing: false,
292    };
293
294    for attr in &input.attrs {
295        let Ok(ident) = attr.meta.path().require_ident() else {
296            return compile_error(&attr, "This attribute is unnamed.".to_string());
297        };
298        if ident.eq("little_endian") {
299            config.big_endian = false;
300        } else if ident.eq("big_endian") {
301            config.big_endian = true;
302        } else if ident.eq("strict_sizing") {
303            config.strict_sizing = true;
304        }
305    }
306
307    let struct_name = &input.ident;
308
309    let Data::Struct(s) = input.data else {
310        return compile_error(&input, "Can only derive on struct type");
311    };
312    let Fields::Named(n) = s.fields else {
313        return compile_error(&s.fields, "Can only derive on named fields.");
314    };
315
316    let mut ts = TokenStream::new();
317    let mut sizing = StructSizing::default();
318    ts.add_ident("impl");
319    ts.add_path(&["irox_structs", "Struct"]);
320    ts.add_ident("for");
321    ts.add_ident(&struct_name.to_string());
322    ts.wrap_braces({
323        let mut ts = TokenStream::new();
324        ts.add_ident("type");
325        ts.add_ident("ImplType");
326        ts.add_punc('=');
327        ts.add_ident(&struct_name.to_string());
328        ts.add_punc(';');
329
330        ts.extend(create_write_to_fn(&n, &config, &mut sizing));
331        ts.extend(create_parse_from_fn(&n, &config));
332        ts
333    });
334    if config.strict_sizing {
335        ts.add_ident("impl");
336        ts.add_ident(&struct_name.to_string());
337        ts.wrap_braces({
338            let mut ts = TokenStream::new();
339            ts.add_ident("pub");
340            ts.add_ident("const");
341            ts.add_ident("STRUCT_SIZE");
342            ts.add_punc(':');
343            ts.add_ident("usize");
344            ts.add_punc('=');
345            ts.add_literal(Literal::usize_suffixed(sizing.size));
346            ts.add_punc(';');
347            ts
348        });
349    }
350    ts
351}