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 <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-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<Self>) -> Box<dyn Any> { self }
+    fn as_any_ref(&self) -> &dyn Any { self }
+}
+
 #[derive(Debug)]
 pub struct SetStmt {
     pub expr: ExprP,

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("<temporary expr value>".to_string())?
+                    .clone();
+                self.compile_expr(&stmt.expr)?;
+                let temp_index = self
+                    .insert_local("<temporary index value>".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

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

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<MapTableWrapper>,
 }
 
 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::<Map>().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::<Map>(),
             other_ptr.borrow().as_any().downcast_ref::<Map>(),
         ) {
-            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__")

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

src/parser.rs

@@ -587,7 +587,13 @@ impl Parser {
             let expr = self.expr()?;
             let stmt: StmtP;
 
-            if expr.as_any_ref().downcast_ref::<GetExpr>().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::<GetExpr>().is_some();
+            let is_index_expr = expr.as_any_ref().downcast_ref::<IndexExpr>().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::<GetExpr>().unwrap();
                 let rhs = self.expr()?;
-                // unpack the GetExpr and turn it into a SetExpr instead
+
-                stmt = Box::new(SetStmt {
+                // unpack the GetExpr or IndexExpr and turn it into a SetExpr instead
-                    expr: expr.expr,
+                if is_get_expr {
-                    name: expr.name,
+                    let expr = expr.as_any().downcast::<GetExpr>().unwrap();
-                    op,
+                    stmt = Box::new(SetStmt {
-                    rhs,
+                        expr: expr.expr,
-                });
+                        name: expr.name,
+                        op,
+                        rhs,
+                    });
+                } else if is_index_expr {
+                    let expr = expr.as_any().downcast::<IndexExpr>().unwrap();
+                    stmt = Box::new(IndexAssignStmt {
+                        expr: expr.expr,
+                        index: expr.index,
+                        op,
+                        rhs,
+                    });
+                } else {
+                    unreachable!()
+                }
             } else {
                 stmt = Box::new(ExprStmt { expr });
             }

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)

tests/map.npp.expect

@@ -2,3 +2,8 @@ to_str
 [:]
 ['a': 1]
 ['b': 4]
+__index_assign__
+['a': 1]
+['a': 2]
+['a': 3]
+

tools/genast.py

@@ -256,6 +256,9 @@ 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(
+        "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<StmtP>", "rbrace: Token"])
     .add_struct("Return", ["return_kw: Token", "expr: Option<ExprP>"])