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

enum_assoc/
lib.rs
