talc/talc-lang/src/vm.rs

use std::{rc::Rc, cmp::Ordering, cell::RefCell, collections::HashMap};

use crate::{chunk::Instruction, value::{Value, Function, Result, throw, Exception}, ast::{BinaryOp, UnaryOp}, symbol::Symbol};

struct TryFrame { idx: usize, stack_len: usize }

#[derive(Default)]
struct CallFrame {
	func: Rc<Function>,
	locals: Vec<Value>,
	try_frames: Vec<TryFrame>,
	ip: usize,
	root: bool,
}


impl CallFrame {
	fn new(func: Rc<Function>, locals: Vec<Value>) -> Self {
		Self {
			func,
			locals,
			..Self::default()
		}
	}

	fn new_root(func: Rc<Function>) -> Self {
		Self {
			func: func.clone(),
			locals: vec![Value::Function(func)],
			root: true,
			..Self::default()
		}
	}
}

pub struct Vm {
	stack: Vec<Value>,
	call_stack: Vec<CallFrame>,
	stack_max: usize,
	globals: HashMap<Symbol, Value>,
}

pub fn binary_op(o: BinaryOp, a: Value, b: Value) -> Result<Value> {
	match o {
		BinaryOp::Add => a + b,
		BinaryOp::Sub => a - b,
		BinaryOp::Mul => a * b,
		BinaryOp::Div => a / b,
		BinaryOp::Mod => a.modulo(b),
		BinaryOp::IntDiv => a.int_div(b),
		BinaryOp::Pow => a.pow(b),
		BinaryOp::Eq => Ok(Value::Bool(a == b)),
		BinaryOp::Ne => Ok(Value::Bool(a != b)),
		BinaryOp::Gt => a.val_cmp(&b).map(|o| Value::Bool(o == Ordering::Greater)),
		BinaryOp::Ge => a.val_cmp(&b).map(|o| Value::Bool(o != Ordering::Less)),
		BinaryOp::Lt => a.val_cmp(&b).map(|o| Value::Bool(o == Ordering::Less)),
		BinaryOp::Le => a.val_cmp(&b).map(|o| Value::Bool(o != Ordering::Greater)),
		BinaryOp::Range => a.range(&b, false),
		BinaryOp::RangeIncl => a.range(&b, true),
		BinaryOp::Concat => a.concat(&b),
		BinaryOp::Append => todo!("append"),
	}
}

pub fn unary_op(o: UnaryOp, a: Value) -> Result<Value> {
	match o {
		UnaryOp::Neg => -a,
		UnaryOp::Not => Ok(Value::Bool(!a.truthy())),
		UnaryOp::RangeEndless => a.range_endless(),
	}
}

impl Vm {
	pub fn new(stack_max: usize) -> Self {
		Self {
			stack: Vec::with_capacity(16),
			call_stack: Vec::with_capacity(16),
			globals: HashMap::with_capacity(16),
			stack_max,
		}
	}

	pub fn set_global(&mut self, name: Symbol, val: Value) {
		self.globals.insert(name, val);
	}

	pub fn get_global(&self, name: Symbol) -> Option<&Value> {
		self.globals.get(&name)
	}

	#[inline]
	fn push(&mut self, v: Value) {
		self.stack.push(v);
	}

	#[inline]
	fn pop(&mut self) -> Value {
		self.stack.pop().expect("temporary stack underflow")
	}

	#[inline]
	fn pop_n(&mut self, n: usize) -> Vec<Value> {
		let res = self.stack.split_off(self.stack.len() - n);
		assert!(res.len() == n, "temporary stack underflow");
		res
	}

	fn debug_instr(&self, frame: &CallFrame, instr: Instruction) {
		let framecode = (Rc::as_ptr(&frame.func) as usize >> 4) & 0xffff;
		println!("({:04x}) {:04}: {instr}", framecode, frame.ip);
	}

	fn run_instr(&mut self, frame: &mut CallFrame, instr: Instruction) -> Result<Option<Value>> {
		use Instruction as I;
		match instr {
			I::Nop => (),
			I::LoadLocal(n)
				=> self.push(frame.locals[usize::from(n)].clone()),
			I::StoreLocal(n)
				=> frame.locals[usize::from(n)] = self.pop(),
			I::NewLocal
				=> frame.locals.push(self.pop()),
			I::DropLocal(n)
				=> frame.locals.truncate(frame.locals.len() - usize::from(n)),
			I::LoadGlobal(s) => {
				let sym = unsafe { s.to_symbol_unchecked() };
				let v = match self.globals.get(&sym) {
					Some(v) => v.clone(),
					None => throw!(name_error, "undefined global {}", sym.name()),
				};
				self.push(v);
			},
			I::StoreGlobal(s) => {
				let sym = unsafe { s.to_symbol_unchecked() };
				let v = self.pop();
				self.globals.insert(sym, v);
			},
			I::Const(n)
				=> self.push(frame.func.chunk.consts[usize::from(n)].clone()),
			I::Nil => self.push(Value::Nil),
			I::Bool(b) => self.push(Value::Bool(b)),
			I::Symbol(n) => {
				let sym = unsafe { Symbol::from_id_unchecked(u32::from(n)) };
				self.push(Value::Symbol(sym));
			},
			I::Int(n) => self.push(Value::Int(i64::from(n))),
			I::Dup => self.push(self.stack[self.stack.len() - 1].clone()),
			I::DupTwo => {
				self.push(self.stack[self.stack.len() - 2].clone());
				self.push(self.stack[self.stack.len() - 2].clone());
			},
			I::Drop(n) => for _ in 0..u32::from(n) { self.pop(); },
			I::Swap => {
				let len = self.stack.len();
				self.stack.swap(len - 1, len - 2);
			},
			I::BinaryOp(op) => {
				let b = self.pop();
				let a = self.pop();
				self.push(binary_op(op, a, b)?);
			},
			I::UnaryOp(op) => {
				let a = self.pop();
				self.push(unary_op(op, a)?);
			},
			I::NewList(n) => {
				let list = self.pop_n(n as usize);
				self.push(Value::List(Rc::new(RefCell::new(list))));
			},
			I::GrowList(n) => {
				let ext = self.pop_n(n as usize);
				let list = self.pop();
				let Value::List(list) = list else { panic!("not a list") };
				list.borrow_mut().extend(ext);
				self.push(Value::List(list));
			},
			I::NewTable(n) => {
				let mut table = HashMap::new();
				for _ in 0..n {
					let v = self.pop();
					let k = self.pop();
					table.insert(k.try_into()?, v);
				}
				self.push(Value::Table(Rc::new(RefCell::new(table))));
			},
			I::GrowTable(n) => {
				let mut ext = self.pop_n(2 * n as usize);
				let table = self.pop();
				let Value::Table(table) = table else { panic!("not a table") };
				let mut table_ref = table.borrow_mut();
				for _ in 0..n {
					// can't panic: pop_n checked that ext would have len 2*n
					let v = ext.pop().unwrap();
					let k = ext.pop().unwrap();
					table_ref.insert(k.try_into()?, v);
				}
				drop(table_ref);
				self.push(Value::Table(table));
			},
			I::Index => {
				let idx = self.pop();
				let ct = self.pop();
				self.push(ct.index(idx)?);
			},
			I::StoreIndex => {
				let v = self.pop();
				let idx = self.pop();
				let ct = self.pop();
				ct.store_index(idx, v.clone())?;
				self.push(v);
			},
			I::Jump(n) => frame.ip = usize::from(n),
			I::JumpTrue(n)  => if self.pop().truthy() { frame.ip = usize::from(n) },
			I::JumpFalse(n) => if !self.pop().truthy() { frame.ip = usize::from(n) },
			I::IterBegin => {
				let iter = self.pop().to_iter_function()?;
				self.push(iter);
			},
			I::IterTest(n) => {
				let v = &self.stack[self.stack.len() - 1];
				if v == &Value::Nil {
					self.pop();
					self.pop();
					frame.ip = usize::from(n);
				}
			},
			I::BeginTry(t) => {
				let tryframe = TryFrame {
					idx: usize::from(t),
					stack_len: self.stack.len()
				};
				frame.try_frames.push(tryframe);
			},
			I::EndTry => {
				frame.try_frames.pop().expect("no try to pop");
			},
			I::Call(n) => {
				let n = usize::from(n);
				let func = &self.stack[self.stack.len() - n - 1];
				if let Value::NativeFunc(nf) = func {
					let nf = nf.clone();
					if nf.arity != n {
						throw!(type_error, "function call with wrong argument count");
					}

					let args = self.pop_n(n);
					let func = self.pop();

					self.call_stack.push(std::mem::take(frame));
					let res = (nf.f)(func, args)?;
					*frame = self.call_stack.pop().expect("no frame left on stack");

					self.stack.push(res);
				} else if let Value::Function(func) = func {
					if func.arity != n {
						throw!(type_error, "function call with wrong argument count");
					}

					if self.call_stack.len() + 1 >= self.stack_max {
						throw!(call_stack_overflow, "call stack overflow")
					}
					self.call_stack.push(std::mem::take(frame));

					let func = func.clone();
					let args = self.pop_n(n + 1);
					*frame = CallFrame::new(func, args);
				} else {
					throw!(type_error, "attempt to call non-function {func}");
				}
			},
			I::Return if frame.root => {
				let v = self.pop();
				self.stack.clear();
				return Ok(Some(v));
			},
			I::Return => {
				*frame = self.call_stack.pop().expect("no root frame");
			},
		}

		Ok(None)
	}

	fn handle_exception(&mut self, frame: &mut CallFrame, exc: Exception) -> Result<()> {
		loop {
			while let Some(try_frame) = frame.try_frames.pop() {
				let table = &frame.func.chunk.try_tables[try_frame.idx];
				for catch in &table.catches {
						if catch.types.is_none() || catch.types.as_ref().unwrap().contains(&exc.ty) {
						frame.ip = catch.addr;
						frame.locals.truncate(table.local_count);
						self.stack.truncate(try_frame.stack_len);
						self.stack.push(Value::Table(exc.to_table()));
						return Ok(())
					}
				}
			}
			if frame.root {
				return Err(exc)
			}
			*frame = self.call_stack.pop().expect("no root frame");
		}
	}

	pub fn run(&mut self, func: Rc<Function>) -> Result<Value> {
		assert!(func.arity == 0, "root function must not take arguments");
		let init_stack_len = self.stack.len();
		let mut frame = CallFrame::new_root(func.clone());

		loop {
			let instr = frame.func.chunk.instrs[frame.ip];
			self.debug_instr(&frame, instr);
			frame.ip += 1;
			match self.run_instr(&mut frame, instr) {
				Ok(None) => (),
				Ok(Some(v)) => {
					self.stack.truncate(init_stack_len);
					return Ok(v)
				}
				Err(e) => {
					if let Err(e) = self.handle_exception(&mut frame, e) {
						self.stack.truncate(init_stack_len);
						return Err(e)
					}
				}
			}
		}
	}

}