Merge pull request #5 from FredTheDino/op-assign

op assign

3 files changed, 98 insertions, 21 deletions

diff --git a/src/compiler.rs b/src/compiler.rs
index b7fdce1..69bf6eb 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -621,9 +621,38 @@ impl Compiler {
         }
     }
 
-    fn assign(&mut self, name: &str, block: &mut Block) {
+    fn assign(&mut self, block: &mut Block) {
+        let name = match self.eat() {
+            Token::Identifier(name) => name,
+            _ => {
+                error!(self, format!("Expected identifier in assignment"));
+                return;
+            }
+        };
+
+        let op = match self.eat() {
+            Token::Equal => None,
+
+            Token::PlusEqual => Some(Op::Add),
+            Token::MinusEqual => Some(Op::Sub),
+            Token::StarEqual => Some(Op::Mul),
+            Token::SlashEqual => Some(Op::Div),
+
+            _ => {
+                error!(self, format!("Expected '=' in assignment"));
+                return;
+            }
+        };
+
         if let Some(var) = self.find_variable(&name) {
-            self.expression(block);
+            if let Some(op) = op {
+                block.add(Op::Copy, self.line());
+                self.expression(block);
+                block.add(op, self.line());
+            } else {
+                self.expression(block);
+            }
+
             if var.upvalue {
                 block.add(Op::AssignUpvalue(var.slot), self.line());
             } else {
@@ -832,14 +861,31 @@ impl Compiler {
                             return Err(());
                         };
 
-                        if self.peek() == Token::Equal {
-                            self.eat();
-                            self.expression(block);
-                            block.add(Op::Set(field), self.line());
-                            return Ok(());
-                        } else {
-                            block.add(Op::Get(field), self.line());
-                        }
+                        let op = match self.peek() {
+                            Token::Equal => {
+                                self.eat();
+                                self.expression(block);
+                                block.add(Op::Set(field), self.line());
+                                return Ok(());
+                            }
+
+                            Token::PlusEqual => Op::Add,
+                            Token::MinusEqual => Op::Sub,
+                            Token::StarEqual => Op::Mul,
+                            Token::SlashEqual => Op::Div,
+
+                            _ => {
+                                block.add(Op::Get(field), self.line());
+                                continue;
+                            }
+                        };
+                        block.add(Op::Copy, self.line());
+                        block.add(Op::Get(field.clone()), self.line());
+                        self.eat();
+                        self.expression(block);
+                        block.add(op, self.line());
+                        block.add(Op::Set(field), self.line());
+                        return Ok(());
                     }
                     Token::LeftParen => {
                         self.call(block);
@@ -867,6 +913,16 @@ impl Compiler {
                 block.add(Op::Print, self.line());
             }
 
+            (Token::Identifier(_), Token::Equal, ..) |
+            (Token::Identifier(_), Token::PlusEqual, ..) |
+            (Token::Identifier(_), Token::MinusEqual, ..) |
+            (Token::Identifier(_), Token::SlashEqual, ..) |
+            (Token::Identifier(_), Token::StarEqual, ..)
+
+                => {
+                self.assign(block);
+            }
+
             (Token::Identifier(_), Token::Dot, ..) => {
                 let block_length = block.ops.len();
                 let token_length = self.curr;
@@ -900,12 +956,6 @@ impl Compiler {
                 self.definition_statement(&name, Type::UnknownType, block);
             }
 
-            (Token::Identifier(name), Token::Equal, ..) => {
-                self.eat();
-                self.eat();
-                self.assign(&name, block);
-            }
-
             (Token::Blob, Token::Identifier(_), ..) => {
                 self.blob_statement(block);
             }
diff --git a/src/lib.rs b/src/lib.rs
index 42f24bd..453d6ce 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -325,6 +325,26 @@ a() <=> 4
 
     test_file!(scoping, "tests/scoping.tdy");
     test_file!(for_, "tests/for.tdy");
+
+    test_multiple!(
+        op_assign,
+        add: "a := 1\na += 1\na <=> 2",
+        sub: "a := 2\na -= 1\na <=> 1",
+        mul: "a := 2\na *= 2\na <=> 4",
+        div: "a := 2\na /= 2\na <=> 1",
+        cluster: "
+blob A { a: int }
+a := A()
+a.a = 0
+a.a += 1
+a.a <=> 1
+a.a *= 2
+a.a <=> 2
+a.a /= 2
+a.a <=> 1
+a.a -= 1
+a.a <=> 0"
+    );
 }
 
 #[derive(Clone)]
@@ -431,6 +451,7 @@ pub enum Op {
 
     Pop,
     PopUpvalue,
+    Copy,
     Constant(Value),
 
     Get(String),
diff --git a/src/vm.rs b/src/vm.rs
index fd8eb73..6ff7b65 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -159,17 +159,23 @@ impl VM {
                 self.pop();
             }
 
-            Op::Yield => {
-                self.frame_mut().ip += 1;
-                return Ok(OpResult::Yield);
-            }
-
             Op::PopUpvalue => {
                 let value = self.pop();
                 let slot = self.stack.len();
                 self.drop_upvalue(slot, value);
             }
 
+            Op::Copy => {
+                let v = self.pop();
+                self.push(v.clone());
+                self.push(v);
+            }
+
+            Op::Yield => {
+                self.frame_mut().ip += 1;
+                return Ok(OpResult::Yield);
+            }
+
             Op::Constant(value) => {
                 let offset = self.frame().stack_offset;
                 let value = match value {