rudus/src/vm.rs

use crate::base::BaseFn;
use crate::compiler::{Chunk, Op};
use crate::parser::Ast;
use crate::spans::Spanned;
use crate::value::{LFn, Value};
use chumsky::prelude::SimpleSpan;
use imbl::{HashMap, Vector};
use num_traits::FromPrimitive;
use std::cell::OnceCell;
use std::cell::RefCell;
use std::fmt;
use std::mem::swap;
use std::rc::Rc;

#[derive(Debug, Clone, PartialEq)]
// pub struct Panic {
//     pub input: &'static str,
//     pub src: &'static str,
//     pub msg: String,
//     pub span: SimpleSpan,
//     pub trace: Vec<Trace>,
//     pub extra: String,
// }
pub enum Panic {
    Str(&'static str),
    String(String),
}

impl fmt::Display for Panic {
    fn fmt(self: &Panic, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Panic::Str(msg) => write!(f, "{msg}"),
            Panic::String(msg) => write!(f, "{msg}"),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Trace {
    pub callee: Spanned<Ast>,
    pub caller: Spanned<Ast>,
    pub function: Value,
    pub arguments: Value,
    pub input: &'static str,
    pub src: &'static str,
}

pub struct CallFrame {
    pub function: Value,
    pub arity: u8,
    pub stack_base: usize,
    pub ip: usize,
}

impl CallFrame {
    pub fn chunk(&self) -> &Chunk {
        let Value::Fn(ref function) = self.function else {
            unreachable!()
        };
        &function.chunk(self.arity)
    }
}

impl fmt::Display for CallFrame {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let Value::Fn(ref function) = self.function else {
            unreachable!()
        };
        write!(
            f,
            "CallFrame: {}/{} @ {}",
            function.name(),
            self.arity,
            self.ip
        )
    }
}

pub struct Vm {
    pub stack: Vec<Value>,
    pub call_stack: Vec<CallFrame>,
    pub frame: CallFrame,
    pub ip: usize,
    pub return_register: [Value; 8],
    pub matches: bool,
    pub match_depth: u8,
    pub result: Option<Result<Value, Panic>>,
}

impl Vm {
    pub fn new(chunk: Chunk) -> Vm {
        let lfn = LFn::Defined {
            name: "user script",
            doc: None,
            chunks: vec![(0, chunk)],
            closed: RefCell::new(vec![]),
        };
        let base_fn = Value::Fn(Rc::new(lfn));
        let base_frame = CallFrame {
            function: base_fn.clone(),
            stack_base: 0,
            ip: 0,
            arity: 0,
        };
        Vm {
            stack: vec![],
            call_stack: Vec::with_capacity(64),
            frame: base_frame,
            ip: 0,
            return_register: [const { Value::Nothing }; 8],
            matches: false,
            match_depth: 0,
            result: None,
        }
    }

    pub fn chunk(&self) -> &Chunk {
        self.frame.chunk()
    }

    pub fn push(&mut self, value: Value) {
        self.stack.push(value);
    }

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

    pub fn peek(&self) -> &Value {
        self.stack.last().unwrap()
    }

    pub fn print_stack(&self) {
        let inner = self
            .stack
            .iter()
            .map(|val| val.to_string())
            .collect::<Vec<_>>()
            .join("|");
        let register = self
            .return_register
            .iter()
            .map(|val| val.to_string())
            .collect::<Vec<_>>()
            .join(",");
        println!("{:04}: [{inner}] ({register})", self.ip);
    }

    fn print_debug(&self) {
        self.print_stack();
        let mut ip = self.ip;
        self.chunk().dissasemble_instr(&mut ip);
    }

    pub fn run(&mut self) -> &Result<Value, Panic> {
        while self.result.is_none() {
            self.interpret();
        }
        self.result.as_ref().unwrap()
    }

    pub fn panic(&mut self, msg: &'static str) {
        self.result = Some(Err(Panic::Str(msg)));
    }

    pub fn panic_with(&mut self, msg: String) {
        self.result = Some(Err(Panic::String(msg)));
    }

    pub fn interpret(&mut self) {
        loop {
            let Some(byte) = self.chunk().bytecode.get(self.ip) else {
                self.result = Some(Ok(self.stack.pop().unwrap()));
                return;
            };
            if crate::DEBUG_RUN {
                self.print_debug();
            }
            let op = Op::from_u8(*byte).unwrap();
            use Op::*;
            match op {
                Noop => {
                    self.ip += 1;
                }
                Nil => {
                    self.push(Value::Nil);
                    self.ip += 1;
                }
                True => {
                    self.push(Value::True);
                    self.ip += 1;
                }
                False => {
                    self.push(Value::False);
                    self.ip += 1;
                }
                Constant => {
                    let const_idx = self.chunk().bytecode[self.ip + 1];
                    let value = self.chunk().constants[const_idx as usize].clone();
                    self.push(value);
                    self.ip += 2;
                }
                Jump => {
                    let jump_len = self.chunk().bytecode[self.ip + 1];
                    self.ip += jump_len as usize + 2;
                }
                JumpBack => {
                    let jump_len = self.chunk().bytecode[self.ip + 1];
                    self.ip -= jump_len as usize;
                }
                JumpIfFalse => {
                    let jump_len = self.chunk().bytecode[self.ip + 1];
                    let cond = self.pop();
                    match cond {
                        Value::Nil | Value::False => {
                            self.ip += jump_len as usize + 2;
                        }
                        _ => {
                            self.ip += 2;
                        }
                    }
                }
                JumpIfTrue => {
                    let jump_len = self.chunk().bytecode[self.ip + 1];
                    let cond = self.pop();
                    match cond {
                        Value::Nil | Value::False => {
                            self.ip += 2;
                        }
                        _ => {
                            self.ip += jump_len as usize + 2;
                        }
                    }
                }
                JumpIfZero => {
                    let jump_len = self.chunk().bytecode[self.ip + 1];
                    let cond = self.pop();
                    match cond {
                        Value::Number(x) if x <= 0.0 => {
                            self.ip += jump_len as usize + 2;
                        }
                        Value::Number(..) => {
                            self.ip += 2;
                        }
                        _ => return self.panic("repeat requires a number"),
                    }
                }
                Pop => {
                    self.pop();
                    self.ip += 1;
                }
                PopN => {
                    let n = self.chunk().bytecode[self.ip + 1] as usize;
                    self.stack.truncate(self.stack.len() - n);
                    self.ip += 2;
                }
                PushBinding => {
                    let binding_idx =
                        self.chunk().bytecode[self.ip + 1] as usize + self.frame.stack_base;
                    let binding_value = self.stack[binding_idx].clone();
                    self.push(binding_value);
                    self.ip += 2;
                }
                Store => {
                    self.return_register[0] = self.pop();
                    self.push(Value::Nothing);
                    self.ip += 1;
                }
                StoreAt => {
                    let i = self.chunk().bytecode[self.ip + 1] as usize;
                    self.return_register[i] = self.pop();
                    self.ip += 2;
                }
                Stash => {
                    self.return_register[0] = self.peek().clone();
                    self.ip += 1;
                }
                Load => {
                    let mut i = 0;
                    while i < 8 && self.return_register[i] != Value::Nothing {
                        let mut value = Value::Nothing;
                        swap(&mut self.return_register[i], &mut value);
                        self.push(value);
                        i += 1;
                    }
                    self.ip += 1;
                }
                ResetMatch => {
                    self.matches = false;
                    self.match_depth = 0;
                    self.ip += 1;
                }
                Match => {
                    self.matches = true;
                    self.ip += 1;
                }
                MatchType => {
                    let as_type = self.pop();
                    let Value::Keyword(as_type) = as_type else {
                        unreachable!()
                    };
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    let val_type = self.stack[idx].type_of();
                    self.matches = val_type == as_type;
                    self.ip += 1;
                }
                MatchNil => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    if self.stack[idx] == Value::Nil {
                        self.matches = true;
                    };
                    self.ip += 1;
                }
                MatchTrue => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    if self.stack[idx] == Value::True {
                        self.matches = true;
                    };
                    self.ip += 1;
                }
                MatchFalse => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    if self.stack[idx] == Value::False {
                        self.matches = true;
                    }
                    self.ip += 1;
                }
                PanicIfNoMatch => {
                    if !self.matches {
                        return self.panic("no match");
                    } else {
                        self.ip += 1;
                    }
                }
                MatchConstant => {
                    let const_idx = self.chunk().bytecode[self.ip + 1];
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    self.matches = self.stack[idx] == self.chunk().constants[const_idx as usize];
                    self.ip += 2;
                }
                MatchTuple => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    let tuple_len = self.chunk().bytecode[self.ip + 1];
                    let scrutinee = self.stack[idx].clone();
                    match scrutinee {
                        Value::Tuple(members) => self.matches = members.len() == tuple_len as usize,
                        _ => self.matches = false,
                    };
                    self.ip += 2;
                }
                PushTuple => {
                    let tuple_len = self.chunk().bytecode[self.ip + 1];
                    let tuple_members = self.stack.split_off(self.stack.len() - tuple_len as usize);
                    let tuple = Value::Tuple(Rc::new(tuple_members));
                    self.push(tuple);
                    self.ip += 2;
                }
                LoadTuple => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    let tuple = self.stack[idx].clone();
                    match tuple {
                        Value::Tuple(members) => {
                            for member in members.iter() {
                                self.push(member.clone());
                            }
                        }
                        _ => return self.panic("internal error: expected tuple"),
                    };
                    self.ip += 1;
                }
                PushList => {
                    self.push(Value::List(Box::new(Vector::new())));
                    self.ip += 1;
                }
                AppendList => {
                    let value = self.pop();
                    let list = self.pop();
                    let Value::List(mut list) = list else {
                        return self.panic("only lists may be splatted into lists");
                    };
                    list.push_back(value);
                    self.push(Value::List(list));
                    self.ip += 1;
                }
                ConcatList => {
                    let splatted = self.pop();
                    let target = self.pop();
                    let Value::List(mut target) = target else {
                        unreachable!()
                    };
                    let Value::List(splatted) = splatted else {
                        return self.panic("only lists may be splatted into lists");
                    };
                    target.append(*splatted);
                    self.push(Value::List(target));
                    self.ip += 1;
                }
                MatchList => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    let tuple_len = self.chunk().bytecode[self.ip + 1];
                    let scrutinee = self.stack[idx].clone();
                    match scrutinee {
                        Value::List(members) => self.matches = members.len() == tuple_len as usize,
                        _ => self.matches = false,
                    };
                    self.ip += 2;
                }
                LoadList => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    let tuple = self.stack[idx].clone();
                    match tuple {
                        Value::List(members) => {
                            for member in members.iter() {
                                self.push(member.clone());
                            }
                        }
                        _ => return self.panic("internal error: expected tuple"),
                    };
                    self.ip += 1;
                }
                // PushDict => {
                //     let dict_len = self.chunk().bytecode[self.ip + 1] as usize * 2;
                //     let dict_members = self.stack.split_off(self.stack.len() - dict_len);
                //     let mut dict = HashMap::new();
                //     let mut dict_iter = dict_members.iter();
                //     while let Some(kw) = dict_iter.next() {
                //         let Value::Keyword(key) = kw else {
                //             unreachable!()
                //         };
                //         let value = dict_iter.next().unwrap();
                //         dict.insert(*key, value.clone());
                //     }
                //     self.push(Value::Dict(Box::new(dict)));
                //     self.ip += 2;
                // }
                PushDict => {
                    self.push(Value::Dict(Box::new(HashMap::new())));
                    self.ip += 1;
                }
                AppendDict => {
                    let value = self.pop();
                    let Value::Keyword(key) = self.pop() else {
                        unreachable!()
                    };
                    let Value::Dict(mut dict) = self.pop() else {
                        unreachable!()
                    };
                    dict.insert(key, value);
                    self.push(Value::Dict(dict));
                    self.ip += 1;
                }
                ConcatDict => {
                    let Value::Dict(splatted) = self.pop() else {
                        return self.panic("only dicts may be splatted into dicts");
                    };
                    let Value::Dict(target) = self.pop() else {
                        unreachable!()
                    };
                    let union = splatted.union(*target);
                    self.push(Value::Dict(Box::new(union)));
                    self.ip += 1;
                }
                LoadDictValue => {
                    let dict_idx = self.chunk().bytecode[self.ip + 1] as usize;
                    let Value::Dict(dict) = self.stack[dict_idx].clone() else {
                        unreachable!("expected dict, got something else")
                    };
                    let Value::Keyword(key) = self.pop() else {
                        unreachable!("expected keyword, got something else")
                    };
                    let value = dict.get(&key).unwrap_or(&Value::Nil);
                    self.push(value.clone());
                    self.ip += 2;
                }
                MatchDict => {
                    let idx = self.stack.len() - self.match_depth as usize - 1;
                    let dict_len = self.chunk().bytecode[self.ip + 1];
                    let scrutinee = self.stack[idx].clone();
                    match scrutinee {
                        Value::Dict(members) => self.matches = members.len() == dict_len as usize,
                        _ => self.matches = false,
                    };
                    self.ip += 2;
                }
                PushBox => {
                    let val = self.pop();
                    self.push(Value::Box(Rc::new(RefCell::new(val))));
                    self.ip += 1;
                }
                GetKey => {
                    let key = self.pop();
                    let Value::Keyword(idx) = key else {
                        unreachable!()
                    };
                    let dict = self.pop();
                    let value = match dict {
                        Value::Dict(d) => d.as_ref().get(&idx).unwrap_or(&Value::Nil).clone(),
                        _ => Value::Nil,
                    };
                    self.push(value);
                    self.ip += 1;
                }
                JumpIfNoMatch => {
                    let jump_len = self.chunk().bytecode[self.ip + 1] as usize;
                    if !self.matches {
                        self.ip += jump_len + 2;
                    } else {
                        self.ip += 2;
                    }
                }
                JumpIfMatch => {
                    let jump_len = self.chunk().bytecode[self.ip + 1] as usize;
                    if self.matches {
                        self.ip += jump_len + 2;
                    } else {
                        self.ip += 2;
                    }
                }
                TypeOf => {
                    let val = self.pop();
                    let type_of = Value::Keyword(val.type_of());
                    self.push(type_of);
                    self.ip += 1;
                }
                Truncate => {
                    let val = self.pop();
                    if let Value::Number(x) = val {
                        self.push(Value::Number(x as usize as f64));
                        self.ip += 1;
                    } else {
                        return self.panic("repeat requires a number");
                    }
                }
                Decrement => {
                    let val = self.pop();
                    if let Value::Number(x) = val {
                        self.push(Value::Number(x - 1.0));
                        self.ip += 1;
                    } else {
                        return self.panic("you may only decrement a number");
                    }
                }
                Duplicate => {
                    self.push(self.peek().clone());
                    self.ip += 1;
                }
                MatchDepth => {
                    self.match_depth = self.chunk().bytecode[self.ip + 1];
                    self.ip += 2;
                }
                PanicNoWhen | PanicNoMatch => {
                    return self.panic("no match");
                }
                Eq => {
                    let first = self.pop();
                    let second = self.pop();
                    if first == second {
                        self.push(Value::True)
                    } else {
                        self.push(Value::False)
                    }
                    self.ip += 1;
                }
                Add => {
                    let first = self.pop();
                    let second = self.pop();
                    if let (Value::Number(x), Value::Number(y)) = (first, second) {
                        self.push(Value::Number(x + y))
                    } else {
                        return self.panic("`add` requires two numbers");
                    }
                    self.ip += 1;
                }
                Sub => {
                    let first = self.pop();
                    let second = self.pop();
                    if let (Value::Number(x), Value::Number(y)) = (first, second) {
                        self.push(Value::Number(y - x))
                    } else {
                        return self.panic("`sub` requires two numbers");
                    }
                    self.ip += 1;
                }
                Mult => {
                    let first = self.pop();
                    let second = self.pop();
                    if let (Value::Number(x), Value::Number(y)) = (first, second) {
                        self.push(Value::Number(x * y))
                    } else {
                        return self.panic("`mult` requires two numbers");
                    }
                    self.ip += 1;
                }
                Div => {
                    let first = self.pop();
                    let second = self.pop();
                    if let (Value::Number(x), Value::Number(y)) = (first, second) {
                        if x == 0.0 {
                            return self.panic("division by 0");
                        }
                        self.push(Value::Number(y / x))
                    } else {
                        return self.panic("`div` requires two numbers");
                    }
                    self.ip += 1;
                }
                Unbox => {
                    let the_box = self.pop();
                    let inner = if let Value::Box(b) = the_box {
                        b.borrow().clone()
                    } else {
                        return self.panic("`unbox` requires a box");
                    };
                    self.push(inner);
                    self.ip += 1;
                }
                BoxStore => {
                    let new_value = self.pop();
                    let the_box = self.pop();
                    if let Value::Box(b) = the_box {
                        b.replace(new_value.clone());
                    } else {
                        return self.panic("`store` requires a box");
                    }
                    self.push(new_value);
                    self.ip += 1;
                }
                Assert => {
                    let value = self.stack.last().unwrap();
                    if let Value::Nil | Value::False = value {
                        return self.panic("asserted falsy value");
                    }
                    self.ip += 1;
                }
                Get => {
                    let key = self.pop();
                    let dict = self.pop();
                    let value = match (key, dict) {
                        (Value::Keyword(k), Value::Dict(d)) => {
                            d.as_ref().get(&k).unwrap_or(&Value::Nil).clone()
                        }
                        (Value::Keyword(_), _) => Value::Nil,
                        _ => return self.panic("keys must be keywords"),
                    };
                    self.push(value);
                    self.ip += 1;
                }
                At => {
                    let idx = self.pop();
                    let ordered = self.pop();
                    let value = match (ordered, idx) {
                        (Value::List(l), Value::Number(i)) => {
                            l.get(i as usize).unwrap_or(&Value::Nil).clone()
                        }
                        (Value::Tuple(t), Value::Number(i)) => {
                            t.get(i as usize).unwrap_or(&Value::Nil).clone()
                        }
                        (_, Value::Number(_)) => Value::Nil,
                        _ => return self.panic("indexes must be numbers"),
                    };
                    self.push(value);
                    self.ip += 1;
                }
                Not => {
                    let value = self.pop();
                    let negated = match value {
                        Value::Nil | Value::False => Value::True,
                        _ => Value::False,
                    };
                    self.push(negated);
                    self.ip += 1;
                }
                Panic => {
                    let msg = self.pop().show();
                    return self.panic_with(msg);
                }
                EmptyString => {
                    self.push(Value::String(Rc::new("".to_string())));
                    self.ip += 1;
                }
                //TODO: don't use the schlemiel's algo here
                ConcatStrings => {
                    let second = self.pop();
                    let first = self.pop();
                    let combined = match (first, second) {
                        (Value::String(first), Value::String(second)) => {
                            let mut new = first.as_ref().clone();
                            new.push_str(second.as_str());
                            Value::String(Rc::new(new))
                        }
                        _ => unreachable!(),
                    };
                    self.push(combined);
                    self.ip += 1;
                }
                Stringify => {
                    let to_stringify = self.pop();
                    let the_string = to_stringify.stringify();
                    let stringified = Value::String(Rc::new(the_string));
                    self.push(stringified);
                    self.ip += 1;
                }
                Call => {
                    let arity = self.chunk().bytecode[self.ip + 1];
                    self.ip += 2;

                    let val = self.pop();

                    match val {
                        Value::Fn(_) => {
                            let mut frame = CallFrame {
                                function: val,
                                arity,
                                stack_base: self.stack.len() - arity as usize,
                                ip: 0,
                            };

                            swap(&mut self.frame, &mut frame);
                            frame.ip = self.ip;

                            self.call_stack.push(frame);
                            self.ip = 0;

                            if crate::DEBUG_RUN {
                                println!("== calling into {} ==", self.frame.function.show());
                            }
                        }
                        Value::BaseFn(base_fn) => {
                            let value = match (arity, base_fn) {
                                (0, BaseFn::Nullary(f)) => f(),
                                (1, BaseFn::Unary(f)) => f(&self.pop()),
                                (2, BaseFn::Binary(f)) => f(&self.pop(), &self.pop()),
                                (3, BaseFn::Ternary(f)) => f(&self.pop(), &self.pop(), &self.pop()),
                                _ => return self.panic("internal ludus error"),
                            };
                            self.push(value);
                        }
                        _ => return self.panic_with(format!("{} is not a function", val.show())),
                    }
                }
                Return => {
                    if crate::DEBUG_RUN {
                        println!("== returning from {} ==", self.frame.function.show())
                    }
                    self.frame = self.call_stack.pop().unwrap();
                    self.ip = self.frame.ip;
                    let mut value = Value::Nothing;
                    swap(&mut self.return_register[0], &mut value);
                    self.push(value);
                }
                Print => {
                    println!("{}", self.pop().show());
                    self.push(Value::Keyword("ok"));
                    self.ip += 1;
                }
                SetUpvalue => {
                    let idx = self.chunk().bytecode[self.ip + 1];
                    self.ip += 2;
                    let closed_over = self.stack[idx as usize].clone();
                    if let Value::Fn(lfn) = self.peek() {
                        lfn.close(closed_over);
                    }
                }
                GetUpvalue => {
                    let idx = self.chunk().bytecode[self.ip + 1];
                    self.ip += 2;
                    if let Value::Fn(ref inner) = self.frame.function {
                        self.push(inner.as_ref().upvalue(idx));
                    } else {
                        unreachable!();
                    }
                }
            }
        }
    }
}