rocket_model_codegen/
lib.rs

1use quote::{ToTokens, quote};
2use syn::{braced, parse_macro_input, parse_quote};
3use syn::{Attribute, Visibility, Ident, Field, Fields, FieldsNamed, Token, ItemStruct, Generics, Type, token};
4use syn::parse::{Parse, ParseStream, Result};
5use syn::punctuated::Punctuated;
6
7enum MetaField {
8    Add(Field),
9    Remove(Ident),
10}
11
12impl Parse for MetaField {
13    fn parse(input: ParseStream) -> Result<Self> {
14        Ok(
15            match input.parse::<Token![-]>() {
16                Ok(_) => MetaField::Remove(input.parse()?),
17                Err(_) => MetaField::Add(input.call(Field::parse_named)?),
18            }
19        )
20    }
21}
22
23struct MetaFields {
24    add_fields: Punctuated<Field, Token![,]>,
25    rem_fields: Vec<Ident>,
26}
27
28impl Parse for MetaFields {
29    fn parse(input: ParseStream) -> Result<Self> {
30        let fields: Punctuated<_, Token![,]> = Punctuated::parse_terminated(input)?;
31
32        Ok(
33            MetaFields {
34                add_fields: fields.iter().filter_map(|f| match f {
35                    MetaField::Add(a) => Some(a.clone()),
36                    _ => None,
37                }).collect(),
38                rem_fields: fields.iter().filter_map(|f| match f {
39                    MetaField::Remove(r) => Some(r.clone()),
40                    _ => None,
41                }).collect(),
42            }
43        )
44    }
45}
46
47struct MetaStruct {
48    attrs: Vec<Attribute>,
49    visibility: Visibility,
50    struct_token: Token![struct],
51    name: Ident,
52    generics: Generics,
53    question_token: Option<Token![?]>,
54    brace_token: token::Brace,
55    fields: MetaFields,
56}
57
58impl Parse for MetaStruct {
59    fn parse(input: ParseStream) -> Result<Self> {
60        let content;
61        Ok(MetaStruct {
62            attrs: input.call(Attribute::parse_outer)?,
63            visibility: input.parse()?,
64            struct_token: input.parse()?,
65            name: input.parse()?,
66            generics: input.parse()?,
67            question_token: input.parse().ok(),
68            brace_token: braced!(content in input),
69            fields: content.parse()?,
70        })
71    }
72}
73
74struct MetaStructs(Vec<MetaStruct>);
75
76impl Parse for MetaStructs {
77    fn parse(input: ParseStream) -> Result<Self> {
78        let mut structs = vec![];
79
80        while !input.is_empty() {
81            structs.push(input.parse()?);
82        }
83
84        Ok(MetaStructs(structs))
85    }
86}
87
88impl ToTokens for MetaStructs {
89    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
90        for (i, meta_struct) in self.0.iter().enumerate() {
91            let mut named = Punctuated::new();
92
93            'outer: for field in &meta_struct.fields.add_fields {
94                for rem_field in &meta_struct.fields.rem_fields {
95                    if field.ident.as_ref().unwrap().eq(rem_field) {
96                        continue 'outer;
97                    }
98                }
99
100                named.push(field.clone());
101            }
102
103            if i != 0 {
104                let parent_fields = &self.0.first().unwrap().fields;
105
106                'parent_outer: for field in &parent_fields.add_fields {
107                    for rem_field in &meta_struct.fields.rem_fields {
108                        if field.ident.as_ref().unwrap().eq(rem_field) {
109                            continue 'parent_outer;
110                        }
111                    }
112    
113                    named.push(field.clone());
114                }
115
116                if meta_struct.question_token.is_some() {
117                    named = named.into_iter().map(|f| {
118                        let orig = f.ty;
119                        Field {
120                            ty: parse_quote! {
121                                Option<#orig>
122                            },
123                            ..f 
124                        }
125                    }).collect();
126                }
127            }
128
129            ItemStruct {
130                attrs: meta_struct.attrs.clone(),
131                vis: meta_struct.visibility.clone(),
132                struct_token: meta_struct.struct_token,
133                ident: meta_struct.name.clone(),
134                generics: meta_struct.generics.clone(),
135                fields: Fields::Named(FieldsNamed {
136                    brace_token: meta_struct.brace_token,
137                    named,
138                }),
139                semi_token: None,
140            }.to_tokens(tokens);
141        }
142    }
143}
144
145#[proc_macro]
146pub fn gen_structs(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
147    let meta_structs = parse_macro_input!(input as MetaStructs);
148
149    let token_stream = quote! {
150        #meta_structs
151    };
152
153    token_stream.into()
154}