feat(jit): fib!

src/vm/env.rs

@@ -4,14 +4,20 @@ use std::rc::Rc;
 use crate::ty::internal::{AttrSet, Value};
 
 #[derive(Debug, Default, Clone)]
-pub struct Env<'jit, 'vm> {
-    pub map: Rc<HashMap<usize, Value<'jit, 'vm>>>,
-    pub last: Option<Rc<Env<'jit, 'vm>>>,
+pub struct LetEnv<'jit, 'vm> {
+    map: Rc<HashMap<usize, Value<'jit, 'vm>>>,
+    last: Option<Rc<LetEnv<'jit, 'vm>>>,
 }
 
-impl<'jit, 'vm> Env<'jit, 'vm> {
-    pub fn empty() -> Self {
-        Env::default()
+#[derive(Debug, Default, Clone)]
+pub struct WithEnv<'jit, 'vm> {
+    map: Rc<AttrSet<'jit, 'vm>>,
+    last: Option<Rc<WithEnv<'jit, 'vm>>>,
+}
+
+impl<'jit, 'vm> LetEnv<'jit, 'vm> {
+    pub fn new(map: Rc<HashMap<usize, Value<'jit, 'vm>>>) -> Self {
+        Self { map, last: None }
     }
 
     pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
@@ -21,46 +27,63 @@ impl<'jit, 'vm> Env<'jit, 'vm> {
         self.last.as_ref().map(|env| env.lookup(symbol)).flatten()
     }
 
-    pub fn insert(&mut self, symbol: usize, value: Value<'jit, 'vm>) {
-        Rc::make_mut(&mut self.map).insert(symbol, value);
-    }
-
-    pub fn enter(self, new: impl Iterator<Item = (usize, Value<'jit, 'vm>)>) -> Self {
+    pub fn enter_let(self, new: impl Iterator<Item = (usize, Value<'jit, 'vm>)>) -> Self {
         let map = Rc::new(new.collect());
-        let last = Some(
-            Env {
-                last: self.last,
-                map: self.map,
-            }
-            .into(),
-        );
-        Env { last, map }
+        let last = Some(self.into());
+        LetEnv { last, map }
     }
 
     pub fn enter_with(self, new: Rc<AttrSet<'jit, 'vm>>) -> Self {
-        let map = Rc::new(
-            new.as_inner()
-                .iter()
-                .map(|(&k, v)| {
-                    (
-                        k,
-                        if let Value::Builtins(weak) = v {
-                            Value::AttrSet(weak.upgrade().unwrap())
-                        } else {
-                            v.clone()
-                        },
-                    )
-                })
-                .collect(),
-        );
-        let last = Some(
-            Env {
-                last: self.last.clone(),
-                map: self.map.clone(),
-            }
-            .into(),
-        );
-        Env { last, map }
+        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<_, _>>()
+            .into();
+        let last = Some(self.into());
+        LetEnv { last, map }
+    }
+
+    pub fn leave(self) -> Self {
+        self.last.unwrap().as_ref().clone()
+    }
+}
+
+impl<'jit, 'vm> WithEnv<'jit, 'vm> {
+    pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
+        if let Some(val) = self.map.select(symbol) {
+            return Some(val);
+        }
+        self.last.as_ref().map(|env| env.lookup(symbol)).flatten()
+    }
+
+    pub fn enter_with(self, new: Rc<AttrSet<'jit, 'vm>>) -> Self {
+        let map = Rc::new(new
+            .as_inner()
+            .iter()
+            .map(|(&k, v)| {
+                (
+                    k,
+                    if let Value::Builtins(weak) = v {
+                        Value::AttrSet(weak.upgrade().unwrap())
+                    } else {
+                        v.clone()
+                    },
+                )
+            })
+            .collect::<HashMap<_, _>>()
+            .into());
+        let last = Some(self.into());
+        WithEnv { last, map }
     }
 
     pub fn leave(self) -> Self {