Skip to content

Commit 9250753

Browse files
Support pointwise addition for arrays and tuples
Fixes #342 We want to support examples like these: ```rust 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, } struct TupleRecursive((i32, u8), [(i32, u8); 10]); ``` Supporting arrays and tuples inside of enums would also be useful, but that's not in this PR.
1 parent 2a001d6 commit 9250753

File tree

3 files changed

+87
-20
lines changed

3 files changed

+87
-20
lines changed

impl/src/add_helpers.rs

Lines changed: 60 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,69 @@
1-
use proc_macro2::TokenStream;
1+
use proc_macro2::{Span, TokenStream};
22
use quote::quote;
3-
use syn::{Field, Ident, Index};
3+
use syn::{Field, Ident, Index, Type, TypeArray, TypeTuple};
44

55
pub fn tuple_exprs(fields: &[&Field], method_ident: &Ident) -> Vec<TokenStream> {
6-
let mut exprs = vec![];
7-
8-
for i in 0..fields.len() {
9-
let i = Index::from(i);
10-
// generates `self.0.add(rhs.0)`
11-
let expr = quote! { self.#i.#method_ident(rhs.#i) };
12-
exprs.push(expr);
13-
}
14-
exprs
6+
let fields: Vec<&Type> = fields.iter().map(|field| &field.ty).collect::<Vec<_>>();
7+
inner_tuple_exprs(0, &quote! {}, &fields, method_ident)
158
}
169

1710
pub fn struct_exprs(fields: &[&Field], method_ident: &Ident) -> Vec<TokenStream> {
18-
let mut exprs = vec![];
11+
fields
12+
.iter()
13+
.map(|field| {
14+
// It's safe to unwrap because struct fields always have an identifier
15+
let field_path = field.ident.as_ref().unwrap();
16+
elem_content(0, &quote! { .#field_path }, &field.ty, method_ident)
17+
})
18+
.collect()
19+
}
20+
21+
pub fn inner_tuple_exprs(
22+
// `depth` is needed for `index_var` generation for nested arrays
23+
depth: usize,
24+
field_path: &TokenStream,
25+
fields: &[&Type],
26+
method_ident: &Ident,
27+
) -> Vec<TokenStream> {
28+
fields
29+
.iter()
30+
.enumerate()
31+
.map(|(i, ty)| {
32+
let i = Index::from(i);
33+
elem_content(depth + 1, &quote! { #field_path.#i }, ty, method_ident)
34+
})
35+
.collect()
36+
}
37+
38+
pub fn elem_content(
39+
depth: usize,
40+
field_path: &TokenStream,
41+
ty: &Type,
42+
method_ident: &Ident,
43+
) -> TokenStream {
44+
match ty {
45+
Type::Array(TypeArray { elem, .. }) => {
46+
let index_var = Ident::new(&format!("i{}", depth), Span::call_site());
47+
let fn_body = elem_content(
48+
depth + 1,
49+
&quote! { #field_path[#index_var] },
50+
elem.as_ref(),
51+
method_ident,
52+
);
1953

20-
for field in fields {
21-
// It's safe to unwrap because struct fields always have an identifier
22-
let field_id = field.ident.as_ref().unwrap();
23-
// generates `x: self.x.add(rhs.x)`
24-
let expr = quote! { self.#field_id.#method_ident(rhs.#field_id) };
25-
exprs.push(expr)
54+
// generates `core::array::from_fn(|i0| self.x[i0].add(rhs.x[i0]))`
55+
quote! { core::array::from_fn(|#index_var| #fn_body) }
56+
}
57+
Type::Tuple(TypeTuple { elems, .. }) => {
58+
let exprs = inner_tuple_exprs(
59+
depth + 1,
60+
field_path,
61+
&elems.iter().collect::<Vec<_>>(),
62+
method_ident,
63+
);
64+
quote! { (#(#exprs),*) }
65+
}
66+
// generates `self.x.add(rhs.x)`
67+
_ => quote! { self #field_path.#method_ident(rhs #field_path) },
2668
}
27-
exprs
2869
}

impl/src/add_like.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ pub fn expand(input: &DeriveInput, trait_name: &str) -> TokenStream {
1818
let generics = add_extra_type_param_bound_op_output(&input.generics, &trait_ident);
1919
let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
2020

21+
// TODO(Matthias): do we add support for arrays here?
2122
let (output_type, block) = match input.data {
2223
Data::Struct(ref data_struct) => match data_struct.fields {
2324
Fields::Unnamed(ref fields) => (
@@ -53,7 +54,7 @@ pub fn expand(input: &DeriveInput, trait_name: &str) -> TokenStream {
5354
}
5455
}
5556

56-
fn tuple_content<T: ToTokens>(
57+
pub(crate) fn tuple_content<T: ToTokens>(
5758
input_type: &T,
5859
fields: &[&Field],
5960
method_ident: &Ident,

tests/add.rs

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,3 +22,28 @@ enum MixedInts {
2222
UnsignedTwo(u32),
2323
Unit,
2424
}
25+
26+
#[derive(Add)]
27+
#[derive(Default)]
28+
struct StructRecursive {
29+
a: i32,
30+
b: [i32; 2],
31+
c: [[i32; 2]; 3],
32+
d: (i32, i32),
33+
e: ((u8, [i32; 3]), i32),
34+
f: ((u8, i32), (u8, ((i32, u64, ((u8, u8), u16)), u8))),
35+
g: i32,
36+
}
37+
38+
#[test]
39+
fn test_sanity() {
40+
let mut a: StructRecursive = Default::default();
41+
let mut b: StructRecursive = Default::default();
42+
a.c[0][1] = 1;
43+
b.c[0][1] = 2;
44+
let c = a + b;
45+
assert_eq!(c.c[0][1], 3);
46+
}
47+
48+
#[derive(Add)]
49+
struct TupleRecursive((i32, u8), [(i32, u8); 10]);

0 commit comments

Comments
 (0)