nixjit/src/vm/mod.rs

use std::cell::RefCell;
use std::rc::Rc;

use hashbrown::{HashMap, HashSet};
use inkwell::context::Context;

use crate::bytecode::{BinOp, Func as F, OpCode, OpCodes, Program, UnOp};
use crate::env::VmEnv;
use crate::error::*;
use crate::jit::{JITContext, JITFunc};
use crate::stack::Stack;
use crate::ty::common::Const;
use crate::ty::internal::*;
use crate::ty::public::{self as p, Symbol};

use derive_more::Constructor;
use ecow::EcoString;

#[cfg(test)]
mod test;

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

struct ContextWrapper(Context);

pub fn run(mut prog: Program) -> Result<p::Value> {
    let jit = Context::create();
    prog.thunks.iter_mut().for_each(|code| code.reverse());
    prog.funcs
        .iter_mut()
        .for_each(|F { opcodes, .. }| opcodes.reverse());
    let vm = VM {
        thunks: prog.thunks,
        funcs: prog.funcs,
        consts: prog.consts,
        symbols: RefCell::new(prog.symbols),
        symmap: RefCell::new(prog.symmap),
        jit: JITContext::new(&jit),
    };
    prog.top_level.reverse();
    Ok(eval(prog.top_level, &vm, VmEnv::new(vec![]))?.to_public(&vm, &mut HashSet::new()))
}

pub fn eval<'gc>(opcodes: Box<[OpCode]>, vm: &VM<'gc>, env: Rc<VmEnv<'gc>>) -> Result<Value<'gc>> {
    let mut opcodes = opcodes.into_vec();
    let mut envs = vec![env];
    let mut stack = Stack::<_, STACK_SIZE>::new();
    while let Some(opcode) = opcodes.pop() {
        let consq = single_op(vm, opcode, &mut stack, envs.last_mut().unwrap())?;
        match consq {
            Consq::NoOp => (),
            Consq::Jmp(step) => opcodes.resize_with(opcodes.len() - step, || unreachable!()),
            Consq::Force => {
                let thunk = stack.tos().as_ref().unwrap_thunk();
                let (code, env) = thunk.suspend()?;
                opcodes.push(OpCode::InsertValue);
                opcodes.push(OpCode::PopEnv);
                opcodes.extend(code);
                envs.push(env);
            }
            Consq::PopEnv => {
                let _ = envs.pop().unwrap();
            }
            Consq::Call => {
                let arg = stack.pop();
                let func = stack.pop().unwrap_func();
                let env = func.env.clone().enter_arg(arg);
                let count = func.count.get();
                func.count.set(count + 1);
                if count > 1 {
                    let compiled = func.compiled.get_or_init(|| vm.compile_func(func.func));
                    let ret = unsafe { compiled(env.as_ref() as *const VmEnv) };
                    stack.push(ret.into())?;
                } else {
                    envs.push(env);
                    opcodes.push(OpCode::PopEnv);
                    opcodes.extend(&func.func.opcodes);
                }
            }
        }
    }
    stack.pop().ok()
}

enum Consq {
    Jmp(usize),
    Call,
    Force,
    PopEnv,
    NoOp,
}

#[inline(always)]
fn single_op<'gc, const CAP: usize>(
    vm: &VM<'gc>,
    opcode: OpCode,
    stack: &mut Stack<Value<'gc>, CAP>,
    env: &mut Rc<VmEnv<'gc>>,
) -> Result<Consq> {
    match opcode {
        OpCode::Illegal => panic!("illegal opcode"),
        OpCode::Const { idx } => stack.push(match vm.get_const(idx) {
            Const::Int(x) => Value::Int(x),
            Const::Float(x) => Value::Float(x),
            Const::Bool(x) => Value::Bool(x),
            Const::String(x) => Value::String(Rc::new(x.into())),
            Const::Null => Value::Null,
        })?,
        OpCode::LoadThunk { idx } => stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx))))?,
        OpCode::LoadValue { idx } => {
            stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx))))?;
            stack.tos().as_ref().unwrap_thunk().capture_env(env.clone());
            return Ok(Consq::Force);
        }
        OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture_env(env.clone()),
        OpCode::ForceValue => {
            if !stack.tos().is_thunk() {
                return Ok(Consq::NoOp);
            }
            let Some(val) = stack.tos().as_ref().unwrap_thunk().get_value() else {
                return Ok(Consq::Force);
            };
            *stack.tos_mut() = val.clone();
        }
        OpCode::InsertValue => {
            let val = stack.pop();
            stack
                .tos()
                .as_ref()
                .unwrap_thunk()
                .insert_value(val.clone());
            *stack.tos_mut() = val;
        }
        OpCode::Jmp { step } => return Ok(Consq::Jmp(step)),
        OpCode::JmpIfFalse { step } => {
            if let Value::Bool(false) = stack.pop() {
                return Ok(Consq::Jmp(step));
            }
        }
        OpCode::Call => {
            let arg = stack.pop();
            let func = stack.tos_mut();
            if func.is_func() {
                let _ = stack.push(arg);
                return Ok(Consq::Call);
            }
            func.call(arg, vm)?;
        }
        OpCode::Func { idx } => {
            let func = vm.get_func(idx);
            stack.push(Value::Func(Rc::new(Func::new(func, env.clone()))))?;
        }
        OpCode::Arg { level } => {
            stack.push(env.lookup_arg(level).clone())?;
        }
        OpCode::UnOp { op } => {
            use UnOp::*;
            let value = stack.tos_mut();
            match op {
                Neg => value.neg(),
                Not => value.not(),
            }
        }
        OpCode::BinOp { op } => {
            use BinOp::*;
            let mut rhs = stack.pop();
            let lhs = stack.tos_mut();
            match op {
                Add => lhs.add(rhs),
                Sub => {
                    rhs.neg();
                    lhs.add(rhs);
                }
                Mul => lhs.mul(rhs),
                Div => lhs.div(rhs)?,
                And => lhs.and(rhs),
                Or => lhs.or(rhs),
                Eq => Value::eq(lhs, rhs),
                Lt => lhs.lt(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 { cap } => {
            stack.push(Value::List(Rc::new(List::with_capacity(cap))))?;
        }
        OpCode::PushElem => {
            let elem = stack.pop();
            stack.tos_mut().push(elem);
        }
        OpCode::AttrSet { cap } => {
            stack.push(Value::AttrSet(Rc::new(AttrSet::with_capacity(cap))))?;
        }
        OpCode::FinalizeLet => {
            let mut list = stack.pop().unwrap_list();
            let map = list.as_ref().clone().into_inner();
            *env = env.clone().enter_let(map);
            Rc::make_mut(&mut list).capture(&Rc::downgrade(env));
        }
        OpCode::PushStaticAttr { name } => {
            let val = stack.pop();
            stack.tos_mut().push_attr(name, val);
        }
        OpCode::PushDynamicAttr => {
            let val = stack.pop();
            let sym = stack.pop();
            let sym = vm.new_sym::<&str>(&sym.unwrap_string());
            stack.tos_mut().push_attr(sym, val);
        }
        OpCode::Select { sym } => {
            stack.tos_mut().select(sym, vm)?;
        }
        OpCode::SelectOrDefault { sym } => {
            let default = stack.pop();
            stack.tos_mut().select_with_default(sym, default)?;
        }
        OpCode::SelectDynamic => {
            let mut val = stack.pop();
            val.coerce_to_string();
            let sym = vm.new_sym::<&str>(&val.unwrap_string());
            stack.tos_mut().select(sym, vm)?;
        }
        OpCode::SelectDynamicOrDefault => {
            let default = stack.pop();
            let mut val = stack.pop();
            val.coerce_to_string();
            let sym = vm.new_sym::<&str>(&val.unwrap_string());
            stack.tos_mut().select_with_default(sym, default)?;
        }
        OpCode::HasAttr { sym } => {
            stack.tos_mut().has_attr(sym);
        }
        OpCode::HasDynamicAttr => {
            let mut val = stack.pop();
            val.coerce_to_string();
            let sym = vm.new_sym::<&str>(&val.unwrap_string());
            stack.tos_mut().has_attr(sym);
        }
        OpCode::LookUp { sym } => {
            stack.push(
                env.lookup_with(&sym)
                    .ok_or_else(|| Error::EvalError(format!("{} not found", vm.get_sym(sym))))?
                    .clone(),
            )?;
        }
        OpCode::LookUpLet { level, idx } => {
            let val = env.lookup_let(level, idx);
            if let Value::Thunk(thunk) = val {
                thunk.upgrade();
            }
            stack.push(val.clone())?;
        }
        OpCode::LeaveEnv => *env = env.leave(),
        OpCode::EnterWithEnv => {
            let mut new = HashMap::new();
            stack
                .pop()
                .unwrap_attr_set()
                .as_inner()
                .iter()
                .map(|(&k, v)| (k, v.clone()))
                .collect_into(&mut new);
            *env = env.clone().enter_with(new.into());
        }
        OpCode::PopEnv => return Ok(Consq::PopEnv),
        OpCode::Assert => {
            if !stack.pop().unwrap_bool() {
                todo!()
            }
        }
    }
    Ok(Consq::NoOp)
}

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

impl<'gc> VM<'gc> {
    pub fn get_thunk(&self, idx: usize) -> &'gc OpCodes {
        unsafe { &*(&self.thunks[idx] as *const _) }
    }

    pub fn get_func(&self, idx: usize) -> &'gc F {
        unsafe { &*(&self.funcs[idx] as *const _) }
    }

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

    pub fn new_sym<T: Into<EcoString>>(&self, sym: T) -> 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 get_const(&self, idx: usize) -> Const {
        self.consts[idx].clone()
    }

    pub fn compile_func(&self, func: &'gc F) -> JITFunc<'gc> {
        self.jit
            .compile_seq(func.opcodes.iter().copied().rev(), self)
            .unwrap()
    }
}