Merge pull request #42 from rustysec/compound_assignment

candidate compound assignment implementation

src/engine.rs

@@ -7,7 +7,7 @@ use std::fmt;
 use parser::{lex, parse, Expr, Stmt, FnDef};
 use fn_register::FnRegister;
 
-use std::ops::{Add, Sub, Mul, Div, Neg};
+use std::ops::{Add, Sub, Mul, Div, Neg, BitAnd, BitOr, BitXor, Shl, Shr};
 use std::cmp::{PartialOrd, PartialEq};
 
 #[derive(Debug)]
@@ -1411,6 +1411,11 @@ impl Engine {
         fn or(x: bool, y: bool)    -> bool { x || y }
         fn not(x: bool)            -> bool { !x }
         fn concat(x: String, y: String) -> String { x + &y }
+        fn binary_and<T: BitAnd>(x: T, y: T)   -> <T as BitAnd>::Output { x & y }
+        fn binary_or<T: BitOr>(x: T, y: T)   -> <T as BitOr>::Output { x | y }
+        fn binary_xor<T: BitXor>(x: T, y: T)   -> <T as BitXor>::Output { x ^ y }
+        fn left_shift<T: Shl<T>>(x: T, y: T)   -> <T as Shl<T>>::Output { x.shl(y) }
+        fn right_shift<T: Shr<T>>(x: T, y: T)   -> <T as Shr<T>>::Output { x.shr(y) }
 
         reg_op!(engine, "+", add, i32, i64, u32, u64, f32, f64);
         reg_op!(engine, "-", sub, i32, i64, u32, u64, f32, f64);
@@ -1426,6 +1431,13 @@ impl Engine {
 
         reg_op!(engine, "||", or, bool);
         reg_op!(engine, "&&", and, bool);
+        reg_op!(engine, "|", binary_or, i32, i64, u32, u64);
+        reg_op!(engine, "|", or, bool);
+        reg_op!(engine, "&", binary_and, i32, i64, u32, u64);
+        reg_op!(engine, "&", and, bool);
+        reg_op!(engine, "^", binary_xor, i32, i64, u32, u64);
+        reg_op!(engine, "<<", left_shift, i32, i64, u32, u64);
+        reg_op!(engine, ">>", right_shift, i32, i64, u32, u64);
 
         reg_un!(engine, "-", neg, i32, i64, f32, f64);
         reg_un!(engine, "!", not, bool);

src/parser.rs

@@ -158,6 +158,16 @@ pub enum Token {
     Return,
     PlusEquals,
     MinusEquals,
+    MultiplyEquals,
+    DivideEquals,
+    LeftShiftEquals,
+    RightShiftEquals,
+    AndEquals,
+    OrEquals,
+    XOrEquals,
+    LeftShift,
+    RightShift,
+    XOr,
     LexErr(LexError),
 }
 
@@ -473,7 +483,15 @@ impl<'a> TokenIterator<'a> {
                         _ => Some(Token::Minus),
                     }
                 },
-                '*' => return Some(Token::Multiply),
+                '*' => {
+                    return match self.char_stream.peek() {
+                        Some(&'=') => {
+                            self.char_stream.next();
+                            Some(Token::MultiplyEquals)
+                        },
+                        _ => Some(Token::Multiply)
+                    }
+                },
                 '/' => {
                     match self.char_stream.peek() {
                         Some(&'/') => {
@@ -501,6 +519,10 @@ impl<'a> TokenIterator<'a> {
                                 }
                             }
                         }
+                        Some(&'=') => {
+                            self.char_stream.next();
+                            return Some(Token::DivideEquals);
+                        }
                         _ => return Some(Token::Divide),
                     }
                 }
@@ -523,6 +545,19 @@ impl<'a> TokenIterator<'a> {
                             self.char_stream.next();
                             return Some(Token::LessThanEqual);
                         }
+                        Some(&'<') => {
+                            self.char_stream.next();
+                            return match self.char_stream.peek() {
+                                Some(&'=') => {
+                                    self.char_stream.next();
+                                    Some(Token::LeftShiftEquals)
+                                },
+                                _ => {
+                                    self.char_stream.next();
+                                    Some(Token::LeftShift)
+                                }
+                            }
+                        }
                         _ => return Some(Token::LessThan),
                     }
                 }
@@ -532,6 +567,19 @@ impl<'a> TokenIterator<'a> {
                             self.char_stream.next();
                             return Some(Token::GreaterThanEqual);
                         }
+                        Some(&'>') => {
+                            self.char_stream.next();
+                            return match self.char_stream.peek() {
+                                Some(&'=') => {
+                                    self.char_stream.next();
+                                    Some(Token::RightShiftEquals)
+                                },
+                                _ => {
+                                    self.char_stream.next();
+                                    Some(Token::RightShift)
+                                }
+                            }
+                        }
                         _ => return Some(Token::GreaterThan),
                     }
                 }
@@ -550,6 +598,10 @@ impl<'a> TokenIterator<'a> {
                             self.char_stream.next();
                             return Some(Token::Or);
                         }
+                        Some(&'=') => {
+                            self.char_stream.next();
+                            return Some(Token::OrEquals);
+                        }
                         _ => return Some(Token::Pipe),
                     }
                 }
@@ -559,9 +611,22 @@ impl<'a> TokenIterator<'a> {
                             self.char_stream.next();
                             return Some(Token::And);
                         }
+                        Some(&'=') => {
+                            self.char_stream.next();
+                            return Some(Token::AndEquals);
+                        }
                         _ => return Some(Token::Ampersand),
                     }
                 }
+                '^' => {
+                    match self.char_stream.peek() {
+                        Some(&'=') => {
+                            self.char_stream.next();
+                            return Some(Token::XOrEquals);
+                        }
+                        _ => return Some(Token::XOr)
+                    }
+                }
                 _x if _x.is_whitespace() => (),
                 _ => return Some(Token::LexErr(LexError::UnexpectedChar)),
             }
@@ -592,8 +657,17 @@ fn get_precedence(token: &Token) -> i32 {
     match *token {
         Token::Equals
         | Token::PlusEquals
-        | Token::MinusEquals => 10,
+        | Token::MinusEquals
-        Token::Or => 11,
+        | Token::MultiplyEquals
+        | Token::DivideEquals
+        | Token::LeftShiftEquals
+        | Token::RightShiftEquals
+        | Token::AndEquals
+        | Token::OrEquals
+        | Token::XOrEquals => 10,
+        Token::Or
+        | Token::XOr
+        | Token::Pipe  => 11,
         Token::And => 12,
         Token::LessThan
         | Token::LessThanEqual
@@ -605,6 +679,8 @@ fn get_precedence(token: &Token) -> i32 {
         | Token::Minus => 20,
         Token::Divide
         | Token::Multiply => 40,
+        Token::LeftShift
+        | Token::RightShift => 50,
         Token::Period => 100,
         _ => -1,
     }
@@ -816,6 +892,61 @@ fn parse_binop<'a>(input: &mut Peekable<TokenIterator<'a>>,
                 }
                 Token::Or => Expr::FnCall("||".to_string(), vec![lhs_curr, rhs]),
                 Token::And => Expr::FnCall("&&".to_string(), vec![lhs_curr, rhs]),
+                Token::XOr => Expr::FnCall("^".to_string(), vec![lhs_curr, rhs]),
+                Token::OrEquals => {
+                    let lhs_copy = lhs_curr.clone();
+                    Expr::Assignment(
+                        Box::new(lhs_curr),
+                        Box::new(Expr::FnCall("|".to_string(), vec![lhs_copy, rhs]))
+                    )
+                },
+                Token::AndEquals => {
+                    let lhs_copy = lhs_curr.clone();
+                    Expr::Assignment(
+                        Box::new(lhs_curr),
+                        Box::new(Expr::FnCall("&".to_string(), vec![lhs_copy, rhs]))
+                    )
+                },
+                Token::XOrEquals => {
+                    let lhs_copy = lhs_curr.clone();
+                    Expr::Assignment(
+                        Box::new(lhs_curr),
+                        Box::new(Expr::FnCall("^".to_string(), vec![lhs_copy, rhs]))
+                    )
+                },
+                Token::MultiplyEquals => {
+                    let lhs_copy = lhs_curr.clone();
+                    Expr::Assignment(
+                        Box::new(lhs_curr),
+                        Box::new(Expr::FnCall("*".to_string(), vec![lhs_copy, rhs]))
+                    )
+                },
+                Token::DivideEquals => {
+                    let lhs_copy = lhs_curr.clone();
+                    Expr::Assignment(
+                        Box::new(lhs_curr),
+                        Box::new(Expr::FnCall("/".to_string(), vec![lhs_copy, rhs]))
+                    )
+                },
+                Token::Pipe => {
+                    Expr::FnCall("|".to_string(), vec![lhs_curr, rhs])
+                },
+                Token::LeftShift => Expr::FnCall("<<".to_string(), vec![lhs_curr, rhs]),
+                Token::RightShift => Expr::FnCall(">>".to_string(), vec![lhs_curr, rhs]),
+                Token::LeftShiftEquals => {
+                    let lhs_copy = lhs_curr.clone();
+                    Expr::Assignment(
+                        Box::new(lhs_curr),
+                        Box::new(Expr::FnCall("<<".to_string(), vec![lhs_copy, rhs]))
+                    )
+                },
+                Token::RightShiftEquals => {
+                    let lhs_copy = lhs_curr.clone();
+                    Expr::Assignment(
+                        Box::new(lhs_curr),
+                        Box::new(Expr::FnCall(">>".to_string(), vec![lhs_copy, rhs]))
+                    )
+                },
                 _ => return Err(ParseError::UnknownOperator),
             };
         }

src/tests.rs

@@ -560,3 +560,132 @@ fn test_decrement() {
         assert!(true);
     }
 }
+
+#[test]
+fn test_or_equals() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("let x = 16; x |= 74; x") {
+        assert_eq!(result, 90);
+    } else {
+        assert!(false);
+    }
+
+    if let Ok(result) = engine.eval::<bool>("let x = true; x |= false; x") {
+        assert!(result);
+    } else {
+        assert!(false);
+    }
+
+    if let Ok(result) = engine.eval::<bool>("let x = false; x |= true; x") {
+        assert!(result);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_and_equals() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("let x = 16; x &= 31; x") {
+        assert_eq!(result, 16);
+    } else {
+        assert!(false);
+    }
+
+    if let Ok(result) = engine.eval::<bool>("let x = true; x &= false; x") {
+        assert_eq!(result, false);
+    } else {
+        assert!(false);
+    }
+
+    if let Ok(result) = engine.eval::<bool>("let x = false; x &= true; x") {
+        assert_eq!(result, false);
+    } else {
+        assert!(false);
+    }
+
+    if let Ok(result) = engine.eval::<bool>("let x = true; x &= true; x") {
+        assert!(result);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_xor_equals() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("let x = 90; x ^= 12; x") {
+        assert_eq!(result, 86);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_multiply_equals() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("let x = 2; x *= 3; x") {
+        assert_eq!(result, 6);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_divide_equals() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("let x = 6; x /= 2; x") {
+        assert_eq!(result, 3);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_left_shift() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("4 << 2") {
+        assert_eq!(result, 16);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_right_shift() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("9 >> 1") {
+        assert_eq!(result, 4);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_left_shift_equals() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("let x = 9; x >>=1; x") {
+        assert_eq!(result, 4);
+    } else {
+        assert!(false);
+    }
+}
+
+#[test]
+fn test_right_shift_equals() {
+    let mut engine = Engine::new();
+
+    if let Ok(result) = engine.eval::<i64>("let x = 4; x<<= 2; x") {
+        assert_eq!(result, 16);
+    } else {
+        assert!(false);
+    }
+}