fix: thunk & function

src/vm/env.rs

@@ -26,11 +26,6 @@ impl Clone for Env<'_> {
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct CapturedEnv<'vm> {
-    env: Rc<Env<'vm>>,
-}
-
 impl<'vm> Env<'vm> {
     pub fn empty() -> Self {
         Env::default()
@@ -78,18 +73,4 @@ impl<'vm> Env<'vm> {
         let map = last.map.borrow().clone();
         *self.map.borrow_mut() = map;
     }
-
-    pub fn captured(self: Rc<Self>) -> CapturedEnv<'vm> {
-        CapturedEnv { env: self }
-    }
-}
-
-impl<'vm> CapturedEnv<'vm> {
-    pub fn lookup(&self, symbol: &Symbol) -> Option<Value<'vm>> {
-        self.env.lookup(symbol)
-    }
-
-    pub fn released(self) -> Rc<Env<'vm>> {
-        Rc::new(self.env.as_ref().clone())
-    }
 }

src/vm/mod.rs

@@ -1,8 +1,7 @@
-use std::cell::OnceCell;
 use std::rc::Rc;
 
 use crate::builtins::env;
-use crate::bytecode::{BinOp, OpCode, OpCodes, Program, UnOp};
+use crate::bytecode::{BinOp, OpCode, OpCodes, Program, UnOp, Func as F};
 use crate::error::*;
 use crate::ty::common::Symbol;
 use crate::ty::internal::*;
@@ -10,7 +9,7 @@ use crate::ty::public as p;
 
 use stack::{STACK_SIZE, Stack};
 
-pub use env::{CapturedEnv, Env};
+pub use env::Env;
 pub use jit::JITContext;
 
 mod env;
@@ -23,15 +22,7 @@ mod test;
 pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
     let vm = VM::new(
         prog.thunks,
-        prog.funcs
-            .into_iter()
-            .map(|f| Func {
-                env: OnceCell::new(),
-                param: f.param,
-                opcodes: f.opcodes,
-                compiled: None,
-            })
-            .collect(),
+        prog.funcs,
         jit
     );
     let env = env();
@@ -40,36 +31,32 @@ pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
     Ok(temp)
 }
 
-pub struct VM<'vm, 'jit> {
-    thunks: Box<[Thunk<'vm>]>,
-    funcs: Box<[Func<'vm>]>,
+pub struct VM<'jit> {
+    thunks: Box<[OpCodes]>,
+    funcs: Box<[F]>,
     jit: JITContext<'jit>,
 }
 
-impl<'vm, 'jit> VM<'vm, 'jit> {
-    fn new(thunks: Box<[OpCodes]>, funcs: Box<[Func<'vm>]>, jit: JITContext<'jit>) -> Self {
-        let thunks = thunks
-            .into_iter()
-            .map(|opcodes| Thunk::new(opcodes))
-            .collect();
+impl<'vm, 'jit: 'vm> VM<'jit> {
+    fn new(thunks: Box<[OpCodes]>, funcs: Box<[F]>, jit: JITContext<'jit>) -> Self {
         VM { thunks, funcs, jit }
     }
 
-    pub fn get_thunk(&self, idx: usize) -> &'vm Thunk<'vm> {
+    pub fn get_thunk(&self, idx: usize) -> &OpCodes {
         // SAFETY: The `idx` is within bounds as `thunks` is initialized with `prog.thunks`
         // and `idx` is expected to be a valid index into this collection.
         // The lifetime of the returned reference is tied to `&self`.
         unsafe { &*(&self.thunks[idx] as *const _) }
     }
 
-    pub fn get_func(&self, idx: usize) -> &'vm Func<'vm> {
+    pub fn get_func(&self, idx: usize) -> &F {
         // SAFETY: The `idx` is within bounds as `funcs` is initialized with `prog.funcs`
         // and `idx` is expected to be a valid index into this collection.
         // The lifetime of the returned reference is tied to `&self`.
         unsafe { &*(&self.funcs[idx] as *const _) }
     }
 
-    pub fn eval(&self, opcodes: OpCodes, env: Rc<Env<'vm>>) -> Result<Value<'vm>> {
+    pub fn eval(&'vm self, opcodes: OpCodes, env: Rc<Env<'vm>>) -> Result<Value<'vm>> {
         let mut stack = Stack::<STACK_SIZE>::new();
         let mut iter = opcodes.into_iter();
         while let Some(opcode) = iter.next() {
@@ -86,7 +73,7 @@ impl<'vm, 'jit> VM<'vm, 'jit> {
 
     #[inline]
     fn single_op<'s, const CAP: usize>(
-        &self,
+        &'vm self,
         opcode: OpCode,
         stack: &'s mut Stack<'vm, CAP>,
         env: &Rc<Env<'vm>>,
@@ -95,12 +82,13 @@ impl<'vm, 'jit> VM<'vm, 'jit> {
             OpCode::Illegal => panic!("illegal opcode"),
             OpCode::Const { value } => stack.push(Value::Const(value))?,
             OpCode::LoadThunk { idx } => {
-                let thunk = self.get_thunk(idx);
-                thunk.capture(env.clone());
-                stack.push(Value::ThunkRef(self.get_thunk(idx)))?
+                stack.push(Value::Thunk(Thunk::new(self.get_thunk(idx))))?
             }
-            OpCode::LoadValue { idx } => {
-                stack.push(self.get_thunk(idx).force(self)?)?;
+            OpCode::CaptureEnv => {
+                match stack.tos()? {
+                    Value::Thunk(thunk) => thunk.capture(env.clone()),
+                    _ => ()
+                }
             }
             OpCode::ForceValue => {
                 stack.tos_mut()?.force(self)?;
@@ -127,8 +115,7 @@ impl<'vm, 'jit> VM<'vm, 'jit> {
             }
             OpCode::Func { idx } => {
                 let func = self.get_func(idx);
-                func.env.get_or_init(|| env.clone().captured());
-                stack.push(Value::Func(func))?;
+                stack.push(Value::Func(Func::new(func, env.clone(), None)))?;
             }
             OpCode::UnOp { op } => {
                 use UnOp::*;
@@ -177,7 +164,7 @@ impl<'vm, 'jit> VM<'vm, 'jit> {
                     .tos_mut()?
                     .as_mut()
                     .unwrap_attr_set()
-                    .capture(env.clone());
+                    .capture(env);
             }
             OpCode::PushStaticAttr { name } => {
                 let val = stack.pop();