feat: ref

src/compile.rs

@@ -1,3 +1,5 @@
+use std::cell::OnceCell;
+
 use crate::bytecode::*;
 use crate::ir;
 use crate::ty::internal::Const;
@@ -19,7 +21,7 @@ pub fn compile(downgraded: ir::Downgraded) -> Program {
             .funcs
             .into_iter()
             .map(|func| Func {
-                env: None,
+                env: OnceCell::new(),
                 param: func.param.into(),
                 opcodes: Compiler::new().compile(*func.body),
             })

src/ty/internal/attrset.rs

@@ -17,7 +17,7 @@ pub struct AttrSet {
 }
 
 impl AttrSet {
-    pub fn empty() -> AttrSet {
+    pub fn empty() -> Self {
         AttrSet {
             data: HashTrieMapSync::new_sync(),
         }

src/ty/internal/func.rs

@@ -1,3 +1,5 @@
+use std::cell::OnceCell;
+
 use ecow::EcoString;
 use itertools::Itertools;
 use rpds::HashTrieMap;
@@ -40,19 +42,19 @@ impl From<ir::Param> for Param {
 
 #[derive(Debug, Clone)]
 pub struct Func {
-    pub env: Option<CapturedEnv>,
+    pub env: OnceCell<CapturedEnv>,
     pub param: Param,
     pub opcodes: OpCodes,
 }
 
 impl Func {
-    pub fn call(self, vm: &VM, arg: Value) -> Result<Value> {
+    pub fn call(&self, vm: &VM, arg: Value) -> Result<Value> {
         use Param::*;
 
-        let env = self.env.unwrap().released();
+        let env = self.env.get().unwrap().clone().released();
 
-        match self.param {
+        match &self.param {
-            Ident(ident) => env.enter(HashTrieMap::new_sync().insert(ident.into(), arg)),
+            Ident(ident) => env.enter(HashTrieMap::new_sync().insert(ident.clone().into(), arg)),
             Formals {
                 formals,
                 ellipsis,
@@ -73,18 +75,18 @@ impl Func {
                 for (formal, default) in formals {
                     let arg = arg
                         .select(formal.clone().into())
-                        .or_else(|| default.map(|idx| Value::Thunk(vm.get_thunk(idx))))
+                        .or_else(|| default.map(Value::ThunkRef))
                         .unwrap();
-                    new.insert_mut(formal.into(), arg);
+                    new.insert_mut(formal.clone().into(), arg);
                 }
                 if let Some(alias) = alias {
-                    new.insert_mut(alias.into(), Value::AttrSet(arg));
+                    new.insert_mut(alias.clone().into(), Value::AttrSet(arg));
                 }
                 env.enter(new);
             }
         }
 
-        vm.eval(self.opcodes, env)
+        vm.eval(self.opcodes.clone(), env)
     }
 }
 

src/ty/internal/mod.rs

@@ -101,48 +101,52 @@ impl PartialEq for Thunk {
 pub enum Value {
     Const(Const),
     Thunk(Thunk),
+    ThunkRef(usize),
     AttrSet(AttrSet),
     RecAttrSet(RecAttrSet),
     List(List),
     Catchable(c::Catchable),
     PrimOp(PrimOp),
     PartialPrimOp(PartialPrimOp),
-    Func(Func),
+    Func(usize),
 }
 
 #[derive(Debug, IsVariant, Unwrap, Clone, PartialEq)]
-pub enum ValueAsRef<'a> {
+pub enum ValueAsRef<'v> {
-    Const(&'a Const),
+    Const(&'v Const),
-    Thunk(&'a Thunk),
+    Thunk(&'v Thunk),
-    AttrSet(&'a AttrSet),
+    ThunkRef(&'v usize),
-    RecAttrSet(&'a RecAttrSet),
+    AttrSet(&'v AttrSet),
-    List(&'a List),
+    RecAttrSet(&'v RecAttrSet),
-    Catchable(&'a c::Catchable),
+    List(&'v List),
-    PrimOp(&'a PrimOp),
+    Catchable(&'v c::Catchable),
-    PartialPrimOp(&'a PartialPrimOp),
+    PrimOp(&'v PrimOp),
-    Func(&'a Func),
+    PartialPrimOp(&'v PartialPrimOp),
+    Func(&'v usize),
 }
 
 #[derive(Debug, IsVariant, Unwrap, PartialEq)]
-pub enum ValueAsMut<'a> {
+pub enum ValueAsMut<'v> {
-    Const(&'a mut Const),
+    Const(&'v mut Const),
-    Thunk(&'a mut Thunk),
+    Thunk(&'v Thunk),
-    AttrSet(&'a mut AttrSet),
+    ThunkRef(&'v usize),
-    RecAttrSet(&'a mut RecAttrSet),
+    AttrSet(&'v mut AttrSet),
-    List(&'a mut List),
+    RecAttrSet(&'v mut RecAttrSet),
-    Catchable(&'a mut c::Catchable),
+    List(&'v mut List),
-    PrimOp(&'a mut PrimOp),
+    Catchable(&'v mut c::Catchable),
-    PartialPrimOp(&'a mut PartialPrimOp),
+    PrimOp(&'v mut PrimOp),
-    Func(&'a mut Func),
+    PartialPrimOp(&'v mut PartialPrimOp),
+    Func(&'v usize),
 }
 
 impl Value {
-    pub fn as_ref(&self) -> ValueAsRef<'_> {
+    pub fn as_ref(&self) -> ValueAsRef {
         use Value::*;
         use ValueAsRef as R;
         match self {
             Const(x) => R::Const(x),
             Thunk(x) => R::Thunk(x),
+            ThunkRef(x) => R::ThunkRef(x),
             AttrSet(x) => R::AttrSet(x),
             RecAttrSet(x) => R::RecAttrSet(x),
             List(x) => R::List(x),
@@ -153,12 +157,13 @@ impl Value {
         }
     }
 
-    pub fn as_mut(&mut self) -> ValueAsMut<'_> {
+    pub fn as_mut(&mut self) -> ValueAsMut {
         use Value::*;
         use ValueAsMut as M;
         match self {
             Const(x) => M::Const(x),
             Thunk(x) => M::Thunk(x),
+            ThunkRef(x) => M::ThunkRef(x),
             AttrSet(x) => M::AttrSet(x),
             RecAttrSet(x) => M::RecAttrSet(x),
             List(x) => M::List(x),
@@ -177,6 +182,7 @@ impl Value {
         match self {
             Const(_) => todo!(),
             Thunk(_) => "thunk",
+            ThunkRef(_) => "thunk",
             AttrSet(_) => "set",
             RecAttrSet(_) => "set",
             List(_) => "list",
@@ -210,7 +216,7 @@ impl Value {
                         PartialPrimOp(func) => {
                             return Ok(func.call(vm, [arg].into_iter().chain(iter).collect()));
                         }
-                        Func(func) => func.call(vm, arg)?,
+                        Func(func) => vm.get_func(func).call(vm, arg)?,
                         _ => todo!(),
                     }
                 }
@@ -361,22 +367,19 @@ impl Value {
     }
 
     pub fn select(&mut self, sym: Symbol) -> &mut Self {
-        if let Value::AttrSet(attrs) = self {
+        let val = match self {
-            let val = attrs.select(sym.clone()).unwrap_or_else(|| {
+            Value::AttrSet(attrs) => attrs.select(sym.clone()).unwrap_or_else(|| {
                 Value::Catchable(c::Catchable::new(Some(format!("{sym:?} not found"))))
-            });
+            }),
-            *self = val;
+            Value::RecAttrSet(attrs) => attrs.select(sym.clone()).unwrap_or_else(|| {
-        } else if let Value::RecAttrSet(attrs) = self {
-            let val = attrs.select(sym.clone()).unwrap_or_else(|| {
                 Value::Catchable(c::Catchable::new(Some(format!("{sym:?} not found"))))
-            });
+            }),
-            *self = val;
+            _ => Value::Catchable(Catchable::new(Some(format!(
-        } else {
-            *self = Value::Catchable(Catchable::new(Some(format!(
                 "cannot select from {:?}",
                 self.typename()
-            ))))
+            )))),
-        }
+        };
+        *self = val;
         self
     }
 
@@ -426,7 +429,7 @@ impl Value {
         match self {
             Value::Thunk(thunk) => {
                 let mut value = thunk.force(vm)?;
-                value.force_deep(vm)?;
+                let _ = value.force_deep(vm)?;
                 *self = value;
             }
             Value::List(list) => list.force_deep(vm)?,
@@ -447,6 +450,7 @@ impl ToPublic for Value {
             Value::Catchable(catchable) => p::Value::Catchable(catchable),
             Value::Const(cnst) => p::Value::Const(cnst.into()),
             Value::Thunk(_) => p::Value::Thunk,
+            Value::ThunkRef(_) => p::Value::Thunk,
             Value::PrimOp(primop) => p::Value::PrimOp(primop.name),
             Value::PartialPrimOp(primop) => p::Value::PartialPrimOp(primop.name),
             Value::Func(_) => p::Value::Func,

src/vm/mod.rs

@@ -37,12 +37,12 @@ impl VM {
         VM { thunks, funcs }
     }
 
-    pub fn get_thunk(&self, idx: usize) -> Thunk {
+    pub fn get_thunk(&self, idx: usize) -> &Thunk {
-        self.thunks[idx].clone()
+        &self.thunks[idx]
     }
 
-    pub fn get_func(&self, idx: usize) -> Func {
+    pub fn get_func(&self, idx: usize) -> &Func {
-        self.funcs[idx].clone()
+        &self.funcs[idx]
     }
 
     pub fn eval(&self, opcodes: OpCodes, env: Arc<Env>) -> Result<Value> {
@@ -101,9 +101,8 @@ impl VM {
                 stack.push(func.call(self, args)?)?;
             }
             OpCode::Func { idx } => {
-                let mut func = self.get_func(idx);
+                self.get_func(idx).env.get_or_init(|| env.captured());
-                func.env = Some(env.captured());
+                stack.push(Value::Func(idx))?;
-                stack.push(Value::Func(func))?;
             }
             OpCode::UnOp { op } => {
                 use UnOp::*;