deimos-lang/dvm-lib/src/vm/mod.rs

use crate::instruction::Instruction;
use crate::platform_function::PlatformFunction;
use crate::vm::constant::Constant;
use crate::vm::function::Function;
use crate::vm::value::Value;
use std::collections::HashMap;
use std::rc::Rc;

pub mod constant;
pub mod function;
pub mod value;

pub struct DvmContext {
    functions: HashMap<Rc<str>, Function>,
    platform_functions: HashMap<Rc<str>, PlatformFunction>,
    constants: HashMap<Rc<str>, Constant>,
}

impl DvmContext {
    pub fn new() -> Self {
        Self {
            functions: HashMap::new(),
            platform_functions: HashMap::new(),
            constants: HashMap::new(),
        }
    }

    pub fn functions(&self) -> &HashMap<Rc<str>, Function> {
        &self.functions
    }

    pub fn functions_mut(&mut self) -> &mut HashMap<Rc<str>, Function> {
        &mut self.functions
    }

    pub fn platform_functions(&self) -> &HashMap<Rc<str>, PlatformFunction> {
        &self.platform_functions
    }

    pub fn platform_functions_mut(&mut self) -> &mut HashMap<Rc<str>, PlatformFunction> {
        &mut self.platform_functions
    }

    #[deprecated]
    pub fn add_function(&mut self, function: Function) {
        self.functions.insert(function.name_owned(), function);
    }

    pub fn constants(&self) -> &HashMap<Rc<str>, Constant> {
        &self.constants
    }

    #[deprecated]
    pub fn add_constant(&mut self, constant: Constant) {
        match &constant {
            Constant::String(string_constant) => {
                self.constants
                    .insert(string_constant.name_owned(), constant);
            }
        }
    }
}

pub struct CallFrame<'a> {
    function_name: Rc<str>,
    instructions: &'a Vec<Instruction>,
    ip: usize,
    stack: Vec<Value>,
    fp: usize,
    return_value: Option<Value>,
}

impl<'a> CallFrame<'a> {
    fn new(function_name: Rc<str>, instructions: &'a Vec<Instruction>) -> Self {
        Self {
            function_name,
            instructions,
            ip: 0,
            stack: vec![],
            fp: 0,
            return_value: None,
        }
    }

    fn function_name(&self) -> &str {
        &self.function_name
    }

    fn instructions(&self) -> &'a [Instruction] {
        self.instructions
    }

    fn ip(&self) -> usize {
        self.ip
    }

    fn increment_ip(&mut self) {
        self.ip += 1;
    }

    fn stack(&self) -> &[Value] {
        &self.stack
    }

    fn stack_mut(&mut self) -> &mut Vec<Value> {
        &mut self.stack
    }

    fn fp(&self) -> usize {
        self.fp
    }

    fn set_fp(&mut self, fp: usize) {
        self.fp = fp;
    }

    fn return_value(&self) -> Option<&Value> {
        self.return_value.as_ref()
    }

    fn set_return_value(&mut self, return_value: Value) {
        self.return_value = Some(return_value);
    }

    fn take_return_value(&mut self) -> Option<Value> {
        self.return_value.take()
    }
}

pub struct CallStack<'a> {
    call_frames: Vec<CallFrame<'a>>,
}

impl<'a> CallStack<'a> {
    pub fn new() -> Self {
        Self {
            call_frames: vec![],
        }
    }

    fn push(&mut self, call_frame: CallFrame<'a>) {
        self.call_frames.push(call_frame);
    }

    fn pop(&mut self) -> CallFrame<'a> {
        self.call_frames.pop().unwrap()
    }

    fn top(&self) -> &CallFrame<'a> {
        &self.call_frames[self.call_frames.len() - 1]
    }

    fn top_mut(&mut self) -> &mut CallFrame<'a> {
        let call_frames_len = self.call_frames.len();
        &mut self.call_frames[call_frames_len - 1] // dang borrow checker
    }

    fn maybe_top(&self) -> Option<&CallFrame<'a>> {
        self.call_frames.last()
    }

    fn maybe_top_mut(&mut self) -> Option<&mut CallFrame<'a>> {
        self.call_frames.last_mut()
    }

    fn penultimate(&self) -> Option<&CallFrame<'a>> {
        self.call_frames.get(self.call_frames.len() - 2)
    }

    fn penultimate_mut(&mut self) -> Option<&mut CallFrame<'a>> {
        let call_frames_len = self.call_frames.len();
        self.call_frames.get_mut(call_frames_len - 2) // dang borrow checker
    }
}

pub fn call<'a>(
    context: &'a DvmContext,
    registers: &mut Vec<Value>,
    call_stack: &mut CallStack<'a>,
    function_name: &str,
    arguments: Vec<Value>,
) -> Option<Value> {
    let function = context
        .functions
        .get(function_name)
        .expect(&format!("Function {} not found", function_name));

    // ensure enough registers
    registers.resize_with(function.register_count(), Default::default);

    // push call frame
    call_stack.push(CallFrame::new(
        function.name_owned(),
        function.instructions(),
    ));

    // put each arg on the stack
    for argument in arguments {
        call_stack.top_mut().stack_mut().push(argument);
    }

    while call_stack.maybe_top().is_some()
        && call_stack.top().ip() < call_stack.top().instructions().len()
    {
        let instruction = &call_stack.top().instructions()[call_stack.top().ip()];
        // println!("{}", instruction);

        match instruction {
            /* Move instructions */
            Instruction::MoveRegister(source, destination) => {
                // copy value from one register to another register
                let value = registers[*source].clone();
                registers[*destination] = value;
                call_stack.top_mut().increment_ip();
            }
            Instruction::MoveInt(value, destination) => {
                registers[*destination] = Value::Int(*value);
                call_stack.top_mut().increment_ip();
            }
            Instruction::MoveStackFrameOffset(offset, destination) => {
                // copy a value offset from the current frame pointer (fp) to a register
                let value_index =
                    call_stack
                        .top()
                        .fp()
                        .checked_add_signed(*offset)
                        .expect(&format!(
                            "Overflow when adding offset {} to fp {}",
                            *offset,
                            call_stack.top().fp()
                        ));
                let value = call_stack.top().stack()[value_index].clone();
                registers[*destination] = value;
                call_stack.top_mut().increment_ip();
            }

            /* Push instructions */
            Instruction::PushRegister(source) => {
                // copy a value from a register to the top of the stack
                let value = registers[*source].clone();
                call_stack.top_mut().stack_mut().push(value);
                call_stack.top_mut().increment_ip();
            }
            Instruction::PushInt(value) => {
                call_stack.top_mut().stack_mut().push(Value::Int(*value));
                call_stack.top_mut().increment_ip();
            }
            Instruction::PushStackFrameOffset(offset) => {
                // copy a value from somewhere on the stack to the top of the stack
                let value_index = call_stack
                    .top()
                    .fp()
                    .checked_add_signed(*offset)
                    .expect("Overflow when adding offset to fp");
                let value = call_stack.top().stack()[value_index].clone();
                call_stack.top_mut().stack_mut().push(value);
                call_stack.top_mut().increment_ip();
            }

            /* Invoke instructions */
            Instruction::InvokeStatic(function_name, arg_count) => {
                let function = context
                    .functions
                    .get(function_name)
                    .expect(&format!("Function {} not found", function_name));

                // get args
                let mut args =
                    call_stack.top().stack()[call_stack.top().stack().len() - arg_count..].to_vec();

                // increment caller frame ip
                call_stack.top_mut().increment_ip();

                // push a new call frame
                call_stack.push(CallFrame::new(
                    function.name_owned(),
                    function.instructions(),
                ));

                // push args
                call_stack.top_mut().stack_mut().append(&mut args);

                // set fp for callee frame
                let callee_frame = call_stack.top_mut();
                callee_frame.set_fp(*arg_count);
            }

            Instruction::InvokePlatformStatic(function_name, arg_count) => {
                let platform_function = context
                    .platform_functions()
                    .get(function_name)
                    .expect(&format!("Platform function {} not found", function_name));

                let args = &call_stack.top().stack()[call_stack.top().stack().len() - arg_count..];

                let result = platform_function(args);
                match result {
                    Ok(return_value) => {
                        // push return value to top of caller stack if it's not null
                        if !matches!(return_value, Value::Null) {
                            call_stack.top_mut().stack_mut().push(return_value);
                        }
                    }
                    Err(error) => {
                        // Eventually we will have some kind of exception handling
                        panic!("{}", error);
                    }
                }
                call_stack.top_mut().increment_ip();
            }

            /* Load constant instructions */
            Instruction::LoadStringConstant(constant_name, destination) => {
                let constant = &context.constants()[constant_name];
                match constant {
                    Constant::String(string_constant) => {
                        registers[*destination] = Value::String(string_constant.content_owned());
                    }
                }
                call_stack.top_mut().increment_ip();
            }

            /* Add instructions */
            Instruction::AddIntInt(lhs, rhs, destination) => {
                let lhs_value = registers[*lhs].unwrap_int();
                let rhs_value = registers[*rhs].unwrap_int();
                let result = lhs_value + rhs_value;
                registers[*destination] = Value::Int(result);
                call_stack.top_mut().increment_ip();
            }
            Instruction::AddStringInt(lhs, rhs, destination) => {
                let lhs_value = registers[*lhs].unwrap_string();
                let rhs_value = registers[*rhs].unwrap_int();
                let formatted = format!("{}{}", lhs_value, rhs_value);
                registers[*destination] = Value::String(Rc::from(formatted));
                call_stack.top_mut().increment_ip();
            }
            Instruction::AddStringString(lhs, rhs, destination) => {
                let lhs_value = registers[*lhs].unwrap_string();
                let rhs_value = registers[*rhs].unwrap_string();
                let formatted = format!("{}{}", lhs_value, rhs_value);
                registers[*destination] = Value::String(Rc::from(formatted));
                call_stack.top_mut().increment_ip();
            }

            /* Pop instructions */
            Instruction::Pop(maybe_register) => {
                let value = call_stack.top_mut().stack_mut().pop().unwrap();
                if let Some(register) = maybe_register {
                    registers[*register] = value;
                }
                call_stack.top_mut().increment_ip();
            }

            /* Return instructions */
            Instruction::SetReturnValue(register) => {
                call_stack
                    .top_mut()
                    .set_return_value(registers[*register].clone());
                call_stack.top_mut().increment_ip();
            }

            Instruction::Return(caller_pop_count) => {
                call_stack.pop(); // pop this frame
                let maybe_caller = call_stack.maybe_top_mut();
                if let Some(caller) = maybe_caller {
                    for _ in 0..*caller_pop_count {
                        caller.stack_mut().pop();
                    }
                    if let Some(return_value) = caller.take_return_value() {
                        caller.stack_mut().push(return_value);
                    }
                }
                // do not increment ip of the caller; this was done above BEFORE the call
            }
        }

        // if let Some(top) = call_stack.maybe_top() {
        //     println!("  stack: {:?}", top.stack());
        //     println!("  registers: {:?}", registers);
        //     println!("  rv: {:?}", top.return_value());
        // }
    }
    None // todo: returning results from main functions
}