From 283eaa1ebee9ffa28b941fb974708d5345ed2731 Mon Sep 17 00:00:00 2001 From: Alek Ratzloff Date: Fri, 18 Oct 2024 22:03:10 -0700 Subject: [PATCH] Add index assignment and augmented assignment This allows for syntax like `foo['a'] = 1` and more complex assignments like `foo.bar()[a() + b()] += 1` Signed-off-by: Alek Ratzloff --- src/ast.rs | 20 ++++++++- src/compiler.rs | 76 ++++++++++++++++++++++++++++++++- src/compiler/visitors.rs | 20 +++++++++ src/obj/map.rs | 92 +++++++++++++++++++++++++--------------- src/obj/ty.rs | 1 + src/parser.rs | 37 ++++++++++++---- tests/map.npp | 12 ++++++ tests/map.npp.expect | 5 +++ tools/genast.py | 3 ++ 9 files changed, 221 insertions(+), 45 deletions(-) diff --git a/src/ast.rs b/src/ast.rs index 73ddbd9..b48f4c2 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-14 20:16:29 +// This file was created at: 2024-10-16 10:33:48 #![allow(dead_code)] use std::fmt::Debug; use std::any::Any; @@ -174,6 +174,7 @@ pub trait StmtVisitor { fn visit_import_stmt(&mut self, stmt: &ImportStmt) -> Result<(), Error>; fn visit_expr_stmt(&mut self, stmt: &ExprStmt) -> Result<(), Error>; fn visit_assign_stmt(&mut self, stmt: &AssignStmt) -> Result<(), Error>; + fn visit_index_assign_stmt(&mut self, stmt: &IndexAssignStmt) -> Result<(), Error>; fn visit_set_stmt(&mut self, stmt: &SetStmt) -> Result<(), Error>; fn visit_block_stmt(&mut self, stmt: &BlockStmt) -> Result<(), Error>; fn visit_return_stmt(&mut self, stmt: &ReturnStmt) -> Result<(), Error>; @@ -234,6 +235,23 @@ impl Stmt for AssignStmt { fn as_any_ref(&self) -> &dyn Any { self } } +#[derive(Debug)] +pub struct IndexAssignStmt { + pub expr: ExprP, + pub index: ExprP, + pub op: Token, + pub rhs: ExprP, +} + +impl Stmt for IndexAssignStmt { + fn accept(&self, visitor: &mut dyn StmtVisitor) -> Result<(), Error> { + visitor.visit_index_assign_stmt(self) + } + + fn as_any(self: Box) -> Box { self } + fn as_any_ref(&self) -> &dyn Any { self } +} + #[derive(Debug)] pub struct SetStmt { pub expr: ExprP, diff --git a/src/compiler.rs b/src/compiler.rs index 3e57712..3abfcef 100644 --- a/src/compiler.rs +++ b/src/compiler.rs @@ -639,6 +639,81 @@ impl StmtVisitor for Compiler<'_> { Ok(()) } + fn visit_index_assign_stmt(&mut self, stmt: &IndexAssignStmt) -> Result<()> { + let line = stmt_line_number(stmt); + let index_assign_constant = self.insert_constant(Str::create("__index_assign__"))?; + + match stmt.op.kind { + TokenKind::Eq => { + self.compile_expr(&stmt.expr)?; + self.emit(line, Op::GetAttr(index_assign_constant)); + self.compile_expr(&stmt.index)?; + self.compile_expr(&stmt.rhs)?; + self.emit(line, Op::Call(2)); + self.emit(line, Op::Pop); + } + TokenKind::PlusEq | TokenKind::MinusEq | TokenKind::StarEq | TokenKind::SlashEq => { + // Augmented index assigns are a bit complex. + // Most cases are going to be something like: + // + // foo[n] += 1 + // + // In the worst case, however, we might have: + // + // foo.bar()[index.decrement()] += 1 + // + // which should desugar to: + // + // temp_base = foo.bar() + // temp_index = index.decrement() + // temp_base[temp_index] = temp_base[temp_index] + 1 + // + // We only want to compile the base expression and the index expression once, in + // case they have side effects. So, we need to store these values in temporary + // variables. + let base_line = expr_line_number(&*stmt.expr); + let index_line = expr_line_number(&*stmt.index); + + self.begin_scope(ScopeKind::Local); + let temp_base = self + .insert_local("".to_string())? + .clone(); + self.compile_expr(&stmt.expr)?; + let temp_index = self + .insert_local("".to_string())? + .clone(); + self.compile_expr(&stmt.index)?; + + self.emit(base_line, Op::GetLocal(temp_base.index)); + self.emit(base_line, Op::GetAttr(index_assign_constant)); + self.emit(index_line, Op::GetLocal(temp_index.index)); + + let index_constant = self.insert_constant(Str::create("__index__"))?.clone(); + self.emit(base_line, Op::GetLocal(temp_base.index)); + self.emit(base_line, Op::GetAttr(index_constant)); + self.emit(index_line, Op::GetLocal(temp_index.index)); + self.emit(base_line, Op::Call(1)); // call __index__ function + + 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(index_line, Op::GetAttr(op_constant)); + self.compile_expr(&stmt.rhs)?; + self.emit(line, Op::Call(1)); // call binary op function + self.emit(line, Op::Call(2)); // index_assign function + + // pop __index_assign__ result (should be nil) + self.emit(line, Op::Pop); + + self.end_scope(line); + } + _ => unreachable!(), + } + + Ok(()) + } + fn visit_set_stmt(&mut self, stmt: &SetStmt) -> Result<()> { let name = self.insert_constant(Str::create(&stmt.name.text))?; let line = stmt_line_number(stmt); @@ -653,7 +728,6 @@ impl StmtVisitor for Compiler<'_> { } // 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 diff --git a/src/compiler/visitors.rs b/src/compiler/visitors.rs index 83cd73f..8fce27b 100644 --- a/src/compiler/visitors.rs +++ b/src/compiler/visitors.rs @@ -47,6 +47,12 @@ impl StmtVisitor for LineNumber { Ok(()) } + fn visit_index_assign_stmt(&mut self, stmt: &IndexAssignStmt) -> Result<()> { + stmt.expr.accept(self).unwrap(); + stmt.rhs.accept(self).unwrap(); + Ok(()) + } + fn visit_set_stmt(&mut self, stmt: &SetStmt) -> Result<()> { stmt.expr.accept(self).unwrap(); stmt.rhs.accept(self).unwrap(); @@ -196,6 +202,13 @@ impl StmtVisitor for LocalAssignCollector { Ok(()) } + fn visit_index_assign_stmt(&mut self, stmt: &IndexAssignStmt) -> Result<()> { + stmt.expr.accept(self)?; + stmt.index.accept(self)?; + stmt.rhs.accept(self)?; + Ok(()) + } + fn visit_set_stmt(&mut self, stmt: &SetStmt) -> Result<()> { stmt.expr.accept(self)?; stmt.rhs.accept(self)?; @@ -340,6 +353,13 @@ impl StmtVisitor for LocalNameCollector { Ok(()) } + fn visit_index_assign_stmt(&mut self, stmt: &IndexAssignStmt) -> Result<()> { + stmt.expr.accept(self)?; + stmt.index.accept(self)?; + stmt.rhs.accept(self)?; + Ok(()) + } + fn visit_set_stmt(&mut self, stmt: &SetStmt) -> Result<()> { stmt.expr.accept(self)?; stmt.rhs.accept(self)?; diff --git a/src/obj/map.rs b/src/obj/map.rs index 4bb6dc8..a236f10 100644 --- a/src/obj/map.rs +++ b/src/obj/map.rs @@ -1,7 +1,8 @@ use std::fmt::{self, Debug, Display}; use std::hash::BuildHasher; +use std::ops::{Deref, DerefMut}; -use gc::{custom_trace, Finalize, Trace}; +use gc::{custom_trace, Finalize, GcCell, Trace}; use hashbrown::{hash_table::Entry, DefaultHashBuilder, HashTable}; use crate::obj::{macros::*, prelude::*}; @@ -34,15 +35,54 @@ fn map_hash_index(vm: &mut Vm, this: &ObjP, index: &ObjP) -> u64 { }) } +//////////////////////////////////////////////////////////////////////////////// +// Hash table wrapper +//////////////////////////////////////////////////////////////////////////////// + +pub type MapTable = HashTable<(ObjP, ObjP)>; + +#[derive(Finalize, Default)] +struct MapTableWrapper(MapTable); + +impl Deref for MapTableWrapper { + type Target = HashTable<(ObjP, ObjP)>; + fn deref(&self) -> &MapTable { + &self.0 + } +} + +impl DerefMut for MapTableWrapper { + fn deref_mut(&mut self) -> &mut MapTable { + &mut self.0 + } +} + +unsafe impl Trace for MapTableWrapper { + custom_trace! { + this, + { + for (k, v) in this.0.iter() { + mark(k); + mark(v); + } + } + } +} + //////////////////////////////////////////////////////////////////////////////// // Map //////////////////////////////////////////////////////////////////////////////// -#[derive(Finalize, Default)] +#[derive(Trace, Finalize, Default)] pub struct Map { base: Obj, + #[unsafe_ignore_trace] hash_builder: DefaultHashBuilder, - table: HashTable<(ObjP, ObjP)>, + // In Map::insert, we will be needing to borrow `table` mutably, but also we need to be able to + // use the `Map` mutably. Normally, we'd use RefCell, but that causes double-root errors. We + // also need to be able to implement Finalize and Trace for the table, so we use a wrapper here + // as well. + table: GcCell, } impl Map { @@ -50,14 +90,6 @@ impl Map { Default::default() } - pub fn table(&self) -> &HashTable<(ObjP, ObjP)> { - &self.table - } - - pub fn table_mut(&mut self) -> &mut HashTable<(ObjP, ObjP)> { - &mut self.table - } - pub fn make_hasher(&self) -> DefaultHasher { self.hash_builder.build_hasher() } @@ -67,7 +99,8 @@ impl Map { let hash = map_hash_index(vm, &this, &index); with_obj_downcast(this.clone(), |map: &Map| { - map.table() + map.table + .borrow() .find(hash, |(key, _)| { let method = index .borrow() @@ -85,18 +118,6 @@ impl Map { impl_create!(); } -unsafe impl Trace for Map { - custom_trace! { - this, - { - for (k, v) in this.table.iter() { - mark(k); - mark(v); - } - } - } -} - impl Display for Map { fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { Debug::fmt(self, fmt) @@ -106,7 +127,7 @@ impl Display for Map { impl Debug for Map { fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { let mut debug_map = fmt.debug_map(); - for (k, v) in self.table.iter() { + for (k, v) in self.table.borrow().iter() { debug_map.entry(&k.borrow(), &v.borrow()); } debug_map.finish() @@ -134,12 +155,13 @@ impl Map { let this = vm.frame_stack()[0].clone(); let this_borrowed = this.borrow(); let map_obj = this_borrowed.as_any().downcast_ref::().unwrap(); - if map_obj.table().len() == 0 { + if map_obj.table.borrow().len() == 0 { return Str::create("[:]"); } let mut repr = "[".to_string(); - let mut iter = map_obj.table().iter(); + let table_borrowed = map_obj.table.borrow(); + let mut iter = table_borrowed.iter(); // first item { @@ -180,7 +202,7 @@ impl Map { // to_list returns the keys of this value let this = vm.frame_stack()[0].clone(); let list: Vec<_> = with_obj_downcast(this, |map: &Map| { - map.table().iter().map(|(k, _)| k.clone()).collect() + map.table.borrow().iter().map(|(k, _)| k.clone()).collect() }); List::create(list) } @@ -193,7 +215,7 @@ impl Map { this.borrow().as_any().downcast_ref::(), other_ptr.borrow().as_any().downcast_ref::(), ) { - if this.table().len() != other.table().len() { + if this.table.borrow().len() != other.table.borrow().len() { return Bool::create(false); } @@ -201,7 +223,7 @@ impl Map { let contains_method = other .get_vtable_attr(other_ptr.clone(), "contains") .unwrap(); - this.table().iter().all(|(key, value)| { + this.table.borrow().iter().all(|(key, value)| { let contains = vm.call(contains_method.clone(), &[key.clone()]); if contains.borrow().is_truthy() { let other_value = vm.call(index_method.clone(), &[key.clone()]); @@ -266,7 +288,7 @@ impl Map { pub(crate) fn len(vm: &mut Vm) -> ObjP { let this = vm.frame_stack()[0].clone(); - let len = with_obj_downcast(this.clone(), |map: &Map| map.table().len()); + let len = with_obj_downcast(this.clone(), |map: &Map| map.table.borrow().len()); Int::create(len as i64) } @@ -277,12 +299,13 @@ impl Map { let hash = map_hash_index(vm, &this, &index); - let old_value = with_obj_downcast_mut(this.clone(), |map: &mut Map| { + let old_value = with_obj_downcast(this.clone(), |map: &Map| { // NOTE: we have to borrow `vm` mutably twice in both of these functions. // This is safe because these closures are not called simultaneously. let vm = vm as *mut Vm; // get the entry - let entry = map.table_mut().entry( + let mut table_borrowed = map.table.borrow_mut(); + let entry = table_borrowed.entry( hash, // eq |(key, _)| { @@ -330,7 +353,8 @@ impl Map { let hash = map_hash_index(vm, &this, &index); let removed = with_obj_downcast_mut(this.clone(), |map: &mut Map| { - let result = map.table_mut().find_entry(hash, |(key, _)| { + let mut table_borrowed = map.table.borrow_mut(); + let result = table_borrowed.find_entry(hash, |(key, _)| { let method = index .borrow() .get_vtable_attr(index.clone(), "__eq__") diff --git a/src/obj/ty.rs b/src/obj/ty.rs index a8d7bc0..9944e28 100644 --- a/src/obj/ty.rs +++ b/src/obj/ty.rs @@ -227,6 +227,7 @@ pub fn init_types() { // Operators __eq__ => BuiltinFunction::create("__eq__", Map::eq, 2), __index__ => BuiltinFunction::create("__index__", Map::index, 2), + __index_assign__ => BuiltinFunction::create("__index_assign__", Map::insert, 3), // Methods contains => BuiltinFunction::create("contains", Map::contains, 2), diff --git a/src/parser.rs b/src/parser.rs index e02b1be..f9723c7 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -587,7 +587,13 @@ impl Parser { let expr = self.expr()?; let stmt: StmtP; - if expr.as_any_ref().downcast_ref::().is_some() + // TODO Parser::stmt_wrapped - complex assign statements could probably be cleaner and + // probably need their own function at this point. + + let is_get_expr = expr.as_any_ref().downcast_ref::().is_some(); + let is_index_expr = expr.as_any_ref().downcast_ref::().is_some(); + + if (is_get_expr || is_index_expr) && mat!( self, TokenKind::Eq, @@ -598,15 +604,28 @@ impl Parser { ) { let op = self.prev.clone().unwrap(); - let expr = expr.as_any().downcast::().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, - }); + + // unpack the GetExpr or IndexExpr and turn it into a SetExpr instead + if is_get_expr { + let expr = expr.as_any().downcast::().unwrap(); + stmt = Box::new(SetStmt { + expr: expr.expr, + name: expr.name, + op, + rhs, + }); + } else if is_index_expr { + let expr = expr.as_any().downcast::().unwrap(); + stmt = Box::new(IndexAssignStmt { + expr: expr.expr, + index: expr.index, + op, + rhs, + }); + } else { + unreachable!() + } } else { stmt = Box::new(ExprStmt { expr }); } diff --git a/tests/map.npp b/tests/map.npp index 6af0b46..9372f54 100644 --- a/tests/map.npp +++ b/tests/map.npp @@ -6,3 +6,15 @@ println("to_str") println(a) println(['a': 1]) println(['b': 2 + 2]) + +println("__index_assign__") +a['a'] = 1 +println(a) +a['a'] += 1 +println(a) + +foo = () { return 'a' } +bar = () { return '' } + +a[foo() + bar()] += 1 +println(a) diff --git a/tests/map.npp.expect b/tests/map.npp.expect index 1aac46a..3c49cf9 100644 --- a/tests/map.npp.expect +++ b/tests/map.npp.expect @@ -2,3 +2,8 @@ to_str [:] ['a': 1] ['b': 4] +__index_assign__ +['a': 1] +['a': 2] +['a': 3] + diff --git a/tools/genast.py b/tools/genast.py index 32ec14c..547e7ae 100755 --- a/tools/genast.py +++ b/tools/genast.py @@ -256,6 +256,9 @@ GENERATE = [ .add_struct("Import", ["import_kw: Token", "what: Vec", "module: Token"]) .add_struct("Expr", ["expr: ExprP"]) .add_struct("Assign", ["lhs: Token", "op: Token", "rhs: ExprP"]) + .add_struct( + "IndexAssign", ["expr: ExprP", "index: ExprP", "op: Token", "rhs: ExprP"] + ) .add_struct("Set", ["expr: ExprP", "name: Token", "op: Token", "rhs: ExprP"]) .add_struct("Block", ["lbrace: Token", "stmts: Vec", "rbrace: Token"]) .add_struct("Return", ["return_kw: Token", "expr: Option"])