Files
nixjit/src/ty/internal/func.rs
2025-05-19 11:33:18 +08:00

108 lines
3.2 KiB
Rust

use std::cell::{Cell, OnceCell};
use derive_more::Constructor;
use inkwell::execution_engine::JitFunction;
use itertools::Itertools;
use crate::bytecode::Func as BFunc;
use crate::error::Result;
use crate::ir;
use crate::jit::JITFunc;
use crate::ty::internal::{Thunk, Value};
use crate::vm::{Env, VM};
#[derive(Debug, Clone)]
pub enum Param {
Ident(usize),
Formals {
formals: Vec<(usize, Option<usize>)>,
ellipsis: bool,
alias: Option<usize>,
},
}
impl From<ir::Param> for Param {
fn from(value: ir::Param) -> Self {
match value {
ir::Param::Ident(ident) => Param::Ident(ident),
ir::Param::Formals {
formals,
ellipsis,
alias,
} => Param::Formals {
formals: formals
.into_iter()
.map(|(sym, default)| (sym, default.map(|default| default.idx)))
.collect(),
ellipsis,
alias,
},
}
}
}
#[derive(Debug, Clone, Constructor)]
pub struct Func<'jit: 'vm, 'vm> {
pub func: &'vm BFunc,
pub env: Env<'jit, 'vm>,
pub compiled: OnceCell<JitFunction<'jit, JITFunc<'jit, 'vm>>>,
pub count: Cell<usize>,
}
impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
pub fn call(&self, vm: &'vm VM<'jit>, arg: Value<'jit, 'vm>) -> Result<Value<'jit, 'vm>> {
use Param::*;
let env = match self.func.param.clone() {
Ident(ident) => self.env.clone().enter([(ident.into(), arg)].into_iter()),
Formals {
formals,
ellipsis,
alias,
} => {
let arg = arg.unwrap_attr_set();
let mut new = Vec::with_capacity(formals.len() + alias.iter().len());
if !ellipsis
&& arg
.as_inner()
.iter()
.map(|(k, _)| k)
.sorted()
.ne(formals.iter().map(|(k, _)| k).sorted())
{
todo!()
}
for (formal, default) in formals {
let formal = formal.clone().into();
let arg = arg
.select(formal)
.or_else(|| {
default.map(|idx| Value::Thunk(Thunk::new(vm.get_thunk(idx)).into()))
})
.unwrap();
new.push((formal, arg));
}
if let Some(alias) = alias {
new.push((alias.clone().into(), Value::AttrSet(arg)));
}
self.env.clone().enter(new.into_iter())
}
};
let count = self.count.get();
self.count.replace(count + 1);
if count >= 1 {
let compiled = self.compiled.get_or_init(|| vm.compile_func(self.func));
let ret = unsafe { compiled.call(vm as *const VM, &env as *const Env) };
return Ok(ret.into());
}
vm.eval(self.func.opcodes.iter().copied(), env)
}
}
impl PartialEq for Func<'_, '_> {
fn eq(&self, _: &Self) -> bool {
false
}
}