feat: less env clone

src/bytecode.rs

@@ -71,10 +71,14 @@ pub enum OpCode {
     SelectDynamic,
     /// [ .. set sym default ] select `sym` from `set` or `default`
     SelectDynamicOrDefault,
-    /// enter the environment of the attribute set at TOS
+    /// enter the let environment of the attribute set at TOS
-    EnterEnv,
+    EnterLetEnv,
-    /// exit current envrironment
+    /// exit current let envrironment
-    LeaveEnv,
+    LeaveLetEnv,
+    /// enter the with environment of the attribute set at TOS
+    EnterWithEnv,
+    /// exit current with envrironment
+    LeaveWithEnv,
 
     /// illegal operation, used as termporary placeholder
     Illegal,

src/compile.rs

@@ -364,18 +364,18 @@ impl Compile for ir::If {
 impl Compile for ir::Let {
     fn compile(self, comp: &mut Compiler) {
         self.attrs.compile(comp);
-        comp.push(OpCode::EnterEnv);
+        comp.push(OpCode::EnterLetEnv);
         self.expr.compile(comp);
-        comp.push(OpCode::LeaveEnv);
+        comp.push(OpCode::LeaveLetEnv);
     }
 }
 
 impl Compile for ir::With {
     fn compile(self, comp: &mut Compiler) {
         self.namespace.compile(comp);
-        comp.push(OpCode::EnterEnv);
+        comp.push(OpCode::EnterWithEnv);
         self.expr.compile(comp);
-        comp.push(OpCode::LeaveEnv);
+        comp.push(OpCode::LeaveWithEnv);
     }
 }
 

src/ty/internal/func.rs

@@ -55,8 +55,9 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
     pub fn call(&self, vm: &'vm VM<'jit>, arg: Value<'jit, 'vm>) -> Result<Value<'jit, 'vm>> {
         use Param::*;
 
-        let env = match self.func.param.clone() {
+        let mut env = self.env.clone();
-            Ident(ident) => self.env.clone().enter_arg(ident, arg),
+        match self.func.param.clone() {
+            Ident(ident) => env.enter_arg(ident, arg),
             Formals {
                 formals,
                 ellipsis,
@@ -86,7 +87,7 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
                 if let Some(alias) = alias {
                     new.insert(alias, Value::AttrSet(arg));
                 }
-                self.env.clone().enter_attrs(new.into())
+                env.enter_let(new.into())
             }
         };
 

src/vm/env.rs

@@ -3,23 +3,30 @@ use std::{hash::Hash, rc::Rc};
 
 use hashbrown::HashMap;
 
-use crate::ty::internal::{AttrSet, Value};
+use crate::ty::internal::Value;
 
-pub struct Env<K: Hash + Eq, V> {
+#[derive(Clone)]
-    map: Node<K, V>,
+pub struct Env<K: Hash + Eq + Clone, V: Clone> {
-    last: Option<Rc<Env<K, V>>>,
+    let_: Rc<LetEnv<K, V>>,
+    with: Rc<With<K, V>>,
 }
 
-impl<K: Clone + Hash + Eq, V: Clone> Clone for Env<K, V> {
+#[derive(Clone)]
-    fn clone(&self) -> Self {
+pub struct LetEnv<K: Hash + Eq + Clone, V: Clone> {
-        Self {
+    map: Let<K, V>,
-            map: self.map.clone(),
+    last: Option<Rc<LetEnv<K, V>>>,
-            last: self.last.clone(),
+}
-        }
+
+impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for Env<K, V> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Env")
+            .field("let_", &self.let_)
+            .field("with", &self.with)
+            .finish()
     }
 }
 
-impl<K: Debug + Hash + Eq, V: Debug> Debug for Env<K, V> {
+impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for LetEnv<K, V> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Env")
             .field("map", &self.map)
@@ -31,13 +38,13 @@ impl<K: Debug + Hash + Eq, V: Debug> Debug for Env<K, V> {
 pub type VmEnv<'jit, 'vm> = Env<usize, Value<'jit, 'vm>>;
 
 #[derive(Debug, Default, Clone)]
-pub struct WithEnv<'jit, 'vm> {
+pub struct With<K: Hash + Eq, V> {
-    map: Rc<AttrSet<'jit, 'vm>>,
+    map: Option<Rc<HashMap<K, V>>>,
-    last: Option<Rc<WithEnv<'jit, 'vm>>>,
+    last: Option<Rc<With<K, V>>>,
 }
 
 #[derive(Debug, Clone)]
-enum Node<K: Hash + Eq, V> {
+enum Let<K: Hash + Eq, V> {
     Let(Rc<HashMap<K, V>>),
     SingleArg(K, V),
     MultiArg(Rc<HashMap<K, V>>),
@@ -50,16 +57,56 @@ pub enum Type {
     With,
 }
 
-impl<K: Hash + Eq, V> Env<K, V> {
+impl<K: Hash + Eq + Clone, V: Clone> Env<K, V> {
     pub fn new(map: Rc<HashMap<K, V>>) -> Self {
         Self {
-            map: Node::Let(map),
+            let_: LetEnv::new(map).into(),
+            with: With {
+                map: None,
+                last: None,
+            }
+            .into(),
+        }
+    }
+
+    pub fn lookup(&self, symbol: &K) -> Option<&V> {
+        if let Some(val) = self.let_.lookup(symbol) {
+            return Some(val);
+        }
+        self.with.lookup(symbol)
+    }
+
+    pub fn enter_arg(self: &mut Rc<Self>, ident: K, val: V) {
+        Rc::make_mut(self).let_.enter_arg(ident, val);
+    }
+
+    pub fn enter_let(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
+        Rc::make_mut(self).let_.enter_let(map);
+    }
+
+    pub fn enter_with(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
+        Rc::make_mut(self).with.enter(map);
+    }
+
+    pub fn leave_let(self: &mut Rc<Self>) {
+        Rc::make_mut(self).let_.leave();
+    }
+
+    pub fn leave_with(self: &mut Rc<Self>) {
+        Rc::make_mut(self).with.leave();
+    }
+}
+
+impl<K: Hash + Eq + Clone, V: Clone> LetEnv<K, V> {
+    pub fn new(map: Rc<HashMap<K, V>>) -> Self {
+        Self {
+            map: Let::Let(map),
             last: None,
         }
     }
 
     pub fn lookup(&self, symbol: &K) -> Option<&V> {
-        use Node::*;
+        use self::Let::*;
         match &self.map {
             Let(map) | MultiArg(map) => {
                 if let Some(val) = map.get(symbol) {
@@ -75,59 +122,45 @@ impl<K: Hash + Eq, V> Env<K, V> {
         self.last.as_ref().and_then(|env| env.lookup(symbol))
     }
 
-    pub fn enter_arg(self: Rc<Self>, ident: K, val: V) -> Rc<Self> {
+    pub fn enter_arg(self: &mut Rc<Self>, ident: K, val: V) {
-        let last = Some(self);
+        let cloned = self.clone();
-        let map = Node::SingleArg(ident, val);
+        let mutref = Rc::make_mut(self);
-        Env { last, map }.into()
+        mutref.last = Some(cloned);
+        mutref.map = Let::SingleArg(ident, val);
     }
 
-    pub fn enter_attrs(self: Rc<Self>, map: Rc<HashMap<K, V>>) -> Rc<Self> {
+    pub fn enter_let(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
-        let last = Some(self);
+        let cloned = self.clone();
-        let map = Node::Let(map);
+        let mutref = Rc::make_mut(self);
-        Env { last, map }.into()
+        mutref.last = Some(cloned);
+        mutref.map = Let::Let(map);
     }
 
-    pub fn enter_with(self: Rc<Self>, map: Rc<HashMap<K, V>>) -> Rc<Self> {
+    pub fn leave(self: &mut Rc<Self>) {
-        let map = Node::Let(map);
+        let refmut = Rc::make_mut(self);
-        let last = Some(self);
+        let last = refmut.last.take().unwrap();
-        Env { last, map }.into()
+        *self = last;
-    }
-
-    pub fn leave(self: Rc<Self>) -> Rc<Self> {
-        self.last.clone().unwrap()
     }
 }
 
-impl<'jit, 'vm> WithEnv<'jit, 'vm> {
+impl<K: Hash + Eq + Clone, V: Clone> With<K, V> {
-    pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
+    pub fn lookup(&self, symbol: &K) -> Option<&V> {
-        if let Some(val) = self.map.select(symbol) {
+        if let Some(val) = self.map.as_ref()?.get(symbol) {
             return Some(val);
         }
         self.last.as_ref().and_then(|env| env.lookup(symbol))
     }
 
-    pub fn enter_with(self, new: Rc<AttrSet<'jit, 'vm>>) -> Self {
+    pub fn enter(self: &mut Rc<Self>, map: Rc<HashMap<K, V>>) {
-        let map = Rc::new(
+        let cloned = self.clone();
-            new.as_inner()
+        let mutref = Rc::make_mut(self);
-                .iter()
+        mutref.last = Some(cloned);
-                .map(|(&k, v)| {
+        mutref.map = Some(map);
-                    (
-                        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 {
+    pub fn leave(self: &mut Rc<Self>) {
-        self.last.unwrap().as_ref().clone()
+        let refmut = Rc::make_mut(self);
+        let last = refmut.last.take().unwrap();
+        *self = last;
     }
 }

src/vm/mod.rs

@@ -189,9 +189,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
                 stack.push(Value::AttrSet(AttrSet::with_capacity(cap).into()))?;
             }
             OpCode::FinalizeRec => {
-                let env = env
+                env.enter_let(stack.tos()?.clone().unwrap_attr_set().into_inner());
-                    .clone()
-                    .enter_attrs(stack.tos()?.clone().unwrap_attr_set().into_inner());
                 stack.tos_mut()?.as_mut().unwrap_attr_set().capture(&env);
             }
             OpCode::PushStaticAttr { name } => {
@@ -251,12 +249,18 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
                         .clone(),
                 )?;
             }
-            OpCode::EnterEnv => match stack.pop() {
+            OpCode::EnterLetEnv => match stack.pop() {
-                Value::AttrSet(attrs) => *env = env.clone().enter_with(attrs.into_inner()),
+                Value::AttrSet(attrs) => env.enter_let(attrs.into_inner()),
 
                 _ => unreachable!(),
             },
-            OpCode::LeaveEnv => *env = env.clone().leave(),
+            OpCode::LeaveLetEnv => env.leave_let(),
+            OpCode::EnterWithEnv => match stack.pop() {
+                Value::AttrSet(attrs) => env.enter_with(attrs.into_inner()),
+
+                _ => unreachable!(),
+            },
+            OpCode::LeaveWithEnv => env.leave_with(),
             OpCode::Assert => {
                 if !stack.pop().unwrap_const().unwrap_bool() {
                     todo!()