optimize(env): single arg

src/vm/env.rs

@@ -3,9 +3,9 @@ use std::rc::Rc;
 
 use crate::ty::internal::{AttrSet, Value};
 
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone)]
 pub struct LetEnv<'jit, 'vm> {
-    map: Rc<HashMap<usize, Value<'jit, 'vm>>>,
+    map: Env<'jit, 'vm>,
     last: Option<Rc<LetEnv<'jit, 'vm>>>,
 }
 
@@ -15,25 +15,51 @@ pub struct WithEnv<'jit, 'vm> {
     last: Option<Rc<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>>),
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum Type {
+    Arg,
+    Let,
+    With
+}
+
 impl<'jit, 'vm> LetEnv<'jit, 'vm> {
-    pub fn new(map: Rc<HashMap<usize, Value<'jit, 'vm>>>) -> Self {
-        Self { map, last: None }
+    pub fn new(map: Rc<AttrSet<'jit, 'vm>>) -> Self {
+        Self { map: Env::Let(map), last: None }
     }
 
     pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
-        if let Some(val) = self.map.get(&symbol).cloned() {
-            return Some(val);
+        use Env::*;
+        match &self.map {
+            Let(map) | MultiArg(map) => if let Some(val) = map.select(symbol) {
+                return Some(val)
+            }
+            SingleArg(sym, val) => if *sym == symbol {
+                return Some(val.clone())
+            }
         }
         self.last.as_ref().map(|env| env.lookup(symbol)).flatten()
     }
 
-    pub fn enter_let(self, new: impl Iterator<Item = (usize, Value<'jit, 'vm>)>) -> Self {
-        let map = Rc::new(new.collect());
-        let last = Some(self.into());
-        LetEnv { last, map }
+    pub fn enter_arg(self: Rc<Self>, ident: usize, val: Value<'jit, 'vm>) -> Rc<Self> {
+        let last = Some(self);
+        let map = Env::SingleArg(ident, val);
+        LetEnv { last, map }.into()
     }
 
-    pub fn enter_with(self, new: Rc<AttrSet<'jit, 'vm>>) -> Self {
+    pub fn enter_attrs(self: Rc<Self>, map: Rc<AttrSet<'jit, 'vm>>) -> Rc<Self> {
+        let last = Some(self);
+        let map = Env::Let(map);
+        LetEnv { last, map }.into()
+    }
+
+    pub fn enter_with(self: Rc<Self>, new: Rc<AttrSet<'jit, 'vm>>) -> Rc<Self> {
         let map = new
             .as_inner()
             .iter()
@@ -47,14 +73,14 @@ impl<'jit, 'vm> LetEnv<'jit, 'vm> {
                     },
                 )
             })
-            .collect::<HashMap<_, _>>()
-            .into();
+            .collect::<HashMap<_, _>>();
+        let map = Env::Let(AttrSet::new(map).into());
         let last = Some(self.into());
-        LetEnv { last, map }
+        LetEnv { last, map }.into()
     }
 
-    pub fn leave(self) -> Self {
-        self.last.unwrap().as_ref().clone()
+    pub fn leave(self: Rc<Self>) -> Rc<Self> {
+        self.last.clone().unwrap()
     }
 }