path: root/matrix-lang/src/vm.rs

use std::{fmt::Display, ops::{IndexMut, Index}, sync::{atomic::{AtomicUsize, Ordering}, Arc}, collections::HashMap};
use crate::prelude::*;

pub struct Stack<T> {
    inner: Vec<T>
}

impl<T> Stack<T> {
    pub fn new() -> Self {
        Self { inner: vec![] }
    }

    pub fn pop(&mut self) -> T {
        self.inner.pop().expect("stack empty")
    }

    pub fn push(&mut self, val: T) {
        self.inner.push(val)
    }

    pub fn split_off(&mut self, len: usize) -> Self {
        let inner = self.inner.split_off(len);
        Self { inner }
    }

    pub fn truncate(&mut self, len: usize) {
        self.inner.truncate(len);
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }
}

impl<T: Display> Display for Stack<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f,"[ ")?;
        for (i, el) in self.inner.iter().enumerate() {
            if i != 0 {
                write!(f, ", ")?;
            }
            write!(f, "{el}")?;
        }
        write!(f, " ]")?;
        Ok(())
    }
}

impl<T> IndexMut<usize> for Stack<T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.inner[index]
    }
}

impl<T> Index<usize> for Stack<T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        &self.inner[index]
    }
}

#[derive(Clone)]
pub struct StackFrame {
    #[allow(dead_code)]
    name: Rc<str>,
    body: Rc<Chunk>,
    ip: usize,
    bp: usize,
    lp: usize,
    depth: usize,
}

struct VmError {
    err: Exception,
    frames: Vec<StackFrame>
}

impl From<Exception> for VmError {
    fn from(err: Exception) -> Self {
        VmError {
            err,
            frames: Vec::new()
        }
    }
}

type VmResult<T> = std::result::Result<T, VmError>;

impl StackFrame {
    fn new(vm: &Vm, name: Rc<str>, body: Rc<Chunk>, arity: usize, depth: usize) -> Self {
        Self {
            name,
            body,
            bp: vm.stack.len() - arity,
            lp: vm.locals.len(),
            ip: 0,
            depth
        }
    }
}

impl Default for StackFrame {
    fn default() -> Self {
        StackFrame {
            name: Rc::from("<root>"),
            body: Rc::new(Chunk::default()),
            bp: 0,
            lp: 0,
            ip: 0,
            depth: 0
        }
    }
}

pub enum Interupt {
    KeyboardInterupt = 1
}

struct TryScope {
    stack_len: usize,
    locals_len: usize,
    frame_depth: usize,
    err_idx: usize,
}

pub struct Vm {
    stack: Stack<Value>,
    locals: Stack<Value>,
    trystack: Vec<TryScope>,
    globals: Rc<RefCell<HashMap<u16, Value>>>,
    names: NamesTable,
    global_names: GlobalsTable,
    interupt: Arc<AtomicUsize>,
}

macro_rules! error {
    ($($arg:tt)*) => {
        exception!(RUNTIME_EXCEPTION, $($arg)*)
    };
}

impl Vm {

    fn push(&mut self, val: Value) {
        self.stack.push(val)
    }

    fn pop(&mut self) -> Value {
        self.stack.pop()
    }

    fn top(&mut self) -> Value {
        self.stack[self.stack.len() - 1].clone()
    }

    pub fn names(&self) -> NamesTable {
        self.names.clone()
    }

    pub fn globals(&self) -> GlobalsTable {
        self.global_names.clone()
    }

    pub fn interupt(&self) -> Arc<AtomicUsize> {
        self.interupt.clone()
    }

    pub fn new() -> Self {
        Self {
            stack: Stack::new(),
            locals: Stack::new(),
            trystack: Vec::new(),
            globals: Rc::new(RefCell::new(HashMap::new())),
            names: Rc::new(RefCell::new(Vec::new())),
            global_names: Rc::new(RefCell::new(Vec::new())),
            interupt: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub fn load_global(&mut self, value: Value, name: &str) {
        let idx = self.global_names.borrow().len();
        self.global_names.borrow_mut().push(Global { idx, name: Rc::from(name), is_const: true });
        self.globals.borrow_mut().insert(idx as u16, value);
    }

    pub fn load_global_fn(&mut self, value: Rc<Function>, name: &str) {
        self.load_global(Value::Function(value), name)
    }

    fn init_frame(&mut self, fun: Rc<Function>) -> Result<StackFrame> {
        let InnerFunction::Compiled(ref compiled_chunk) = fun.fun else {
            return Err(error!("cannot create stack frame on builtin function"))
        };
        Ok(StackFrame::new(
            self,
            fun.name.clone(),
            compiled_chunk.clone(),
            fun.arity,
            0
        ))
    }

    fn check_interupt(&self) -> bool {
        self.interupt.load(Ordering::Relaxed) != 0
    }

    fn exec_ins(&mut self, frame: &mut StackFrame, ins: Instruction) -> VmResult<Option<Value>> {
        use Instruction as I;
        match ins {
            I::NoOp => {},
            I::CreateLocal => self.locals.push(self.stack.pop()),
            I::LoadLocal(idx) => {self.stack.push(self.locals[frame.lp + idx as usize].clone());},
            I::StoreLocal(idx) => self.locals[frame.bp + idx as usize] = self.pop(),
            I::DiscardLocals(count) => {self.locals.truncate(self.locals.len() - count as usize)},
            I::LoadGlobal(idx) => {
                let val = self.globals
                    .borrow_mut()
                    .get(&idx)
                    .unwrap()
                    .clone();
                self.stack.push(val);
            },
            I::StoreGlobal(idx) => {
                let val = self.pop();
                self.globals.borrow_mut().insert(idx, val);
            },
            I::Const(idx) => self.push(frame.body.constants[idx as usize].clone()),
            I::Int(i) => self.push(Value::Int(i as i64)),
            I::True => self.push(Value::Bool(true)),
            I::False => self.push(Value::Bool(false)),
            I::Nil => self.push(Value::Nil),
            I::Dup => self.push(self.stack[self.stack.len() - 1].clone()),
            I::Discard(count) => {self.stack.truncate(self.stack.len() - count as usize)},
            I::UnaryOp(op) => {
                let val = self.pop();
                self.push(Value::unary_op(op, val));
            },
            I::BinaryOp(op) => {
                let rhs = self.pop();
                let lhs = self.pop();
                self.push(Value::binary_op(op, lhs, rhs)?);
            },
            I::NewList(items) => {
                let list = self.stack.split_off(self.stack.len() - items as usize);
                self.push(Value::List(list.inner.into()));
            },
            I::NewTable(items) => {
                let mut table = ValueMap::new();
                for _ in 0..items {
                    let value = self.pop();
                    let key = self.pop();
                    table.insert(key, value)?;
                }
                self.push(Value::Table(table.into()))
            },
            I::NewMatrix(items, domain) => {
                let values = self.stack.split_off(self.stack.len() - items as usize).inner;
                let domain = domain as usize;
                let codomain = values.len() / domain;
                self.push(Value::Matrix(Gc::new(Matrix::new(domain, codomain, values))));
            }
            I::Jump(idx) => {
                if self.check_interupt() {
                    return Ok(Some(Value::Nil))
                }
                frame.ip = idx as usize;
            }
            I::JumpTrue(idx) => {
                if self.check_interupt() {
                    return Ok(Some(Value::Nil))
                }
                if !!self.pop() {
                    frame.ip = idx as usize;
                }
            },
            I::JumpFalse(idx) => {
                if self.check_interupt() {
                    return Ok(Some(Value::Nil))
                }
                if !self.pop() {
                    frame.ip = idx as usize;
                }
            },
            I::JumpNil(idx) => {
                if self.check_interupt() {
                    return Ok(Some(Value::Nil))
                }
                if self.pop() == Value::Nil {
                    frame.ip = idx as usize;
                }
            },
            I::Call(arity) => {
                let arity = arity as usize;

                let fun = self.pop();
                let Value::Function(fun) = fun else {
                    Err(error!("cannot call {fun} (not a function)"))?
                };

                if !fun.variadic && arity > fun.arity {
                    let needs = fun.arity;
                    Err(error!("function {fun} takes {needs} args, given {arity}"))?
                }

                let mut params = self.stack.split_off(self.stack.len() - arity).inner;

                if params.len() < fun.arity {
                    let new_arity = fun.arity - arity;
                    let val = native!("<partial>", new_arity, false, move |(vm, vm_frame), args| {
                        let mut params = params.clone();
                        params.extend(args);
                        vm.exec_fn(vm_frame, fun.clone(), params).map_err(|e| e.err)
                    });
                    self.stack.push(val);
                    return Ok(None);
                }

                if fun.variadic {
                    if arity == fun.arity {
                        params.push(Value::List(vec![].into()));
                    } else {
                        let count = arity - fun.arity;
                        let varadic_params = params.split_off(params.len() - count);
                        params.push(Value::List(varadic_params.into()));
                    }
                }

                let res = self.exec_fn(frame, fun, params)?;
                self.push(res);
            },
            I::Return => {
                if self.check_interupt() {
                    return Ok(Some(Value::Nil))
                }
                let ret = self.pop();
                self.stack.truncate(frame.bp);
                self.locals.truncate(frame.lp);
                return Ok(Some(ret))
            },
            I::Field(name_idx) => {
                let expr = self.pop();
                let name = self.names.borrow()[name_idx as usize].clone();
                self.push(expr.field_access(&name)?);
            },
            I::StoreField(name_idx) => {
                let mut expr = self.pop();
                let value = self.pop();
                let name = self.names.borrow()[name_idx as usize].clone();
                expr.store_field_access(&name, value)?;
            },
            I::Index(count) => {
                let index = self.stack.split_off(self.stack.len() - count as usize);
                let collection = self.pop();
                let value = collection.index(&index.inner)?;
                self.stack.push(value);
            },
            I::StoreIndex(count) => {
                let index = self.stack.split_off(self.stack.len() - count as usize);
                let mut collection = self.pop();
                let value = self.pop();
                collection.store_index(&index.inner, value)?;
            },
            I::IterCreate => {
                let val = self.pop().into_iter();
                self.push(val?);
            },
            I::IterNext => {
                let Value::Iter(iter) = self.top() else { panic!("bypassed iter check"); };
                let val = self.exec_fn(frame, iter, vec![])?;
                self.push(val);
            },
            I::Try(idx) => {
                let scope = TryScope {
                    err_idx: idx as usize,
                    frame_depth: frame.depth,
                    stack_len: self.stack.len(),
                    locals_len: self.locals.len(),
                };
                self.trystack.push(scope);
            },
            I::TryEnd => {
                self.trystack.pop().unwrap();
            },
        };

        Ok(None)
    }

    fn stack_trace(&mut self, frames: Vec<StackFrame>, e: Exception) -> Exception {
        let mut e = e;
        for frame in frames {
            let ip = frame.ip - 1;
            let pos = frame.body.pos[ip];
            e = e.trace(frame.name, pos);
        }
        e
    }

    fn exec_fn(&mut self, frame: &mut StackFrame, fun: Rc<Function>, params: Vec<Value>) -> VmResult<Value> {
        if self.check_interupt() {
            return Ok(Value::Nil)
        }
        let name = fun.name.clone();
        let params_len = params.len();
        match &fun.fun {
            InnerFunction::Compiled(body) => {
                for param in params {
                    self.stack.push(param);
                }
                let mut new_frame = StackFrame::new(self, name, body.clone(), params_len, frame.depth + 1);
                let res = self.exec(&mut new_frame);
                res
            },
            InnerFunction::Native(native_fun) => {
                Ok(native_fun((self, frame), params)?)
            }
        }
    }

    fn exec(&mut self, frame: &mut StackFrame) -> VmResult<Value> {
        loop {
            let ins = frame.body.code[frame.ip].clone();
            frame.ip += 1;
            match self.exec_ins(frame, ins) {
                Ok(Some(val)) => return Ok(val),
                Ok(None) => {},
                Err(err) => {
                    if let Some(catch) = self.trystack.pop() {
                        if frame.depth != catch.frame_depth {
                            self.trystack.push(catch);
                            return Err(err);
                        }
                        self.stack.truncate(catch.stack_len);
                        self.locals.truncate(catch.locals_len);
                        frame.ip = catch.err_idx;
                        self.stack.push(Value::Exception(err.err));
                    } else {
                        let mut err = err;
                        err.frames.push(frame.clone());
                        return Err(err)
                    }
                },
            }
        }
    }

    pub fn run_fn(&mut self, frame: &mut StackFrame, fun: Rc<Function>, params: Vec<Value>) -> Result<Value> {
        self.exec_fn(frame, fun, params).map_err(|e| e.err)
    }

    pub fn run(&mut self, fun: Rc<Function>) -> Result<Value> {
        self.interupt.store(0, Ordering::SeqCst);
        self.stack = Stack::new();
        self.locals = Stack::new();
        let mut frame = self.init_frame(fun)?;
        self.exec(&mut frame).map_err(|e| {
            self.stack_trace(e.frames, e.err)
        })
    }
}