ruff/crates/ruff_macros/src/cache_key.rs
2023-05-11 18:00:00 -04:00

103 lines
3.5 KiB
Rust

use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::spanned::Spanned;
use syn::{Data, DeriveInput, Error, Fields};
pub(crate) fn derive_cache_key(item: &DeriveInput) -> syn::Result<TokenStream> {
let fields = match &item.data {
Data::Enum(item_enum) => {
let arms = item_enum.variants.iter().enumerate().map(|(i, variant)| {
let variant_name = &variant.ident;
match &variant.fields {
Fields::Named(fields) => {
let field_names: Vec<_> = fields
.named
.iter()
.map(|field| field.ident.clone().unwrap())
.collect();
let fields_code = field_names
.iter()
.map(|field| quote!(#field.cache_key(key);));
quote! {
Self::#variant_name{#(#field_names),*} => {
key.write_usize(#i);
#(#fields_code)*
}
}
}
Fields::Unnamed(fields) => {
let field_names: Vec<_> = fields
.unnamed
.iter()
.enumerate()
.map(|(i, _)| format_ident!("field_{i}"))
.collect();
let fields_code = field_names
.iter()
.map(|field| quote!(#field.cache_key(key);));
quote! {
Self::#variant_name(#(#field_names),*) => {
key.write_usize(#i);
#(#fields_code)*
}
}
}
Fields::Unit => {
quote! {
Self::#variant_name => {
key.write_usize(#i);
}
}
}
}
});
quote! {
match self {
#(#arms)*
}
}
}
Data::Struct(item_struct) => {
let fields = item_struct.fields.iter().enumerate().map(|(i, field)| {
let field_attr = match &field.ident {
Some(ident) => quote!(self.#ident),
None => {
let index = syn::Index::from(i);
quote!(self.#index)
}
};
quote!(#field_attr.cache_key(key);)
});
quote! {#(#fields)*}
}
Data::Union(_) => {
return Err(Error::new(
item.span(),
"CacheKey does not support unions. Only structs and enums are supported",
))
}
};
let name = &item.ident;
let (impl_generics, ty_generics, where_clause) = &item.generics.split_for_impl();
Ok(quote!(
impl #impl_generics ruff_cache::CacheKey for #name #ty_generics #where_clause {
fn cache_key(&self, key: &mut ruff_cache::CacheKeyHasher) {
use std::hash::Hasher;
use ruff_cache::CacheKey;
#fields
}
}
))
}