optimize(env): single arg

src/ty/internal/attrset.rs

@@ -1,3 +1,5 @@
+use std::rc::Rc;
+
 use hashbrown::{HashMap, HashSet};
 
 use derive_more::Constructor;
@@ -47,7 +49,7 @@ impl<'jit: 'vm, 'vm> AttrSet<'jit, 'vm> {
         self.data.get(&sym).is_some()
     }
 
-    pub fn capture(&mut self, env: &LetEnv<'jit, 'vm>) {
+    pub fn capture(&mut self, env: &Rc<LetEnv<'jit, 'vm>>) {
         self.data.iter().for_each(|(_, v)| match v.clone() {
             Value::Thunk(ref thunk) => {
                 thunk.capture(env.clone());

src/ty/internal/func.rs

@@ -1,6 +1,8 @@
 use std::cell::{Cell, OnceCell};
+use std::rc::Rc;
 
 use derive_more::Constructor;
+use hashbrown::HashMap;
 use inkwell::execution_engine::JitFunction;
 use itertools::Itertools;
 
@@ -8,7 +10,7 @@ use crate::bytecode::Func as BFunc;
 use crate::error::Result;
 use crate::ir;
 use crate::jit::JITFunc;
-use crate::ty::internal::{Thunk, Value};
+use crate::ty::internal::{AttrSet, Thunk, Value};
 use crate::vm::{LetEnv, VM};
 
 #[derive(Debug, Clone)]
@@ -44,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: LetEnv<'jit, 'vm>,
+    pub env: Rc<LetEnv<'jit, 'vm>>,
     pub compiled: OnceCell<JitFunction<'jit, JITFunc<'jit, 'vm>>>,
     pub count: Cell<usize>,
 }
@@ -57,14 +59,14 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
             Ident(ident) => self
                 .env
                 .clone()
-                .enter_let([(ident.into(), arg)].into_iter()),
+                .enter_arg(ident.into(), arg),
             Formals {
                 formals,
                 ellipsis,
                 alias,
             } => {
                 let arg = arg.unwrap_attr_set();
-                let mut new = Vec::with_capacity(formals.len() + alias.iter().len());
+                let mut new = HashMap::with_capacity(formals.len() + alias.iter().len());
                 if !ellipsis
                     && arg
                         .as_inner()
@@ -83,20 +85,24 @@ impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
                             default.map(|idx| Value::Thunk(Thunk::new(vm.get_thunk(idx)).into()))
                         })
                         .unwrap();
-                    new.push((formal, arg));
+                    new.insert(formal, arg);
                 }
                 if let Some(alias) = alias {
-                    new.push((alias.clone().into(), Value::AttrSet(arg)));
+                    new.insert(alias.clone().into(), Value::AttrSet(arg));
                 }
-                self.env.clone().enter_let(new.into_iter())
+                self.env.clone().enter_attrs(AttrSet::new(new).into())
             }
-        };
+        }.into();
 
         let count = self.count.get();
         self.count.replace(count + 1);
         if count >= 1 {
             let compiled = self.compiled.get_or_init(|| vm.compile_func(self.func));
-            let ret = unsafe { compiled.call(vm as *const VM, &env as *const LetEnv) };
+            let env = Rc::into_raw(env);
+            let ret = unsafe { compiled.call(vm as *const VM, env) };
+            unsafe {
+                Rc::decrement_strong_count(env);
+            }
             return Ok(ret.into());
         }
         vm.eval(self.func.opcodes.iter().copied(), env)

src/ty/internal/mod.rs

@@ -485,7 +485,7 @@ pub struct Thunk<'jit, 'vm> {
 
 #[derive(Debug, IsVariant, Unwrap, Clone)]
 pub enum _Thunk<'jit, 'vm> {
-    Code(&'vm OpCodes, OnceCell<LetEnv<'jit, 'vm>>),
+    Code(&'vm OpCodes, OnceCell<Rc<LetEnv<'jit, 'vm>>>),
     SuspendedFrom(*const Thunk<'jit, 'vm>),
     Value(Value<'jit, 'vm>),
 }
@@ -497,7 +497,7 @@ impl<'jit, 'vm> Thunk<'jit, 'vm> {
         }
     }
 
-    pub fn capture(&self, env: LetEnv<'jit, 'vm>) {
+    pub fn capture(&self, env: Rc<LetEnv<'jit, 'vm>>) {
         if let _Thunk::Code(_, envcell) = &*self.thunk.borrow() {
             envcell.get_or_init(|| env);
         }