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

feat: bumpalo

Browse Source
This commit is contained in:
imxyy_soope_
2025-05-23 12:09:53 +08:00
parent 53cbb37b00
commit a47a08b051
12 changed files with 130 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
Cargo.lock generated
View File

@@ -45,6 +45,15 @@ version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd"
[[package]]
name = "bumpalo"
version = "3.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1628fb46dfa0b37568d12e5edd512553eccf6a22a78e8bde00bb4aed84d5bdbf"
dependencies = [
"allocator-api2",
]
[[package]]
name = "cc"
version = "1.2.21"
@@ -307,6 +316,7 @@ dependencies = [
name = "nixjit"
version = "0.0.0"
dependencies = [
"bumpalo",
"derive_more",
"ecow",
"hashbrown 0.15.3",

4
Cargo.toml
View File

@@ -25,11 +25,11 @@ rnix = "0.12"
thiserror = "2.0"
itertools = "0.14"
rpds = "1.1"
derive_more = { version = "2.0", features = [ "full" ] }
derive_more = { version = "2.0", features = ["full"] }
ecow = "0.2"
regex = "1.11"
hashbrown = "0.15"
inkwell = { version = "0.6.0", features = ["llvm18-1"] }
bumpalo = { version = "3.17", features = ["allocator-api2"] }
rustyline = { version = "15.0", optional = true }

14
src/builtins/mod.rs
View File

@@ -1,15 +1,12 @@
use hashbrown::HashMap;
use std::rc::Rc;
use hashbrown::HashMap;
use crate::ty::common::Const;
use crate::ty::internal::{AttrSet, PrimOp, Value};
use crate::vm::{VM, VmEnv};
pub fn env<'jit, 'vm>(vm: &'vm VM<'jit>) -> VmEnv<'jit, 'vm> {
let mut env_map = HashMap::new();
env_map.insert(vm.new_sym("true"), Value::Const(Const::Bool(true)));
env_map.insert(vm.new_sym("false"), Value::Const(Const::Bool(false)));
let primops = [
PrimOp::new("add", 2, |_, args| {
let [mut first, second]: [Value; 2] = args.try_into().unwrap();
@@ -49,6 +46,11 @@ pub fn env<'jit, 'vm>(vm: &'vm VM<'jit>) -> VmEnv<'jit, 'vm> {
}),
];
let mut env_map = HashMap::new_in(&vm.bump);
env_map.insert(vm.new_sym("true"), Value::Const(Const::Bool(true)));
env_map.insert(vm.new_sym("false"), Value::Const(Const::Bool(false)));
let mut map = HashMap::new();
for primop in primops {
let primop = Rc::new(primop);
@@ -66,5 +68,5 @@ pub fn env<'jit, 'vm>(vm: &'vm VM<'jit>) -> VmEnv<'jit, 'vm> {
let builtins = Value::AttrSet(attrs);
env_map.insert(sym, builtins);
VmEnv::new(env_map.into())
VmEnv::new(vm.bump.alloc(env_map), &vm.bump)
}

7
src/jit/helpers.rs
View File

@@ -1,5 +1,3 @@
use std::rc::Rc;
use inkwell::AddressSpace;
use inkwell::context::Context;
use inkwell::execution_engine::ExecutionEngine;
@@ -209,9 +207,8 @@ extern "C" fn helper_debug(value: JITValue) {
extern "C" fn helper_capture_env(thunk: JITValue, env: *const VmEnv) {
let thunk = unsafe { (thunk.data.ptr as *const Thunk).as_ref().unwrap() };
let env = unsafe { Rc::from_raw(env) };
thunk.capture(env.clone());
std::mem::forget(env);
let env = unsafe { env.as_ref() }.unwrap();
thunk.capture(env);
}
extern "C" fn helper_neg(rhs: JITValue, _env: *const VmEnv) -> JITValue {

4
src/jit/mod.rs
View File

@@ -165,10 +165,10 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
.unwrap()
}
pub fn compile_function(
pub fn compile_function<'bump>(
&self,
func: &Func,
vm: &'vm VM<'_>,
vm: &'vm VM<'ctx>,
) -> Result<JitFunction<'ctx, JITFunc<'ctx, 'vm>>> {
let mut stack = Stack::<_, STACK_SIZE>::new();
let mut iter = func.opcodes.iter().copied();

4
src/jit/test.rs
View File

@@ -2,6 +2,7 @@
extern crate test;
use bumpalo::Bump;
use hashbrown::{HashMap, HashSet};
use inkwell::context::Context;
@@ -30,11 +31,12 @@ fn test_expr(expr: &str, expected: Value) {
prog.symbols.into(),
prog.symmap.into(),
prog.consts,
Bump::new(),
jit,
);
let env = env(&vm);
let value = vm
.eval(prog.top_level.into_iter(), env.into())
.eval(prog.top_level.into_iter(), vm.bump.alloc(env))
.unwrap()
.to_public(&vm, &mut HashSet::new());
assert_eq!(value, expected);

1
src/lib.rs
View File

@@ -1,5 +1,6 @@
#![cfg_attr(test, feature(test))]
#![allow(dead_code)]
#![feature(iter_collect_into)]
mod builtins;
mod bytecode;

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

@@ -52,10 +52,10 @@ impl<'jit: 'vm, 'vm> AttrSet<'jit, 'vm> {
self.data.get(&sym).is_some()
}
pub fn capture(&mut self, env: &Rc<VmEnv<'jit, 'vm>>) {
pub fn capture(&mut self, env: &'vm VmEnv<'jit, 'vm>) {
self.data.iter().for_each(|(_, v)| {
if let Value::Thunk(ref thunk) = v.clone() {
thunk.capture(Rc::clone(env));
thunk.capture(env);
}
})
}

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

@@ -1,5 +1,4 @@
use std::cell::{Cell, OnceCell};
use std::rc::Rc;
use derive_more::Constructor;
use hashbrown::HashMap;
@@ -46,7 +45,7 @@ impl From<ir::Param> for Param {
#[derive(Debug, Clone, Constructor)]
pub struct Func<'jit: 'vm, 'vm> {
pub func: &'vm BFunc,
pub env: Rc<VmEnv<'jit, 'vm>>,
pub env: &'vm VmEnv<'jit, 'vm>,
pub compiled: OnceCell<JitFunction<'jit, JITFunc<'jit, 'vm>>>,
pub count: Cell<usize>,
}
@@ -55,16 +54,16 @@ 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 mut env = self.env.clone();
match self.func.param.clone() {
Ident(ident) => env.enter_arg(ident, arg),
let mut env = self.env;
env = match self.func.param.clone() {
Ident(ident) => env.enter_arg(ident, arg, &vm.bump),
Formals {
formals,
ellipsis,
alias,
} => {
let arg = arg.unwrap_attr_set();
let mut new = HashMap::with_capacity(formals.len() + alias.iter().len());
let mut new = HashMap::with_capacity_in(formals.len() + alias.iter().len(), &vm.bump);
if !ellipsis
&& arg
.as_inner()
@@ -87,7 +86,7 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
if let Some(alias) = alias {
new.insert(alias, Value::AttrSet(arg));
}
env.enter_let(new.into())
env.enter_let(vm.bump.alloc(new), &vm.bump)
}
};
@@ -95,14 +94,11 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
self.count.replace(count + 1);
if count >= 1 {
let compiled = self.compiled.get_or_init(|| vm.compile_func(self.func));
let env = Rc::into_raw(env);
let env = env as *const _;
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)
vm.eval(self.func.opcodes.iter().copied(), vm.bump.alloc(env))
}
}

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

@@ -467,17 +467,11 @@ pub struct Thunk<'jit, 'vm> {
#[derive(Debug, IsVariant, Unwrap, Clone)]
pub enum _Thunk<'jit, 'vm> {
Code(&'vm OpCodes, OnceCell<EnvRef<'jit, 'vm>>),
Code(&'vm OpCodes, OnceCell<&'vm VmEnv<'jit, 'vm>>),
SuspendedFrom(*const Thunk<'jit, 'vm>),
Value(Value<'jit, 'vm>),
}
#[derive(Debug, IsVariant, Unwrap, Clone)]
pub enum EnvRef<'jit, 'vm> {
Strong(Rc<VmEnv<'jit, 'vm>>),
Weak(Weak<VmEnv<'jit, 'vm>>),
}
impl<'jit, 'vm> Thunk<'jit, 'vm> {
pub fn new(opcodes: &'vm OpCodes) -> Self {
Thunk {
@@ -485,15 +479,9 @@ impl<'jit, 'vm> Thunk<'jit, 'vm> {
}
}
pub fn capture(&self, env: Rc<VmEnv<'jit, 'vm>>) {
pub fn capture(&self, env: &'vm VmEnv<'jit, 'vm>) {
if let _Thunk::Code(_, envcell) = &*self.thunk.borrow() {
envcell.get_or_init(|| EnvRef::Strong(env));
}
}
pub fn capture_weak(&self, env: Weak<VmEnv<'jit, 'vm>>) {
if let _Thunk::Code(_, envcell) = &*self.thunk.borrow() {
envcell.get_or_init(|| EnvRef::Weak(env));
envcell.get_or_init(|| env);
}
}
@@ -519,10 +507,7 @@ impl<'jit, 'vm> Thunk<'jit, 'vm> {
_Thunk::SuspendedFrom(self as *const Thunk),
)
.unwrap_code();
let env = match env.get().unwrap() {
EnvRef::Strong(env) => env.clone(),
EnvRef::Weak(env) => env.upgrade().unwrap(),
};
let env = env.get().unwrap();
let value = vm.eval(opcodes.iter().copied(), env)?;
let _ = std::mem::replace(&mut *self.thunk.borrow_mut(), _Thunk::Value(value));
Ok(match unsafe { &*(&*self.thunk.borrow() as *const _) } {

135
src/vm/env.rs
View File

@@ -1,53 +1,58 @@
use std::fmt::Debug;
use std::{hash::Hash, rc::Rc};
use std::hash::Hash;
use hashbrown::HashMap;
use hashbrown::{DefaultHashBuilder, HashMap};
use bumpalo::Bump;
use crate::ty::internal::Value;
type Map<'bump, K, V> = HashMap<K, V, DefaultHashBuilder, &'bump Bump>;
#[derive(Clone)]
pub struct Env<K: Hash + Eq + Clone, V: Clone> {
let_: Rc<LetEnv<K, V>>,
with: Rc<With<K, V>>,
pub struct Env<'bump, K: Hash + Eq + Clone, V: Clone> {
let_: &'bump LetEnv<'bump, K, V>,
with: &'bump With<'bump, K, V>,
last: Option<&'bump Env<'bump, K, V>>,
}
#[derive(Clone)]
pub struct LetEnv<K: Hash + Eq + Clone, V: Clone> {
map: LetNode<K, V>,
last: Option<Rc<LetEnv<K, V>>>,
pub struct LetEnv<'bump, K: Hash + Eq + Clone, V: Clone> {
map: LetNode<'bump, K, V>,
last: Option<&'bump LetEnv<'bump, K, V>>,
}
impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for Env<K, V> {
impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for Env<'_, K, V> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Env")
.field("let_", &self.let_)
.field("with", &self.with)
.field("last", &self.last)
.finish()
}
}
impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for LetEnv<K, V> {
impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for LetEnv<'_, K, V> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Env")
f.debug_struct("LetEnv")
.field("map", &self.map)
.field("last", &self.last)
.finish()
}
}
pub type VmEnv<'jit, 'vm> = Env<usize, Value<'jit, 'vm>>;
pub type VmEnv<'jit, 'vm> = Env<'vm, usize, Value<'jit, 'vm>>;
#[derive(Debug, Default, Clone)]
pub struct With<K: Hash + Eq, V> {
map: Option<Rc<HashMap<K, V>>>,
last: Option<Rc<With<K, V>>>,
pub struct With<'bump, K: Hash + Eq, V> {
map: Option<&'bump Map<'bump, K, V>>,
last: Option<&'bump With<'bump, K, V>>,
}
#[derive(Debug, Clone)]
enum LetNode<K: Hash + Eq, V> {
Let(Rc<HashMap<K, V>>),
enum LetNode<'bump, K: Hash + Eq + Clone, V: Clone> {
Let(&'bump Map<'bump, K, V>),
SingleArg(K, V),
MultiArg(Rc<HashMap<K, V>>),
MultiArg(&'bump Map<'bump, K, V>),
}
#[derive(Debug, Clone, Copy)]
@@ -57,15 +62,15 @@ pub enum Type {
With,
}
impl<K: Hash + Eq + Clone, V: Clone> Env<K, V> {
pub fn new(map: Rc<HashMap<K, V>>) -> Self {
impl<'bump, K: Hash + Eq + Clone, V: Clone> Env<'bump, K, V> {
pub fn new(map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> Self {
Self {
let_: LetEnv::new(map).into(),
with: With {
let_: &*bump.alloc(LetEnv::new(map)),
with: &*bump.alloc(With {
map: None,
last: None,
}
.into(),
}),
last: None
}
}
@@ -76,29 +81,37 @@ impl<K: Hash + Eq + Clone, V: Clone> Env<K, V> {
self.with.lookup(symbol)
}
pub fn enter_arg(self: &mut Rc<Self>, ident: K, val: V) {
Rc::make_mut(self).let_.enter_arg(ident, val);
pub fn enter_arg(&'bump self, ident: K, val: V, bump: &'bump Bump) -> &'bump Self {
&*bump.alloc(Env {
let_: self.let_.enter_arg(ident, val, bump),
with: self.with,
last: Some(self)
})
}
pub fn enter_let(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
Rc::make_mut(self).let_.enter_let(map);
pub fn enter_let(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
&*bump.alloc(Env {
let_: self.let_.enter_let(map, bump),
with: self.with,
last: Some(self)
})
}
pub fn enter_with(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
Rc::make_mut(self).with.enter(map);
pub fn enter_with(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
&*bump.alloc(Env {
let_: self.let_,
with: self.with.enter(map, bump),
last: Some(self)
})
}
pub fn leave_let(self: &mut Rc<Self>) {
Rc::make_mut(self).let_.leave();
}
pub fn leave_with(self: &mut Rc<Self>) {
Rc::make_mut(self).with.leave();
pub fn leave(&self) -> &Self {
self.last.unwrap()
}
}
impl<K: Hash + Eq + Clone, V: Clone> LetEnv<K, V> {
pub fn new(map: Rc<HashMap<K, V>>) -> Self {
impl<'bump, K: Hash + Eq + Clone, V: Clone> LetEnv<'bump, K, V> {
pub fn new(map: &'bump HashMap<K, V, DefaultHashBuilder, &Bump>) -> Self {
Self {
map: LetNode::Let(map),
last: None,
@@ -122,45 +135,33 @@ impl<K: Hash + Eq + Clone, V: Clone> LetEnv<K, V> {
self.last.as_ref().and_then(|env| env.lookup(symbol))
}
pub fn enter_arg(self: &mut Rc<Self>, ident: K, val: V) {
let cloned = self.clone();
let mutref = Rc::make_mut(self);
mutref.last = Some(cloned);
mutref.map = LetNode::SingleArg(ident, val);
pub fn enter_arg(&'bump self, ident: K, val: V, bump: &'bump Bump) -> &'bump Self {
&*bump.alloc(Self {
map: LetNode::SingleArg(ident, val),
last: Some(self)
})
}
pub fn enter_let(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
let cloned = self.clone();
let mutref = Rc::make_mut(self);
mutref.last = Some(cloned);
mutref.map = LetNode::Let(map);
}
pub fn leave(self: &mut Rc<Self>) {
let refmut = Rc::make_mut(self);
let last = refmut.last.take().unwrap();
*self = last;
pub fn enter_let(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
&*bump.alloc(Self {
map: LetNode::Let(map),
last: Some(self)
})
}
}
impl<K: Hash + Eq + Clone, V: Clone> With<K, V> {
pub fn lookup(&self, symbol: &K) -> Option<&V> {
impl<'bump, K: Hash + Eq + Clone, V: Clone> With<'bump, K, V> {
pub fn lookup(&'bump self, symbol: &K) -> Option<&'bump V> {
if let Some(val) = self.map.as_ref()?.get(symbol) {
return Some(val);
}
self.last.as_ref().and_then(|env| env.lookup(symbol))
}
pub fn enter(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
let cloned = self.clone();
let mutref = Rc::make_mut(self);
mutref.last = Some(cloned);
mutref.map = Some(map);
}
pub fn leave(self: &mut Rc<Self>) {
let refmut = Rc::make_mut(self);
let last = refmut.last.take().unwrap();
*self = last;
pub fn enter(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
&*bump.alloc(Self {
map: Some(map),
last: Some(self)
})
}
}

31
src/vm/mod.rs
View File

@@ -1,7 +1,7 @@
use bumpalo::Bump;
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};
@@ -30,12 +30,13 @@ pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
symbols: RefCell::new(prog.symbols),
symmap: RefCell::new(prog.symmap),
consts: prog.consts,
bump: Bump::new(),
jit,
};
let env = env(&vm);
let mut seen = HashSet::new();
let value = vm
.eval(prog.top_level.into_iter(), env.into())?
.eval(prog.top_level.into_iter(), vm.bump.alloc(env))?
.to_public(&vm, &mut seen);
Ok(value)
}
@@ -47,6 +48,7 @@ pub struct VM<'jit> {
symbols: RefCell<Vec<EcoString>>,
symmap: RefCell<HashMap<EcoString, usize>>,
consts: Box<[Const]>,
pub bump: Bump,
jit: JITContext<'jit>,
}
@@ -83,7 +85,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
pub fn eval(
&'vm self,
opcodes: impl Iterator<Item = OpCode>,
mut env: Rc<VmEnv<'jit, 'vm>>,
mut env: &'vm VmEnv<'jit, 'vm>,
) -> Result<Value<'jit, 'vm>> {
let mut stack = Stack::<_, STACK_SIZE>::new();
let mut iter = opcodes.into_iter();
@@ -104,7 +106,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
&'vm self,
opcode: OpCode,
stack: &'s mut Stack<Value<'jit, 'vm>, CAP>,
env: &mut Rc<VmEnv<'jit, 'vm>>,
env: &mut &'vm VmEnv<'jit, 'vm>,
) -> Result<usize> {
match opcode {
OpCode::Illegal => panic!("illegal opcode"),
@@ -112,7 +114,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
OpCode::LoadThunk { idx } => {
stack.push(Value::Thunk(Thunk::new(self.get_thunk(idx)).into()))?
}
OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture(env.clone()),
OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture(*env),
OpCode::ForceValue => {
stack.tos_mut().force(self)?;
}
@@ -131,7 +133,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
OpCode::Func { idx } => {
let func = self.get_func(idx);
stack.push(Value::Func(
Func::new(func, env.clone(), OnceCell::new(), Cell::new(0)).into(),
Func::new(func, env, OnceCell::new(), Cell::new(0)).into(),
))?;
}
OpCode::UnOp { op } => {
@@ -184,7 +186,8 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
stack.push(Value::AttrSet(AttrSet::with_capacity(cap).into()))?;
}
OpCode::FinalizeRec => {
env.enter_let(stack.tos().clone().unwrap_attr_set().into_inner());
let new = self.bump.alloc(HashMap::new_in(&self.bump));
*env = env.enter_let(stack.tos().as_ref().unwrap_attr_set().as_inner().iter().map(|(&k, v)| (k, v.clone())).collect_into(new), &self.bump);
stack.tos_mut().as_mut().unwrap_attr_set().capture(env);
}
OpCode::PushStaticAttr { name } => {
@@ -244,10 +247,16 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
.clone(),
)?;
}
OpCode::EnterLetEnv => env.enter_let(stack.pop().unwrap_attr_set().into_inner()),
OpCode::LeaveLetEnv => env.leave_let(),
OpCode::EnterWithEnv => env.enter_with(stack.pop().unwrap_attr_set().into_inner()),
OpCode::LeaveWithEnv => env.leave_with(),
OpCode::EnterLetEnv => {
let new = self.bump.alloc(HashMap::new_in(&self.bump));
*env = env.enter_let(stack.pop().unwrap_attr_set().into_inner().iter().map(|(&k, v)| (k, v.clone())).collect_into(new), &self.bump);
}
OpCode::LeaveLetEnv => *env = env.leave(),
OpCode::EnterWithEnv => {
let new = self.bump.alloc(HashMap::new_in(&self.bump));
*env = env.enter_with(stack.pop().unwrap_attr_set().into_inner().iter().map(|(&k, v)| (k, v.clone())).collect_into(new), &self.bump);
}
OpCode::LeaveWithEnv => *env = env.leave(),
OpCode::Assert => {
if !stack.pop().unwrap_const().unwrap_bool() {
todo!()