nixjit/src/vm/mod.rs

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::pin::Pin;

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

use crate::stack::Stack;

use derive_more::Constructor;
use ecow::EcoString;
pub use env::Env;
pub use jit::JITContext;

mod env;
mod jit;

#[cfg(test)]
mod test;

pub const STACK_SIZE: usize = 8 * 1024 / size_of::<Value>();

pub fn run(prog: Program, jit: Pin<Box<JITContext<'_>>>) -> Result<p::Value> {
    let vm = VM::new(
        prog.thunks,
        prog.funcs,
        RefCell::new(prog.symbols),
        RefCell::new(prog.symmap),
        prog.consts,
        jit,
    );
    let env = env(&vm);
    let mut seen = HashSet::new();
    let value = vm
        .eval(prog.top_level.into_iter(), env)?
        .to_public(&vm, &mut seen);
    Ok(value)
}

#[derive(Constructor)]
pub struct VM<'jit> {
    thunks: Box<[OpCodes]>,
    funcs: Box<[F]>,
    symbols: RefCell<Vec<EcoString>>,
    symmap: RefCell<HashMap<EcoString, usize>>,
    consts: Box<[Const]>,
    jit: Pin<Box<JITContext<'jit>>>,
}

impl<'vm, 'jit: 'vm> VM<'jit> {
    pub fn get_thunk(&self, idx: usize) -> &OpCodes {
        &self.thunks[idx]
    }

    pub fn get_func(&self, idx: usize) -> &F {
        &self.funcs[idx]
    }

    pub fn get_sym(&self, idx: usize) -> Symbol {
        self.symbols.borrow()[idx].clone().into()
    }

    pub fn new_sym(&self, sym: impl Into<EcoString>) -> usize {
        let sym = sym.into();
        if let Some(&idx) = self.symmap.borrow().get(&sym) {
            idx
        } else {
            self.symmap
                .borrow_mut()
                .insert(sym.clone(), self.symbols.borrow().len());
            self.symbols.borrow_mut().push(sym);
            self.symbols.borrow().len() - 1
        }
    }

    pub fn eval(
        &'vm self,
        opcodes: impl Iterator<Item = OpCode>,
        mut env: Env<'vm>,
    ) -> Result<Value<'vm>> {
        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, &mut env)?;
            for _ in 0..jmp {
                iter.next().unwrap();
            }
        }
        assert_eq!(stack.len(), 1);
        let mut ret = stack.pop();
        ret.force(self)?;
        Ok(ret)
    }

    #[inline(always)]
    fn single_op<'s, const CAP: usize>(
        &'vm self,
        opcode: OpCode,
        stack: &'s mut Stack<Value<'vm>, CAP>,
        env: &mut Env<'vm>,
    ) -> Result<usize> {
        match opcode {
            OpCode::Illegal => panic!("illegal opcode"),
            OpCode::Const { idx } => stack.push(Value::Const(self.consts[idx].clone()))?,
            OpCode::LoadThunk { idx } => {
                stack.push(Value::Thunk(Thunk::new(self.get_thunk(idx)).into()))?
            }
            OpCode::CaptureEnv => match stack.tos()? {
                Value::Thunk(thunk) => thunk.capture(env.clone()),
                _ => (),
            },
            OpCode::ForceValue => {
                stack.tos_mut()?.force(self)?;
            }
            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 mut func = stack.pop();
                func.force(self)?;
                stack.push(func.call(self, args)?)?;
            }
            OpCode::Func { idx } => {
                let func = self.get_func(idx);
                stack.push(Value::Func(Func::new(func, env.clone(), None).into()))?;
            }
            OpCode::UnOp { op } => {
                use UnOp::*;
                let mut value = stack.pop();
                value.force(self)?;
                stack.push(match op {
                    Neg => value.neg(),
                    Not => value.not(),
                })?;
            }
            OpCode::BinOp { op } => {
                use BinOp::*;
                let mut rhs = stack.pop();
                let mut lhs = stack.pop();
                lhs.force(self)?;
                rhs.force(self)?;
                stack.push(match op {
                    Add => lhs.add(rhs),
                    Sub => lhs.add(rhs.neg()),
                    Mul => lhs.mul(rhs),
                    Div => lhs.div(rhs)?,
                    And => lhs.and(rhs),
                    Or => lhs.or(rhs),
                    Eq => lhs.eq(rhs),
                    Lt => lhs.lt(rhs),
                    Con => lhs.concat(rhs),
                    Upd => lhs.update(rhs),
                })?;
            }
            OpCode::ConcatString => {
                let mut rhs = stack.pop();
                rhs.force(self)?;
                stack.tos_mut()?.concat_string(rhs);
            }
            OpCode::Path => {
                todo!()
            }
            OpCode::List => {
                stack.push(Value::List(List::empty().into()))?;
            }
            OpCode::PushElem => {
                let elem = stack.pop();
                stack.tos_mut()?.push(elem);
            }
            OpCode::AttrSet => {
                stack.push(Value::AttrSet(AttrSet::empty().into()))?;
            }
            OpCode::FinalizeRec => {
                let env = env.clone().enter(
                    stack
                        .tos()?
                        .clone()
                        .unwrap_attr_set()
                        .as_inner()
                        .iter()
                        .map(|(k, v)| (k.clone(), v.clone())),
                );
                stack.tos_mut()?.as_mut().unwrap_attr_set().capture(&env);
            }
            OpCode::PushStaticAttr { name } => {
                let val = stack.pop();
                stack.tos_mut()?.push_attr(name, val);
            }
            OpCode::PushDynamicAttr => {
                let val = stack.pop();
                let mut sym = stack.pop();
                sym.force(self)?.coerce_to_string();
                let sym = self.new_sym(sym.unwrap_const().unwrap_string());
                stack.tos_mut()?.push_attr(sym, val);
            }
            OpCode::Select { sym } => {
                stack.tos_mut()?.force(self)?.select(sym, self)?;
            }
            OpCode::SelectOrDefault { sym } => {
                let default = stack.pop();
                stack
                    .tos_mut()?
                    .force(self)?
                    .select_with_default(sym, default)?;
            }
            OpCode::SelectDynamic => {
                let mut val = stack.pop();
                val.force(self)?;
                val.coerce_to_string();
                let sym = self.new_sym(val.unwrap_const().unwrap_string());
                stack.tos_mut()?.force(self)?.select(sym, self)?;
            }
            OpCode::SelectDynamicOrDefault => {
                let default = stack.pop();
                let mut val = stack.pop();
                val.force(self)?;
                val.coerce_to_string();
                let sym = self.new_sym(val.unwrap_const().unwrap_string());
                stack
                    .tos_mut()?
                    .force(self)?
                    .select_with_default(sym, default)?;
            }
            OpCode::HasAttr { sym } => {
                stack.tos_mut()?.force(self)?.has_attr(sym);
            }
            OpCode::HasDynamicAttr => {
                let mut val = stack.pop();
                val.coerce_to_string();
                let sym = self.new_sym(val.unwrap_const().unwrap_string());
                stack.tos_mut()?.force(self)?.has_attr(sym);
            }
            OpCode::LookUp { sym } => {
                stack.push(env.lookup(sym).ok_or_else(|| {
                    Error::EvalError(format!("{} not found", self.get_sym(sym)))
                })?)?;
            }
            OpCode::EnterEnv => match stack.pop() {
                Value::AttrSet(attrs) => *env = env.clone().enter_with(attrs),

                _ => unreachable!(),
            },
            OpCode::LeaveEnv => *env = env.clone().leave(),
            OpCode::Assert => {
                if !stack.pop().unwrap_const().unwrap_bool() {
                    todo!()
                }
            }
        }
        Ok(0)
    }
}