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

feat: generalize env

Browse Source
This commit is contained in:
imxyy_soope_
2025-05-21 09:33:43 +08:00
parent 36f29a9cac
commit 177acfabcf
8 changed files with 77 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/builtins/mod.rs
View File

@@ -3,9 +3,9 @@ use std::rc::Rc;
use crate::ty::common::Const;
use crate::ty::internal::{AttrSet, PrimOp, Value};
use crate::vm::{LetEnv, VM};
use crate::vm::{VmEnv, VM};
pub fn env<'jit, 'vm>(vm: &'vm VM<'jit>) -> LetEnv<'jit, 'vm> {
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)));
@@ -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(AttrSet::new(env_map).into())
VmEnv::new(env_map.into())
}

14
src/jit/helpers.rs
View File

@@ -9,7 +9,7 @@ use inkwell::values::{BasicValueEnum, FunctionValue};
use crate::jit::JITValueData;
use crate::ty::internal::{Thunk, Value};
use crate::vm::{LetEnv, VM};
use crate::vm::{VmEnv, VM};
use super::{JITValue, ValueTag};
@@ -210,7 +210,7 @@ extern "C" fn helper_debug(value: JITValue) {
}
#[unsafe(no_mangle)]
extern "C" fn helper_capture_env(thunk: JITValue, env: *const LetEnv) {
extern "C" fn helper_capture_env(thunk: JITValue, env: *const VmEnv) {
let thunk: &Thunk = unsafe { std::mem::transmute(thunk.data.ptr.as_ref().unwrap()) };
let env = unsafe { Rc::from_raw(env) };
thunk.capture(env.clone());
@@ -218,7 +218,7 @@ extern "C" fn helper_capture_env(thunk: JITValue, env: *const LetEnv) {
}
#[unsafe(no_mangle)]
extern "C" fn helper_neg(rhs: JITValue, _env: *const LetEnv) -> JITValue {
extern "C" fn helper_neg(rhs: JITValue, _env: *const VmEnv) -> JITValue {
use ValueTag::*;
match rhs.tag {
Int => JITValue {
@@ -238,7 +238,7 @@ extern "C" fn helper_neg(rhs: JITValue, _env: *const LetEnv) -> JITValue {
}
#[unsafe(no_mangle)]
extern "C" fn helper_not(rhs: JITValue, _env: *const LetEnv) -> JITValue {
extern "C" fn helper_not(rhs: JITValue, _env: *const VmEnv) -> JITValue {
use ValueTag::*;
match rhs.tag {
Bool => JITValue {
@@ -347,10 +347,10 @@ extern "C" fn helper_call<'jit>(
}
#[unsafe(no_mangle)]
extern "C" fn helper_lookup<'jit, 'vm>(sym: usize, env: *const LetEnv<'jit, 'vm>) -> JITValue {
extern "C" fn helper_lookup<'jit, 'vm>(sym: usize, env: *const VmEnv<'jit, 'vm>) -> JITValue {
let env = unsafe { env.as_ref() }.unwrap();
let val = env.lookup(sym);
val.unwrap().into()
let val = env.lookup(&sym);
val.cloned().unwrap().into()
}
#[unsafe(no_mangle)]

4
src/jit/mod.rs
View File

@@ -12,7 +12,7 @@ use crate::error::*;
use crate::stack::Stack;
use crate::ty::common::Const;
use crate::ty::internal::{Thunk, Value};
use crate::vm::{LetEnv, VM};
use crate::vm::{VmEnv, VM};
mod helpers;
@@ -92,7 +92,7 @@ impl From<Value<'_, '_>> for JITValue {
}
pub type JITFunc<'jit, 'vm> =
unsafe extern "C" fn(*const VM<'jit>, *const LetEnv<'jit, 'vm>) -> JITValue;
unsafe extern "C" fn(*const VM<'jit>, *const VmEnv<'jit, 'vm>) -> JITValue;
pub struct JITContext<'ctx> {
context: &'ctx Context,

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

@@ -6,12 +6,12 @@ use derive_more::Constructor;
use itertools::Itertools;
use crate::error::Result;
use crate::vm::{LetEnv, VM};
use crate::vm::{VmEnv, VM};
use super::super::public as p;
use super::Value;
#[repr(C)]
#[repr(transparent)]
#[derive(Debug, Default, Constructor, Clone, PartialEq)]
pub struct AttrSet<'jit: 'vm, 'vm> {
data: HashMap<usize, Value<'jit, 'vm>>,
@@ -52,7 +52,7 @@ impl<'jit: 'vm, 'vm> AttrSet<'jit, 'vm> {
self.data.get(&sym).is_some()
}
pub fn capture(&mut self, env: &Rc<LetEnv<'jit, 'vm>>) {
pub fn capture(&mut self, env: &Rc<VmEnv<'jit, 'vm>>) {
self.data.iter().for_each(|(_, v)| if let Value::Thunk(ref thunk) = v.clone() {
thunk.capture(env.clone());
})
@@ -68,6 +68,10 @@ impl<'jit: 'vm, 'vm> AttrSet<'jit, 'vm> {
&self.data
}
pub fn into_inner(self: Rc<Self>) -> Rc<HashMap<usize, Value<'jit, 'vm>>> {
unsafe { std::mem::transmute(self) }
}
pub fn from_inner(data: HashMap<usize, Value<'jit, 'vm>>) -> Self {
Self { data }
}

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

@@ -10,8 +10,8 @@ use crate::bytecode::Func as BFunc;
use crate::error::Result;
use crate::ir;
use crate::jit::JITFunc;
use crate::ty::internal::{AttrSet, Thunk, Value};
use crate::vm::{LetEnv, VM};
use crate::ty::internal::{Thunk, Value};
use crate::vm::{VmEnv, VM};
#[derive(Debug, Clone)]
pub enum Param {
@@ -46,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: Rc<LetEnv<'jit, 'vm>>,
pub env: Rc<VmEnv<'jit, 'vm>>,
pub compiled: OnceCell<JitFunction<'jit, JITFunc<'jit, 'vm>>>,
pub count: Cell<usize>,
}
@@ -86,7 +86,7 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
if let Some(alias) = alias {
new.insert(alias, Value::AttrSet(arg));
}
self.env.clone().enter_attrs(AttrSet::new(new).into())
self.env.clone().enter_attrs(new.into())
}
};

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

@@ -11,7 +11,7 @@ use super::public as p;
use crate::bytecode::OpCodes;
use crate::error::*;
use crate::vm::{LetEnv, VM};
use crate::vm::{VmEnv, VM};
mod attrset;
mod func;
@@ -477,7 +477,7 @@ pub struct Thunk<'jit, 'vm> {
#[derive(Debug, IsVariant, Unwrap, Clone)]
pub enum _Thunk<'jit, 'vm> {
Code(&'vm OpCodes, OnceCell<Rc<LetEnv<'jit, 'vm>>>),
Code(&'vm OpCodes, OnceCell<Rc<VmEnv<'jit, 'vm>>>),
SuspendedFrom(*const Thunk<'jit, 'vm>),
Value(Value<'jit, 'vm>),
}
@@ -489,7 +489,7 @@ impl<'jit, 'vm> Thunk<'jit, 'vm> {
}
}
pub fn capture(&self, env: Rc<LetEnv<'jit, 'vm>>) {
pub fn capture(&self, env: Rc<VmEnv<'jit, 'vm>>) {
if let _Thunk::Code(_, envcell) = &*self.thunk.borrow() {
envcell.get_or_init(|| env);
}

85
src/vm/env.rs
View File

@@ -1,14 +1,32 @@
use std::{hash::Hash, rc::Rc};
use std::fmt::Debug;
use hashbrown::HashMap;
use std::rc::Rc;
use crate::ty::internal::{AttrSet, Value};
#[derive(Debug, Clone)]
pub struct LetEnv<'jit, 'vm> {
map: Env<'jit, 'vm>,
last: Option<Rc<LetEnv<'jit, 'vm>>>,
pub struct Env<K: Hash + Eq, V> {
map: Node<K, V>,
last: Option<Rc<Env<K, V>>>
}
impl<K: Clone + Hash + Eq, V: Clone> Clone for Env<K, V> {
fn clone(&self) -> Self {
Self {
map: self.map.clone(),
last: self.last.clone()
}
}
}
impl<K: Debug + Hash + Eq, V: Debug> Debug for Env<K, V> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Env").field("map", &self.map).field("last", &self.last).finish()
}
}
pub type VmEnv<'jit, 'vm> = Env<usize, Value<'jit, 'vm>>;
#[derive(Debug, Default, Clone)]
pub struct WithEnv<'jit, 'vm> {
map: Rc<AttrSet<'jit, 'vm>>,
@@ -16,10 +34,10 @@ pub struct 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>>),
enum Node<K: Hash + Eq, V> {
Let(Rc<HashMap<K, V>>),
SingleArg(K, V),
MultiArg(Rc<HashMap<K, V>>),
}
#[derive(Debug, Clone, Copy)]
@@ -29,61 +47,46 @@ pub enum Type {
With,
}
impl<'jit, 'vm> LetEnv<'jit, 'vm> {
pub fn new(map: Rc<AttrSet<'jit, 'vm>>) -> Self {
impl<K: Hash + Eq, V> Env<K, V> {
pub fn new(map: Rc<HashMap<K, V>>) -> Self {
Self {
map: Env::Let(map),
map: Node::Let(map),
last: None,
}
}
pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
use Env::*;
pub fn lookup(&self, symbol: &K) -> Option<&V> {
use Node::*;
match &self.map {
Let(map) | MultiArg(map) => {
if let Some(val) = map.select(symbol) {
if let Some(val) = map.get(symbol) {
return Some(val);
}
}
SingleArg(sym, val) => {
if *sym == symbol {
return Some(val.clone());
if sym == symbol {
return Some(val);
}
}
}
self.last.as_ref().and_then(|env| env.lookup(symbol))
}
pub fn enter_arg(self: Rc<Self>, ident: usize, val: Value<'jit, 'vm>) -> Rc<Self> {
pub fn enter_arg(self: Rc<Self>, ident: K, val: V) -> Rc<Self> {
let last = Some(self);
let map = Env::SingleArg(ident, val);
LetEnv { last, map }.into()
let map = Node::SingleArg(ident, val);
Env { last, map }.into()
}
pub fn enter_attrs(self: Rc<Self>, map: Rc<AttrSet<'jit, 'vm>>) -> Rc<Self> {
pub fn enter_attrs(self: Rc<Self>, map: Rc<HashMap<K, V>>) -> Rc<Self> {
let last = Some(self);
let map = Env::Let(map);
LetEnv { last, map }.into()
let map = Node::Let(map);
Env { last, map }.into()
}
pub fn enter_with(self: Rc<Self>, new: Rc<AttrSet<'jit, 'vm>>) -> Rc<Self> {
let map = new
.as_inner()
.iter()
.map(|(&k, v)| {
(
k,
if let Value::Builtins(weak) = v {
Value::AttrSet(weak.upgrade().unwrap())
} else {
v.clone()
},
)
})
.collect::<HashMap<_, _>>();
let map = Env::Let(AttrSet::new(map).into());
let last = Some(self);
LetEnv { last, map }.into()
pub fn enter_with(self: Rc<Self>, map: Rc<HashMap<K, V>>) -> Rc<Self> {
let map = Node::Let(map);
let last = Some(self);Env { last, map }.into()
}
pub fn leave(self: Rc<Self>) -> Rc<Self> {

14
src/vm/mod.rs
View File

@@ -14,7 +14,7 @@ use crate::ty::public::{self as p, Symbol};
use derive_more::Constructor;
use ecow::EcoString;
pub use env::LetEnv;
pub use env::VmEnv;
mod env;
@@ -83,7 +83,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
pub fn eval(
&'vm self,
opcodes: impl Iterator<Item = OpCode>,
mut env: Rc<LetEnv<'jit, 'vm>>,
mut env: Rc<VmEnv<'jit, 'vm>>,
) -> Result<Value<'jit, 'vm>> {
let mut stack = Stack::<_, STACK_SIZE>::new();
let mut iter = opcodes.into_iter();
@@ -104,7 +104,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
&'vm self,
opcode: OpCode,
stack: &'s mut Stack<Value<'jit, 'vm>, CAP>,
env: &mut Rc<LetEnv<'jit, 'vm>>,
env: &mut Rc<VmEnv<'jit, 'vm>>,
) -> Result<usize> {
match opcode {
OpCode::Illegal => panic!("illegal opcode"),
@@ -191,7 +191,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
OpCode::FinalizeRec => {
let env = env
.clone()
.enter_attrs(stack.tos()?.clone().unwrap_attr_set());
.enter_attrs(stack.tos()?.clone().unwrap_attr_set().into_inner());
stack.tos_mut()?.as_mut().unwrap_attr_set().capture(&env);
}
OpCode::PushStaticAttr { name } => {
@@ -243,12 +243,12 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
stack.tos_mut()?.force(self)?.has_attr(sym);
}
OpCode::LookUp { sym } => {
stack.push(env.lookup(sym).ok_or_else(|| {
stack.push(env.lookup(&sym).ok_or_else(|| {
Error::EvalError(format!("{} not found", self.get_sym(sym)))
})?)?;
})?.clone())?;
}
OpCode::EnterEnv => match stack.pop() {
Value::AttrSet(attrs) => *env = env.clone().enter_with(attrs),
Value::AttrSet(attrs) => *env = env.clone().enter_with(attrs.into_inner()),
_ => unreachable!(),
},