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>
This commit is contained in:
2024-10-18 22:03:10 -07:00
parent 43891bd2a3
commit 283eaa1ebe
9 changed files with 221 additions and 45 deletions

View File

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

View File

@@ -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

View File

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

View File

@@ -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__")

View File

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

View File

@@ -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
// unpack the GetExpr or IndexExpr and turn it into a SetExpr instead
if is_get_expr {
let expr = expr.as_any().downcast::<GetExpr>().unwrap();
stmt = Box::new(SetStmt {
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 });
}

View File

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

View File

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

View File

@@ -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>"])