Add augmented assignment operators for set statements

This allows us to do things like

    self.foo += 5

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

src/compiler.rs

@@ -39,6 +39,40 @@ fn unescape(s: &str) -> String {
         .replace("\\\\", "\\")
 }
 
+static BIN_OP_NAMES: LazyLock<HashMap<TokenKind, &'static str>> = LazyLock::new(|| {
+    hash_map! {
+        TokenKind::Plus => "__add__",
+        TokenKind::Minus => "__sub__",
+        TokenKind::Star => "__mul__",
+        TokenKind::Slash => "__div__",
+        TokenKind::And => "__and__",
+        TokenKind::Or => "__or__",
+        TokenKind::BangEq => "__ne__",
+        TokenKind::EqEq => "__eq__",
+        TokenKind::Greater => "__gt__",
+        TokenKind::GreaterEq => "__ge__",
+        TokenKind::Less => "__lt__",
+        TokenKind::LessEq => "__le__",
+    }
+});
+
+static UNARY_OP_NAMES: LazyLock<HashMap<TokenKind, &'static str>> = LazyLock::new(|| {
+    hash_map! {
+        TokenKind::Plus => "__pos__",
+        TokenKind::Minus => "__neg__",
+        TokenKind::Bang => "__not__",
+    }
+});
+
+static AUG_ASSIGN_NAMES: LazyLock<HashMap<TokenKind, &'static str>> = LazyLock::new(|| {
+    hash_map! {
+        TokenKind::PlusEq => "__add__",
+        TokenKind::MinusEq => "__sub__",
+        TokenKind::StarEq => "__mul__",
+        TokenKind::SlashEq => "__div__",
+    }
+});
+
 ////////////////////////////////////////////////////////////////////////////////
 // Scope
 ////////////////////////////////////////////////////////////////////////////////
@@ -573,16 +607,8 @@ impl StmtVisitor for Compiler<'_> {
             }
             // augmented assignment
             TokenKind::PlusEq | TokenKind::MinusEq | TokenKind::StarEq | TokenKind::SlashEq => {
-                static OP_NAMES: LazyLock<HashMap<TokenKind, &'static str>> = LazyLock::new(|| {
-                    hash_map! {
-                        TokenKind::PlusEq => "__add__",
-                        TokenKind::MinusEq => "__sub__",
-                        TokenKind::StarEq => "__mul__",
-                        TokenKind::SlashEq => "__div__",
-                    }
-                });
                 self.emit_lookup(line, name)?;
-                let op_name = OP_NAMES
+                let op_name = AUG_ASSIGN_NAMES
                     .get(&stmt.op.kind)
                     .expect("invalid augmented assignment operator");
                 let op_constant = self.insert_constant(Str::create(op_name))?;
@@ -598,10 +624,35 @@ impl StmtVisitor for Compiler<'_> {
     }
 
     fn visit_set_stmt(&mut self, stmt: &SetStmt) -> Result<()> {
-        self.compile_expr(&stmt.expr)?;
         let name = self.insert_constant(Str::create(&stmt.name.text))?;
+        let line = stmt_line_number(stmt);
+
+        self.compile_expr(&stmt.expr)?;
+
+        match stmt.op.kind {
+            // normal assignment
+            TokenKind::Eq => {
+                self.compile_expr(&stmt.rhs)?;
+                self.emit(line, Op::SetAttr(name));
+            }
+            // augmented assignment
+            TokenKind::PlusEq | TokenKind::MinusEq | TokenKind::StarEq | TokenKind::SlashEq => {
+                //
+                self.emit(line, Op::Dup);
+                self.emit(line, Op::GetAttr(name));
+                let op = AUG_ASSIGN_NAMES
+                    .get(&stmt.op.kind)
+                    .expect("invalid augmented assign operator");
+                let op_constant = self.insert_constant(Str::create(op))?;
+                self.emit(line, Op::GetAttr(op_constant));
+                self.compile_expr(&stmt.rhs)?;
+                self.emit(line, Op::Call(1));
+                self.emit(line, Op::SetAttr(name));
+            }
+            _ => unreachable!(),
+        }
+
         self.compile_expr(&stmt.rhs)?;
-        self.emit(stmt_line_number(stmt), Op::SetAttr(name));
         Ok(())
     }
 
@@ -677,23 +728,6 @@ impl StmtVisitor for Compiler<'_> {
 
 impl ExprVisitor for Compiler<'_> {
     fn visit_binary_expr(&mut self, expr: &BinaryExpr) -> Result<()> {
-        static OP_NAMES: LazyLock<HashMap<TokenKind, &'static str>> = LazyLock::new(|| {
-            hash_map! {
-                TokenKind::Plus => "__add__",
-                TokenKind::Minus => "__sub__",
-                TokenKind::Star => "__mul__",
-                TokenKind::Slash => "__div__",
-                TokenKind::And => "__and__",
-                TokenKind::Or => "__or__",
-                TokenKind::BangEq => "__ne__",
-                TokenKind::EqEq => "__eq__",
-                TokenKind::Greater => "__gt__",
-                TokenKind::GreaterEq => "__ge__",
-                TokenKind::Less => "__lt__",
-                TokenKind::LessEq => "__le__",
-            }
-        });
-
         self.compile_expr(&expr.lhs)?;
 
         // short-circuit setup
@@ -711,7 +745,7 @@ impl ExprVisitor for Compiler<'_> {
             }
         }
 
-        let name = OP_NAMES
+        let name = BIN_OP_NAMES
             .get(&expr.op.kind)
             .expect("invalid binary operator");
         let constant_id = self.insert_constant(Str::create(name))?;
@@ -750,15 +784,10 @@ impl ExprVisitor for Compiler<'_> {
     }
 
     fn visit_unary_expr(&mut self, expr: &UnaryExpr) -> Result<()> {
-        static OP_NAMES: LazyLock<HashMap<TokenKind, &'static str>> = LazyLock::new(|| {
-            hash_map! {
-                TokenKind::Plus => "__pos__",
-                TokenKind::Minus => "__neg__",
-                TokenKind::Bang => "__not__",
-            }
-        });
         self.compile_expr(&expr.expr)?;
-        let name = OP_NAMES.get(&expr.op.kind).expect("invalid unary operator");
+        let name = UNARY_OP_NAMES
+            .get(&expr.op.kind)
+            .expect("invalid unary operator");
         let constant_id = self.insert_constant(Str::create(name))?;
         self.emit(expr_line_number(expr), Op::GetAttr(constant_id));
         self.emit(expr_line_number(expr), Op::Call(0));