rudus/src/vm.rs

use crate::compiler::{Chunk, Op};
use crate::parser::Ast;
use crate::spans::Spanned;
use crate::value::Value;
use chumsky::prelude::SimpleSpan;
use imbl::{HashMap, Vector};
use num_traits::FromPrimitive;
use std::cell::RefCell;
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 struct Panic(&'static str);

#[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 Vm<'a> {
    pub stack: Vec<Value>,
    pub chunk: &'a Chunk,
    pub ip: usize,
    pub return_register: [Value; 8],
    pub matches: bool,
    pub match_depth: u8,
    pub result: Option<Result<Value, Panic>>,
}

impl<'a> Vm<'a> {
    pub fn new(chunk: &'a Chunk) -> Vm<'a> {
        Vm {
            chunk,
            stack: vec![],
            ip: 0,
            return_register: [
                Value::Nothing,
                Value::Nothing,
                Value::Nothing,
                Value::Nothing,
                Value::Nothing,
                Value::Nothing,
                Value::Nothing,
                Value::Nothing,
            ],
            matches: false,
            match_depth: 0,
            result: None,
        }
    }

    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();
        self.chunk.dissasemble_instr(self.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(msg)));
    }

    pub fn interpret(&mut self) {
        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;
                    }
                }
            }
            JumpIfZero => {
                let jump_len = self.chunk.bytecode[self.ip + 1];
                let cond = self.pop();
                match cond {
                    Value::Number(0.0) => {
                        self.ip += jump_len as usize + 2;
                        self.interpret()
                    }
                    Value::Number(..) => {
                        self.ip += 2;
                        self.interpret()
                    }
                    _ => 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;
                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;
            }
            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;
            }
            MatchWord => {
                self.matches = true;
                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;
                self.interpret()
            }
            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 {
                    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.stack.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());
                        }
                    }
                    _ => self.panic("internal error: expected tuple"),
                };
                self.ip += 1;
            }
            PushList => {
                let list_len = self.chunk.bytecode[self.ip + 1];
                let list_members = self.stack.split_off(self.stack.len() - list_len as usize);
                let list = Value::List(Box::new(Vector::from(list_members)));
                self.stack.push(list);
                self.ip += 2;
            }
            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());
                        }
                    }
                    _ => 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.stack.push(Value::Dict(Box::new(dict)));
                self.ip += 2;
            }
            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.stack.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 = self.chunk.kw_from(val.type_of()).unwrap();
                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 {
                    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;
                    self.interpret()
                } else {
                    self.panic("you may only decrement a number");
                }
            }
            Duplicate => {
                self.push(self.peek().clone());
                self.ip += 1;
                self.interpret()
            }
            MatchDepth => {
                self.match_depth = self.chunk.bytecode[self.ip + 1];
                self.ip += 2;
                self.interpret()
            }
            PanicNoWhen | PanicNoMatch => self.panic("no match"),

            Eq => {
                let first = self.stack.pop().unwrap();
                let second = self.stack.pop().unwrap();
                if first == second {
                    self.stack.push(Value::True)
                } else {
                    self.stack.push(Value::False)
                }
                self.ip += 1;
            }
            Add => {
                let first = self.stack.pop().unwrap();
                let second = self.stack.pop().unwrap();
                if let (Value::Number(x), Value::Number(y)) = (first, second) {
                    self.stack.push(Value::Number(x + y))
                } else {
                    self.panic("`add` requires two numbers")
                }
                self.ip += 1;
            }
            Sub => {
                let first = self.stack.pop().unwrap();
                let second = self.stack.pop().unwrap();
                if let (Value::Number(x), Value::Number(y)) = (first, second) {
                    self.stack.push(Value::Number(y - x))
                } else {
                    self.panic("`sub` requires two numbers")
                }
                self.ip += 1;
            }
            Mult => {
                let first = self.stack.pop().unwrap();
                let second = self.stack.pop().unwrap();
                if let (Value::Number(x), Value::Number(y)) = (first, second) {
                    self.stack.push(Value::Number(x * y))
                } else {
                    self.panic("`mult` requires two numbers")
                }
                self.ip += 1;
            }
            Div => {
                let first = self.stack.pop().unwrap();
                let second = self.stack.pop().unwrap();
                if let (Value::Number(x), Value::Number(y)) = (first, second) {
                    if x == 0.0 {
                        self.panic("division by 0")
                    }
                    self.stack.push(Value::Number(y / x))
                } else {
                    self.panic("`div` requires two numbers")
                }
                self.ip += 1;
            }
            Unbox => {
                let the_box = self.stack.pop().unwrap();
                let inner = if let Value::Box(b) = the_box {
                    b.borrow().clone()
                } else {
                    return self.panic("`unbox` requires a box");
                };
                self.stack.push(inner);
                self.ip += 1;
            }
            BoxStore => {
                let new_value = self.stack.pop().unwrap();
                let the_box = self.stack.pop().unwrap();
                if let Value::Box(b) = the_box {
                    b.replace(new_value.clone());
                } else {
                    return self.panic("`store` requires a box");
                }
                self.stack.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.stack.pop().unwrap();
                let dict = self.stack.pop().unwrap();
                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.stack.push(value);
                self.ip += 1;
            }
            At => {
                let idx = self.stack.pop().unwrap();
                let ordered = self.stack.pop().unwrap();
                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.stack.push(value);
                self.ip += 1;
            }
        }
    }
}