talc/talc-lang/src/vm.rs

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

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

use anyhow::{Result, anyhow, bail};

struct CallFrame {
	locals: Vec<Value>,
	func: Rc<Function>,
	ip: usize,
}

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

	pub fn run(&mut self, func: Rc<Function>) -> Result<Value> {
		use Instruction as I;

		assert!(func.arity == 0, "root function must not take arguments");
		let mut frame = CallFrame {
			func: func.clone(),
			locals: Vec::with_capacity(16),
			ip: 0,
		};
		frame.locals.push(Value::Function(func));

		loop {
			let instr = frame.func.chunk.instrs[frame.ip];
			self.debug_instr(&frame, instr);
			frame.ip += 1;
			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 = self.globals.get(&sym)
						.ok_or_else(|| anyhow!("global not defined"))?.clone();
					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().try_into()?;
						table.insert(k, 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") };
					for _ in 0..n {
						// can't panic:
						let v = ext.pop().unwrap();
						let k = ext.pop().unwrap().try_into()?;
						table.borrow_mut().insert(k, v);
					}
					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 v = self.pop().to_iter_function()?;
					self.push(v);
				},
				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::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();
						let args = self.pop_n(n);
						let func = self.pop();
						if nf.arity != n {
							 bail!("function call with wrong argument count");
						}
						let res = (nf.f)(func, args)?;
						self.stack.push(res);
					} else if let Value::Function(func) = func {
						if func.arity != n {
							bail!("function call with wrong argument count");
						}
						if self.call_stack.len() + 1 >= self.stack_max {
							bail!("call stack overflow")
						}
						self.call_stack.push(frame);
						let func = func.clone();
						let args = self.pop_n(n + 1);
						frame = CallFrame {
							func,
							ip: 0,
							locals: args,
						};
					} else {
						bail!("attempt to call non-function {}", func);
					}
				},
				I::Return if self.call_stack.is_empty() => {
					let v = self.pop();
					self.stack.clear();
					return Ok(v);
				},
				I::Return => {
					// can't panic: checked that call_stack wasn't empty
					frame = self.call_stack.pop().unwrap();
				},
			}
		}
	}

}