use proc_macro::TokenStream;
use proc_macro2::{Ident, TokenStream as TokenStream2};
use quote::{format_ident, quote};
use syn::{
parse_macro_input, AngleBracketedGenericArguments, Data, DataStruct, DeriveInput, FieldsNamed,
GenericArgument, PathArguments, Type,
};
#[proc_macro_derive(Builder)]
pub fn derive_builder(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let ts = do_expand(&input);
ts.into()
}
fn do_expand(st: &DeriveInput) -> TokenStream2 {
let ident = &st.ident;
let name = format_ident!("{}Builder", ident);
let tokens = parse_data(&st.data);
let data: Vec<_> = tokens.iter().map(|token| &token.field).collect();
let content: Vec<_> = tokens.iter().map(|token| &token.value).collect();
let impl_block: Vec<_> = tokens.iter().map(|token| &token.impl_block).collect();
let build_fields: Vec<_> = tokens.iter().map(|token| &token.build_field).collect();
quote!(
pub struct #name {
#(#data),*
}
impl #ident {
pub fn builder () -> #name {
#name {
#(#content),*
}
}
}
impl #name {
#(#impl_block)*
}
impl #name {
pub fn build(self) -> #ident {
#ident {
#(#build_fields),*
}
}
}
)
}
fn parse_data(data: &Data) -> Vec<Token> {
match data {
Data::Struct(data) => parse_struct(data),
Data::Enum(_) => todo!(),
Data::Union(_) => todo!(),
}
}
struct Token {
field: TokenStream2,
value: TokenStream2,
impl_block: TokenStream2,
build_field: TokenStream2,
}
fn parse_struct(data: &DataStruct) -> Vec<Token> {
let fields = &data.fields;
match fields {
syn::Fields::Named(name) => parse_named(name),
syn::Fields::Unnamed(_) => todo!(),
syn::Fields::Unit => todo!(),
}
}
fn get_inner_option_type(ty: &Type) -> Option<&Type> {
if let Type::Path(ty) = ty {
let seg = ty.path.segments.last();
if let Some(seg) = seg {
let ident = seg.ident.to_string();
if ident == "Option" {
if let PathArguments::AngleBracketed(AngleBracketedGenericArguments {
args, ..
}) = &seg.arguments
{
if let Some(GenericArgument::Type(ty)) = args.last() {
return Some(ty);
}
}
}
}
None
} else {
None
}
}
fn get_type_name(ty: &Type) -> Option<String> {
match ty {
Type::Path(ty) => ty.path.segments.last().map(|seg| seg.ident.to_string()),
_ => None,
}
}
fn generate_impl_block(ident: &Option<Ident>, ty_name: &Option<String>, ty: &Type) -> TokenStream2 {
match ty_name {
Some(name) if name == "String" => {
quote!(
pub fn #ident<T>(&mut self, #ident: T) ->&mut Self where T:AsRef<str> {
self.#ident = Some(#ident.as_ref().to_owned());
self
}
)
}
_ => {
quote!(
pub fn #ident(&mut self, #ident: #ty) ->&mut Self {
self.#ident = Some(#ident);
self
}
)
}
}
}
fn generate_build_field(inner_ty: &Option<&Type>, ident: &Option<Ident>) -> TokenStream2 {
if inner_ty.is_some() {
quote!(
#ident: self.#ident
)
} else {
quote!(
#ident: self.#ident.expect("#ident is missing")
)
}
}
fn parse_named(name: &FieldsNamed) -> Vec<Token> {
name.named
.iter()
.map(|field| {
let ident = field.ident.to_owned();
let old_ty = &field.ty;
let inner_ty = get_inner_option_type(old_ty);
let ty = match inner_ty {
Some(ty) => ty,
None => old_ty,
};
let ty_name = get_type_name(ty);
let field = quote!(
#ident: Option<#ty>
);
let value = quote!(
#ident: None
);
let impl_block = generate_impl_block(&ident, &ty_name, ty);
let build_field = generate_build_field(&inner_ty, &ident);
Token {
field,
value,
impl_block,
build_field,
}
})
.collect::<Vec<_>>()
}