nixjit/src/vm/mod.rs

use std::cell::RefCell;

use gc_arena::{Arena, Collect, Gc, Mutation, Rootable};
use hashbrown::{HashMap, HashSet};
use inkwell::context::Context;

use crate::builtins::vm_env;
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>();

#[derive(Collect)]
#[collect(require_static)]
struct ContextWrapper(Context);

pub fn run(mut prog: Program) -> Result<p::Value> {
    let mut arena: Arena<Rootable![GcRoot<'_>]> = Arena::new(|mc| {
        let jit = Gc::new(mc, ContextWrapper(Context::create()));
        let mut thunks = std::mem::take(&mut prog.thunks);
        thunks.iter_mut().for_each(|code| code.reverse());
        let mut funcs = std::mem::take(&mut prog.funcs);
        funcs
            .iter_mut()
            .for_each(|F { opcodes, .. }| opcodes.reverse());
        let symbols = std::mem::take(&mut prog.symbols);
        let symmap = std::mem::take(&mut prog.symmap);
        let consts = std::mem::take(&mut prog.consts);
        let vm = VM {
            thunks,
            funcs,
            consts,
            symbols: RefCell::new(symbols),
            symmap: RefCell::new(symmap),
            jit: JITContext::new(&jit.as_ref().0),
        };
        let vm = Gc::new(mc, vm);
        GcRoot {
            vm,
            jit,
            stack: Stack::new(),
            envs: vec![vm_env(&vm, mc)],
        }
    });
    prog.top_level.reverse();
    eval(prog.top_level, &mut arena, |val, root, _| {
        Ok(val.to_public(&root.vm, &mut HashSet::new()))
    })
}
pub fn eval<T, F: for<'gc> FnOnce(Value<'gc>, &mut GcRoot<'gc>, &Mutation<'gc>) -> Result<T>>(
    opcodes: Box<[OpCode]>,
    arena: &mut Arena<impl for<'gc> Rootable<'gc, Root = GcRoot<'gc>>>,
    f: F,
) -> Result<T> {
    let mut opcodes = opcodes.into_vec();
    while let Some(opcode) = opcodes.pop() {
        arena.mutate_root(|mc, root| {
            let consq = single_op(
                &root.vm,
                opcode,
                &mut root.stack,
                root.envs.last_mut().unwrap(),
                mc,
            )?;
            match consq {
                Consq::NoOp => (),
                Consq::Jmp(step) => opcodes.resize_with(opcodes.len() - step, || unreachable!()),
                Consq::Force => {
                    let thunk = root.stack.tos().as_ref().unwrap_thunk();
                    let (code, env) = thunk.suspend(mc)?;
                    opcodes.push(OpCode::InsertValue);
                    opcodes.push(OpCode::PopEnv);
                    opcodes.extend(code);
                    root.envs.push(env);
                }
                Consq::PopEnv => _ = root.envs.pop().unwrap(),
                Consq::Call => {
                    let arg = root.stack.pop();
                    let func = root.stack.pop().unwrap_func();
                    let env = func.env.enter_arg(arg, mc);
                    let count = func.count.get();
                    func.count.set(count + 1);
                    if count >= 1 {
                        let compiled = func
                            .compiled
                            .borrow_mut(mc)
                            .get_or_insert_with(|| root.vm.compile_func(func.func))
                            .clone();
                        let ret =
                            unsafe { compiled.call(env.as_ref() as *const VmEnv, mc as *const _) };
                        root.stack.push(ret.into())?;
                    } else {
                        root.envs.push(env);
                        opcodes.push(OpCode::PopEnv);
                        opcodes.extend(&func.func.opcodes);
                    }
                }
            }
            Result::Ok(())
        })?;
        arena.finish_cycle();
    }
    arena.collect_debt();
    arena.mutate_root(|mc, root| {
        assert_eq!(root.stack.len(), 1);
        let ret = root.stack.pop();
        f(ret, root, mc)
    })
}

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

#[inline(always)]
fn single_op<'gc, const CAP: usize>(
    vm: &'gc VM<'gc>,
    opcode: OpCode,
    stack: &mut Stack<Value<'gc>, CAP>,
    env: &mut Gc<'gc, VmEnv<'gc>>,
    mc: &'gc Mutation<'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(CoW::new(x.into(), mc)),
            Const::Null => Value::Null,
        })?,
        OpCode::LoadThunk { idx } => stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx), mc)))?,
        OpCode::LoadValue { idx } => {
            stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx), mc)))?;
            stack.tos().as_ref().unwrap_thunk().capture_env(*env, mc);
            return Ok(Consq::Force);
        }
        OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture_env(*env, mc),
        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;
        }
        OpCode::InsertValue => {
            let val = stack.pop();
            stack.tos().as_ref().unwrap_thunk().insert_value(val.clone(), mc);
            *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, mc)?;
        }
        OpCode::Func { idx } => {
            let func = vm.get_func(idx);
            stack.push(Value::Func(Gc::new(mc, Func::new(func, *env, mc))))?;
        }
        OpCode::Arg { level } => {
            stack.push(env.lookup_arg(level))?;
        }
        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, mc),
                Sub => {
                    rhs.neg();
                    lhs.add(rhs, mc);
                }
                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, mc),
                Upd => lhs.update(rhs, mc),
            }
        }
        OpCode::ConcatString => {
            let rhs = stack.pop();
            stack.tos_mut().concat_string(rhs, mc);
        }
        OpCode::Path => {
            todo!()
        }
        OpCode::List { cap } => {
            stack.push(Value::List(CoW::new(List::with_capacity(cap), mc)))?;
        }
        OpCode::PushElem => {
            let elem = stack.pop();
            stack.tos_mut().push(elem, mc);
        }
        OpCode::AttrSet { cap } => {
            stack.push(Value::AttrSet(CoW::new(AttrSet::with_capacity(cap), mc)))?;
        }
        OpCode::FinalizeLet => {
            let mut list = stack.pop().unwrap_list();
            let map = list
                .as_ref()
                .clone()
                .into_inner();
            *env = env.enter_let(map, mc);
            list.make_mut(|list| list.capture(*env, mc), mc);
        }
        OpCode::PushStaticAttr { name } => {
            let val = stack.pop();
            stack.tos_mut().push_attr(name, val, mc);
        }
        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, mc);
        }
        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 } => {
            stack.push(env.lookup_let(level, idx))?;
        }
        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.enter_with(Gc::new(mc, new), mc);
        }
        OpCode::PopEnv => return Ok(Consq::PopEnv),
        OpCode::Assert => {
            if !stack.pop().unwrap_bool() {
                todo!()
            }
        }
    }
    Ok(Consq::NoOp)
}

#[derive(Collect)]
#[collect(no_drop)]
pub struct GcRoot<'gc, const CAP: usize = STACK_SIZE> {
    vm: Gc<'gc, VM<'gc>>,
    jit: Gc<'gc, ContextWrapper>,

    stack: Stack<Value<'gc>, CAP>,
    envs: Vec<Gc<'gc, VmEnv<'gc>>>,
}

#[derive(Constructor, Collect)]
#[collect(no_drop)]
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) -> &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<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(&'gc self, func: &'gc F) -> JITFunc<'gc> {
        self.jit
            .compile_seq(func.opcodes.iter().copied().rev(), self)
            .unwrap()
    }
}