imxyy1soope1/nixjit
1
0
Fork 0
Code Issues Pull Requests Actions Projects 1 Releases Activity

optimize(env): single arg

Browse Source
This commit is contained in:
imxyy_soope_
2025-05-20 09:47:30 +08:00
parent b4db46d48a
commit d0298ce2a6
8 changed files with 76 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/builtins/mod.rs
View File

@@ -60,5 +60,5 @@ pub fn env<'jit, 'vm>(vm: &'vm VM<'jit>) -> LetEnv<'jit, 'vm> {
let builtins = Value::AttrSet(attrs);
env_map.insert(sym, builtins);
LetEnv::new(env_map.into())
LetEnv::new(AttrSet::new(env_map).into())
}

8
src/jit/helpers.rs
View File

@@ -1,3 +1,5 @@
use std::rc::Rc;
use inkwell::AddressSpace;
use inkwell::context::Context;
use inkwell::execution_engine::ExecutionEngine;
@@ -205,7 +207,11 @@ extern "C" fn helper_debug(value: JITValue) {
#[unsafe(no_mangle)]
extern "C" fn helper_capture_env(thunk: JITValue, env: *const LetEnv) {
let thunk: &Thunk = unsafe { std::mem::transmute(thunk.data.ptr.as_ref().unwrap()) };
thunk.capture(unsafe { env.as_ref().unwrap() }.clone());
let env = unsafe {
Rc::from_raw(env)
};
thunk.capture(env.clone());
std::mem::forget(env);
}
#[unsafe(no_mangle)]

2
src/jit/test.rs
View File

@@ -32,7 +32,7 @@ fn test_expr(expr: &str, expected: Value) {
);
let env = env(&vm);
let value = vm
.eval(prog.top_level.into_iter(), env)
.eval(prog.top_level.into_iter(), env.into())
.unwrap()
.to_public(&vm, &mut HashSet::new());
assert_eq!(value, expected);

4
src/ty/internal/attrset.rs
View File

@@ -1,3 +1,5 @@
use std::rc::Rc;
use hashbrown::{HashMap, HashSet};
use derive_more::Constructor;
@@ -47,7 +49,7 @@ impl<'jit: 'vm, 'vm> AttrSet<'jit, 'vm> {
self.data.get(&sym).is_some()
}
pub fn capture(&mut self, env: &LetEnv<'jit, 'vm>) {
pub fn capture(&mut self, env: &Rc<LetEnv<'jit, 'vm>>) {
self.data.iter().for_each(|(_, v)| match v.clone() {
Value::Thunk(ref thunk) => {
thunk.capture(env.clone());

24
src/ty/internal/func.rs
View File

@@ -1,6 +1,8 @@
use std::cell::{Cell, OnceCell};
use std::rc::Rc;
use derive_more::Constructor;
use hashbrown::HashMap;
use inkwell::execution_engine::JitFunction;
use itertools::Itertools;
@@ -8,7 +10,7 @@ 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::ty::internal::{AttrSet, Thunk, Value};
use crate::vm::{LetEnv, VM};
#[derive(Debug, Clone)]
@@ -44,7 +46,7 @@ impl From<ir::Param> for Param {
#[derive(Debug, Clone, Constructor)]
pub struct Func<'jit: 'vm, 'vm> {
pub func: &'vm BFunc,
pub env: LetEnv<'jit, 'vm>,
pub env: Rc<LetEnv<'jit, 'vm>>,
pub compiled: OnceCell<JitFunction<'jit, JITFunc<'jit, 'vm>>>,
pub count: Cell<usize>,
}
@@ -57,14 +59,14 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
Ident(ident) => self
.env
.clone()
.enter_let([(ident.into(), arg)].into_iter()),
.enter_arg(ident.into(), arg),
Formals {
formals,
ellipsis,
alias,
} => {
let arg = arg.unwrap_attr_set();
let mut new = Vec::with_capacity(formals.len() + alias.iter().len());
let mut new = HashMap::with_capacity(formals.len() + alias.iter().len());
if !ellipsis
&& arg
.as_inner()
@@ -83,20 +85,24 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
default.map(|idx| Value::Thunk(Thunk::new(vm.get_thunk(idx)).into()))
})
.unwrap();
new.push((formal, arg));
new.insert(formal, arg);
}
if let Some(alias) = alias {
new.push((alias.clone().into(), Value::AttrSet(arg)));
new.insert(alias.clone().into(), Value::AttrSet(arg));
}
self.env.clone().enter_let(new.into_iter())
self.env.clone().enter_attrs(AttrSet::new(new).into())
}
};
}.into();
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 LetEnv) };
let env = Rc::into_raw(env);
let ret = unsafe { compiled.call(vm as *const VM, env) };
unsafe {
Rc::decrement_strong_count(env);
}
return Ok(ret.into());
}
vm.eval(self.func.opcodes.iter().copied(), env)

4
src/ty/internal/mod.rs
View File

@@ -485,7 +485,7 @@ pub struct Thunk<'jit, 'vm> {
#[derive(Debug, IsVariant, Unwrap, Clone)]
pub enum _Thunk<'jit, 'vm> {
Code(&'vm OpCodes, OnceCell<LetEnv<'jit, 'vm>>),
Code(&'vm OpCodes, OnceCell<Rc<LetEnv<'jit, 'vm>>>),
SuspendedFrom(*const Thunk<'jit, 'vm>),
Value(Value<'jit, 'vm>),
}
@@ -497,7 +497,7 @@ impl<'jit, 'vm> Thunk<'jit, 'vm> {
}
}
pub fn capture(&self, env: LetEnv<'jit, 'vm>) {
pub fn capture(&self, env: Rc<LetEnv<'jit, 'vm>>) {
if let _Thunk::Code(_, envcell) = &*self.thunk.borrow() {
envcell.get_or_init(|| env);
}

58
src/vm/env.rs
View File

@@ -3,9 +3,9 @@ use std::rc::Rc;
use crate::ty::internal::{AttrSet, Value};
#[derive(Debug, Default, Clone)]
#[derive(Debug, Clone)]
pub struct LetEnv<'jit, 'vm> {
map: Rc<HashMap<usize, Value<'jit, 'vm>>>,
map: Env<'jit, 'vm>,
last: Option<Rc<LetEnv<'jit, 'vm>>>,
}
@@ -15,25 +15,51 @@ pub struct WithEnv<'jit, 'vm> {
last: Option<Rc<WithEnv<'jit, 'vm>>>,
}
#[derive(Debug, Clone)]
enum Env<'jit, 'vm> {
Let(Rc<AttrSet<'jit, 'vm>>),
SingleArg(usize, Value<'jit, 'vm>),
MultiArg(Rc<AttrSet<'jit, 'vm>>),
}
#[derive(Debug, Clone, Copy)]
pub enum Type {
Arg,
Let,
With
}
impl<'jit, 'vm> LetEnv<'jit, 'vm> {
pub fn new(map: Rc<HashMap<usize, Value<'jit, 'vm>>>) -> Self {
Self { map, last: None }
pub fn new(map: Rc<AttrSet<'jit, 'vm>>) -> Self {
Self { map: Env::Let(map), last: None }
}
pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
if let Some(val) = self.map.get(&symbol).cloned() {
return Some(val);
use Env::*;
match &self.map {
Let(map) | MultiArg(map) => if let Some(val) = map.select(symbol) {
return Some(val)
}
SingleArg(sym, val) => if *sym == symbol {
return Some(val.clone())
}
}
self.last.as_ref().map(|env| env.lookup(symbol)).flatten()
}
pub fn enter_let(self, new: impl Iterator<Item = (usize, Value<'jit, 'vm>)>) -> Self {
let map = Rc::new(new.collect());
let last = Some(self.into());
LetEnv { last, map }
pub fn enter_arg(self: Rc<Self>, ident: usize, val: Value<'jit, 'vm>) -> Rc<Self> {
let last = Some(self);
let map = Env::SingleArg(ident, val);
LetEnv { last, map }.into()
}
pub fn enter_with(self, new: Rc<AttrSet<'jit, 'vm>>) -> Self {
pub fn enter_attrs(self: Rc<Self>, map: Rc<AttrSet<'jit, 'vm>>) -> Rc<Self> {
let last = Some(self);
let map = Env::Let(map);
LetEnv { last, map }.into()
}
pub fn enter_with(self: Rc<Self>, new: Rc<AttrSet<'jit, 'vm>>) -> Rc<Self> {
let map = new
.as_inner()
.iter()
@@ -47,14 +73,14 @@ impl<'jit, 'vm> LetEnv<'jit, 'vm> {
},
)
})
.collect::<HashMap<_, _>>()
.into();
.collect::<HashMap<_, _>>();
let map = Env::Let(AttrSet::new(map).into());
let last = Some(self.into());
LetEnv { last, map }
LetEnv { last, map }.into()
}
pub fn leave(self) -> Self {
self.last.unwrap().as_ref().clone()
pub fn leave(self: Rc<Self>) -> Rc<Self> {
self.last.clone().unwrap()
}
}

12
src/vm/mod.rs
View File

@@ -1,6 +1,7 @@
use hashbrown::{HashMap, HashSet};
use inkwell::execution_engine::JitFunction;
use std::cell::{Cell, OnceCell, RefCell};
use std::rc::Rc;
use crate::builtins::env;
use crate::bytecode::{BinOp, Func as F, OpCode, OpCodes, Program, UnOp};
@@ -34,7 +35,7 @@ pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
let env = env(&vm);
let mut seen = HashSet::new();
let value = vm
.eval(prog.top_level.into_iter(), env)?
.eval(prog.top_level.into_iter(), env.into())?
.to_public(&vm, &mut seen);
Ok(value)
}
@@ -82,7 +83,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
pub fn eval(
&'vm self,
opcodes: impl Iterator<Item = OpCode>,
mut env: LetEnv<'jit, 'vm>,
mut env: Rc<LetEnv<'jit, 'vm>>,
) -> Result<Value<'jit, 'vm>> {
let mut stack = Stack::<_, STACK_SIZE>::new();
let mut iter = opcodes.into_iter();
@@ -103,7 +104,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
&'vm self,
opcode: OpCode,
stack: &'s mut Stack<Value<'jit, 'vm>, CAP>,
env: &mut LetEnv<'jit, 'vm>,
env: &mut Rc<LetEnv<'jit, 'vm>>,
) -> Result<usize> {
match opcode {
OpCode::Illegal => panic!("illegal opcode"),
@@ -188,14 +189,11 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
stack.push(Value::AttrSet(AttrSet::with_capacity(cap).into()))?;
}
OpCode::FinalizeRec => {
let env = env.clone().enter_let(
let env = env.clone().enter_attrs(
stack
.tos()?
.clone()
.unwrap_attr_set()
.as_inner()
.iter()
.map(|(k, v)| (k.clone(), v.clone())),
);
stack.tos_mut()?.as_mut().unwrap_attr_set().capture(&env);
}