nek-lang/src/vm.rs

use crate::{bytecode::op::*, interpreter::Value};

#[derive(Debug, Default)]
pub struct Vm {
    prog: Vec<u32>,
    ip: usize,
    stack: Vec<Value>,

    /// This isn't actually a heap. It's actually still more of a f*cked up stack
    heap: Vec<Value>,
}

macro_rules! binop_stack {
    ($self:ident, $op:tt) => {
        {
            let rhs = $self.stack.pop().unwrap();
            let lhs = $self.stack.last_mut().unwrap();
            match (lhs, rhs) {
                (Value::I64(lhs), Value::I64(rhs)) => *lhs = *lhs $op rhs,
                _ => panic!("Invalid data for add"),
            }
        }
    };
}

impl Vm {
    pub fn new(prog: Vec<u32>) -> Self {
        Self {
            prog,
            ..Default::default()
        }
    }

    pub fn run(&mut self) {
        while let Some(op) = self.prog.get(self.ip).copied() {
            self.ip += 1;

            match op & 0xff {
                PUSH => {
                    let val = self.read_i64();
                    self.stack.push(Value::I64(val));
                }
                POP => {
                    self.stack.pop();
                }
                LOAD => {
                    // let addr = self.read_i64() as usize;
                    let addr = (op >> 8) as usize;

                    if let Some(val) = self.heap.get(addr) {
                        self.stack.push(val.clone());
                    } else {
                        panic!("Trying to load from uninitialized heap");
                    }
                }
                STORE => {
                    let val = self
                        .stack
                        .pop()
                        .expect("Trying to pop value from stack for storing");
                    // let addr = self.read_i64() as usize;
                    let addr = (op >> 8) as usize;

                    if self.heap.len() == addr {
                        self.heap.push(val);
                    } else {
                        self.heap[addr] = val;
                    }
                }
                PRINT => {
                    let val = self
                        .stack
                        .pop()
                        .expect("Trying to pop value from stack for printing");
                    print!("{}", val);
                }
                DBG_PRINT => {
                    let val = self
                        .stack
                        .pop()
                        .expect("Trying to pop value from stack for printing");
                    print!("{:?}", val);
                }

                ADD => {
                    binop_stack!(self, +);
                    // self.stack.push(Value::I64(vals.0 + vals.1))
                }
                SUB => {
                    binop_stack!(self, -);
                    // let vals = self.pop2_i64();
                    // self.stack.push(Value::I64(vals.0 - vals.1))
                }
                MUL => {
                    binop_stack!(self, *);
                    // let vals = self.pop2_i64();
                    // self.stack.push(Value::I64(vals.0 * vals.1))
                }
                DIV => {
                    binop_stack!(self, /);
                    // let vals = self.pop2_i64();
                    // self.stack.push(Value::I64(vals.0 / vals.1))
                }
                MOD => {
                    binop_stack!(self, %);
                    // let vals = self.pop2_i64();
                    // self.stack.push(Value::I64(vals.0 % vals.1))
                }
                EQ => {
                    let vals = self.pop2_i64();
                    self.stack
                        .push(Value::I64(if vals.0 == vals.1 { 1 } else { 0 }))
                }
                NEQ => {
                    let vals = self.pop2_i64();
                    self.stack
                        .push(Value::I64(if vals.0 != vals.1 { 1 } else { 0 }))
                }
                GT => {
                    let vals = self.pop2_i64();
                    self.stack
                        .push(Value::I64(if vals.0 > vals.1 { 1 } else { 0 }))
                }
                GE => {
                    let vals = self.pop2_i64();
                    self.stack
                        .push(Value::I64(if vals.0 >= vals.1 { 1 } else { 0 }))
                }
                LT => {
                    let vals = self.pop2_i64();
                    self.stack
                        .push(Value::I64(if vals.0 < vals.1 { 1 } else { 0 }))
                }
                LE => {
                    let vals = self.pop2_i64();
                    self.stack
                        .push(Value::I64(if vals.0 <= vals.1 { 1 } else { 0 }))
                }
                BOR => {
                    let vals = self.pop2_i64();
                    self.stack.push(Value::I64(vals.0 | vals.1))
                }
                BAND => {
                    let vals = self.pop2_i64();
                    self.stack.push(Value::I64(vals.0 & vals.1))
                }
                BXOR => {
                    let vals = self.pop2_i64();
                    self.stack.push(Value::I64(vals.0 ^ vals.1))
                }
                SHL => {
                    let vals = self.pop2_i64();
                    self.stack.push(Value::I64(vals.0 << vals.1))
                }
                SHR => {
                    let vals = self.pop2_i64();
                    self.stack.push(Value::I64(vals.0 >> vals.1))
                }
                JUMP => {
                    self.ip = self.read_i64() as usize;
                }
                JUMP_TRUE => {
                    let jmp_target = self.read_i64() as usize;
                    if !matches!(self.stack.pop(), Some(Value::I64(0))) {
                        self.ip = jmp_target;
                    }
                }
                JUMP_FALSE => {
                    let jmp_target = self.read_i64() as usize;
                    if matches!(self.stack.pop(), Some(Value::I64(0))) {
                        self.ip = jmp_target;
                    }
                }
                _ => panic!("Invalid opcode")
            }
        }
    }

    fn pop2_i64(&mut self) -> (i64, i64) {
        let rhs = self.stack.pop();
        let lhs = self.stack.pop();
        match (lhs, rhs) {
            (Some(Value::I64(lhs)), Some(Value::I64(rhs))) => (lhs, rhs),
            _ => panic!("Invalid data for add"),
        }
    }

    fn read_i64(&mut self) -> i64 {
        let mut val = if let Some(val) = self.prog.get(self.ip).copied() {
            val
        } else {
            panic!("Expected Value as next OP")
        } as i64;

        self.ip += 1;

        val |= (if let Some(val) = self.prog.get(self.ip).copied() {
            val
        } else {
            panic!("Expected Value as next OP")
        } as i64)
            << 32;

        self.ip += 1;

        val
    }
}