fix: thunk & function

src/ty/internal/mod.rs

@@ -1,3 +1,4 @@
+use std::cell::OnceCell;
 use std::cell::RefCell;
 use std::rc::Rc;
 
@@ -36,7 +37,7 @@ pub enum Value<'vm> {
     Catchable(c::Catchable),
     PrimOp(PrimOp<'vm>),
     PartialPrimOp(PartialPrimOp<'vm>),
-    Func(&'vm Func<'vm>),
+    Func(Func<'vm>),
 }
 
 #[derive(Debug, IsVariant, Unwrap, Clone, PartialEq)]
@@ -49,7 +50,7 @@ pub enum ValueAsRef<'v, 'vm: 'v> {
     Catchable(&'v c::Catchable),
     PrimOp(&'v PrimOp<'vm>),
     PartialPrimOp(&'v PartialPrimOp<'vm>),
-    Func(&'vm Func<'vm>),
+    Func(&'v Func<'vm>),
 }
 
 #[derive(Debug, IsVariant, Unwrap, PartialEq)]
@@ -62,7 +63,7 @@ pub enum ValueAsMut<'v, 'vm: 'v> {
     Catchable(&'v mut c::Catchable),
     PrimOp(&'v mut PrimOp<'vm>),
     PartialPrimOp(&'v mut PartialPrimOp<'vm>),
-    Func(&'vm Func<'vm>),
+    Func(&'v Func<'vm>),
 }
 
 impl<'v, 'vm: 'v> Value<'vm> {
@@ -127,7 +128,7 @@ impl<'vm> Value<'vm> {
         }
     }
 
-    pub fn call(self, vm: &VM<'vm, '_>, args: Vec<Value<'vm>>) -> Result<Value<'vm>> {
+    pub fn call(self, vm: &'vm VM<'_>, args: Vec<Value<'vm>>) -> Result<Value<'vm>> {
         use Value::*;
         Ok(match self {
             PrimOp(func) => func.call(vm, args),
@@ -365,7 +366,7 @@ impl<'vm> Value<'vm> {
         self
     }
 
-    pub fn force(&mut self, vm: &VM<'vm, '_>) -> Result<&mut Self> {
+    pub fn force(&mut self, vm: &'vm VM<'_>) -> Result<&mut Self> {
         if let Value::Thunk(thunk) = self {
             let value = thunk.force(vm)?;
             *self = value
@@ -376,7 +377,7 @@ impl<'vm> Value<'vm> {
         Ok(self)
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<&mut Self> {
+    pub fn force_deep(&mut self, vm: &'vm VM<'_>) -> Result<&mut Self> {
         match self {
             Value::Thunk(thunk) => {
                 let mut value = thunk.force(vm)?;
@@ -420,34 +421,30 @@ pub trait ToPublic {
 
 #[derive(Debug, Clone)]
 pub struct Thunk<'vm> {
-    pub thunk: RefCell<_Thunk<'vm>>,
-    pub env: RefCell<Option<Rc<Env<'vm>>>>,
+    pub thunk: Rc<RefCell<_Thunk<'vm>>>,
 }
 
 #[derive(Debug, IsVariant, Unwrap, Clone)]
 pub enum _Thunk<'vm> {
-    Code(OpCodes),
+    Code(&'vm OpCodes, OnceCell<Rc<Env<'vm>>>),
     SuspendedFrom(*const Thunk<'vm>),
     Value(Box<Value<'vm>>),
 }
 
 impl<'vm> Thunk<'vm> {
-    pub fn new(opcodes: OpCodes) -> Self {
+    pub fn new(opcodes: &'vm OpCodes) -> Self {
         Thunk {
-            thunk: RefCell::new(_Thunk::Code(opcodes)),
-            env: RefCell::new(None),
+            thunk: Rc::new(RefCell::new(_Thunk::Code(opcodes, OnceCell::new()))),
         }
     }
 
-    pub fn unwrap_code(&self) -> OpCodes {
-        self.thunk.borrow().clone().unwrap_code()
-    }
-
     pub fn capture(&self, env: Rc<Env<'vm>>) {
-        *self.env.borrow_mut() = Some(env);
+        if let _Thunk::Code(_, envcell) = &*self.thunk.borrow() {
+            envcell.get_or_init(|| env);
+        }
     }
 
-    pub fn force(&self, vm: &VM<'vm, '_>) -> Result<Value<'vm>> {
+    pub fn force(&self, vm: &'vm VM<'_>) -> Result<Value<'vm>> {
         match &*self.thunk.borrow() {
             _Thunk::Value(value) => return Ok(value.as_ref().clone()),
             _Thunk::SuspendedFrom(from) => {
@@ -456,14 +453,14 @@ impl<'vm> Thunk<'vm> {
                     self as *const Thunk
                 )));
             }
-            _Thunk::Code(_) => (),
+            _Thunk::Code(..) => (),
         }
-        let opcodes = std::mem::replace(
+        let (opcodes, env) = std::mem::replace(
             &mut *self.thunk.borrow_mut(),
             _Thunk::SuspendedFrom(self as *const Thunk),
         )
         .unwrap_code();
-        let value = vm.eval(opcodes, self.env.borrow().clone().unwrap())?;
+        let value = vm.eval(opcodes.clone(), env.get().unwrap().clone())?;
         let _ = std::mem::replace(
             &mut *self.thunk.borrow_mut(),
             _Thunk::Value(value.clone().into()),