Skip to main content

ops_derive2/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Literal;
3use quote::{ToTokens, quote};
4use syn::{self, Data, DeriveInput, Ident, parse_macro_input};
5
6enum Ops {
7    Add,
8    Sub,
9    Mul,
10    Div,
11    AddAssign,
12    SubAssign,
13    MulAssign,
14    DivAssign,
15}
16
17fn parse(ast: DeriveInput) -> (Ident, Vec<Box<dyn ToTokens>>) {
18    let name = ast.ident;
19    let struct_fields = if let Data::Struct(struct_data) = ast.data {
20        struct_data.fields
21    } else {
22        panic!("Expected a struct!")
23    };
24
25    let mut field_names: Vec<Box<dyn ToTokens>> = vec![];
26
27    for (id, field) in struct_fields.into_iter().enumerate() {
28        if let Some(name) = field.ident {
29            field_names.push(Box::new(name));
30        } else {
31            // Tuple struct
32            field_names.push(Box::new(Literal::usize_unsuffixed(id)))
33        }
34    }
35
36    (name, field_names)
37}
38
39#[proc_macro_derive(AutoAdd)]
40pub fn derive_add(input: TokenStream) -> TokenStream {
41    complete(input, Ops::Add)
42}
43
44#[proc_macro_derive(AutoSub)]
45pub fn derive_sub(input: TokenStream) -> TokenStream {
46    complete(input, Ops::Sub)
47}
48
49#[proc_macro_derive(AutoMul)]
50pub fn derive_mul(input: TokenStream) -> TokenStream {
51    complete(input, Ops::Mul)
52}
53
54#[proc_macro_derive(AutoDiv)]
55pub fn derive_div(input: TokenStream) -> TokenStream {
56    complete(input, Ops::Div)
57}
58
59#[proc_macro_derive(AutoAddAssign)]
60pub fn derive_add_assign(input: TokenStream) -> TokenStream {
61    complete(input, Ops::AddAssign)
62}
63
64#[proc_macro_derive(AutoSubAssign)]
65pub fn derive_sub_assign(input: TokenStream) -> TokenStream {
66    complete(input, Ops::SubAssign)
67}
68
69#[proc_macro_derive(AutoMulAssign)]
70pub fn derive_mul_assign(input: TokenStream) -> TokenStream {
71    complete(input, Ops::MulAssign)
72}
73
74#[proc_macro_derive(AutoDivAssign)]
75pub fn derive_div_assign(input: TokenStream) -> TokenStream {
76    complete(input, Ops::DivAssign)
77}
78
79#[proc_macro_derive(AutoNeg)]
80pub fn derive_neg(input: TokenStream) -> TokenStream {
81    let (name, fields) = parse(parse_macro_input!(input as DeriveInput));
82
83    let ret = quote! {
84        impl ::std::ops::Neg for #name {
85            type Output = Self;
86
87            fn neg(self) -> Self::Output {
88                #name {
89                    #(
90                        #fields: -self.#fields,
91                    )*
92                }
93            }
94        }
95    };
96    ret.into()
97}
98
99#[proc_macro_derive(AutoAll)]
100pub fn derive_all(input: TokenStream) -> TokenStream {
101    let (name, fields) = parse(parse_macro_input!(input as DeriveInput));
102
103    let add: proc_macro2::TokenStream = complete_internal(&name, &fields, Ops::Add).into();
104    let sub: proc_macro2::TokenStream = complete_internal(&name, &fields, Ops::Sub).into();
105    let mul: proc_macro2::TokenStream = complete_internal(&name, &fields, Ops::Mul).into();
106    let div: proc_macro2::TokenStream = complete_internal(&name, &fields, Ops::Div).into();
107    let add_assign: proc_macro2::TokenStream =
108        complete_internal(&name, &fields, Ops::AddAssign).into();
109    let sub_assign: proc_macro2::TokenStream =
110        complete_internal(&name, &fields, Ops::SubAssign).into();
111    let mul_assign: proc_macro2::TokenStream =
112        complete_internal(&name, &fields, Ops::MulAssign).into();
113    let div_assign: proc_macro2::TokenStream =
114        complete_internal(&name, &fields, Ops::DivAssign).into();
115    let ret = quote! {
116        #add
117        #sub
118        #mul
119        #div
120        #add_assign
121        #sub_assign
122        #mul_assign
123        #div_assign
124
125        impl ::std::ops::Neg for #name {
126            type Output = Self;
127
128            fn neg(self) -> Self::Output {
129                #name {
130                    #(
131                        #fields: -self.#fields,
132                    )*
133                }
134            }
135        }
136    };
137    ret.into()
138}
139
140fn complete(input: TokenStream, trait_: Ops) -> TokenStream {
141    let (name, fields) = parse(parse_macro_input!(input as DeriveInput));
142    complete_internal(&name, &fields, trait_)
143}
144
145fn complete_internal(name: &Ident, fields: &[Box<dyn ToTokens>], trait_: Ops) -> TokenStream {
146    let (trait_name, func_header, operation) = match trait_ {
147        Ops::Add => (
148            quote! {Add},
149            quote! {add(self, rhs: T) -> Self::Output},
150            quote! {+},
151        ),
152        Ops::Sub => (
153            quote! {Sub},
154            quote! {sub(self, rhs: T) -> Self::Output},
155            quote! {-},
156        ),
157        Ops::Mul => (
158            quote! {Mul},
159            quote! {mul(self, rhs: T) -> Self::Output},
160            quote! {*},
161        ),
162        Ops::Div => (
163            quote! {Div},
164            quote! {div(self, rhs: T) -> Self::Output},
165            quote! {/},
166        ),
167        Ops::AddAssign => (
168            quote! {AddAssign},
169            quote! {add_assign(&mut self, rhs: T)},
170            quote! {+},
171        ),
172        Ops::SubAssign => (
173            quote! {SubAssign},
174            quote! {sub_assign(&mut self, rhs: T)},
175            quote! {-},
176        ),
177        Ops::MulAssign => (
178            quote! {MulAssign},
179            quote! {mul_assign(&mut self, rhs: T)},
180            quote! {*},
181        ),
182        Ops::DivAssign => (
183            quote! {DivAssign},
184            quote! {div_assign(&mut self, rhs: T)},
185            quote! {/},
186        ),
187    };
188
189    let output = match trait_ {
190        Ops::Add | Ops::Sub | Ops::Mul | Ops::Div => quote! {type Output = Self;},
191        _ => quote! {},
192    };
193
194    let deref = match trait_ {
195        Ops::AddAssign | Ops::SubAssign | Ops::MulAssign | Ops::DivAssign => quote! {*self = },
196        _ => quote! {},
197    };
198
199    let ret = quote! {
200        impl<T: Into<#name>> ::std::ops::#trait_name<T> for #name {
201            #output
202
203            fn #func_header {
204                let rhs = rhs.into();
205
206                #deref #name {
207                    #(
208                        #fields: self.#fields #operation rhs.#fields,
209                    )*
210                }
211            }
212        }
213    };
214    ret.into()
215}