Skip to content

Commit

Permalink
simplify some stuff
Browse files Browse the repository at this point in the history
  • Loading branch information
ArthurZucker committed Jun 12, 2024
1 parent f87bb97 commit c4b4f3c
Showing 1 changed file with 58 additions and 182 deletions.
240 changes: 58 additions & 182 deletions tokenizers/display_derive/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,213 +10,89 @@ use vendored::FmtAttribute;
pub fn display_derive(input: TokenStream) -> TokenStream {
// Parse the parsed_input tokens into a syntax tree
let parsed_input = parse_macro_input!(input as DeriveInput);
let attr_name = "display";
let attrs = syn::parse::<FmtAttribute>(input).unwrap();
// 1. If the attrs are not None, then we defer to this.
// Meaning we juste return quote!{ format!(#fmt, #attr)}
let trait_ident = format_ident!("display");
let ident = &parsed_input.ident;

let ctx = (&attrs, ident, &trait_ident, &trait_ident);
// 2. We automatically parse
let body = match &parsed_input.data {
syn::Data::Struct(s) => expand_struct(s, ctx),
syn::Data::Enum(e) => expand_enum(e, ctx),
syn::Data::Struct(s) => generate_fmt_impl_for_struct(s, ident),
syn::Data::Enum(e) => generate_fmt_impl_for_enum(e, ident),
syn::Data::Union(u) => {
let error = syn::Error::new_spanned(u.union_token, "Unions are not supported");
return proc_macro::TokenStream::from(error.into_compile_error());
}
};

quote! {
impl std::fmt::Display for #ident{
let expanded = quote! {

Check failure on line 29 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.9)

the trait bound `proc_macro::TokenStream: ToTokens` is not satisfied

Check failure on line 29 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.8)

the trait bound `proc_macro::TokenStream: ToTokens` is not satisfied
impl std::fmt::Display for #ident {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
#body
}
}
}.into()
}

/// Type alias for an expansion context:
/// - [`FmtAttribute`].
/// - Struct/enum/union [`syn::Ident`].
/// - Derived trait [`syn::Ident`].
/// - Attribute name [`syn::Ident`].
///
/// [`syn::Ident`]: struct@syn::Ident
type ExpansionCtx<'a> = (
&'a FmtAttribute,
&'a syn::Ident,
&'a syn::Ident,
&'a syn::Ident,
);

/// Expands a [`fmt::Display`]-like derive macro for the provided struct.
fn expand_struct(
s: &syn::DataStruct,
(attrs, ident, trait_ident, _): ExpansionCtx<'_>,
) -> syn::Result<(Vec<syn::WherePredicate>, TokenStream)> {
let s = Expansion {
attrs,
fields: &s.fields,
trait_ident,
ident,
};
let body = s.generate_body()?;

let vars = s.fields.iter().enumerate().map(|(i, f)| {
let var = f.ident.clone().unwrap_or_else(|| format_ident!("_{i}"));
let member = f
.ident
.clone()
.map_or_else(|| syn::Member::Unnamed(i.into()), syn::Member::Named);
TokenStream::from(expanded)
}

fn generate_fmt_impl_for_struct(data_struct: &syn::DataStruct, ident: &syn::Ident) -> TokenStream {
let fields = &data_struct.fields;
let field_fmts = fields.iter().enumerate().map(|(i, field)| {
let field_name = match &field.ident {
Some(ident) => ident,
None => {
// If the field doesn't have a name, we generate a name based on its index
let index = syn::Index::from(i);
quote! { #index }

Check failure on line 48 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.9)

`match` arms have incompatible types

Check failure on line 48 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.8)

`match` arms have incompatible types
}
};
quote! {
let #var = &self.#member;
write!(f, "{}: {}", stringify!(#field_name), self.#field_name)?;
}
});

let body = quote! {
#( #vars )*
#body
};

Ok(body)
// Collect the mapped tokens into a TokenStream
field_fmts

Check failure on line 56 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.9)

mismatched types

Check failure on line 56 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.8)

mismatched types
}

/// Expands a [`fmt`]-like derive macro for the provided enum.
fn expand_enum(
e: &syn::DataEnum,
(attrs, _, trait_ident, attr_name): ExpansionCtx<'_>,
) -> syn::Result<(Vec<syn::WherePredicate>, TokenStream)> {
if attrs.fmt.is_some() {
todo!("https://github.com/JelteF/derive_more/issues/142");
}

let match_arms = e.variants.iter().try_fold(
(Vec::new(), TokenStream::new()),
|mut arms, variant| {
let attrs = FmtAttribute::parse_attrs(&variant.attrs, attr_name)?
.unwrap_or_default();
let ident = &variant.ident;

if attrs.fmt.is_none()
&& variant.fields.is_empty()
&& attr_name != "display"
{
return Err(syn::Error::new(
e.variants.span(),
format!(
"implicit formatting of unit enum variant is supported only for `Display` \
macro, use `#[{attr_name}(\"...\")]` to explicitly specify the formatting",
),
));
fn generate_fmt_impl_for_enum(data_enum: &syn::DataEnum, ident: &syn::Ident) -> TokenStream {
let arms = data_enum.variants.iter().map(|variant| {
let variant_name = &variant.ident;
let variant_fmt = match &variant.fields {
syn::Fields::Unit => {
// If the variant has no fields, we just print its name
quote! { write!(f, "{}", stringify!(#variant_name))?; }
}

let v = Expansion {
attrs: &attrs,
fields: &variant.fields,
trait_ident,
ident,
};
let arm_body = v.generate_body()?;

let fields_idents =
variant.fields.iter().enumerate().map(|(i, f)| {
f.ident.clone().unwrap_or_else(|| format_ident!("_{i}"))
syn::Fields::Named(fields) => {
// If the variant has named fields, we print each field's name and value
let field_fmts = fields.named.iter().map(|field| {
let field_name = field.ident.as_ref().unwrap();
quote! {
write!(f, "{}: {:?}", stringify!(#field_name), self.#field_name)?;
}
});
let matcher = match variant.fields {
syn::Fields::Named(_) => {
quote! { Self::#ident { #( #fields_idents ),* } }
quote! {
write!(f, "{} {{ ", stringify!(#variant_name))?;
#( #field_fmts )*
write!(f, " }}")?;
}
syn::Fields::Unnamed(_) => {
quote! { Self::#ident ( #( #fields_idents ),* ) }
}
syn::Fields::Unit => quote! { Self::#ident },
};

arms.extend([quote! { #matcher => { #arm_body }, }]);

Ok::<_, syn::Error>(arms)
},
)?;

let body = match_arms
.is_empty()
.then(|| quote! { match *self {} })
.unwrap_or_else(|| quote! { match self { #match_arms } });

Ok(body)
}


/// Helper struct to generate [`Display::fmt()`] implementation body and trait
/// bounds for a struct or an enum variant.
///
/// [`Display::fmt()`]: fmt::Display::fmt()
#[derive(Debug)]
struct Expansion<'a> {
/// Derive macro [`FmtAttribute`].
attrs: &'a FmtAttribute,

/// Struct or enum [`syn::Ident`].
///
/// [`syn::Ident`]: struct@syn::Ident
ident: &'a syn::Ident,

/// Struct or enum [`syn::Fields`].
fields: &'a syn::Fields,

/// [`fmt`] trait [`syn::Ident`].
///
/// [`syn::Ident`]: struct@syn::Ident
trait_ident: &'a syn::Ident,
}

impl<'a> Expansion<'a> {
/// Generates [`Display::fmt()`] implementation for a struct or an enum variant.
///
/// # Errors
///
/// In case [`FmtAttribute`] is [`None`] and [`syn::Fields`] length is
/// greater than 1.
///
/// [`Display::fmt()`]: fmt::Display::fmt()
/// [`FmtAttribute`]: super::FmtAttribute
fn generate_body(&self) -> syn::Result<TokenStream> {
match &self.attrs.fmt {
Some(fmt) => {
Ok(if let Some((expr, trait_ident)) = fmt.transparent_call() {
quote! { core::fmt::#trait_ident::fmt(&(#expr), __derive_more_f) }
} else {
quote! { core::write!(__derive_more_f, #fmt) }
})
}
None if self.fields.is_empty() => {
let ident_str = self.ident.to_string();

Ok(quote! {
core::write!(__derive_more_f, #ident_str)
})
}
None if self.fields.len() == 1 => {
let field = self
.fields
.iter()
.next()
.unwrap_or_else(|| unreachable!("count() == 1"));
let ident = field.ident.clone().unwrap_or_else(|| format_ident!("_0"));
let trait_ident = self.trait_ident;

Ok(quote! {
core::fmt::#trait_ident::fmt(#ident, __derive_more_f)
})
syn::Fields::Unnamed(fields) => {
// If the variant has unnamed fields, we print each field's value without names
let field_fmts = fields.unnamed.iter().map(|field| {
quote! {
write!(f, "{:?}, ", self.#field)?;
}
});
quote! {
write!(f, "{}(", stringify!(#variant_name))?;
#( #field_fmts )*
write!(f, ")")?;
}
}
_ => Err(syn::Error::new(
self.fields.span(),
format!(
"TODO ARTHUR! struct or enum variant with more than 1 field must have \
`#[{}(\"...\", ...)]` attribute",
self.trait_ident,
),
)),
}
}
};
quote! { #ident::#variant_name => { #variant_fmt } }
});
arms

Check failure on line 97 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.9)

mismatched types

Check failure on line 97 in tokenizers/display_derive/src/lib.rs

View workflow job for this annotation

GitHub Actions / Check it builds for Windows 32-bit (3.8)

mismatched types
}

0 comments on commit c4b4f3c

Please sign in to comment.