scylladb_parse_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3
4#[proc_macro_derive(ParseFromStr, attributes(parse_via))]
5pub fn parse_from_str_derive(input: TokenStream) -> TokenStream {
6    let syn::DeriveInput {
7        attrs,
8        vis: _,
9        ident,
10        generics,
11        data: _,
12    } = syn::parse_macro_input!(input as syn::DeriveInput);
13    let (imp, ty, wher) = generics.split_for_impl();
14    let mut res = quote! {
15        impl #imp FromStr for #ident #ty #wher {
16            type Err = anyhow::Error;
17            fn from_str(s: &str) -> anyhow::Result<Self> {
18                StatementStream::new(s).parse()
19            }
20        }
21    };
22    if let Some(a) = attrs.iter().find(|a| a.path.is_ident("parse_via")) {
23        let via = match a.parse_meta().expect("Invalid parse_via attribute") {
24            syn::Meta::List(l) => {
25                if l.nested.len() != 1 {
26                    panic!("parse_via attribute must have exactly one argument");
27                }
28                match l.nested.into_iter().next().unwrap() {
29                    syn::NestedMeta::Meta(syn::Meta::Path(p)) => p,
30                    _ => panic!("parse_via attribute must contain a path to the via type"),
31                }
32            }
33            _ => panic!("parse_via attribute must be a list"),
34        };
35        res.extend(quote! {
36            impl #imp Parse for #ident #ty #wher {
37                type Output = Self;
38                fn parse(s: &mut StatementStream<'_>) -> anyhow::Result<Self::Output> {
39                    s.parse::<#via>()?.try_into()
40                }
41            }
42        });
43    }
44    res.into()
45}
46
47fn is_wrappable(ty: &syn::Type) -> bool {
48    match ty {
49        syn::Type::Path(syn::TypePath { qself: None, path }) => {
50            if let Some(seg) = path.segments.last() {
51                match seg.ident.to_string().as_str() {
52                    "Vec" | "Option" | "HashMap" | "Box" | "BTreeMap" | "BTreeSet" | "String" => return true,
53                    _ => (),
54                }
55            }
56        }
57        _ => (),
58    }
59    false
60}
61
62#[proc_macro_derive(ToTokens, attributes(wrap, tokenize_as))]
63pub fn to_tokens_derive(input: TokenStream) -> TokenStream {
64    let syn::DeriveInput {
65        attrs,
66        vis: _,
67        ident,
68        generics,
69        data,
70    } = syn::parse_macro_input!(input as syn::DeriveInput);
71    let mut imp_c = generics.clone();
72    imp_c.params.push(syn::parse_quote!('a));
73    let imp_ct = imp_c.split_for_impl().0;
74    let (imp, ty, wher) = generics.split_for_impl();
75    let mut tokenized: syn::Path = syn::parse_quote!(#ident);
76    if let Some(a) = attrs.iter().find(|a| a.path.is_ident("tokenize_as")) {
77        tokenized = match a.parse_meta().expect("Invalid tokenize_as attribute") {
78            syn::Meta::List(l) => {
79                if l.nested.len() != 1 {
80                    panic!("tokenize_as attribute must have exactly one argument");
81                }
82                match l.nested.into_iter().next().unwrap() {
83                    syn::NestedMeta::Meta(syn::Meta::Path(p)) => p,
84                    _ => panic!("tokenize_as attribute must contain a path to the tokenize type"),
85                }
86            }
87            _ => panic!("tokenize_as attribute must be a list"),
88        };
89    }
90    let res = match data {
91        syn::Data::Struct(s) => {
92            let (destr, restr) = match s.fields {
93                syn::Fields::Named(f) => {
94                    let names = f.named.iter().map(|f| f.ident.as_ref().unwrap()).collect::<Vec<_>>();
95                    let assigns = names.iter().map(|n| quote!(#n: ##n));
96                    let wrapped = f
97                        .named
98                        .iter()
99                        .map(|f| f.attrs.iter().find(|a| a.path.is_ident("wrap")).is_some() || is_wrappable(&f.ty))
100                        .zip(names.iter())
101                        .map(|(w, n)| {
102                            if w {
103                                quote!(TokenWrapper(&self.#n))
104                            } else {
105                                quote!(&self.#n)
106                            }
107                        });
108                    (
109                        quote! {
110                            let (#(#names),*) = (#(#wrapped),*);
111                        },
112                        quote! {
113                            #tokenized { #(#assigns),* }
114                        },
115                    )
116                }
117                syn::Fields::Unnamed(f) => {
118                    let names = f
119                        .unnamed
120                        .iter()
121                        .enumerate()
122                        .map(|(i, _)| {
123                            let idx = syn::Index::from(i);
124                            quote::format_ident!("f_{}", idx)
125                        })
126                        .collect::<Vec<_>>();
127                    let assigns = names.iter().map(|n| quote!(##n));
128                    let wrapped = f
129                        .unnamed
130                        .iter()
131                        .map(|f| f.attrs.iter().find(|a| a.path.is_ident("wrap")).is_some() || is_wrappable(&f.ty))
132                        .enumerate()
133                        .map(|(i, w)| {
134                            let idx = syn::Index::from(i);
135                            if w {
136                                quote!(TokenWrapper(&self.#idx))
137                            } else {
138                                quote!(&self.#idx)
139                            }
140                        });
141                    (
142                        quote! {
143                            let (#(#names),*) = (#(#wrapped),*);
144                        },
145                        quote! {
146                            #tokenized ( #(#assigns),* )
147                        },
148                    )
149                }
150                syn::Fields::Unit => (quote!(), quote!( #tokenized )),
151            };
152            quote! {
153                impl #imp_ct CustomToTokens<'a> for #ident #ty #wher {
154                    fn to_tokens(&'a self, tokens: &mut quote::__private::TokenStream) {
155                        #destr
156                        tokens.extend(quote::quote!(#restr));
157                    }
158                }
159
160                impl #imp quote::ToTokens for #ident #ty #wher {
161                    fn to_tokens(&self, tokens: &mut quote::__private::TokenStream) {
162                        CustomToTokens::to_tokens(self, tokens);
163                    }
164                }
165            }
166        }
167        syn::Data::Enum(e) => {
168            let variants = e.variants.into_iter().map(|v| {
169                let var_id = &v.ident;
170                let (destr, restr) = match v.fields {
171                    syn::Fields::Named(f) => {
172                        let names = f.named.iter().map(|f| f.ident.as_ref().unwrap()).collect::<Vec<_>>();
173                        let assigns = names.iter().map(|n| quote!(#n: ##n));
174                        let wrapped = f
175                            .named
176                            .iter()
177                            .map(|f| f.attrs.iter().find(|a| a.path.is_ident("wrap")).is_some() || is_wrappable(&f.ty))
178                            .zip(names.iter())
179                            .map(|(w, n)| if w { quote!(TokenWrapper(#n)) } else { quote!(#n) });
180                        (
181                            quote! {
182                                { #(#names),* }
183                            },
184                            quote! {
185                                {
186                                    let (#(#names),*) = (#(#wrapped),*);
187                                    quote::quote!(#tokenized::#var_id { #(#assigns),* })
188                                }
189                            },
190                        )
191                    }
192                    syn::Fields::Unnamed(f) => {
193                        let names = f
194                            .unnamed
195                            .iter()
196                            .enumerate()
197                            .map(|(i, _)| {
198                                let idx = syn::Index::from(i);
199                                quote::format_ident!("f_{}", idx)
200                            })
201                            .collect::<Vec<_>>();
202                        let assigns = names.iter().map(|n| quote!(##n));
203                        let wrapped = f
204                            .unnamed
205                            .iter()
206                            .map(|f| f.attrs.iter().find(|a| a.path.is_ident("wrap")).is_some() || is_wrappable(&f.ty))
207                            .zip(names.iter())
208                            .map(|(w, n)| if w { quote!(TokenWrapper(#n)) } else { quote!(#n) });
209                        (
210                            quote! {
211                                ( #(#names),* )
212                            },
213                            quote! {
214                                {
215                                    let (#(#names),*) = (#(#wrapped),*);
216                                    quote::quote!(#tokenized::#var_id ( #(#assigns),* ))
217                                }
218                            },
219                        )
220                    }
221                    syn::Fields::Unit => (quote!(), quote!(quote::quote!(#tokenized::#var_id))),
222                };
223                quote! {
224                    #ident::#var_id #destr => #restr
225                }
226            });
227            quote! {
228                impl #imp_ct CustomToTokens<'a> for #ident #ty #wher {
229                    fn to_tokens(&'a self, tokens: &mut quote::__private::TokenStream) {
230                        tokens.extend(match self {
231                            #(#variants),*
232                        })
233                    }
234                }
235
236                impl #imp quote::ToTokens for #ident #ty #wher {
237                    fn to_tokens(&self, tokens: &mut quote::__private::TokenStream) {
238                        CustomToTokens::to_tokens(self, tokens);
239                    }
240                }
241            }
242        }
243        syn::Data::Union(_) => panic!("Unions not supported!"),
244    };
245    res.into()
246}