144 lines
4.4 KiB
Rust
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,
|
|
}
|