macromath/
lib.rs

1//! Macromath provides macros which allow convenient calculations with `checked_*`, `wrapping_*` or
2//! `saturating_*` behavior.
3
4extern crate proc_macro;
5
6use proc_macro::TokenStream;
7use quote::quote;
8use syn::{
9  fold::Fold,
10  parse_macro_input,
11  punctuated::Punctuated,
12  token::{Dot, Paren, Question},
13  Expr, ExprMethodCall, ExprParen, ExprTry,
14};
15
16/// Replaces all supported operators in the given expression with their `checked_*` equivalent.
17///
18/// Each replaced operation is followed by the `?` operator to allow operator chaining.
19///
20/// # Supported operators
21///
22/// | Operator | Replacement         |
23/// |----------|---------------------|
24/// | `a + b`  | `a.checked_add(b)?` |
25/// | `a - b`  | `a.checked_sub(b)?` |
26/// | `a * b`  | `a.checked_mul(b)?` |
27/// | `a / b`  | `a.checked_div(b)?` |
28/// | `a % b`  | `a.checked_rem(b)?` |
29/// | `a >> b` | `a.checked_shr(b)?` |
30/// | `a << b` | `a.checked_shl(b)?` |
31/// | `-a`     | `a.checked_neg()?`  |
32///
33/// # Returns
34///
35/// The given expression is wrapped with [Option] and returns [Some] if all checked operations in
36/// replaced expression return [Some] and [None] otherwise.
37///
38/// # Note
39///
40/// The unary `-` (`.neg()`) operator will not be replaced for numeric literals. This allows to
41/// write negative literals as before (`-128i8` is valid, but `-(128i8)` is not).
42///
43/// # Example
44///
45/// ```
46/// use macromath::checked;
47///
48/// fn calc() -> Option<i8> {
49///   Some(100)
50/// }
51///
52/// assert_eq!(Some(1), checked!(1));
53/// assert_eq!(Some(1), checked!(0u32 + 1));
54/// assert_eq!(Some(50), checked!(calc()? / 2));
55/// assert_eq!(None, checked!(i16::MIN - 1));
56/// assert_eq!(Some(0), checked!(1u32 - 1));
57/// assert_eq!(None, checked!(0u32 - 1));
58///
59/// assert_eq!(None, checked!(1u32 << 40));
60/// assert_eq!(None, checked!(1u32 >> 40));
61///
62/// assert_eq!(Some(0), checked!(13u32 / 40));
63/// assert_eq!(Some(13), checked!(13u32 % 40));
64///
65/// assert_eq!(None, checked!(3u32 - 9u32 / 4 * 2));
66/// assert_eq!(Some(1), checked!(3u32 - 9u32 / 4 * 1));
67///
68/// assert_eq!(Some(-128), checked!(-128i8));
69/// assert_eq!(None, checked!(-(-128i8)));
70/// ```
71#[proc_macro]
72pub fn checked(input: TokenStream) -> TokenStream {
73  let expr = parse_macro_input!(input as Expr);
74  let output = checked_impl::Replace.fold_expr(expr);
75  quote!(
76    (|| -> Option<_> { Some(#output) })()
77  )
78  .into()
79}
80
81/// Replaces all supported operators in the given expression with their `wrapping_*` equivalent.
82///
83/// # Returns
84///
85/// The same type as the given expression.
86///
87/// # Supported operators
88///
89/// | Operator | Replacement         |
90/// |----------|---------------------|
91/// | `a + b`  | `a.wrapping_add(b)` |
92/// | `a - b`  | `a.wrapping_sub(b)` |
93/// | `a * b`  | `a.wrapping_mul(b)` |
94/// | `a / b`  | `a.wrapping_div(b)` |
95/// | `a % b`  | `a.wrapping_rem(b)` |
96/// | `a >> b` | `a.wrapping_shr(b)` |
97/// | `a << b` | `a.wrapping_shl(b)` |
98/// | `-a`     | `a.wrapping_neg()`  |
99///
100/// # Note
101///
102/// The unary `-` (`.neg()`) operator will not be replaced for numeric literals. This allows to
103/// write negative literals as before (`-128i8` is valid, but `-(128i8)` is not).
104///
105/// # Example
106///
107/// ```
108/// use macromath::wrapping;
109///
110/// assert_eq!(1, wrapping!(1));
111/// assert_eq!(1, wrapping!(0u32 + 1));
112/// assert_eq!(255, wrapping!(1u8 - 2));
113/// assert_eq!(127, wrapping!(i8::MIN - 1));
114///
115/// assert_eq!(0, wrapping!(1u32 << 35));
116/// assert_eq!(8, wrapping!(1u32 >> 35));
117///
118/// assert_eq!(-2, wrapping!(127i8 * 2));
119/// assert_eq!(-50, wrapping!(-100i8 / 2));
120/// assert_eq!(-128, wrapping!(-128i8 / -1));
121///
122/// assert_eq!(-128, wrapping!(-128i8));
123/// assert_eq!(-128, wrapping!(-(-128i8)));
124/// ```
125#[proc_macro]
126pub fn wrapping(input: TokenStream) -> TokenStream {
127  let expr = parse_macro_input!(input as Expr);
128  let output = wrapping_impl::Replace.fold_expr(expr);
129  quote!(#output).into()
130}
131
132/// Replaces all supported operators in the given expression with their `saturating_*` equivalent.
133///
134/// # Returns
135///
136/// The same type as the given expression.
137///
138/// # Supported operators
139///
140/// | Operator | Replacement           |
141/// |----------|-----------------------|
142/// | `a + b`  | `a.saturating_add(b)` |
143/// | `a - b`  | `a.saturating_sub(b)` |
144/// | `a * b`  | `a.saturating_mul(b)` |
145/// | `a / b`  | `a.saturating_div(b)` |
146/// | `a % b`  | `a.saturating_rem(b)` |
147/// | `-a`     | `a.saturating_neg()`  |
148///
149/// # Note
150///
151/// The unary `-` (`.neg()`) operator will not be replaced for numeric literals. This allows to
152/// write negative literals as before (`-128i8` is valid, but `-(128i8)` is not).
153///
154/// # Example
155///
156/// ```
157/// use macromath::saturating;
158///
159/// assert_eq!(1, saturating!(0u32 + 1));
160/// assert_eq!(0, saturating!(0u32 - 1));
161/// assert_eq!(-128, saturating!(-128i8));
162/// assert_eq!(127, saturating!(-(-128i8)));
163/// assert_eq!(-128, saturating!(i8::MIN - 1));
164///
165/// assert_eq!(1, saturating!(1u8));
166/// assert_eq!(-128, saturating!(-128i8));
167/// assert_eq!(127, saturating!(-(-128i8)));
168/// assert_eq!(0, saturating!(1u8 - 2));
169/// assert_eq!(127, saturating!(100i8 * 2));
170/// ```
171#[proc_macro]
172pub fn saturating(input: TokenStream) -> TokenStream {
173  let expr = parse_macro_input!(input as Expr);
174  let output = saturating_impl::Replace.fold_expr(expr);
175  quote!(#output).into()
176}
177
178mod checked_impl {
179  use crate::{try_op_binary, try_op_unary};
180  use syn::{
181    fold::{self, Fold},
182    BinOp, Expr, ExprBinary, ExprUnary, UnOp,
183  };
184
185  #[derive(Debug)]
186  pub(super) struct Replace;
187
188  impl Fold for Replace {
189    fn fold_expr(&mut self, e: Expr) -> Expr {
190      fold::fold_expr(
191        self,
192        match e {
193          Expr::Binary(ExprBinary {
194            attrs,
195            left,
196            op,
197            right,
198          }) => match op {
199            BinOp::Add(_) => try_op_binary("checked_add", left, *right),
200            BinOp::Sub(_) => try_op_binary("checked_sub", left, *right),
201            BinOp::Mul(_) => try_op_binary("checked_mul", left, *right),
202            BinOp::Div(_) => try_op_binary("checked_div", left, *right),
203            BinOp::Rem(_) => try_op_binary("checked_rem", left, *right),
204            BinOp::Shl(_) => try_op_binary("checked_shr", left, *right),
205            BinOp::Shr(_) => try_op_binary("checked_shl", left, *right),
206            _ => Expr::Binary(ExprBinary {
207              attrs,
208              left,
209              op,
210              right,
211            }),
212          },
213          Expr::Unary(ExprUnary { attrs, op, expr }) => match (op, &*expr) {
214            (_, Expr::Lit(_)) => Expr::Unary(ExprUnary { attrs, op, expr }),
215            (UnOp::Neg(_), _) => try_op_unary("checked_neg", expr),
216            _ => Expr::Unary(ExprUnary { attrs, op, expr }),
217          },
218          _ => e,
219        },
220      )
221    }
222  }
223}
224
225mod wrapping_impl {
226  use crate::{op_binary, op_unary};
227  use syn::{
228    fold::{self, Fold},
229    BinOp, Expr, ExprBinary, ExprUnary, UnOp,
230  };
231
232  #[derive(Debug)]
233  pub(super) struct Replace;
234
235  impl Fold for Replace {
236    fn fold_expr(&mut self, e: Expr) -> Expr {
237      fold::fold_expr(
238        self,
239        match e {
240          Expr::Binary(ExprBinary {
241            attrs,
242            left,
243            op,
244            right,
245          }) => match op {
246            BinOp::Add(_) => op_binary("wrapping_add", left, *right),
247            BinOp::Sub(_) => op_binary("wrapping_sub", left, *right),
248            BinOp::Mul(_) => op_binary("wrapping_mul", left, *right),
249            BinOp::Div(_) => op_binary("wrapping_div", left, *right),
250            BinOp::Rem(_) => op_binary("wrapping_rem", left, *right),
251            BinOp::Shl(_) => op_binary("wrapping_shr", left, *right),
252            BinOp::Shr(_) => op_binary("wrapping_shl", left, *right),
253            _ => Expr::Binary(ExprBinary {
254              attrs,
255              left,
256              op,
257              right,
258            }),
259          },
260          Expr::Unary(ExprUnary { attrs, op, expr }) => match (op, &*expr) {
261            (_, Expr::Lit(_)) => Expr::Unary(ExprUnary { attrs, op, expr }),
262            (UnOp::Neg(_), _) => op_unary("wrapping_neg", expr),
263            _ => Expr::Unary(ExprUnary { attrs, op, expr }),
264          },
265          _ => e,
266        },
267      )
268    }
269  }
270}
271
272mod saturating_impl {
273  use crate::{op_binary, op_unary};
274  use syn::{
275    fold::{self, Fold},
276    BinOp, Expr, ExprBinary, ExprUnary, UnOp,
277  };
278
279  #[derive(Debug)]
280  pub(super) struct Replace;
281
282  impl Fold for Replace {
283    fn fold_expr(&mut self, e: Expr) -> Expr {
284      fold::fold_expr(
285        self,
286        match e {
287          Expr::Binary(ExprBinary {
288            attrs,
289            left,
290            op,
291            right,
292          }) => match op {
293            BinOp::Add(_) => op_binary("saturating_add", left, *right),
294            BinOp::Sub(_) => op_binary("saturating_sub", left, *right),
295            BinOp::Mul(_) => op_binary("saturating_mul", left, *right),
296            BinOp::Div(_) => op_binary("saturating_div", left, *right),
297            BinOp::Rem(_) => op_binary("saturating_rem", left, *right),
298            _ => Expr::Binary(ExprBinary {
299              attrs,
300              left,
301              op,
302              right,
303            }),
304          },
305          Expr::Unary(ExprUnary { attrs, op, expr }) => match (op, &*expr) {
306            (_, Expr::Lit(_)) => Expr::Unary(ExprUnary { attrs, op, expr }),
307            (UnOp::Neg(_), _) => op_unary("saturating_neg", expr),
308            _ => Expr::Unary(ExprUnary { attrs, op, expr }),
309          },
310          _ => e,
311        },
312      )
313    }
314  }
315}
316
317fn op_binary(op: &str, left: Box<Expr>, right: Expr) -> Expr {
318  use proc_macro2::{Ident, Span};
319
320  Expr::MethodCall(ExprMethodCall {
321    attrs: vec![],
322    receiver: Box::new(Expr::Paren(ExprParen {
323      attrs: vec![],
324      paren_token: Paren::default(),
325      expr: left,
326    })),
327    dot_token: Dot::default(),
328    method: Ident::new(op, Span::call_site()),
329    turbofish: None,
330    paren_token: Paren::default(),
331    args: Punctuated::from_iter([right]),
332  })
333}
334
335fn op_unary(op: &str, expr: Box<Expr>) -> Expr {
336  use proc_macro2::{Ident, Span};
337
338  Expr::MethodCall(ExprMethodCall {
339    attrs: vec![],
340    receiver: Box::new(Expr::Paren(ExprParen {
341      attrs: vec![],
342      paren_token: Paren::default(),
343      expr,
344    })),
345    dot_token: Dot::default(),
346    method: Ident::new(op, Span::call_site()),
347    turbofish: None,
348    paren_token: Paren::default(),
349    args: Punctuated::default(),
350  })
351}
352
353fn try_op_binary(op: &str, left: Box<Expr>, right: Expr) -> Expr {
354  use proc_macro2::{Ident, Span};
355
356  Expr::Try(ExprTry {
357    attrs: vec![],
358    expr: Box::new(Expr::MethodCall(ExprMethodCall {
359      attrs: vec![],
360      receiver: Box::new(Expr::Paren(ExprParen {
361        attrs: vec![],
362        paren_token: Paren::default(),
363        expr: left,
364      })),
365      dot_token: Dot::default(),
366      method: Ident::new(op, Span::call_site()),
367      turbofish: None,
368      paren_token: Paren::default(),
369      args: Punctuated::from_iter([right]),
370    })),
371    question_token: Question::default(),
372  })
373}
374
375fn try_op_unary(op: &str, expr: Box<Expr>) -> Expr {
376  use proc_macro2::{Ident, Span};
377
378  Expr::Try(ExprTry {
379    attrs: vec![],
380    expr: Box::new(Expr::MethodCall(ExprMethodCall {
381      attrs: vec![],
382      receiver: Box::new(Expr::Paren(ExprParen {
383        attrs: vec![],
384        paren_token: Paren::default(),
385        expr,
386      })),
387      dot_token: Dot::default(),
388      method: Ident::new(op, Span::call_site()),
389      turbofish: None,
390      paren_token: Paren::default(),
391      args: Punctuated::default(),
392    })),
393    question_token: Question::default(),
394  })
395}