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/ast.rs

@@ -1,5 +1,5 @@
 // This is an auto-generated file. Any changes made to this file may be overwritten.
-// This file was created at: 2024-10-07 09:59:23
+// This file was created at: 2024-10-07 10:41:44
 #![allow(dead_code)]
 use std::fmt::Debug;
 use std::any::Any;
@@ -221,6 +221,7 @@ impl Stmt for AssignStmt {
 pub struct SetStmt {
     pub expr: ExprP,
     pub name: Token,
+    pub op: Token,
     pub rhs: ExprP,
 }
 

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));

src/parser.rs

@@ -587,13 +587,24 @@ impl Parser {
             let expr = self.expr()?;
             let stmt: StmtP;
 
-            if expr.as_any_ref().downcast_ref::<GetExpr>().is_some() && self.mat(TokenKind::Eq)? {
+            if expr.as_any_ref().downcast_ref::<GetExpr>().is_some()
+                && mat!(
+                    self,
+                    TokenKind::Eq,
+                    TokenKind::PlusEq,
+                    TokenKind::MinusEq,
+                    TokenKind::StarEq,
+                    TokenKind::SlashEq
+                )
+            {
+                let op = self.prev.clone().unwrap();
                 let expr = expr.as_any().downcast::<GetExpr>().unwrap();
                 let rhs = self.expr()?;
                 // unpack the GetExpr and turn it into a SetExpr instead
                 stmt = Box::new(SetStmt {
                     expr: expr.expr,
                     name: expr.name,
+                    op,
                     rhs,
                 });
             } else {

tools/genast.py

@@ -235,7 +235,7 @@ GENERATE = [
     .add_struct("Import", ["import_kw: Token", "what: Vec<Token>", "module: Token"])
     .add_struct("Expr", ["expr: ExprP"])
     .add_struct("Assign", ["lhs: Token", "op: Token", "rhs: ExprP"])
-    .add_struct("Set", ["expr: ExprP", "name: Token", "rhs: ExprP"])
+    .add_struct("Set", ["expr: ExprP", "name: Token", "op: Token", "rhs: ExprP"])
     .add_struct("Block", ["lbrace: Token", "stmts: Vec<StmtP>", "rbrace: Token"])
     .add_struct("Return", ["return_kw: Token", "expr: Option<ExprP>"])
     .add_struct(