not-python-rust/src/vm.rs

use std::collections::HashMap;
use std::rc::Rc;

use crate::obj::function::*;
use crate::obj::*;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Op {
    // Stack functions
    Pop,
    PushConstant(LongOpArg),

    // Variables
    GetLocal(LocalIndex),
    SetLocal(LocalIndex),
    GetGlobal(GlobalId),
    SetGlobal(GlobalId),

    // Attributes
    GetAttr(ConstantId),
    SetAttr(ConstantId),

    // Jumps
    Jump(JumpOpArg),
    JumpFalse(JumpOpArg),
    JumpTrue(JumpOpArg),

    // Functions
    Call(Argc),
    Return,
    CloseOver { depth: ShortOpArg, slot: ShortOpArg },

    // VM control
    Nop,
    Halt,
}

pub type LineRange = (usize, usize);

type ShortOpArg = u16;
type LongOpArg = u32;

pub type JumpOpArg = i32;
pub type LocalIndex = LongOpArg;
pub type LocalSlot = ShortOpArg;
pub type ConstantId = LongOpArg;
pub type GlobalId = LongOpArg;
pub type Argc = LongOpArg;
pub type FrameDepth = ShortOpArg;

#[derive(Debug, Clone)]
pub struct Local {
    pub(crate) slot: LocalSlot,
    pub(crate) index: LocalIndex,
    pub(crate) name: String,
}

#[derive(Debug, Default, Clone)]
pub struct Chunk {
    pub(crate) code: Vec<Op>,
    pub(crate) lines: Vec<LineRange>,
    pub(crate) locals: Vec<Local>,
}

#[derive(Debug)]
pub(crate) enum Function {
    Chunk(Rc<Chunk>),
    Builtin(BuiltinFunctionPtr, FunctionState),
}

#[derive(Debug)]
pub struct Frame {
    pub(crate) name: Rc<String>,
    pub(crate) function: Function,
    pub(crate) ip: usize,
    pub(crate) stack_base: usize,
}

impl Frame {
    pub fn new(name: Rc<String>, function: Function, stack_base: usize) -> Self {
        Self {
            name,
            function,
            ip: 0,
            stack_base,
        }
    }
}

pub struct Vm {
    constants: Vec<ObjP>,
    //global_names: Vec<String>,
    globals: Vec<ObjP>,
    stack: Vec<ObjP>,
    frames: Vec<Frame>,
    builtins: HashMap<String, ObjP>,
}

impl Vm {
    /// Create a new virtual machine with the given chunk, constants, and global names.
    pub fn new(
        chunk: Rc<Chunk>,
        constants: Vec<ObjP>,
        global_names: Vec<String>,
        builtins: HashMap<String, ObjP>,
    ) -> Self {
        // set up globals
        let nil = builtins.create_nil();
        let mut globals: Vec<_> = global_names.iter().map(|_| ObjP::clone(&nil)).collect();

        let mut register_global = |name: &str, value: ObjP| {
            let index = global_names
                .iter()
                .position(|global| global == name)
                .expect("could not find global");
            globals[index] = value;
        };

        for (name, builtin) in builtins.iter() {
            register_global(&name, ObjP::clone(&builtin));
        }

        // stack and frames
        let stack = Vec::new();
        let frames = vec![Frame::new(
            "__main__".to_string().into(),
            Function::Chunk(chunk),
            0,
        )];

        Vm {
            constants,
            //global_names,
            globals,
            stack,
            frames,
            builtins,
        }
    }

    /// Get the stack.
    pub fn stack(&self) -> &Vec<ObjP> {
        &self.stack
    }

    /// Get the stack, mutably.
    pub fn stack_mut(&mut self) -> &mut Vec<ObjP> {
        &mut self.stack
    }

    /// Gets the current stack, starting at the frame's stack base.
    pub fn frame_stack(&self) -> &[ObjP] {
        &self.stack()[self.frame().stack_base..]
    }

    /// Current stack frame.
    pub fn frame(&self) -> &Frame {
        self.frames.last().unwrap()
    }

    /// Current stack frame, mutably.
    pub fn frame_mut(&mut self) -> &mut Frame {
        self.frames.last_mut().unwrap()
    }

    /// Push a new stack frame.
    pub fn push_frame(&mut self, frame: Frame) {
        self.frames.push(frame);
    }

    /// Pop the current stack frame.
    pub fn pop_frame(&mut self) -> Frame {
        self.frames.pop().expect("no frame")
    }

    /// Gets the chunk of the currently executing frame.
    pub fn chunk(&self) -> Option<&Chunk> {
        if let Function::Chunk(chunk) = &self.frame().function {
            Some(chunk)
        } else {
            None
        }
    }

    /// Instruction pointer of the current frame.
    pub fn ip(&self) -> usize {
        self.frame().ip
    }

    /// Update the current instruction pointer.
    pub fn set_ip(&mut self, ip: usize) {
        self.frame_mut().ip = ip;
    }

    /*
    /// Gets the line of the current instruction.
    fn line(&self, offset: isize) -> LineRange {
        let index = (((self.ip() as isize) + offset).max(0) as usize).min(self.chunk().lines.len());
        self.chunk().lines[index]
    }
    */

    /// Get the current instruction and advance the IP.
    ///
    /// This may actually end up calling a function on top of the stack, if it's a builtin function.
    fn dispatch(&mut self) -> Op {
        let mut ip = self.ip();
        let op = match &self.frame().function {
            Function::Chunk(chunk) => {
                let op = chunk.code[ip];
                ip += 1;
                op
            }
            Function::Builtin(function, state) => {
                // keep track of where the current frame index is in case we need to yield
                let frame_index = self.frames.len() - 1;
                let result = (function)(self, *state);
                match result {
                    FunctionResult::ReturnPush(value) => {
                        // push value to the stack and let the VM handle return protocols
                        self.push(value);
                        Op::Return
                    }
                    // value is already on top of the stack, let the VM handle return protocols
                    FunctionResult::Return => Op::Return,
                    // new stack frame has been pushed, yield control while keeping track of the
                    // old state
                    FunctionResult::Yield(resume_state) => {
                        // update the current state
                        if let Function::Builtin(_function, ref mut state) =
                            &mut self.frames[frame_index].function
                        {
                            *state = FunctionState::Resume(resume_state)
                        } else {
                            panic!("function stack got really messed up - function stack was changed under us");
                        }
                        // inject a no-op so the VM can load a new instruction or dispatch a new
                        // instruction.
                        Op::Nop
                    }
                }
            }
        };
        self.set_ip(ip);
        op
    }

    /// Pop a value from the stack.
    pub fn pop(&mut self) -> ObjP {
        self.stack.pop().expect("stack empty")
    }

    /// Peek the top value of the stack.
    pub fn peek(&self) -> ObjP {
        self.stack.last().map(Ptr::clone).expect("stack empty")
    }

    /// Push a value to the stack.
    pub fn push(&mut self, value: ObjP) {
        self.stack.push(value);
    }

    pub fn run(&mut self) {
        // Cached pointers that we just always want to have on hand
        let method_type = self.builtins().get("Method").unwrap().clone();

        loop {
            match self.dispatch() {
                Op::Pop => {
                    self.pop();
                }
                Op::PushConstant(constant_id) => {
                    let constant = Ptr::clone(&self.constants[constant_id as usize]);
                    self.push(constant);
                }
                Op::GetLocal(local_index) => {
                    let local = &self.chunk().expect("no chunk").locals[local_index as usize];
                    let value =
                        Ptr::clone(&self.stack[self.frame().stack_base + local.slot as usize]);
                    self.push(value);
                }
                Op::SetLocal(local_index) => {
                    let value = self.pop();
                    let local = &self.chunk().expect("no chunk").locals[local_index as usize];
                    let index = self.frame().stack_base + local.slot as usize;
                    self.stack[index] = value;
                }
                Op::GetGlobal(global_index) => {
                    let value = Ptr::clone(&self.globals[global_index as usize]);
                    self.push(value);
                }
                Op::SetGlobal(global_index) => {
                    let value = self.pop();
                    self.globals[global_index as usize] = value;
                }
                Op::GetAttr(constant_id) => {
                    // need both declarations to borrow cell value
                    let name_obj = Ptr::clone(&self.constants[constant_id as usize]);
                    let name =
                        with_obj_downcast(name_obj, |name: &StrInst| Rc::clone(&name.str_value()));
                    let owner = self.pop();
                    let value =
                        owner
                            .borrow_mut()
                            .get_attr_lazy(owner.clone(), method_type.clone(), &name);
                    if let Some(value) = value {
                        self.push(value);
                    } else {
                        // TODO Vm::run, Op::GetAttr - throw an exception when the attribute
                        //                             doesn't exist
                        // BLOCKED-ON: exceptions
                        todo!(
                            "throw an error because we couldn't read attr '{}' on '{}'",
                            name,
                            owner.borrow(),
                        );
                    }
                }
                Op::SetAttr(constant_id) => {
                    let name_obj = Ptr::clone(&self.constants[constant_id as usize]);
                    let name =
                        with_obj_downcast(name_obj, |name: &StrInst| Rc::clone(&name.str_value()));
                    let value = self.pop();
                    let target = self.pop();

                    target.borrow_mut().set_attr(&name, value);
                }
                Op::Jump(offset) => {
                    let base = (self.ip() - 1) as JumpOpArg;
                    assert!(base + offset > 0, "tried to jump to negative IP");
                    self.set_ip((base + offset) as usize);
                }
                Op::JumpFalse(offset) => {
                    let base = (self.ip() - 1) as JumpOpArg;
                    let value = self.peek();
                    if !value.borrow().is_truthy() {
                        self.set_ip((base + offset) as usize);
                    }
                }
                Op::JumpTrue(offset) => {
                    let base = (self.ip() - 1) as JumpOpArg;
                    let value = self.peek();
                    if value.borrow().is_truthy() {
                        self.set_ip((base + offset) as usize);
                    }
                }
                Op::Call(argc) => {
                    let argc = argc;
                    let index = self.stack.len() - (argc as usize) - 1;
                    let fun_ptr = Ptr::clone(&self.stack[index]);

                    let arity = if let Some(arity) = fun_ptr.borrow().arity() {
                        arity
                    } else {
                        // TODO Vm::run, Op::Call - throw an exception when the value isn't
                        //                          callable
                        // BLOCKED-ON: exceptions
                        todo!(
                            "throw an error because we couldn't call {}",
                            fun_ptr.borrow()
                        );
                    };

                    // Check arity
                    if arity != argc {
                        // TODO Vm::run, Op::Call - throw an exception when the number of arguments
                        //                          does not match the function's arity
                        // BLOCKED-ON: exceptions
                        todo!(
                            "throw an error because we passed the wrong number of arguments to {}",
                            fun_ptr.borrow()
                        );
                    }
                    fun_ptr.borrow().call(self, argc as Argc);
                }
                Op::Return => {
                    let return_value = self.pop();
                    let old_frame = self.frames.pop().unwrap();
                    // stack_base is always going to be <= current stack size
                    self.stack
                        .resize_with(old_frame.stack_base, || unreachable!());
                    // also pop the function object off of the stack
                    self.stack.pop();
                    self.push(return_value);
                }
                Op::CloseOver { depth, slot } => {
                    // since we're closing over a value, and functions ultimately come from
                    // constants, we want to deep-clone this object so we don't alter any live
                    // objects.
                    // there is some room for optimization here so we aren't cloning the entire
                    // UserFunctionInst for every individual capture in a function.
                    let fun_ptr = self.pop();
                    let mut fun: UserFunctionInst =
                        with_obj_downcast(fun_ptr, UserFunctionInst::clone);
                    let frame_index = self.frames.len() - (depth as usize) - 1;
                    let stack_base = self.frames[frame_index].stack_base;
                    let value = Ptr::clone(&self.stack[stack_base + (slot as usize)]);
                    fun.push_capture(value);
                    self.push(make_ptr(fun));
                }
                Op::Nop => {
                    continue;
                }
                Op::Halt => {
                    break;
                }
            }
        }
    }
}

impl ObjFactory for Vm {
    fn builtins(&self) -> &HashMap<String, ObjP> {
        &self.builtins
    }
}