enum_assoc/
lib.rs

1#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
2
3use proc_macro::TokenStream;
4use quote::{ToTokens, quote};
5use syn::{
6    Error, FnArg, Result, Token, Variant, parenthesized, parse::Parser, punctuated::Punctuated,
7    spanned::Spanned,
8};
9
10const FUNC_ATTR: &str = "func";
11const ASSOC_ATTR: &str = "assoc";
12
13#[proc_macro_derive(Assoc, attributes(func, assoc))]
14pub fn derive_assoc(input: TokenStream) -> TokenStream {
15    impl_macro(&syn::parse(input).expect("Failed to parse macro input"))
16        //.map(|t| {println!("{}", quote!(#t)); t})
17        .unwrap_or_else(syn::Error::into_compile_error)
18        .into()
19}
20
21fn impl_macro(ast: &syn::DeriveInput) -> Result<proc_macro2::TokenStream> {
22    let name = &ast.ident;
23    let generics = &ast.generics;
24    let generic_params = &generics.params;
25    let fns = ast
26        .attrs
27        .iter()
28        .filter(|attr| attr.path().is_ident(FUNC_ATTR))
29        .map(|attr| syn::parse2::<DeriveFuncs>(attr.meta.to_token_stream()))
30        .collect::<Result<Vec<DeriveFuncs>>>()?;
31    let variants: Vec<&Variant> = if let syn::Data::Enum(data) = &ast.data {
32        data.variants.iter().collect()
33    } else {
34        panic!("#[derive(Assoc)] only applicable to enums")
35    };
36    let functions: Vec<proc_macro2::TokenStream> = fns
37        .into_iter()
38        .flat_map(|DeriveFuncs(funcs)| {
39            funcs
40                .iter()
41                .map(|func| build_function(&variants, func, funcs.clone()))
42                .collect::<Vec<_>>()
43        })
44        .collect::<Result<Vec<proc_macro2::TokenStream>>>()?;
45    Ok(quote! {
46        #[allow(clippy::used_underscore_binding)]
47        impl <#generic_params> #name #generics
48        {
49            #(#functions)*
50        }
51    })
52}
53
54fn build_function(
55    variants: &[&Variant],
56    func: &DeriveFunc,
57    associated_funcs: Vec<DeriveFunc>,
58) -> Result<proc_macro2::TokenStream> {
59    let vis = &func.vis;
60    let sig = &func.sig;
61    // has_self determines whether or not this a reverse assoc
62    let has_self = match func.sig.inputs.first() {
63        Some(FnArg::Receiver(_)) => true,
64        Some(FnArg::Typed(pat_type)) => {
65            let pat = &pat_type.pat;
66            quote!(#pat).to_string().trim() == "self"
67        }
68        None => false,
69    };
70    let is_option = if let syn::ReturnType::Type(_, ty) = &func.sig.output {
71        let s = quote!(#ty).to_string();
72        let trimmed = s.trim();
73        trimmed.starts_with("Option") && trimmed.len() > 6 && trimmed[6..].trim().starts_with("<")
74    } else {
75        false
76    };
77    let mut arms = variants
78        .iter()
79        .map(|variant| {
80            build_variant_arm(
81                variant,
82                &func.sig.ident,
83                associated_funcs.iter().map(|func| func.sig.ident.clone()),
84                is_option,
85                has_self,
86                &func.def,
87            )
88        })
89        .collect::<Result<Vec<(proc_macro2::TokenStream, Wildcard)>>>()?;
90    if is_option
91        && !arms
92            .iter()
93            .any(|(_, wildcard)| matches!(wildcard, Wildcard::True))
94    {
95        arms.push((quote!(_ => None,), Wildcard::True))
96    }
97    // make sure wildcards are last
98    if !has_self {
99        arms.sort_by(|(_, wildcard1), (_, wildcard2)| wildcard1.cmp(wildcard2));
100    }
101    let arms = arms.into_iter().map(|(toks, _)| toks);
102    let match_on = if has_self {
103        quote!(self)
104    } else if func.sig.inputs.is_empty() {
105        return Err(syn::Error::new(func.span, "Missing parameter"));
106    } else {
107        let mut result = quote!();
108        for input in &func.sig.inputs {
109            match input {
110                FnArg::Receiver(_) => {
111                    result = quote!(self);
112                    break;
113                }
114                FnArg::Typed(pat_type) => {
115                    let pat = &pat_type.pat;
116                    result = if result.is_empty() {
117                        quote!(#pat)
118                    } else {
119                        quote!(#result, #pat)
120                    };
121                }
122            }
123        }
124        if func.sig.inputs.len() > 1 {
125            result = quote!((#result));
126        }
127        result
128    };
129    Ok(quote! {
130        #vis #sig
131        {
132            match #match_on
133            {
134                #(#arms)*
135            }
136        }
137    })
138}
139
140fn build_variant_arm(
141    variant: &Variant,
142    func: &syn::Ident,
143    mut assoc_funcs: impl Iterator<Item = syn::Ident>,
144    is_option: bool,
145    has_self: bool,
146    def: &Option<proc_macro2::TokenStream>,
147) -> Result<(proc_macro2::TokenStream, Wildcard)> {
148    // Partially parse associations
149    let assocs = Association::get_variant_assocs(variant, !has_self).filter(|assoc| {
150        assoc.func == *func || assoc_funcs.any(|assoc_func| assoc_func == assoc.func)
151    });
152    if has_self {
153        build_fwd_assoc(assocs, variant, is_option, func, def)
154    } else {
155        build_rev_assoc(assocs, variant, is_option)
156    }
157}
158
159fn build_fwd_assoc(
160    assocs: impl Iterator<Item = Association>,
161    variant: &Variant,
162    is_option: bool,
163    func_ident: &syn::Ident,
164    def: &Option<proc_macro2::TokenStream>,
165) -> Result<(proc_macro2::TokenStream, Wildcard)> {
166    let var_ident = &variant.ident;
167    let fields = match &variant.fields {
168        syn::Fields::Named(fields) => {
169            let named = fields
170                .named
171                .iter()
172                .map(|f| {
173                    let ident = &f.ident;
174                    let val: &Option<proc_macro2::Ident> = &f.ident.as_ref().map(|s| {
175                        proc_macro2::Ident::new(&("_".to_string() + &s.to_string()), f.span())
176                    });
177                    quote!(#ident: #val)
178                })
179                .collect::<Vec<proc_macro2::TokenStream>>();
180            quote!({#(#named),*})
181        }
182        syn::Fields::Unnamed(fields) => {
183            let unnamed = fields
184                .unnamed
185                .iter()
186                .enumerate()
187                .map(|(i, f)| {
188                    let ident =
189                        proc_macro2::Ident::new(&("_".to_string() + &i.to_string()), f.span());
190                    quote!(#ident)
191                })
192                .collect::<Vec<proc_macro2::TokenStream>>();
193            quote!((#(#unnamed),*))
194        }
195        _ => quote!(),
196    };
197    let assocs = assocs
198        .filter_map(|assoc| {
199            if let AssociationType::Forward(expr) = assoc.assoc {
200                Some(Ok(expr))
201            } else {
202                None
203            }
204        })
205        .collect::<Result<Vec<syn::Expr>>>()?;
206    match assocs.len() {
207        0 => {
208            if let Some(tokens) = def {
209                Ok(quote! { Self::#var_ident #fields => #tokens, })
210            } else if is_option {
211                Ok(quote! { Self::#var_ident #fields => None, })
212            } else {
213                Err(Error::new_spanned(
214                    variant,
215                    format!("Missing `assoc` attribute for {}", func_ident),
216                ))
217            }
218        }
219        1 => {
220            let val = &assocs[0];
221            if is_option {
222                if quote!(#val).to_string().trim() == "None" {
223                    Ok(quote! { Self::#var_ident #fields => #val, })
224                } else {
225                    Ok(quote! { Self::#var_ident #fields => Some(#val), })
226                }
227            } else {
228                Ok(quote! { Self::#var_ident #fields => #val, })
229            }
230        }
231        _ => Err(Error::new_spanned(
232            variant,
233            format!("Too many `assoc` attributes for {}", func_ident),
234        )),
235    }
236    .map(|toks| (toks, Wildcard::None))
237}
238
239fn build_rev_assoc(
240    assocs: impl Iterator<Item = Association>,
241    variant: &Variant,
242    is_option: bool,
243) -> Result<(proc_macro2::TokenStream, Wildcard)> {
244    let var_ident = &variant.ident;
245    let assocs = assocs
246        .filter_map(|assoc| {
247            if let AssociationType::Reverse(pat) = assoc.assoc {
248                Some(Ok(pat))
249            } else {
250                None
251            }
252        })
253        .collect::<Result<Vec<syn::Pat>>>()?;
254    let mut concrete_pats: Vec<proc_macro2::TokenStream> = Vec::new();
255    let mut wildcard_pat: Option<proc_macro2::TokenStream> = None;
256    let mut wildcard_status = Wildcard::False;
257    for pat in assocs.iter() {
258        if !matches!(variant.fields, syn::Fields::Unit) {
259            return Err(Error::new_spanned(
260                variant,
261                "Reverse associations not allowed for tuple or struct-like variants",
262            ));
263        }
264        let arm = if is_option {
265            quote!(#pat => Some(Self::#var_ident),)
266        } else {
267            quote!(#pat => Self::#var_ident,)
268        };
269        if matches!(pat, syn::Pat::Wild(_)) {
270            if wildcard_pat.is_some() {
271                return Err(syn::Error::new_spanned(
272                    pat,
273                    "Only 1 wildcard allowed per reverse association",
274                ));
275            }
276            wildcard_status = Wildcard::True;
277            wildcard_pat = Some(arm);
278        } else {
279            concrete_pats.push(arm);
280        }
281    }
282    if let Some(wildcard_pat) = wildcard_pat {
283        concrete_pats.push(wildcard_pat)
284    }
285    Ok((quote!(#(#concrete_pats) *), wildcard_status))
286}
287
288/// A container for a function parsed within a `func` attribute. Note that the
289/// span of the `func` atribute is included because the syn nodes were
290/// manipulated as a string and have lost therr own span information.
291#[derive(Clone)]
292struct DeriveFunc {
293    vis: syn::Visibility,
294    sig: syn::Signature,
295    span: proc_macro2::Span,
296    def: Option<proc_macro2::TokenStream>,
297}
298
299/// An association. Contains a function ident as well as the actual tokens of
300/// the VALUE (not the variant) of the association.
301struct Association {
302    func: syn::Ident,
303    assoc: AssociationType,
304}
305
306enum AssociationType {
307    Forward(syn::Expr),
308    Reverse(syn::Pat),
309}
310
311/// For reverse associations, this enum keeps track of wldcard patterns. For
312/// forward associations, the value is always set to "None". This is also used
313/// to sort reverse associations appropriately. If more complex sorting is to
314/// be implemented, updating this enum would be the best way to start.
315#[derive(PartialEq, Eq, PartialOrd, Ord)]
316enum Wildcard {
317    False = 0,
318    None = 1,
319    True = 2,
320}
321
322impl syn::parse::Parse for DeriveFunc {
323    /// Parse a function signature from an attribute
324    fn parse(input: syn::parse::ParseStream) -> Result<Self> {
325        let vis = input.parse::<syn::Visibility>()?;
326        let sig = input.parse::<syn::Signature>()?;
327        let def = if let Ok(block) = input.parse::<syn::Block>() {
328            Some(proc_macro2::TokenStream::from(ToTokens::into_token_stream(
329                block,
330            )))
331        } else {
332            None
333        };
334        Ok(DeriveFunc {
335            vis,
336            sig,
337            span: input.span(),
338            def,
339        })
340    }
341}
342
343struct DeriveFuncs(Vec<DeriveFunc>);
344impl syn::parse::Parse for DeriveFuncs {
345    /// Parse a list of function signatures form an attribute
346    fn parse(input: syn::parse::ParseStream) -> Result<Self> {
347        input.step(|cursor| {
348            if let Some((_, next)) = cursor.token_tree() {
349                Ok(((), next))
350            } else {
351                Err(cursor.error("Missing function signature"))
352            }
353        })?;
354        let content;
355        parenthesized!(content in input);
356        Ok(Self(
357            content
358                .parse_terminated(DeriveFunc::parse, Token!(,))
359                .map(|parsed| parsed.into_iter().collect())?,
360        ))
361    }
362}
363
364/// Used to parse forward associations, which are of form Ident = Expr
365struct ForwardAssocTokens(syn::Ident, syn::Expr);
366impl syn::parse::Parse for ForwardAssocTokens {
367    fn parse(input: syn::parse::ParseStream) -> Result<Self> {
368        let ident = input.parse()?;
369        input.parse::<syn::Token!(=)>()?;
370        let expr = input.parse()?;
371        Ok(Self(ident, expr))
372    }
373}
374
375/// Used to parse reverse associations, which are of form Ident = Pat
376struct ReverseAssocTokens(syn::Ident, syn::Pat);
377impl syn::parse::Parse for ReverseAssocTokens {
378    fn parse(input: syn::parse::ParseStream) -> Result<Self> {
379        let ident = input.parse()?;
380        input.parse::<syn::Token!(=)>()?;
381        let pat = syn::Pat::parse_multi_with_leading_vert(input)?;
382        Ok(Self(ident, pat))
383    }
384}
385
386impl From<ForwardAssocTokens> for Association {
387    fn from(val: ForwardAssocTokens) -> Self {
388        Association {
389            func: val.0,
390            assoc: AssociationType::Forward(val.1),
391        }
392    }
393}
394
395impl From<ReverseAssocTokens> for Association {
396    fn from(val: ReverseAssocTokens) -> Self {
397        Association {
398            func: val.0,
399            assoc: AssociationType::Reverse(val.1),
400        }
401    }
402}
403
404impl Association {
405    fn get_variant_assocs(variant: &Variant, is_reverse: bool) -> impl Iterator<Item = Self> + '_ {
406        variant
407            .attrs
408            .iter()
409            .filter(|assoc_attr| assoc_attr.path().is_ident(ASSOC_ATTR))
410            .filter_map(move |attr| {
411                if let syn::Meta::List(meta_list) = &attr.meta {
412                    if is_reverse {
413                        let parser = Punctuated::<ReverseAssocTokens, Token![,]>::parse_terminated;
414                        parser
415                            .parse2(meta_list.tokens.clone())
416                            .map(|tokens| {
417                                tokens
418                                    .into_iter()
419                                    .map(|tokens| tokens.into())
420                                    .collect::<Vec<Self>>()
421                            })
422                            .ok()
423                    } else {
424                        let parser = Punctuated::<ForwardAssocTokens, Token![,]>::parse_terminated;
425                        parser
426                            .parse2(meta_list.tokens.clone())
427                            .map(|tokens| {
428                                tokens
429                                    .into_iter()
430                                    .map(|tokens| tokens.into())
431                                    .collect::<Vec<Self>>()
432                            })
433                            .ok()
434                    }
435                } else {
436                    None
437                }
438            })
439            .flatten()
440    }
441}
enum_assoc/lib.rs

enum_assoc/
lib.rs
