use proc_macro2::TokenStream;
use quote::quote;
use syn::{
parse_macro_input, Data, DataEnum, DataStruct, DeriveInput, Fields, FieldsNamed, FieldsUnnamed,
Ident, Index,
};
fn get_struct_body(out_type: &Ident, data: &DataStruct) -> TokenStream {
match &data.fields {
Fields::Named(FieldsNamed { named, .. }) => {
let mut init_toks = TokenStream::new();
let mut fields = TokenStream::new();
for f in named.iter() {
let ident = &f.ident;
init_toks.extend(quote! { let #ident = self.#ident.to_tok_stream(); });
fields.extend(quote! { #ident: ##ident, });
}
quote! {
#init_toks
let element = rustifact::internal::quote! {
#out_type {
#fields
}
};
toks.extend(element);
}
}
Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => {
let mut init_toks = TokenStream::new();
let mut fields = TokenStream::new();
for i in 0..unnamed.len() {
let index = Index::from(i);
let ident = Ident::new(&format!("ident{}", i), out_type.span());
init_toks.extend(quote! { let #ident = self.#index.to_tok_stream(); });
fields.extend(quote! { ##ident, });
}
quote! {
#init_toks
let element = rustifact::internal::quote! {
#out_type ( #fields )
};
toks.extend(element);
}
}
Fields::Unit => {
quote! { () }
}
}
}
fn get_enum_body(out_type: &Ident, data: &DataEnum) -> TokenStream {
let mut arms = TokenStream::new();
for v in &data.variants {
let ident = &v.ident;
let toks = match &v.fields {
Fields::Unnamed(fields_unnamed) => {
let mut init_toks = TokenStream::new();
let mut fields = TokenStream::new();
let mut fields_out = TokenStream::new();
for i in 0..fields_unnamed.unnamed.len() {
let id = Ident::new(&format!("ident{}", i), out_type.span());
let id_toks = Ident::new(&format!("ident{}_toks", i), out_type.span());
init_toks.extend(quote! { let #id_toks = #id.to_tok_stream(); });
fields.extend(quote! { #id, });
fields_out.extend(quote! { ##id_toks, });
}
if fields.is_empty() {
quote! {
#out_type::#ident => rustifact::internal::quote! { #out_type::#ident },
}
} else {
quote! {
#out_type::#ident( #fields ) => {
#init_toks
rustifact::internal::quote! { #out_type::#ident( #fields_out ) }
},
}
}
}
Fields::Named(_) => {
panic!("Named fields are not yet supported");
}
Fields::Unit => {
quote! { #out_type::#ident => rustifact::internal::quote! { #out_type::#ident }, }
}
};
arms.extend(toks);
}
quote! {
let element = match self {
#arms
};
toks.extend(element);
}
}
#[proc_macro_derive(ToTokenStream, attributes(OutType))]
pub fn derive_token_stream(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
let ast = parse_macro_input!(input as DeriveInput);
let in_type = &ast.ident;
let mut out_type: Ident = in_type.clone();
for attr in &ast.attrs {
if attr.path().is_ident("OutType") {
if let Ok(id) = attr.parse_args::<Ident>() {
out_type = id;
}
}
}
let body = match &ast.data {
Data::Struct(data) => get_struct_body(&out_type, data),
Data::Enum(data) => get_enum_body(&out_type, data),
Data::Union(_) => {
panic!("Unions are not yet supported");
}
};
let generics = &ast.generics;
let gen_where = &generics.where_clause;
quote! {
impl #generics rustifact::ToTokenStream for #in_type #generics #gen_where {
fn to_toks(&self, toks: &mut rustifact::internal::TokenStream) {
#body
}
}
}
.into()
}