separable_derive/
lib.rs

1use proc_macro::{TokenStream};
2use proc_macro2::{Span};
3use syn::{parse_macro_input, DeriveInput, Data, Ident, Fields, spanned::Spanned};
4use quote::quote;
5
6use self::error::Error;
7mod error;
8
9/// Derive macro for the `Separable` trait
10///
11/// See the main page of the documentation for a quick example
12#[proc_macro_derive(Separable)]
13pub fn separable_derive(input: TokenStream) -> TokenStream {
14    // We first parse de input
15    let input = parse_macro_input!(input as DeriveInput);
16
17    separable_derive_wrapper(input).unwrap_or_else(|e| e.to_syn_error().to_compile_error().into())
18}
19
20fn separable_derive_wrapper(input: DeriveInput) -> Result<TokenStream, Error> {
21    // We quickly check if the derive candidate is a struct
22    let enum_data = match input.data {
23        Data::Enum(enum_data) => enum_data,
24        Data::Struct(s) => return Err(Error::custom("derive macro works only with enums", s.struct_token.span)),
25        Data::Union(u) => return Err(Error::custom("derive macro works only with enums", u.union_token.span))
26    };
27
28    // Name of the enum
29    let name = &input.ident;
30
31    // Check that all the variants are tupples with one element
32    let types = enum_data.variants.iter().map(|variant| {
33        match &variant.fields {
34            Fields::Unnamed(unnamed) => {
35                if unnamed.unnamed.len() != 1 {
36                    Err(Error::custom("tupple variants may contain only 1 element", unnamed.span()))
37                } else {
38                    let field = unnamed.unnamed.first().unwrap();
39                    let field_type = &field.ty;
40                    //Ok(quote! {Vec<#field_type>})
41                    Ok(field_type)
42                }
43            },
44            _ => Err(Error::custom("unsupported variant, this macro only works with unnamed single-value tuple variants", variant.ident.span()))
45        }
46    }).collect::<Result<Vec<_>, _>>()?;
47
48    // We create the referenced types, and mutable referenced types
49    let vecced_types = types.iter().map(|ty| quote! {Vec<#ty>}).collect::<Vec<_>>();
50    let ref_vecced_types = types.iter().map(|ty| quote! {Vec<&'a #ty>}).collect::<Vec<_>>();
51    let ref_mut_vecced_types = types.iter().map(|ty| quote! {Vec<&'a mut #ty>}).collect::<Vec<_>>();
52
53    let collections = enum_data.variants.iter().enumerate().map(|(idx, _)| {
54        Ident::new(&format!("collection_{}", idx), Span::call_site())
55    }).collect::<Vec<_>>();
56
57    let collection_creation = enum_data.variants.iter().enumerate().map(|(idx, _)| {
58        let collection_name = Ident::new(&format!("collection_{}", idx), Span::call_site());
59        quote!{let mut #collection_name = Vec::new();}
60    }).collect::<Vec<_>>();
61
62    let variants = enum_data.variants.iter().enumerate().map(|(idx, variant)| {
63        let variant_name = &variant.ident;
64        let collection_name = Ident::new(&format!("collection_{}", idx), Span::call_site());
65        quote!{#name::#variant_name(v) => #collection_name.push(v)}
66    }).collect::<Vec<_>>();
67
68    let expanded = quote! {
69        impl FromIterator<#name> for (#(#vecced_types),*, ) {
70            fn from_iter<I: IntoIterator<Item=#name>>(iter: I) -> Self {
71                #(#collection_creation);*
72
73                for i in iter {
74                    match i {
75                        #(#variants),*
76                    }
77                }
78        
79                (#(#collections),*, )
80            }
81        }
82
83        // With mutable reference implementation
84        impl<'a> FromIterator<&'a #name> for (#(#ref_vecced_types),*, ) {
85            fn from_iter<I: IntoIterator<Item=&'a #name>>(iter: I) -> Self {
86                #(#collection_creation);*
87
88                for i in iter {
89                    match i {
90                        #(#variants),*
91                    }
92                }
93        
94                (#(#collections),*, )
95            }
96        }
97
98        impl<'a> FromIterator<&'a mut #name> for (#(#ref_mut_vecced_types),*, ) {
99            fn from_iter<I: IntoIterator<Item=&'a mut #name>>(iter: I) -> Self {
100                #(#collection_creation);*
101
102                for i in iter {
103                    match i {
104                        #(#variants),*
105                    }
106                }
107        
108                (#(#collections),*, )
109            }
110        }
111
112        impl separable::Separable for #name {
113            type Target = (#(#vecced_types),*, );
114        }
115
116        impl<'a> separable::Separable for &'a #name {
117            type Target = (#(#ref_vecced_types),*, );
118        }
119
120        impl<'a> separable::Separable for &'a mut #name {
121            type Target = (#(#ref_mut_vecced_types),*, );
122        }
123    };
124
125    Ok(TokenStream::from(expanded))
126}