diff --git a/src/ast.rs b/src/ast.rs index d2dbabc..6ec7767 100644 --- a/src/ast.rs +++ b/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, } diff --git a/src/compiler.rs b/src/compiler.rs index fefcc6d..20ff9b0 100644 --- a/src/compiler.rs +++ b/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 - // 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::() { - fun.set_name(Rc::new(name.to_string())); - } - } + let line = stmt_line_number(stmt); - // compile LHS - self.emit_assign(stmt_line_number(stmt), name)?; + match stmt.op.kind { + // 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::() + { + 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> = 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 - if let Some(local) = self.get_local(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)); - } + let line = expr_line_number(expr); + self.emit_lookup(line, name)?; } TokenKind::Number => { let obj = if expr.token.text.contains('.') { diff --git a/src/parser.rs b/src/parser.rs index 62cd976..a026e03 100644 --- a/src/parser.rs +++ b/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) - } 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, })) } diff --git a/src/token.rs b/src/token.rs index 1c45555..8cb8387 100644 --- a/src/token.rs +++ b/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, diff --git a/tools/genast.py b/tools/genast.py index 45a0990..df0ded9 100755 --- a/tools/genast.py +++ b/tools/genast.py @@ -234,7 +234,7 @@ GENERATE = [ GenerateGroup("Stmt") .add_struct("Import", ["import_kw: Token", "what: Vec", "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", "rbrace: Token"]) .add_struct("Return", ["return_kw: Token", "expr: Option"])