1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88
extern crate proc_macro; #[macro_use] extern crate quote; extern crate syn; use proc_macro::TokenStream; use syn::{Body, DeriveInput, Ident, VariantData}; #[proc_macro_derive(EnumIter)] pub fn enum_iter(input: TokenStream) -> TokenStream { let source = input.to_string(); let ast = syn::parse_derive_input(&source).unwrap(); let expanded = enum_iter_impl(&ast); expanded.parse().unwrap() } fn enum_iter_impl(ast: &DeriveInput) -> quote::Tokens { let name = &ast.ident; let variants = match ast.body { Body::Enum(ref variants) => variants, _ => panic!( "`EnumIter` can be applied only to the enums, {} is not an enum", &name ), }; let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl(); let iter_impl = Ident::from(format!("_EnumIterator{}", name)); let vis = &ast.vis; let match_variants = variants.iter().enumerate().map(|(idx, variant)| { let id = &variant.ident; match variant.data { VariantData::Unit => quote!{ #idx => #name::#id, }, VariantData::Tuple(ref fields) => { let types = fields.iter().map(|f| &f.ty); quote! { #idx => #name::#id( #(#types::default(),)* ), } } VariantData::Struct(ref fields) => { let items = fields.iter().map(|f| { let ident = &f.ident; let ty = &f.ty; quote! { #ident: #ty::default() } }); quote! { #idx => #name::#id{ #(#items,)* }, } } } }); quote!{ #[derive(Default)] #vis struct #iter_impl { count: usize, } impl ::std::iter::Iterator for #iter_impl { type Item = #name; fn next(&mut self) -> Option<Self::Item> { let item = match self.count { #(#match_variants)* _ => return None, }; self.count += 1; Some(item) } } impl #impl_generics #name #ty_generics #where_clause { pub fn iter() -> #iter_impl { #iter_impl::default() } } } }