smooth_operator_impl/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::TokenStream as TokenStream2;
3use quote::{quote, ToTokens};
4use syn::visit_mut::VisitMut;
5
6/// Convert arithmetic operators within the given expression to their checked
7/// variants and provide detailed error strings about which operator has failed
8/// for diagnostic.
9#[proc_macro]
10pub fn checked(expression: TokenStream) -> TokenStream {
11    let (result, expression_str) = checked_inner(expression.into());
12
13    let crate_name = {
14        let this_crate_without_impl = env!("CARGO_PKG_NAME").trim_end_matches("-impl");
15
16        if std::env::var("CARGO_PKG_NAME").unwrap() == this_crate_without_impl {
17            quote!(crate)
18        } else {
19            let ident = this_crate_without_impl.replace('-', "_");
20            let ident = syn::Ident::new(&ident, proc_macro2::Span::call_site());
21            quote!(::#ident)
22        }
23    };
24
25    quote! {
26        (|| -> ::core::result::Result<_, #crate_name::Error> {
27            type Err = #crate_name::Error;
28            const ORIGINAL_EXPR: &'static str = #expression_str;
29
30            Ok(
31                #[allow(clippy::needless_question_mark)]
32                #[allow(unused_parens)]
33                {
34                    #result
35                }
36            )
37        })()
38    }
39    .into()
40}
41
42#[inline]
43fn checked_inner(expression: TokenStream2) -> (TokenStream2, String) {
44    let mut expr: syn::Expr =
45        syn::parse2(expression).expect("Failed to parse arithmetic expression");
46    let original_expr = expr.to_token_stream().to_string();
47    CheckedArith.visit_expr_mut(&mut expr);
48    (expr.to_token_stream(), original_expr)
49}
50
51struct CheckedArith;
52
53impl VisitMut for CheckedArith {
54    fn visit_expr_mut(&mut self, node: &mut syn::Expr) {
55        match node {
56            syn::Expr::Binary(syn::ExprBinary {
57                left, right, op, ..
58            }) => {
59                let op_len = op.to_token_stream().to_string().len();
60                let op_ix = {
61                    let left_len = left.to_token_stream().to_string().len();
62
63                    left_len
64                        + 1 // Add 1 for whitespace
65                        + op_len
66                };
67
68                self.visit_expr_mut(left);
69                self.visit_expr_mut(right);
70
71                match op {
72                    syn::BinOp::Add(_) => {
73                        *node = syn::parse2::<syn::Expr>(quote! {
74                            #left.checked_add(#right).ok_or(Err {
75                                expr: ORIGINAL_EXPR,
76                                __op_ix: #op_ix,
77                                __op_len: #op_len,
78                            })?
79                        })
80                        .unwrap();
81                    }
82                    syn::BinOp::AddAssign(_) => {
83                        *node = syn::parse2::<syn::Expr>(quote! {
84                            #left = #left.checked_add(#right).ok_or(Err {
85                                expr: ORIGINAL_EXPR,
86                                __op_ix: #op_ix,
87                                __op_len: #op_len,
88                            })?
89                        })
90                        .unwrap();
91                    }
92                    syn::BinOp::Sub(_) => {
93                        *node = syn::parse2::<syn::Expr>(quote! {
94                            #left.checked_sub(#right).ok_or(Err {
95                                expr: ORIGINAL_EXPR,
96                                __op_ix: #op_ix,
97                                __op_len: #op_len,
98                            })?
99                        })
100                        .unwrap();
101                    }
102                    syn::BinOp::SubAssign(_) => {
103                        *node = syn::parse2::<syn::Expr>(quote! {
104                            #left = #left.checked_sub(#right).ok_or(Err {
105                                expr: ORIGINAL_EXPR,
106                                __op_ix: #op_ix,
107                                __op_len: #op_len,
108                            })?
109                        })
110                        .unwrap();
111                    }
112                    syn::BinOp::Div(_) => {
113                        *node = syn::parse2::<syn::Expr>(quote! {
114                            #left.checked_div(#right).ok_or(Err {
115                                expr: ORIGINAL_EXPR,
116                                __op_ix: #op_ix,
117                                __op_len: #op_len,
118                            })?
119                        })
120                        .unwrap();
121                    }
122                    syn::BinOp::DivAssign(_) => {
123                        *node = syn::parse2::<syn::Expr>(quote! {
124                            #left = #left.checked_div(#right).ok_or(Err {
125                                expr: ORIGINAL_EXPR,
126                                __op_ix: #op_ix,
127                                __op_len: #op_len,
128                            })?
129                        })
130                        .unwrap();
131                    }
132                    syn::BinOp::Mul(_) => {
133                        *node = syn::parse2::<syn::Expr>(quote! {
134                            #left.checked_mul(#right).ok_or(Err {
135                                expr: ORIGINAL_EXPR,
136                                __op_ix: #op_ix,
137                                __op_len: #op_len,
138                            })?
139                        })
140                        .unwrap();
141                    }
142                    syn::BinOp::MulAssign(_) => {
143                        *node = syn::parse2::<syn::Expr>(quote! {
144                            #left = #left.checked_mul(#right).ok_or(Err {
145                                expr: ORIGINAL_EXPR,
146                                __op_ix: #op_ix,
147                                __op_len: #op_len,
148                            })?
149                        })
150                        .unwrap();
151                    }
152                    syn::BinOp::Rem(_) => {
153                        *node = syn::parse2::<syn::Expr>(quote! {
154                            #left.checked_rem(#right).ok_or(Err {
155                                expr: ORIGINAL_EXPR,
156                                __op_ix: #op_ix,
157                                __op_len: #op_len,
158                            })?
159                        })
160                        .unwrap();
161                    }
162                    syn::BinOp::RemAssign(_) => {
163                        *node = syn::parse2::<syn::Expr>(quote! {
164                            #left = #left.checked_rem(#right).ok_or(Err {
165                                expr: ORIGINAL_EXPR,
166                                __op_ix: #op_ix,
167                                __op_len: #op_len,
168                            })?
169                        })
170                        .unwrap();
171                    }
172                    syn::BinOp::BitXor(_) => {
173                        *node = syn::parse2::<syn::Expr>(quote! {
174                            #left.checked_pow(#right).ok_or(Err {
175                                expr: ORIGINAL_EXPR,
176                                __op_ix: #op_ix,
177                                __op_len: #op_len,
178                            })?
179                        })
180                        .unwrap();
181                    }
182                    syn::BinOp::BitXorAssign(_) => {
183                        *node = syn::parse2::<syn::Expr>(quote! {
184                            #left = #left.checked_pow(#right).ok_or(Err {
185                                expr: ORIGINAL_EXPR,
186                                __op_ix: #op_ix,
187                                __op_len: #op_len,
188                            })?
189                        })
190                        .unwrap();
191                    }
192                    syn::BinOp::Shl(_) => {
193                        *node = syn::parse2::<syn::Expr>(quote! {
194                            #left.checked_shl(#right).ok_or(Err {
195                                expr: ORIGINAL_EXPR,
196                                __op_ix: #op_ix,
197                                __op_len: #op_len,
198                            })?
199                        })
200                        .unwrap();
201                    }
202                    syn::BinOp::ShlAssign(_) => {
203                        *node = syn::parse2::<syn::Expr>(quote! {
204                            #left = #left.checked_shl(#right).ok_or(Err {
205                                expr: ORIGINAL_EXPR,
206                                __op_ix: #op_ix,
207                                __op_len: #op_len,
208                            })?
209                        })
210                        .unwrap();
211                    }
212                    syn::BinOp::Shr(_) => {
213                        *node = syn::parse2::<syn::Expr>(quote! {
214                            #left.checked_shr(#right).ok_or(Err {
215                                expr: ORIGINAL_EXPR,
216                                __op_ix: #op_ix,
217                                __op_len: #op_len,
218                            })?
219                        })
220                        .unwrap();
221                    }
222                    syn::BinOp::ShrAssign(_) => {
223                        *node = syn::parse2::<syn::Expr>(quote! {
224                            #left = #left.checked_shr(#right).ok_or(Err {
225                                expr: ORIGINAL_EXPR,
226                                __op_ix: #op_ix,
227                                __op_len: #op_len,
228                            })?
229                        })
230                        .unwrap();
231                    }
232                    _ => {}
233                }
234            }
235            syn::Expr::Unary(syn::ExprUnary { op, expr, .. }) => {
236                self.visit_expr_mut(expr);
237
238                if let syn::UnOp::Neg(_) = op {
239                    *node = syn::parse2::<syn::Expr>(quote! {
240                        #expr.checked_neg().ok_or(Err {
241                            expr: ORIGINAL_EXPR,
242                            __op_len: 1,
243                            __op_ix: 0, // Negation comes first
244                        })?
245                    })
246                    .unwrap();
247                }
248            }
249            syn::Expr::Paren(expr) => {
250                self.visit_expr_paren_mut(expr);
251            }
252            syn::Expr::Call(expr) => {
253                self.visit_expr_call_mut(expr);
254            }
255            syn::Expr::MethodCall(expr) => {
256                self.visit_expr_method_call_mut(expr);
257            }
258            syn::Expr::Path(_) | syn::Expr::Lit(_) => {}
259            _ => {}
260        }
261    }
262}