diff --git a/impl/src/add_helpers.rs b/impl/src/add_helpers.rs index b9cee1a6..77d4bf1b 100644 --- a/impl/src/add_helpers.rs +++ b/impl/src/add_helpers.rs @@ -1,28 +1,69 @@ -use proc_macro2::TokenStream; +use proc_macro2::{Span, TokenStream}; use quote::quote; -use syn::{Field, Ident, Index}; +use syn::{Field, Ident, Index, Type, TypeArray, TypeTuple}; pub fn tuple_exprs(fields: &[&Field], method_ident: &Ident) -> Vec { - let mut exprs = vec![]; - - for i in 0..fields.len() { - let i = Index::from(i); - // generates `self.0.add(rhs.0)` - let expr = quote! { self.#i.#method_ident(rhs.#i) }; - exprs.push(expr); - } - exprs + let fields: Vec<&Type> = fields.iter().map(|field| &field.ty).collect::>(); + inner_tuple_exprs(0, "e! {}, &fields, method_ident) } pub fn struct_exprs(fields: &[&Field], method_ident: &Ident) -> Vec { - let mut exprs = vec![]; + fields + .iter() + .map(|field| { + // It's safe to unwrap because struct fields always have an identifier + let field_path = field.ident.as_ref().unwrap(); + elem_content(0, "e! { .#field_path }, &field.ty, method_ident) + }) + .collect() +} + +pub fn inner_tuple_exprs( + // `depth` is needed for `index_var` generation for nested arrays + depth: usize, + field_path: &TokenStream, + fields: &[&Type], + method_ident: &Ident, +) -> Vec { + fields + .iter() + .enumerate() + .map(|(i, ty)| { + let i = Index::from(i); + elem_content(depth + 1, "e! { #field_path.#i }, ty, method_ident) + }) + .collect() +} + +pub fn elem_content( + depth: usize, + field_path: &TokenStream, + ty: &Type, + method_ident: &Ident, +) -> TokenStream { + match ty { + Type::Array(TypeArray { elem, .. }) => { + let index_var = Ident::new(&format!("i{}", depth), Span::call_site()); + let fn_body = elem_content( + depth + 1, + "e! { #field_path[#index_var] }, + elem.as_ref(), + method_ident, + ); - for field in fields { - // It's safe to unwrap because struct fields always have an identifier - let field_id = field.ident.as_ref().unwrap(); - // generates `x: self.x.add(rhs.x)` - let expr = quote! { self.#field_id.#method_ident(rhs.#field_id) }; - exprs.push(expr) + // generates `core::array::from_fn(|i0| self.x[i0].add(rhs.x[i0]))` + quote! { core::array::from_fn(|#index_var| #fn_body) } + } + Type::Tuple(TypeTuple { elems, .. }) => { + let exprs = inner_tuple_exprs( + depth + 1, + field_path, + &elems.iter().collect::>(), + method_ident, + ); + quote! { (#(#exprs),*) } + } + // generates `self.x.add(rhs.x)` + _ => quote! { self #field_path.#method_ident(rhs #field_path) }, } - exprs } diff --git a/impl/src/add_like.rs b/impl/src/add_like.rs index 8cc1a9c4..573ef1a3 100644 --- a/impl/src/add_like.rs +++ b/impl/src/add_like.rs @@ -18,6 +18,7 @@ pub fn expand(input: &DeriveInput, trait_name: &str) -> TokenStream { let generics = add_extra_type_param_bound_op_output(&input.generics, &trait_ident); let (impl_generics, ty_generics, where_clause) = generics.split_for_impl(); + // TODO(Matthias): do we add support for arrays here? let (output_type, block) = match input.data { Data::Struct(ref data_struct) => match data_struct.fields { Fields::Unnamed(ref fields) => ( @@ -53,7 +54,7 @@ pub fn expand(input: &DeriveInput, trait_name: &str) -> TokenStream { } } -fn tuple_content( +pub(crate) fn tuple_content( input_type: &T, fields: &[&Field], method_ident: &Ident, diff --git a/tests/add.rs b/tests/add.rs index 4c38fdae..50d71ca7 100644 --- a/tests/add.rs +++ b/tests/add.rs @@ -22,3 +22,28 @@ enum MixedInts { UnsignedTwo(u32), Unit, } + +#[derive(Add)] +#[derive(Default)] +struct StructRecursive { + a: i32, + b: [i32; 2], + c: [[i32; 2]; 3], + d: (i32, i32), + e: ((u8, [i32; 3]), i32), + f: ((u8, i32), (u8, ((i32, u64, ((u8, u8), u16)), u8))), + g: i32, +} + +#[test] +fn test_sanity() { + let mut a: StructRecursive = Default::default(); + let mut b: StructRecursive = Default::default(); + a.c[0][1] = 1; + b.c[0][1] = 2; + let c = a + b; + assert_eq!(c.c[0][1], 3); +} + +#[derive(Add)] +struct TupleRecursive((i32, u8), [(i32, u8); 10]);