extern crate proc_macro;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::{
parse::{self, Parse, ParseStream},
parse_macro_input,
};
struct UnzipN {
n: usize,
generic_types: Vec<syn::Ident>,
collections: Vec<syn::Ident>,
trait_name: syn::Ident,
}
impl Parse for UnzipN {
fn parse(input: ParseStream) -> parse::Result<Self> {
let n: usize = input.parse::<syn::LitInt>()?.base10_parse()?;
let generic_types: Vec<_> = (0..n).map(|id| format_ident!("Type_{}", id)).collect();
let collections: Vec<_> = (0..n)
.map(|id| format_ident!("Collection_{}", id))
.collect();
let trait_name = format_ident!("Unzip{}", n);
Ok(UnzipN {
n,
generic_types,
collections,
trait_name,
})
}
}
impl UnzipN {
fn make_trait(&self, no_std: bool) -> TokenStream {
let trait_doc = format!(
"Extension trait for unzipping iterators over tuples of size {}.",
self.n
);
let generic_types = &self.generic_types;
let collections = &self.collections;
let trait_name = &self.trait_name;
let unzip_n_vec = if no_std {
TokenStream::new()
} else {
quote!(
fn unzip_n_vec(self) -> ( #( Vec< #generic_types >, )* )
where
Self: Sized,
{
self.unzip_n()
}
)
};
quote!(
#[doc = #trait_doc]
trait #trait_name < #( #generic_types, )* > {
fn unzip_n<#(#collections,)*>(self) -> ( #(#collections,)* )
where
#( #collections: Default + Extend< #generic_types >, )*
;
#unzip_n_vec
}
)
}
fn make_implementation(&self) -> TokenStream {
let generic_types = &self.generic_types;
let collections = &self.collections;
let trait_name = &self.trait_name;
let containers: Vec<_> = (0..self.n)
.map(|id| format_ident!("container_{}", id))
.collect();
let values: Vec<_> = (0..self.n).map(|id| format_ident!("val_{}", id)).collect();
quote!(
impl<Iter, #(#generic_types,)*> #trait_name <#(#generic_types,)*> for Iter
where
Iter: Iterator<Item = (#(#generic_types,)*)>,
{
fn unzip_n<#(#collections,)*>(self) -> ( #(#collections,)* )
where
#( #collections: Default + Extend< #generic_types >, )*
{
#( let mut #containers = #collections :: default() ;)*
self.for_each(|( #( #values, )* )| {
#( #containers.extend(Some(#values)) ;)*
});
(#( #containers, )* )
}
}
)
}
pub fn generate(&self, no_std: bool) -> TokenStream {
let trait_decl = self.make_trait(no_std);
let impl_block = self.make_implementation();
quote!( #trait_decl #impl_block )
}
}
#[proc_macro]
pub fn unzip_n(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
parse_macro_input!(input as UnzipN).generate(false).into()
}
#[proc_macro]
pub fn unzip_n_nostd(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
parse_macro_input!(input as UnzipN).generate(true).into()
}