Files
nixjit/src/engine/mod.rs
2025-07-13 19:26:58 +08:00

144 lines
4.4 KiB
Rust

use std::cell::OnceCell;
use core::mem::MaybeUninit;
use std::rc::Rc;
use hashbrown::{HashMap, HashSet};
use priority_queue::PriorityQueue;
use crate::env::Env;
use crate::error::Result;
use crate::eval::jit::{JITCompiler, JITFunc};
use crate::eval::Evaluate;
use crate::ir::{Dep, Downgraded, Ir, SccNode};
use crate::ty::internal as i;
use crate::ty::public::Value;
#[cfg(test)]
mod test;
type ThunkIdx = usize;
type EnvIdx = usize;
pub struct Engine {
pub thunks: Box<[Ir]>,
pub funcs: Box<[Ir]>,
pub func_deps: Vec<HashMap<Dep, usize>>,
jit: JITCompiler,
compiled: Box<[OnceCell<JITFunc>]>,
tasks: PriorityQueue<CompileTask, usize>,
}
pub fn eval(downgraded: Downgraded) -> Result<Value> {
let mut engine = Engine::new(
downgraded.thunks,
downgraded.funcs,
downgraded.func_deps,
JITCompiler::new()
);
engine.eval(downgraded.graph)
}
impl Engine {
pub fn new(thunks: Box<[Ir]>, funcs: Box<[Ir]>, func_deps: Vec<HashMap<Dep, usize>>, jit: JITCompiler) -> Self {
Self {
compiled: (0..thunks.len() + funcs.len()).map(|_| OnceCell::new()).collect(),
tasks: PriorityQueue::new(),
thunks,
funcs,
func_deps,
jit,
}
}
pub fn eval(&mut self, graph: Vec<SccNode>) -> Result<Value> {
let mut env = Env::new();
let last = graph.last().unwrap().members[0];
for SccNode { members, .. } in graph.into_iter() {
if members.len() == 1 {
for member in members.into_iter() {
let engine = unsafe { &mut *(self as *mut Self) };
let mut val = self.thunks[member].eval(engine, &mut env)?;
val.force(engine, &mut env)?;
env.insert_cache(member, val);
}
} else {
todo!();
for member in members.into_iter() {
let val = self.thunks[member].clone().eval(self, &mut env)?;
env.insert_cache(member, val);
}
}
}
env.lookup_cache(last, |_| unreachable!()).map(|mut val| {
Ok(val
.force(self, &mut env)?
.to_public(self, &mut HashSet::new()))
})?
}
pub fn eval_thunk(&mut self, idx: usize, env: &mut Env) -> Result<i::Value> {
let engine = self as *const Engine;
let func = self.compiled[idx].get_or_init(|| self.jit.compile(&self.thunks[idx], idx));
let mut ret: MaybeUninit<i::Value> = MaybeUninit::uninit();
unsafe {
func(engine, env, ret.as_mut_ptr());
Ok(ret.assume_init())
}
}
pub fn call_func(&mut self, idx: usize, env: &mut Env) -> Result<i::Value> {
let engine = self as *const Engine;
let func = self.compiled[idx + self.thunks.len()].get_or_init(|| self.jit.compile(&self.funcs[idx], idx));
let mut ret: MaybeUninit<i::Value> = MaybeUninit::uninit();
unsafe {
func(engine, env, ret.as_mut_ptr());
Ok(ret.assume_init())
}
}
pub fn eval_func_deps(&mut self, idx: usize, env: &mut Env) -> Result<()> {
for (&dep, _) in
unsafe { &*(&self.func_deps[idx] as *const HashMap<Dep, usize>) }.iter()
{
match dep {
Dep::Arg(idx) => {
if let i::Value::Thunk(idx) = env.lookup_arg(idx) {
env.insert_cache_lazy(idx, |env| {
let engine = unsafe { &mut *(self as *mut Self) };
let mut val = self.thunks[idx].eval(engine, env)?;
val.force(engine, env)?;
val.ok()
})?;
}
}
Dep::Thunk(idx) => {
env.insert_cache_lazy(idx, |env| {
let engine = unsafe { &mut *(self as *mut Self) };
let mut val = self.thunks[idx].eval(engine, env)?;
val.force(engine, env)?;
val.ok()
})?;
}
}
}
Ok(())
}
}
enum Thunk {
Expr(Ir),
Compiling,
Compiled(fn(Rc<Env>) -> Value),
}
impl Thunk {
fn new(ir: Ir) -> Thunk {
Thunk::Expr(ir)
}
}
#[derive(Hash, PartialEq, Eq)]
struct CompileTask {
idx: usize,
}