nixjit/src/vm/vm.rs

use anyhow::{Result, anyhow};

use crate::builtins::env;
use crate::bytecode::{self, BinOp, OpCode, OpCodes, Program, Thunks, UnOp};
use crate::ty::common::Symbol;
use crate::ty::internal::*;
use crate::ty::public as p;

use super::env::Env;
use super::stack::{STACK_SIZE, Stack};
use super::vmthunk::*;

pub fn run(prog: Program) -> Result<p::Value> {
    let vm = VM::new(prog.thunks);
    Ok(vm.eval(prog.top_level, &mut env())?.to_public(&vm))
}

pub struct VM {
    thunks: Box<[VmThunk]>,
}

impl VM {
    fn new(thunks: Thunks) -> Self {
        let thunks = thunks
            .into_iter()
            .map(|bytecode::Thunk { opcodes }| VmThunk::new(opcodes))
            .collect();
        VM { thunks }
    }

    pub fn get_thunk_value(&self, idx: usize, env: &mut Env) -> Result<Value> {
        self.thunks.get(idx).unwrap().force(self, env)
    }

    pub fn eval(&self, opcodes: OpCodes, env: &mut Env) -> Result<Value> {
        let mut stack = Stack::<STACK_SIZE>::new();
        let mut iter = opcodes.into_iter();
        while let Some(opcode) = iter.next() {
            let jmp = self.single_op(opcode, &mut stack, env)?;
            for _ in 0..jmp {
                iter.next().unwrap();
            }
        }
        assert_eq!(stack.len(), 1);
        stack.pop()
    }

    #[inline]
    fn single_op<const CAP: usize>(
        &self,
        opcode: OpCode,
        stack: &mut Stack<CAP>,
        env: &mut Env,
    ) -> Result<usize> {
        match opcode {
            OpCode::NoOp => (),
            OpCode::Const { value } => stack.push(Value::Const(value))?,
            OpCode::LoadThunk { idx } => stack.push(Value::Thunk(Thunk::new(idx)))?,
            OpCode::LoadValue { idx } => {
                stack.push(self.get_thunk_value(idx, env)?)?;
            }
            OpCode::ForceValue => {
                stack.tos_mut()?.force(self, env)?;
            }
            OpCode::Jmp { step } => return Ok(step),
            OpCode::JmpIfTrue { step } => {
                if let Value::Const(Const::Bool(true)) = stack.pop()? {
                    return Ok(step);
                }
            }
            OpCode::JmpIfFalse { step } => {
                if let Value::Const(Const::Bool(false)) = stack.pop()? {
                    return Ok(step);
                }
            }
            OpCode::Call { arity } => {
                let mut args = Vec::with_capacity(arity);
                for _ in 0..arity {
                    args.insert(0, stack.pop()?);
                }
                let func = stack.pop()?;
                stack.push(func.call(self, args))?;
            }
            OpCode::Func { idx } => {
                stack.push(Value::Func(Func::new(env.locked(), self.thunks[idx].unwrap_code())))?;
            }
            OpCode::PushIdentParam { sym } => {
                stack
                    .tos_mut()?
                    .as_mut()
                    .unwrap_func()
                    .push_ident_param(sym);
            }
            OpCode::PushFormalParam { sym } => {
                stack
                    .tos_mut()?
                    .as_mut()
                    .unwrap_func()
                    .push_formal_param(sym);
            }
            OpCode::PushDefaultParam { idx } => {
                stack
                    .tos_mut()?
                    .as_mut()
                    .unwrap_func()
                    .push_default_param(idx);
            }
            OpCode::SetEllipsis => {
                stack.tos_mut()?.as_mut().unwrap_func().set_ellipsis();
            }
            OpCode::SetAlias { sym } => {
                stack.tos_mut()?.as_mut().unwrap_func().set_alias(sym);
            }
            OpCode::Ret => {
                todo!()
            }
            OpCode::UnOp { op } => {
                use UnOp::*;
                let value = stack.pop()?;
                stack.push(match op {
                    Not => value.not(),
                })?;
            }
            OpCode::BinOp { op } => {
                use BinOp::*;
                let rhs = stack.pop()?;
                let lhs = stack.pop()?;
                stack.push(match op {
                    Add => lhs.add(rhs),
                    And => lhs.and(rhs),
                    Or => lhs.or(rhs),
                    Eq => lhs.eq(rhs),
                    Con => lhs.concat(rhs),
                    Upd => lhs.update(rhs),
                })?;
            }
            OpCode::ConcatString => {
                let rhs = stack.pop()?;
                stack.tos_mut()?.concat_string(rhs);
            }
            OpCode::Path => {
                todo!()
            }
            OpCode::List => {
                stack.push(Value::List(List::empty()))?;
            }
            OpCode::PushElem => {
                let elem = stack.pop()?;
                stack.tos_mut()?.push(elem);
            }
            OpCode::AttrSet => {
                stack.push(Value::AttrSet(AttrSet::empty()))?;
            }
            OpCode::PushStaticAttr { name } => {
                let val = stack.pop()?;
                stack.tos_mut()?.push_attr(Symbol::new(name), val);
            }
            OpCode::PushDynamicAttr => {
                let val = stack.pop()?;
                let mut sym = stack.pop().unwrap();
                sym.coerce_to_string();
                let sym = sym.unwrap_const().unwrap_string().into();
                stack.tos_mut()?.push_attr(sym, val);
            }
            OpCode::Select { sym } => {
                stack.tos_mut()?.select(Symbol::new(sym)).force(self, env)?;
            }
            OpCode::SelectOrDefault { sym } => {
                let default = stack.pop()?;
                stack
                    .tos_mut()?
                    .select_with_default(Symbol::new(sym), default);
            }
            OpCode::SelectDynamic => {
                let mut val = stack.pop().unwrap();
                val.coerce_to_string();
                let sym = val.unwrap_const().unwrap_string().into();
                stack.tos_mut()?.select(sym);
            }
            OpCode::SelectDynamicOrDefault => {
                let default = stack.pop()?;
                let mut val = stack.pop().unwrap();
                val.coerce_to_string();
                let sym = val.unwrap_const().unwrap_string().into();
                stack.tos_mut()?.select_with_default(sym, default);
            }
            OpCode::HasAttr { sym } => {
                stack.tos_mut()?.has_attr(Symbol::new(sym));
            }
            OpCode::HasDynamicAttr => {
                let mut val = stack.pop().unwrap();
                val.coerce_to_string();
                let sym = val.unwrap_const().unwrap_string().into();
                stack.tos_mut()?.has_attr(sym);
            }
            OpCode::LookUp { sym } => {
                stack.push(
                    env.lookup(Symbol::new(sym.clone()))
                        .ok_or(anyhow!(r#""{sym}" not found"#))?,
                )?;
            }
            OpCode::EnterEnv => {
                env.enter(stack.pop()?.unwrap_attr_set().to_data());
            }
            OpCode::LeaveEnv => {
                env.leave();
            }
            OpCode::Assert => {
                if !stack.pop()?.unwrap_const().unwrap_bool() {
                    todo!()
                }
            }
        }
        Ok(0)
    }
}