feat: generalize 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 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> {