diff --git a/Cargo.lock b/Cargo.lock index 091d8b8..4335382 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", diff --git a/Cargo.toml b/Cargo.toml index 5465aa4..586ff2c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 } diff --git a/src/builtins/mod.rs b/src/builtins/mod.rs index 9c0825a..c7c21a2 100644 --- a/src/builtins/mod.rs +++ b/src/builtins/mod.rs @@ -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) } diff --git a/src/jit/helpers.rs b/src/jit/helpers.rs index 48893a6..609bc2f 100644 --- a/src/jit/helpers.rs +++ b/src/jit/helpers.rs @@ -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 { diff --git a/src/jit/mod.rs b/src/jit/mod.rs index 0410d9c..a8bf29a 100644 --- a/src/jit/mod.rs +++ b/src/jit/mod.rs @@ -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>> { let mut stack = Stack::<_, STACK_SIZE>::new(); let mut iter = func.opcodes.iter().copied(); diff --git a/src/jit/test.rs b/src/jit/test.rs index 908e031..fd92ca4 100644 --- a/src/jit/test.rs +++ b/src/jit/test.rs @@ -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); diff --git a/src/lib.rs b/src/lib.rs index 1b27b13..2b780a5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ #![cfg_attr(test, feature(test))] #![allow(dead_code)] +#![feature(iter_collect_into)] mod builtins; mod bytecode; diff --git a/src/ty/internal/attrset.rs b/src/ty/internal/attrset.rs index b1c6c8c..74e4105 100644 --- a/src/ty/internal/attrset.rs +++ b/src/ty/internal/attrset.rs @@ -52,10 +52,10 @@ impl<'jit: 'vm, 'vm> AttrSet<'jit, 'vm> { self.data.get(&sym).is_some() } - pub fn capture(&mut self, env: &Rc>) { + 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); } }) } diff --git a/src/ty/internal/func.rs b/src/ty/internal/func.rs index aaa804e..f97e69f 100644 --- a/src/ty/internal/func.rs +++ b/src/ty/internal/func.rs @@ -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 for Param { #[derive(Debug, Clone, Constructor)] pub struct Func<'jit: 'vm, 'vm> { pub func: &'vm BFunc, - pub env: Rc>, + pub env: &'vm VmEnv<'jit, 'vm>, pub compiled: OnceCell>>, pub count: Cell, } @@ -55,16 +54,16 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> { pub fn call(&self, vm: &'vm VM<'jit>, arg: Value<'jit, 'vm>) -> Result> { 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)) } } diff --git a/src/ty/internal/mod.rs b/src/ty/internal/mod.rs index 13db9fd..e4b4597 100644 --- a/src/ty/internal/mod.rs +++ b/src/ty/internal/mod.rs @@ -467,17 +467,11 @@ pub struct Thunk<'jit, 'vm> { #[derive(Debug, IsVariant, Unwrap, Clone)] pub enum _Thunk<'jit, 'vm> { - Code(&'vm OpCodes, OnceCell>), + 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>), - Weak(Weak>), -} - 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>) { + 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>) { - 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 _) } { diff --git a/src/vm/env.rs b/src/vm/env.rs index 2d6b9bc..b8f27fc 100644 --- a/src/vm/env.rs +++ b/src/vm/env.rs @@ -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; + #[derive(Clone)] -pub struct Env { - let_: Rc>, - with: Rc>, +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 { - map: LetNode, - last: Option>>, +pub struct LetEnv<'bump, K: Hash + Eq + Clone, V: Clone> { + map: LetNode<'bump, K, V>, + last: Option<&'bump LetEnv<'bump, K, V>>, } -impl Debug for Env { +impl 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 Debug for LetEnv { +impl 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>; +pub type VmEnv<'jit, 'vm> = Env<'vm, usize, Value<'jit, 'vm>>; #[derive(Debug, Default, Clone)] -pub struct With { - map: Option>>, - last: Option>>, +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 { - Let(Rc>), +enum LetNode<'bump, K: Hash + Eq + Clone, V: Clone> { + Let(&'bump Map<'bump, K, V>), SingleArg(K, V), - MultiArg(Rc>), + MultiArg(&'bump Map<'bump, K, V>), } #[derive(Debug, Clone, Copy)] @@ -57,15 +62,15 @@ pub enum Type { With, } -impl Env { - pub fn new(map: Rc>) -> 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 Env { self.with.lookup(symbol) } - pub fn enter_arg(self: &mut Rc, 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, map: Rc>) { - 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, map: Rc>) { - 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) { - Rc::make_mut(self).let_.leave(); - } - - pub fn leave_with(self: &mut Rc) { - Rc::make_mut(self).with.leave(); + pub fn leave(&self) -> &Self { + self.last.unwrap() } } -impl LetEnv { - pub fn new(map: Rc>) -> Self { +impl<'bump, K: Hash + Eq + Clone, V: Clone> LetEnv<'bump, K, V> { + pub fn new(map: &'bump HashMap) -> Self { Self { map: LetNode::Let(map), last: None, @@ -122,45 +135,33 @@ impl LetEnv { self.last.as_ref().and_then(|env| env.lookup(symbol)) } - pub fn enter_arg(self: &mut Rc, 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, map: Rc>) { - 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) { - 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 With { - 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, map: Rc>) { - let cloned = self.clone(); - let mutref = Rc::make_mut(self); - mutref.last = Some(cloned); - mutref.map = Some(map); - } - - pub fn leave(self: &mut Rc) { - 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) + }) } } diff --git a/src/vm/mod.rs b/src/vm/mod.rs index 42b61d2..812b2c4 100644 --- a/src/vm/mod.rs +++ b/src/vm/mod.rs @@ -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 { 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>, symmap: RefCell>, 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, - mut env: Rc>, + mut env: &'vm VmEnv<'jit, 'vm>, ) -> Result> { 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, CAP>, - env: &mut Rc>, + env: &mut &'vm VmEnv<'jit, 'vm>, ) -> Result { 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!()