From ea4143f8062da0dc87bf39f509e6f603fef2afb0 Mon Sep 17 00:00:00 2001 From: Manish Goregaokar Date: Mon, 26 Jun 2023 17:43:49 +0200 Subject: [PATCH] Add support for custom varule fields in #[make_varule] --- utils/zerovec/derive/examples/make_var.rs | 10 ++++ utils/zerovec/derive/src/make_varule.rs | 59 +++++++++++++++++------ utils/zerovec/derive/src/utils.rs | 28 ++++++++++- utils/zerovec/src/lib.rs | 3 +- 4 files changed, 84 insertions(+), 16 deletions(-) diff --git a/utils/zerovec/derive/examples/make_var.rs b/utils/zerovec/derive/examples/make_var.rs index 7dedfd3e9bd..0c9a2a7f70d 100644 --- a/utils/zerovec/derive/examples/make_var.rs +++ b/utils/zerovec/derive/examples/make_var.rs @@ -59,6 +59,16 @@ struct MultiFieldStruct<'a> { f: char, } +#[make_varule(CustomVarFieldULE)] +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, serde::Serialize, serde::Deserialize)] +#[zerovec::derive(Serialize, Deserialize, Debug)] +struct CustomVarField<'a> { + #[zerovec::varule(MultiFieldStructULE)] + #[serde(borrow)] + a: MultiFieldStruct<'a>, + b: u32, +} + #[make_varule(MultiFieldTupleULE)] #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, serde::Serialize, serde::Deserialize)] #[zerovec::derive(Serialize, Deserialize, Debug)] diff --git a/utils/zerovec/derive/src/make_varule.rs b/utils/zerovec/derive/src/make_varule.rs index 7859ae15f5d..56aa0b4f325 100644 --- a/utils/zerovec/derive/src/make_varule.rs +++ b/utils/zerovec/derive/src/make_varule.rs @@ -9,7 +9,7 @@ use quote::{quote, ToTokens}; use syn::spanned::Spanned; use syn::{ parse_quote, Data, DeriveInput, Error, Field, Fields, GenericArgument, Ident, Lifetime, - PathArguments, Type, + PathArguments, Type, TypePath, }; pub fn make_varule_impl(ule_name: Ident, mut input: DeriveInput) -> TokenStream2 { @@ -45,9 +45,10 @@ pub fn make_varule_impl(ule_name: Ident, mut input: DeriveInput) -> TokenStream2 let lt = lt.map(|l| &l.lifetime); let name = &input.ident; + let input_span = input.span(); let fields = match input.data { - Data::Struct(ref s) => &s.fields, + Data::Struct(ref mut s) => &mut s.fields, _ => { return Error::new(input.span(), "#[make_varule] must be applied to a struct") .to_compile_error(); @@ -65,16 +66,32 @@ pub fn make_varule_impl(ule_name: Ident, mut input: DeriveInput) -> TokenStream2 let mut sized_fields = vec![]; let mut unsized_fields = vec![]; + let mut custom_varule_idents = vec![]; + + for field in fields.iter_mut() { + match utils::extract_field_attributes(&mut field.attrs) { + Ok(i) => custom_varule_idents.push(i), + Err(e) => return e.to_compile_error(), + } + } + for (i, field) in fields.iter().enumerate() { - match UnsizedField::new(field, i) { + match UnsizedField::new(field, i, custom_varule_idents[i].clone()) { Ok(o) => unsized_fields.push(o), Err(_) => sized_fields.push(FieldInfo::new_for_field(field, i)), } } if unsized_fields.is_empty() { + let last_field_index = fields.len() - 1; let last_field = fields.iter().next_back().unwrap(); - let e = UnsizedField::new(last_field, fields.len() - 1).unwrap_err(); + + let e = UnsizedField::new( + last_field, + last_field_index, + custom_varule_idents[last_field_index].clone(), + ) + .unwrap_err(); return Error::new(last_field.span(), e).to_compile_error(); } @@ -131,7 +148,7 @@ pub fn make_varule_impl(ule_name: Ident, mut input: DeriveInput) -> TokenStream2 name, &ule_name, lt, - input.span(), + input_span, ); let eq_impl = quote!( @@ -402,11 +419,13 @@ enum OwnULETy<'a> { } /// Represents the type of the last field of the struct -#[derive(Copy, Clone, Debug)] +#[derive(Clone, Debug)] enum UnsizedFieldKind<'a> { Cow(OwnULETy<'a>), ZeroVec(&'a Type), VarZeroVec(&'a Type), + /// Custom VarULE type, and the identifier corresponding to the VarULE type + Custom(&'a TypePath, Ident), // Generally you should be using the above ones for maximum zero-copy, but these will still work Growable(OwnULETy<'a>), @@ -561,9 +580,13 @@ impl<'a> UnsizedFields<'a> { } impl<'a> UnsizedField<'a> { - fn new(field: &'a Field, index: usize) -> Result { + fn new( + field: &'a Field, + index: usize, + custom_varule_ident: Option, + ) -> Result { Ok(UnsizedField { - kind: UnsizedFieldKind::new(&field.ty)?, + kind: UnsizedFieldKind::new(&field.ty, custom_varule_ident)?, field: FieldInfo::new_for_field(field, index), }) } @@ -589,7 +612,10 @@ impl<'a> UnsizedField<'a> { impl<'a> UnsizedFieldKind<'a> { /// Construct a UnsizedFieldKind for the type of a UnsizedFieldKind if possible - fn new(ty: &'a Type) -> Result, String> { + fn new( + ty: &'a Type, + custom_varule_ident: Option, + ) -> Result, String> { static PATH_TYPE_IDENTITY_ERROR: &str = "Can only automatically detect corresponding VarULE types for path types \ that are Cow, ZeroVec, VarZeroVec, Box, String, or Vec"; @@ -600,6 +626,9 @@ impl<'a> UnsizedFieldKind<'a> { match *ty { Type::Reference(ref tyref) => OwnULETy::new(&tyref.elem, "reference").map(UnsizedFieldKind::Ref), Type::Path(ref typath) => { + if let Some(custom_varule_ident) = custom_varule_ident { + return Ok(UnsizedFieldKind::Custom(typath, custom_varule_ident)); + } if typath.path.segments.len() != 1 { return Err("Can only automatically detect corresponding VarULE types for \ path types with a single path segment".into()); @@ -671,6 +700,7 @@ impl<'a> UnsizedFieldKind<'a> { let inner_ule = inner.varule_ty(); quote!(#inner_ule) } + Self::Custom(_, ref name) => quote!(#name), Self::ZeroVec(ref inner) => quote!(zerovec::ZeroSlice<#inner>), Self::VarZeroVec(ref inner) => quote!(zerovec::VarZeroSlice<#inner>), } @@ -681,8 +711,8 @@ impl<'a> UnsizedFieldKind<'a> { match *self { Self::Ref(_) | Self::Cow(_) | Self::Growable(_) | Self::Boxed(_) => quote!(&*#value), - Self::ZeroVec(_) => quote!(&*#value), - Self::VarZeroVec(_) => quote!(&*#value), + Self::Custom(..) => quote!(&#value), + Self::ZeroVec(_) | Self::VarZeroVec(_) => quote!(&*#value), } } @@ -694,15 +724,16 @@ impl<'a> UnsizedFieldKind<'a> { | Self::Growable(ref inner) | Self::Boxed(ref inner) => inner.varule_ty(), - Self::ZeroVec(ty) => quote!(zerovec::ZeroSlice<#ty>), - Self::VarZeroVec(ty) => quote!(zerovec::VarZeroSlice<#ty>), + Self::Custom(ref path, _) => quote!(#path), + Self::ZeroVec(ref ty) => quote!(zerovec::ZeroSlice<#ty>), + Self::VarZeroVec(ref ty) => quote!(zerovec::VarZeroSlice<#ty>), } } fn has_zf(&self) -> bool { matches!( *self, - Self::Ref(_) | Self::Cow(_) | Self::ZeroVec(_) | Self::VarZeroVec(_) + Self::Ref(_) | Self::Cow(_) | Self::ZeroVec(_) | Self::VarZeroVec(_) | Self::Custom(..) ) } } diff --git a/utils/zerovec/derive/src/utils.rs b/utils/zerovec/derive/src/utils.rs index 8863540c2e4..1b8916ccb41 100644 --- a/utils/zerovec/derive/src/utils.rs +++ b/utils/zerovec/derive/src/utils.rs @@ -200,6 +200,32 @@ pub fn extract_zerovec_attributes(attrs: &mut Vec) -> Vec ret } +/// Extract attributes from field, and return them +/// +/// Only current field attribute is `zerovec::varule(VarUleType)` +pub fn extract_field_attributes(attrs: &mut Vec) -> Result> { + let mut zerovec_attrs = extract_zerovec_attributes(attrs); + let varule = extract_parenthetical_zerovec_attrs(&mut zerovec_attrs, "varule")?; + + if varule.len() > 1 { + return Err(Error::new( + attr.span(), + format!("Found multiple #[zerovec::varule()] on one field"), + )); + } + + if !zerovec_attrs.is_empty() { + return Err(Error::new( + zerovec_attrs[1].span(), + format!( + "Found unusable #[zerovec::] attrs on field, only #[zerovec::varule()] supported" + ), + )); + } + + Ok(varule.get(0).cloned()) +} + #[derive(Default, Copy, Clone)] pub struct ZeroVecAttrs { pub skip_kv: bool, @@ -210,7 +236,7 @@ pub struct ZeroVecAttrs { pub hash: bool, } -/// Removes all known zerovec:: attributes from attrs and validates them +/// Removes all known zerovec:: attributes from struct attrs and validates them pub fn extract_attributes_common( attrs: &mut Vec, span: Span, diff --git a/utils/zerovec/src/lib.rs b/utils/zerovec/src/lib.rs index 455e2c09e73..84a51906b39 100644 --- a/utils/zerovec/src/lib.rs +++ b/utils/zerovec/src/lib.rs @@ -441,7 +441,8 @@ pub use zerovec_derive::make_ule; /// /// This can be attached to structs containing only [`AsULE`] types with the last fields being /// [`Cow<'a, str>`](alloc::borrow::Cow), [`ZeroSlice`], or [`VarZeroSlice`]. If there is more than one such field, it will be represented -/// using [`MultiFieldsULE`](crate::ule::MultiFieldsULE) and getters will be generated. +/// using [`MultiFieldsULE`](crate::ule::MultiFieldsULE) and getters will be generated. Other VarULE fields will be detected if they are +/// tagged with `#[zerovec::varule(NameOfVarULETy)]`. /// /// The type must be [`PartialEq`] and [`Eq`]. ///