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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

extern crate proc_macro;

use proc_macro::TokenStream;
use quote::quote;
use syn::fold::{self, Fold};
use syn::parse::{Parse, ParseStream, Result};
use syn::{parse_macro_input, parse_quote, Attribute, BinOp, Expr, ItemFn};

/// Dummy struct to use for token parsing.
#[derive(Debug)]
struct Args {}

impl Parse for Args {
    fn parse(_input: ParseStream) -> Result<Self> {
        Ok(Args {})
    }
}

impl Args {
    /// Produces an expression that is a binary operation.
    /// TODO: do we need attrs?
    fn binary_op(&mut self, _attrs: Vec<Attribute>, left: Expr, op: &BinOp, right: Expr) -> Expr {
        let left = match left {
            Expr::Binary(e) => self.binary_op(e.attrs, *e.left, &e.op, *e.right),
            _ => fold::fold_expr(self, left),
        };
        let right = match right {
            Expr::Binary(e) => self.binary_op(e.attrs, *e.left, &e.op, *e.right),
            _ => fold::fold_expr(self, right),
        };
        match op {
            BinOp::Add(_op) => {
                return parse_quote!(
                    #left.wrapping_add(#right)
                );
            }
            BinOp::Sub(_op) => {
                return parse_quote!(
                    #left.wrapping_sub(#right)
                );
            }
            BinOp::Mul(_op) => {
                return parse_quote!(
                    #left.wrapping_mul(#right)
                );
            }
            _ => {
                return parse_quote!(
                    #left #op #right
                );
            }
        }
    }

    /// Produces an expression that is an assign operation.
    /// TODO: do we need attrs?
    fn assign_op(&mut self, _attrs: Vec<Attribute>, left: Expr, op: &BinOp, right: Expr) -> Expr {
        let left = match left {
            Expr::Binary(e) => self.binary_op(e.attrs, *e.left, &e.op, *e.right),
            _ => fold::fold_expr(self, left),
        };
        let right = match right {
            Expr::Binary(e) => self.binary_op(e.attrs, *e.left, &e.op, *e.right),
            _ => fold::fold_expr(self, right),
        };
        match op {
            BinOp::AddEq(_op) => {
                return parse_quote!(
                    #left = #left.wrapping_add(#right)
                );
            }
            BinOp::SubEq(_op) => {
                return parse_quote!(
                    #left = #left.wrapping_sub(#right)
                );
            }
            BinOp::MulEq(_op) => {
                return parse_quote!(
                    #left = #left.wrapping_mul(#right)
                );
            }
            _ => {
                return parse_quote!(
                    #left #op #right
                );
            }
        }
    }
}

/// The `Fold` trait is a way to traverse an owned syntax tree and replace some
/// of its nodes.
///
/// Syn provides two other syntax tree traversal traits: `Visit` which walks a
/// shared borrow of a syntax tree, and `VisitMut` which walks an exclusive
/// borrow of a syntax tree and can mutate it in place.
///
/// All three traits have a method corresponding to each type of node in Syn's
/// syntax tree. All of these methods have default no-op implementations that
/// simply recurse on any child nodes. We can override only those methods for
/// which we want non-default behavior. In this case the traversal needs to
/// transform `Expr` nodes.
impl Fold for Args {
    fn fold_expr(&mut self, e: Expr) -> Expr {
        match e {
            Expr::Binary(e) => self.binary_op(e.attrs, *e.left, &e.op, *e.right),
            Expr::AssignOp(e) => self.assign_op(e.attrs, *e.left, &e.op, *e.right),
            _ => fold::fold_expr(self, e),
        }
    }
}

/// Make addition and multiplication wrapping in the annotated function.
///
/// # Example
///
/// ```rust,ignore
/// #[wrapit]
/// fn oops() -> bool {
///     let a: u32 = std::u32::MAX;
///     let b: u32 = 2;
///     let r = a + b;
///     r == 1
/// }
/// ```
#[proc_macro_attribute]
pub fn wrappit(args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);

    // Get something to fold on.
    // TODO: is there a nicer way to do this?
    let mut args = parse_macro_input!(args as Args);

    // Use a syntax tree traversal to transform the function body.
    let output = args.fold_item_fn(input);

    // Hand the resulting function body back to the compiler.
    TokenStream::from(quote!(#output))
}