1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
//! Crate providing a procedural macro to make normal operations into overflow checked operations easily
extern crate proc_macro;

use proc_macro2::{Ident, Span, TokenStream, TokenTree};
use quote::quote;
use syn::{parse_macro_input, BinOp, Expr, ExprBinary, ExprMethodCall, ExprUnary, UnOp};

/// Procedural macro to convert normal integer operations into overflow-checked operations.
/// The expression will evaluate to `Some(result)` if all operations succed or to `None` if any of
/// the operations fails.
///
/// This Macro supports checked functions for these operators:
/// `-` (`neg`),`+` (`add`), `-` (`sub`), `*` (`mul`), `/` (`div`), `%` (`rem`), `<<` (`shl`), `>>` (`shr`)
/// And it also substitutes these methods with their checked variants:
/// `abs`, `pow`
///
/// # Examples
/// ```
/// use checked_expr::checked_expr;
/// assert_eq!(checked_expr!(254_u8 + 1), Some(255));
/// assert_eq!(checked_expr!(255_u8 + 1), None);
///
/// assert_eq!(checked_expr!(-(-127_i8)), Some(127));
/// assert_eq!(checked_expr!(-(-128_i8)), None);
///
/// // also works with `pow` and `abs`
/// assert_eq!(checked_expr!((-128_i8).abs()), None);
/// assert_eq!(checked_expr!(10_u16.pow(100)), None);
///
/// // you can also arbitrarily nest expressions although you sometimes need to be very
/// // explicit with the types on literals on the left hand side of operations
/// assert_eq!(checked_expr!((10_i32 - 8) * (3_i32 + 6) % 8), Some(2));
/// ```
#[proc_macro]
pub fn checked_expr(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    // the idea is to convert something like this
    // 1 + 2 * 3
    // into
    // (|| Some(1.checked_add(2.checked_mul(3)?)?) )()
    let input = parse_macro_input!(input as Expr);

    let expr = expr(&input);
    let tokens = quote! {
        (|| Some(#expr) )()
    };
    tokens.into()
}

fn expr(expr: &Expr) -> TokenStream {
    match expr {
        Expr::Unary(expr) => unary_expr(expr),
        Expr::Binary(expr) => binary_expr(expr),
        Expr::Paren(expr) => crate::expr(&expr.expr),
        Expr::MethodCall(expr) => method_call(expr),
        _ => quote! {
            (#expr)
        },
    }
}

fn binary_expr(expr: &ExprBinary) -> TokenStream {
    let function_name = match expr.op {
        BinOp::Add(_) => "checked_add",
        BinOp::Sub(_) => "checked_sub",
        BinOp::Mul(_) => "checked_mul",
        BinOp::Div(_) => "checked_div",
        BinOp::Rem(_) => "checked_rem",
        BinOp::Shl(_) => "checked_shl",
        BinOp::Shr(_) => "checked_shr",
        _ => {
            return quote! {
                (#expr)
            }
        }
    };
    let fn_ident = TokenTree::Ident(Ident::new(function_name, Span::call_site()));

    let left = crate::expr(&expr.left);
    let right = crate::expr(&expr.right);
    quote! {
        #left.#fn_ident(#right)?
    }
}

fn unary_expr(expr: &ExprUnary) -> TokenStream {
    let function_name = match expr.op {
        UnOp::Neg(_) => "checked_neg",
        _ => {
            return quote! {
                (#expr)
            }
        }
    };
    let fn_ident = TokenTree::Ident(Ident::new(function_name, Span::call_site()));

    let inner = crate::expr(&expr.expr);
    quote! {
        #inner.#fn_ident()?
    }
}

fn method_call(expr: &ExprMethodCall) -> TokenStream {
    let method_name = match expr.method.to_string().as_str() {
        "abs" => "checked_abs",
        "pow" => "checked_pow",
        _ => {
            return quote! {
                (#expr)
            }
        }
    };
    let method_ident = TokenTree::Ident(Ident::new(method_name, Span::call_site()));
    let left = crate::expr(&expr.receiver);
    let args = &expr.args;
    quote! {
        #left.#method_ident(#args)?
    }
}