Implement augmented assignment operators

Add support for +=, -=, *=, and /= operators. This is basically just
syntactic sugar, but it's still nice to have

    a += 1

compiles to the equivalent of

    a = a + 1

with all the same implications of scoping rules.

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-03 12:06:28
+// This file was created at: 2024-10-07 09:59:23
 #![allow(dead_code)]
 use std::fmt::Debug;
 use std::any::Any;
@@ -204,6 +204,7 @@ impl Stmt for ExprStmt {
 #[derive(Debug)]
 pub struct AssignStmt {
     pub lhs: Token,
+    pub op: Token,
     pub rhs: ExprP,
 }
 

src/compiler.rs

@@ -417,6 +417,29 @@ impl<'c> Compiler<'c> {
         Ok(())
     }
 
+    /// Emit a name lookup. This will either emit `Op::GetLocal` or `Op::GetGlobal`, depending on
+    /// the context. If the local or global is not found in the current scope, we error out
+    /// instead.
+    fn emit_lookup(&mut self, line: LineRange, name: &str) -> Result<()> {
+        // check if there's a local with this name, otherwise check globals
+        if let Some(local) = self.get_local(name) {
+            self.emit(line, Op::GetLocal(local.index));
+        } else {
+            let global = self.get_global(name).ok_or_else(|| CompileError::Error {
+                line: Some(line),
+                source_path: self.path.display().to_string(),
+                message: if self.is_global_scope() {
+                    format!("unknown global {}", name)
+                } else {
+                    format!("unknown local {}", name)
+                },
+            })?;
+            self.emit(line, Op::GetGlobal(global));
+        }
+
+        Ok(())
+    }
+
     fn search_dir(&self) -> &Path {
         self.path.parent().unwrap()
     }
@@ -524,23 +547,52 @@ impl StmtVisitor for Compiler<'_> {
     }
 
     fn visit_assign_stmt(&mut self, stmt: &AssignStmt) -> Result<()> {
-        // compile RHS
-        self.compile_expr(&stmt.rhs)?;
-
         let name = &stmt.lhs.text;
 
-        // If the last value that was assigned to is a function, set its name here
+        let line = stmt_line_number(stmt);
-        // TODO - maybe this would be smarter to set up in the AST. I'm 99% sure that the last
-        // object created, if it were a function object, will be what we're assigning it to, but I
-        // want to be 100% sure instead of 99%.
-        if let Some(obj) = self.constants.last() {
-            if let Some(fun) = obj.borrow_mut().as_any_mut().downcast_mut::<UserFunction>() {
-                fun.set_name(Rc::new(name.to_string()));
-            }
-        }
 
-        // compile LHS
+        match stmt.op.kind {
-        self.emit_assign(stmt_line_number(stmt), name)?;
+            // normal assignment
+            TokenKind::Eq => {
+                // compile RHS
+                self.compile_expr(&stmt.rhs)?;
+
+                // If the last value that was assigned to is a function, set its name here
+                // TODO - maybe this would be smarter to set up in the AST. I'm 99% sure that the last
+                // object created, if it were a function object, will be what we're assigning it to, but I
+                // want to be 100% sure instead of 99%.
+                if let Some(obj) = self.constants.last() {
+                    if let Some(fun) = obj.borrow_mut().as_any_mut().downcast_mut::<UserFunction>()
+                    {
+                        fun.set_name(Rc::new(name.to_string()));
+                    }
+                }
+
+                // compile LHS
+                self.emit_assign(line, name)?;
+            }
+            // 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
+                    .get(&stmt.op.kind)
+                    .expect("invalid augmented assignment operator");
+                let op_constant = self.insert_constant(Str::create(op_name))?;
+                self.emit(line, Op::GetAttr(op_constant));
+                self.compile_expr(&stmt.rhs)?;
+                self.emit(line, Op::Call(1));
+                self.emit_assign(line, name)?;
+            }
+            _ => unreachable!(),
+        }
 
         Ok(())
     }
@@ -750,21 +802,8 @@ impl ExprVisitor for Compiler<'_> {
         match expr.token.kind {
             TokenKind::Name => {
                 let name = &expr.token.text;
-                // check if there's a local with this name, otherwise check globals
+                let line = expr_line_number(expr);
-                if let Some(local) = self.get_local(name) {
+                self.emit_lookup(line, name)?;
-                    self.emit(expr_line_number(expr), Op::GetLocal(local.index));
-                } else {
-                    let global = self.get_global(name).ok_or_else(|| CompileError::Error {
-                        line: Some(expr_line_number(expr)),
-                        source_path: self.path.display().to_string(),
-                        message: if self.is_global_scope() {
-                            format!("unknown global {}", name)
-                        } else {
-                            format!("unknown local {}", name)
-                        },
-                    })?;
-                    self.emit(expr_line_number(expr), Op::GetGlobal(global));
-                }
             }
             TokenKind::Number => {
                 let obj = if expr.token.text.contains('.') {

src/parser.rs

@@ -178,17 +178,31 @@ impl Lexer {
         } else if STRING_START_CHARS.contains(self.current()) {
             return self.string();
         } else if self.mat('+') {
-            return Ok(self.make_token(TokenKind::Plus));
+            if self.mat('=') {
+                return Ok(self.make_token(TokenKind::PlusEq));
+            } else {
+                return Ok(self.make_token(TokenKind::Plus));
+            }
         } else if self.mat('-') {
             if self.mat('>') {
                 return Ok(self.make_token(TokenKind::Arrow));
+            } else if self.mat('=') {
+                return Ok(self.make_token(TokenKind::MinusEq));
             } else {
                 return Ok(self.make_token(TokenKind::Minus));
             }
         } else if self.mat('*') {
-            return Ok(self.make_token(TokenKind::Star));
+            if self.mat('=') {
+                return Ok(self.make_token(TokenKind::StarEq));
+            } else {
+                return Ok(self.make_token(TokenKind::Star));
+            }
         } else if self.mat('/') {
-            return Ok(self.make_token(TokenKind::Slash));
+            if self.mat('=') {
+                return Ok(self.make_token(TokenKind::SlashEq));
+            } else {
+                return Ok(self.make_token(TokenKind::Slash));
+            }
         } else if self.mat('&') {
             if self.mat('&') {
                 return Ok(self.make_token(TokenKind::And));
@@ -561,7 +575,13 @@ impl Parser {
                 stmts,
                 rbrace,
             }) as Box<dyn Stmt + 'static>)
-        } else if self.current.kind == TokenKind::Name && self.next.kind == TokenKind::Eq {
+        } else if self.current.kind == TokenKind::Name
+            && (self.next.kind == TokenKind::Eq
+                || self.next.kind == TokenKind::PlusEq
+                || self.next.kind == TokenKind::MinusEq
+                || self.next.kind == TokenKind::StarEq
+                || self.next.kind == TokenKind::SlashEq)
+        {
             self.assign_stmt()
         } else {
             let expr = self.expr()?;
@@ -745,7 +765,15 @@ impl Parser {
         let name = self
             .expect("expect name for assign statement", TokenKind::Name)?
             .clone();
-        self.expect("expect '=' after name", TokenKind::Eq)?;
+        let op = expect!(
+            self,
+            "expected '=' or augmented assign after name",
+            TokenKind::Eq,
+            TokenKind::PlusEq,
+            TokenKind::MinusEq,
+            TokenKind::StarEq,
+            TokenKind::SlashEq
+        )?;
         let expr = self.expr()?;
         if !self.check(TokenKind::RBrace) {
             expect!(
@@ -757,6 +785,7 @@ impl Parser {
         }
         Ok(Box::new(AssignStmt {
             lhs: name,
+            op,
             rhs: expr,
         }))
     }

src/token.rs

@@ -21,6 +21,12 @@ pub enum TokenKind {
     Star,
     Slash,
 
+    // Augmented assignment operators
+    PlusEq,
+    MinusEq,
+    StarEq,
+    SlashEq,
+
     // Unary operators (not already covered)
     Bang,
 

tools/genast.py

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