oden/fine/src/vm.rs

use std::rc::Rc;

use crate::compiler::{Export, Function, Instruction, Module};
use crate::semantics::Type;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum VMErrorCode {
    #[error("code panic (syntax or semantic error)")]
    Panic,
    #[error("internal error: stack underflow")]
    StackUnderflow,
    #[error("internal error: stack type mismatch: {0:?} is not {1:?}")]
    StackTypeMismatch(StackValue, Type),

    // TODO: This one is *not* like the others!
    #[error("divide by zero")]
    DivideByZero,

    #[error("internal error: argument {0} out of range")]
    ArgumentOutOfRange(usize),
    #[error("internal error: global {0} out of range")]
    GlobalOutOfRange(usize),
    #[error("internal error: local {0} out of range")]
    LocalOutOfRange(usize),
    #[error("internal error: string {0} out of range")]
    StringOutOfRange(usize),
    #[error("internal error: function {0} out of range")]
    FunctionOutOfRange(usize),
    #[error("internal error: stack type mismatch ({0:?} is not function)")]
    StackExpectedFunction(StackValue),
}

#[derive(Debug)]
pub struct VMError {
    pub code: VMErrorCode,
    pub stack: Box<[Frame]>,
}

type Result<T> = std::result::Result<T, VMErrorCode>;

#[derive(Clone, Debug)]
pub enum StackValue {
    Nothing,
    Bool(bool),
    Float(f64),
    String(Rc<str>),
    Function(Rc<Function>),
    ExternFunction(usize),
}

enum FuncValue {
    Function(Rc<Function>),
    ExternFunction(usize),
}

#[derive(Debug)]
pub struct Frame {
    func: Rc<Function>,
    args: Vec<StackValue>,
    locals: Vec<StackValue>,
    stack: Vec<StackValue>,
    pc: usize,
}

impl Frame {
    fn from_function(func: Rc<Function>, args: Vec<StackValue>) -> Self {
        let mut locals = Vec::new();
        locals.resize(func.locals(), StackValue::Nothing);

        Frame {
            func,
            args,
            locals,
            stack: Vec::new(),
            pc: 0,
        }
    }

    fn pop_value(&mut self) -> Result<StackValue> {
        self.stack.pop().ok_or_else(|| VMErrorCode::StackUnderflow)
    }

    fn pop_bool(&mut self) -> Result<bool> {
        match self.pop_value()? {
            StackValue::Bool(v) => Ok(v),
            v => Err(VMErrorCode::StackTypeMismatch(v, Type::Bool)),
        }
    }

    fn pop_float(&mut self) -> Result<f64> {
        match self.pop_value()? {
            StackValue::Float(v) => Ok(v),
            v => Err(VMErrorCode::StackTypeMismatch(v, Type::F64)),
        }
    }

    fn push_value(&mut self, v: StackValue) {
        self.stack.push(v)
    }

    fn push_bool(&mut self, value: bool) {
        self.push_value(StackValue::Bool(value));
    }

    fn push_float(&mut self, value: f64) {
        self.push_value(StackValue::Float(value));
    }

    fn push_nothing(&mut self) {
        self.push_value(StackValue::Nothing);
    }

    fn push_string(&mut self, v: Rc<str>) {
        self.push_value(StackValue::String(v))
    }

    fn push_function(&mut self, v: Rc<Function>) {
        self.push_value(StackValue::Function(v));
    }

    fn push_extern_function(&mut self, v: usize) {
        self.push_value(StackValue::ExternFunction(v));
    }

    fn get_argument(&self, i: usize) -> Result<StackValue> {
        self.args
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::ArgumentOutOfRange(i))
    }

    fn get_local(&self, i: usize) -> Result<StackValue> {
        self.locals
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::LocalOutOfRange(i))
    }

    fn get_string(&self, i: usize) -> Result<Rc<str>> {
        let strings = self.func.strings();
        strings
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::StringOutOfRange(i))
    }

    fn store_local(&mut self, i: usize, v: StackValue) -> Result<()> {
        if i >= self.locals.len() {
            Err(VMErrorCode::LocalOutOfRange(i))
        } else {
            self.locals[i] = v;
            Ok(())
        }
    }

    fn pop_function(&mut self) -> Result<FuncValue> {
        match self.pop_value()? {
            StackValue::Function(f) => Ok(FuncValue::Function(f)),
            StackValue::ExternFunction(i) => Ok(FuncValue::ExternFunction(i)),
            v => Err(VMErrorCode::StackExpectedFunction(v)),
        }
    }
}

pub struct Context {
    module: Rc<Module>,
    globals: Vec<StackValue>,
}

impl Context {
    pub fn new(module: Rc<Module>) -> Context {
        let mut globals = Vec::new();
        globals.resize(module.globals, StackValue::Nothing);
        Context { module, globals }
    }

    pub fn init(&mut self) -> std::result::Result<(), VMError> {
        let init_index = self.module.init;
        let init_fn = self.module.functions[init_index].clone();
        eval(self, init_fn, vec![])?;
        Ok(())
    }

    fn get_global(&self, i: usize) -> Result<StackValue> {
        self.globals
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::GlobalOutOfRange(i)) // TODO: Test
    }

    fn set_global(&mut self, i: usize, v: StackValue) -> Result<()> {
        if i >= self.globals.len() {
            Err(VMErrorCode::GlobalOutOfRange(i)) // TODO: Test
        } else {
            self.globals[i] = v;
            Ok(())
        }
    }

    fn get_function(&self, i: usize) -> Result<Rc<Function>> {
        let functions = self.module.functions();
        functions
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::FunctionOutOfRange(i)) // TODO: Test
    }
}

enum Flow {
    Break,
    Continue,
}

#[inline(always)]
fn eval_one(
    instruction: Instruction,
    index: &mut usize,
    c: &mut Context,
    f: &mut Frame,
    stack: &mut Vec<Frame>,
) -> Result<Flow> {
    match instruction {
        Instruction::Panic => return Err(VMErrorCode::Panic),
        Instruction::BoolNot => {
            let value = f.pop_bool()?;
            f.push_bool(!value);
        }
        Instruction::Discard => {
            f.pop_value()?;
        }
        Instruction::FloatAdd => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_float(x + y);
        }
        Instruction::FloatDivide => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            if y == 0. {
                return Err(VMErrorCode::DivideByZero);
            }
            f.push_float(x / y);
        }
        Instruction::FloatMultiply => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_float(x * y);
        }
        Instruction::FloatSubtract => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_float(x - y);
        }
        Instruction::Jump(i) => {
            *index = i;
        }
        Instruction::JumpFalse(i) => {
            if !(f.pop_bool()?) {
                *index = i;
            }
        }
        Instruction::JumpTrue(i) => {
            if f.pop_bool()? {
                *index = i;
            }
        }
        Instruction::LoadArgument(i) => {
            let v = f.get_argument(i)?;
            f.push_value(v);
        }
        Instruction::LoadLocal(i) => {
            let v = f.get_local(i)?;
            f.push_value(v);
        }
        Instruction::LoadModule(i) => {
            let v = c.get_global(i)?;
            f.push_value(v);
        }
        Instruction::PushFalse => {
            f.push_bool(false);
        }
        Instruction::PushFloat(v) => {
            f.push_float(v);
        }
        Instruction::PushNothing => {
            f.push_nothing();
        }
        Instruction::PushString(s) => {
            let v = f.get_string(s)?;
            f.push_string(v);
        }
        Instruction::PushTrue => {
            f.push_bool(true);
        }
        Instruction::StoreLocal(i) => {
            let v = f.pop_value()?;
            f.store_local(i, v)?;
        }
        Instruction::StoreModule(i) => {
            let v = f.pop_value()?;
            c.set_global(i, v)?;
        }
        Instruction::LoadFunction(i) => {
            let v = c.get_function(i)?;
            f.push_function(v);
        }
        Instruction::LoadExternFunction(i) => {
            f.push_extern_function(i);
        }
        Instruction::Call(arg_count) => {
            let function = f.pop_function()?;
            let mut args = Vec::new();
            for _ in 0..arg_count {
                args.push(f.pop_value()?);
            }
            match function {
                FuncValue::Function(func) => {
                    let mut frame = Frame::from_function(func, args);
                    std::mem::swap(&mut frame, f);
                    frame.pc = *index;
                    stack.push(frame);
                    *index = 0;
                }
                FuncValue::ExternFunction(_) => todo!(),
            }
        }
        Instruction::Return => match stack.pop() {
            Some(mut frame) => {
                // The return value is at the top of the stack already.
                let retval = f.pop_value()?;
                std::mem::swap(&mut frame, f);
                *index = f.pc;
                f.push_value(retval);
            }
            None => return Ok(Flow::Break),
        },
    }

    Ok(Flow::Continue)
}

pub fn eval(
    c: &mut Context,
    function: Rc<Function>,
    args: Vec<StackValue>,
) -> std::result::Result<StackValue, VMError> {
    let mut stack = Vec::new();
    let mut f = Frame::from_function(function, args);

    let mut index = 0;
    loop {
        let instructions = f.func.instructions();
        let instruction = instructions[index];

        // {
        //     eprint!("{index}: {instruction:?}  [");
        //     for val in f.stack.iter().rev().take(3) {
        //         eprint!("{val:?} ");
        //     }
        //     if f.stack.len() > 3 {
        //         eprint!("...");
        //     }
        //     eprintln!("]");
        // }

        index += 1;

        match eval_one(instruction, &mut index, c, &mut f, &mut stack) {
            Ok(Flow::Break) => match f.pop_value() {
                Ok(v) => return Ok(v),
                Err(e) => {
                    f.pc = index;
                    stack.push(f);
                    return Err(VMError {
                        code: e,
                        stack: stack.into(),
                    });
                }
            },
            Ok(Flow::Continue) => (),
            Err(e) => {
                f.pc = index;
                stack.push(f);
                return Err(VMError {
                    code: e,
                    stack: stack.into(),
                });
            }
        };
    }
}

pub fn eval_export_fn(
    c: &mut Context,
    name: &str,
    args: &[StackValue],
) -> std::result::Result<StackValue, VMError> {
    let export = match c.module.exports.get(name) {
        Some(Export::Function(id)) => id,
        Some(_) => todo!(),
        None => todo!(),
    };

    let function = c.module.functions[*export].clone();
    let args = args.iter().map(|a| a.clone()).collect();
    eval(c, function, args)
}