use quote::{format_ident, quote};
use syn::{
parse::{Parse, ParseStream, Result},
Ident, Token,
};
use crate::rich::{common::Visibility, Field, FieldType};
#[derive(Debug)]
pub(crate) enum WriteFnStyle {
Compose,
Builder,
}
#[derive(Debug)]
pub(crate) struct WriteFn {
attributes: Vec<syn::Attribute>,
rename: Option<Ident>,
style: Option<WriteFnStyle>,
take: bool,
visibility: Option<Visibility>,
option: bool,
}
impl Default for WriteFn {
fn default() -> Self {
Self {
attributes: vec![],
rename: None,
style: None,
take: false,
visibility: None,
option: false,
}
}
}
impl WriteFn {
pub fn get_fn(&self, field: &Field) -> proc_macro2::TokenStream {
let field_name = &field.name;
let write_fn_name = match (&field.ty, self.option, &self.style) {
(_, _, Some(WriteFnStyle::Compose)) => format_ident!("and_{}", &field.name),
(FieldType::Option(_), true, _) => format_ident!("try_set_{}", &field.name),
_ => format_ident!("set_{}", &field.name),
};
let write_fn_name = self.rename.as_ref().unwrap_or(&write_fn_name);
let visibility = self
.visibility
.as_ref()
.map(|vis| quote! { #vis })
.unwrap_or(quote! { pub });
let attributes = &self.attributes;
let (arg_self, ret_self) = if self.take {
(quote! { mut self }, quote! { Self })
} else {
(quote! { &mut self }, quote! { &mut Self })
};
let builder_fn = || match (&field.ty, self.option) {
(FieldType::Normal(ty), _) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, value: impl Into<#ty>) -> #ret_self {
self.#field_name = value.into();
self
}
},
(FieldType::Option(ty), false) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, value: impl Into<#ty>) -> #ret_self {
self.#field_name = Some(value.into());
self
}
},
(FieldType::Option(ty), true) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, value: Option<impl Into<#ty>>) -> #ret_self {
self.#field_name = value.map(|v| v.into());
self
}
},
};
let compose_fn = || match (&field.ty, self.take, self.option) {
(FieldType::Normal(ty), false, _) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, set_value: impl Fn(&mut #ty) -> &mut #ty) -> #ret_self {
set_value(&mut self.#field_name);
self
}
},
(FieldType::Option(ty), false, false) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, set_value: impl Fn(&mut #ty) -> &mut #ty) -> #ret_self
where #ty: Default
{
set_value(self.#field_name.get_or_insert(#ty::default()));
self
}
},
(FieldType::Option(ty), false, true) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, set_value: impl Fn(Option<&mut #ty>) -> Option<&mut #ty>) -> #ret_self {
set_value(self.#field_name.as_ref_mut());
self
}
},
(FieldType::Normal(ty), true, _) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, set_value: impl Fn(#ty) -> #ty) -> #ret_self {
self.#field_name = set_value(self.#field_name);
self
}
},
(FieldType::Option(ty), true, false) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, set_value: impl Fn(#ty) -> #ty) -> #ret_self
where #ty: Default
{
self.#field_name = Some(set_value(self.#field_name.unwrap_or(#ty::default())));
self
}
},
(FieldType::Option(ty), true, true) => quote! {
#( #attributes )*
#visibility fn #write_fn_name(#arg_self, set_value: impl Fn(Option<#ty>) -> Option<#ty>) -> #ret_self {
self.#field_name = set_value(self.#field_name.as_ref_mut());
self
}
},
};
match self.style {
Some(WriteFnStyle::Builder) | None => builder_fn(),
Some(WriteFnStyle::Compose) => compose_fn(),
}
}
}
impl Parse for WriteFn {
fn parse(input: ParseStream) -> Result<Self> {
let attr_name = "write";
if input.parse::<Ident>()? != attr_name {
return Err(input.error(format!("Expected `{}`", attr_name)));
}
if input.peek(Token![,]) || input.is_empty() {
return Ok(WriteFn::default());
}
let mut write_fn = WriteFn::default();
let inner;
syn::parenthesized!(inner in input);
write_fn.attributes = inner.call(syn::Attribute::parse_outer)?;
while !inner.is_empty() {
match inner.parse::<Ident>()?.to_string().as_ref() {
"rename" => {
let _ = inner.parse::<Token![=]>()?;
match write_fn.rename {
Some(_) => {
return Err(inner
.error("`write` attribute have been renamed more than once"))
}
None => write_fn.rename = Some(inner.parse::<Ident>()?),
};
}
"style" => {
if let Some(_) = write_fn.style {
return Err(inner.error("`write` have been styled more than once"));
}
let _ = inner.parse::<Token![=]>()?;
write_fn.style = match inner.parse::<Ident>()?.to_string().as_ref() {
"builder" => Some(WriteFnStyle::Builder),
"compose" => Some(WriteFnStyle::Compose),
_ => Err(input.error(
"`write` attribute only accept `builder` or `compse` as `style`",
))?,
};
}
"take" => {
if write_fn.take {
return Err(
inner.error("`write` have received `take` flag more than once")
);
}
write_fn.take = true;
}
"option" => {
if write_fn.option {
return Err(
inner.error("`write` have received `option` flag more than once")
);
}
write_fn.option = true;
}
"vis" => {
if let Some(_) = write_fn.visibility {
return Err(inner.error(
"`write` attribute have been received visibility value more than once",
))
}
let _ = inner.parse::<Token![=]>()?;
let visibility = inner.parse::<Visibility>()?;
write_fn.visibility = Some(visibility);
}
_ => Err(inner.error(
"`write` attribute only accept `rename`, `style`, `take`, `option` and `visibility` attributes",
))?,
}
if inner.peek(Token![,]) {
inner.parse::<Token![,]>()?;
}
}
Ok(write_fn)
}
}