feat: generalize env

src/builtins/mod.rs

@@ -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())
 }

src/jit/helpers.rs

@@ -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);
+    let val = env.lookup(&sym);
-    val.unwrap().into()
+    val.cloned().unwrap().into()
 }
 
 #[unsafe(no_mangle)]

src/jit/mod.rs

@@ -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,

src/ty/internal/attrset.rs

@@ -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 }
     }

src/ty/internal/func.rs

@@ -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::ty::internal::{Thunk, Value};
-use crate::vm::{LetEnv, VM};
+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())
             }
         };
 

src/ty/internal/mod.rs

@@ -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);
         }

src/vm/env.rs

@@ -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 Env<K: Hash + Eq, V> {
-pub struct LetEnv<'jit, 'vm> {
+    map: Node<K, V>,
-    map: Env<'jit, 'vm>,
+    last: Option<Rc<Env<K, V>>>
-    last: Option<Rc<LetEnv<'jit, 'vm>>>,
 }
 
+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> {
+enum Node<K: Hash + Eq, V> {
-    Let(Rc<AttrSet<'jit, 'vm>>),
+    Let(Rc<HashMap<K, V>>),
-    SingleArg(usize, Value<'jit, 'vm>),
+    SingleArg(K, V),
-    MultiArg(Rc<AttrSet<'jit, 'vm>>),
+    MultiArg(Rc<HashMap<K, V>>),
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -29,61 +47,46 @@ pub enum Type {
     With,
 }
 
-impl<'jit, 'vm> LetEnv<'jit, 'vm> {
+impl<K: Hash + Eq, V> Env<K, V> {
-    pub fn new(map: Rc<AttrSet<'jit, 'vm>>) -> Self {
+    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>> {
+    pub fn lookup(&self, symbol: &K) -> Option<&V> {
-        use Env::*;
+        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 {
+                if sym == symbol {
-                    return Some(val.clone());
+                    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);
+        let map = Node::SingleArg(ident, val);
-        LetEnv { last, map }.into()
+        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);
+        let map = Node::Let(map);
-        LetEnv { last, map }.into()
+        Env { last, map }.into()
     }
 
-    pub fn enter_with(self: Rc<Self>, new: Rc<AttrSet<'jit, 'vm>>) -> Rc<Self> {
+    pub fn enter_with(self: Rc<Self>, map: Rc<HashMap<K, V>>) -> Rc<Self> {
-        let map = new
+        let map = Node::Let(map);
-            .as_inner()
+        let last = Some(self);Env { last, map }.into()
-            .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 leave(self: Rc<Self>) -> Rc<Self> {

src/vm/mod.rs

@@ -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!(),
             },