fix: thunk & function

src/builtins/mod.rs

@@ -54,10 +54,8 @@ pub fn env<'vm>() -> Rc<Env<'vm>> {
     }
     let builtins = Value::RecAttrSet(RecAttrSet::from_inner(map.clone()));
     let thunk = Thunk {
-        thunk: RefCell::new(_Thunk::Value(Box::new(builtins.clone()))),
+        thunk: Rc::new(RefCell::new(_Thunk::Value(Box::new(builtins.clone())))),
-        env: RefCell::default(),
     };
-    thunk.capture(builtins_env);
     map.insert(Symbol::from("builtins"), Value::Thunk(thunk));
 
     env.insert(Symbol::from("builtins"), builtins);

src/bytecode.rs

@@ -14,8 +14,8 @@ pub enum OpCode {
     LookUp { sym: EcoString },
     /// load a thunk lazily onto stack
     LoadThunk { idx: usize },
-    /// load a thunk onto stack and force its value
+    /// TODO:
-    LoadValue { idx: usize },
+    CaptureEnv,
     /// force TOS to value
     ForceValue,
 

src/compile.rs

@@ -98,11 +98,17 @@ impl Compile for ir::Attrs {
         comp.push(OpCode::AttrSet);
         for stc in self.stcs {
             stc.1.compile(comp);
+            if !self.rec {
+                comp.push(OpCode::CaptureEnv);
+            }
             comp.push(OpCode::PushStaticAttr { name: stc.0 });
         }
         for dynamic in self.dyns {
             dynamic.0.compile(comp);
             dynamic.1.compile(comp);
+            if !self.rec {
+                comp.push(OpCode::CaptureEnv);
+            }
             comp.push(OpCode::PushDynamicAttr)
         }
         if self.rec {

src/ty/internal/attrset.rs

@@ -42,23 +42,18 @@ impl<'vm> AttrSet<'vm> {
         self.data.get(&sym).is_some()
     }
 
-    pub fn capture(&mut self, env: Rc<Env<'vm>>) {
+    pub fn capture(&mut self, env: &Rc<Env<'vm>>) {
-        self.data = self
+        self
             .data
-            .into_iter()
+            .iter()
-            .map(|(k, v)| {
+            .for_each(|(_, v)| {
-                (
-                    k.clone(),
                     match v.clone() {
-                        x @ Value::ThunkRef(thunk) => {
+                        Value::Thunk(ref thunk) => {
                             thunk.capture(env.clone());
-                            x
                         }
-                        x => x,
+                        _ => ()
-                    },
+                    }
-                )
             })
-            .collect();
     }
 
     pub fn update(mut self, other: AttrSet<'vm>) -> AttrSet<'vm> {
@@ -83,7 +78,7 @@ impl<'vm> AttrSet<'vm> {
         &self.data
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> {
+    pub fn force_deep(&mut self, vm: &'vm VM<'_>) -> Result<()> {
         let mut map: Vec<_> = self
             .data
             .into_iter()
@@ -170,7 +165,7 @@ impl<'vm> RecAttrSet<'vm> {
         RecAttrSet { data }
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> {
+    pub fn force_deep(&mut self, vm: &'vm VM<'_>) -> Result<()> {
         let mut map: Vec<_> = self
             .data
             .map

src/ty/internal/func.rs

@@ -1,14 +1,15 @@
-use std::cell::OnceCell;
+use std::rc::Rc;
 
 use ecow::EcoString;
 use itertools::Itertools;
 use rpds::HashTrieMap;
+use derive_more::Constructor;
 
-use crate::bytecode::OpCodes;
+use crate::bytecode::Func as BFunc;
 use crate::error::Result;
 use crate::ir;
-use crate::ty::internal::Value;
+use crate::ty::internal::{Thunk, Value};
-use crate::vm::{CapturedEnv, Env, VM};
+use crate::vm::{Env, VM};
 
 #[derive(Debug, Clone)]
 pub enum Param {
@@ -40,24 +41,23 @@ impl From<ir::Param> for Param {
     }
 }
 
-pub type JITFunc<'vm> = unsafe extern "C" fn(vm: *mut VM<'vm, '_>, *mut Env<'vm>, *mut Value<'vm>) -> Value<'vm>;
+pub type JITFunc<'vm> = unsafe extern "C" fn(vm: *mut VM<'_>, *mut Env<'vm>, *mut Value<'vm>) -> Value<'vm>;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Constructor)]
 pub struct Func<'vm> {
-    pub env: OnceCell<CapturedEnv<'vm>>,
+    pub func: &'vm BFunc,
-    pub param: Param,
+    pub env: Rc<Env<'vm>>,
-    pub opcodes: OpCodes,
     pub compiled: Option<JITFunc<'vm>>
 }
 
 impl<'vm> Func<'vm> {
-    pub fn call(&'vm self, vm: &VM<'vm, '_>, arg: Value<'vm>) -> Result<Value<'vm>> {
+    pub fn call(self, vm: &'vm VM<'_>, arg: Value<'vm>) -> Result<Value<'vm>> {
         use Param::*;
 
-        let env = self.env.get().unwrap().clone().released();
+        let env = Rc::new(self.env.as_ref().clone());
 
-        match &self.param {
+        match self.func.param.clone() {
-            Ident(ident) => env.enter(HashTrieMap::new().insert(ident.clone().into(), arg)),
+            Ident(ident) => env.enter(HashTrieMap::new().insert(ident.into(), arg)),
             Formals {
                 formals,
                 ellipsis,
@@ -79,7 +79,7 @@ impl<'vm> Func<'vm> {
                     let formal = formal.clone().into();
                     let arg = arg
                         .select(&formal)
-                        .or_else(|| default.map(|idx| Value::ThunkRef(vm.get_thunk(idx))))
+                        .or_else(|| default.map(|idx| Value::Thunk(Thunk::new(vm.get_thunk(idx)))))
                         .unwrap();
                     new.insert_mut(formal, arg);
                 }
@@ -90,7 +90,7 @@ impl<'vm> Func<'vm> {
             }
         }
 
-        vm.eval(self.opcodes.clone(), env)
+        vm.eval(self.func.opcodes.clone(), env)
     }
 }
 

src/ty/internal/list.rs

@@ -30,7 +30,7 @@ impl<'vm> List<'vm> {
         self
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> {
+    pub fn force_deep(&mut self, vm: &'vm VM<'_>) -> Result<()> {
         let mut vec: Vec<_> = self.data.iter().cloned().collect();
         vec.iter_mut()
             .map(|v| v.force_deep(vm).map(|_| ()))

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 thunk: Rc<RefCell<_Thunk<'vm>>>,
-    pub env: RefCell<Option<Rc<Env<'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)),
+            thunk: Rc::new(RefCell::new(_Thunk::Code(opcodes, OnceCell::new()))),
-            env: RefCell::new(None),
         }
     }
 
-    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()),

src/ty/internal/primop.rs

@@ -8,7 +8,7 @@ use super::Value;
 pub struct PrimOp<'vm> {
     pub name: &'static str,
     arity: u8,
-    func: fn(&VM<'vm, '_>, Vec<Value<'vm>>) -> Value<'vm>,
+    func: fn(&'vm VM<'_>, Vec<Value<'vm>>) -> Value<'vm>,
 }
 
 impl PartialEq for PrimOp<'_> {
@@ -18,7 +18,7 @@ impl PartialEq for PrimOp<'_> {
 }
 
 impl<'vm> PrimOp<'vm> {
-    pub fn call(self, vm: &VM<'vm, '_>, args: Vec<Value<'vm>>) -> Value<'vm> {
+    pub fn call(self, vm: &'vm VM<'_>, args: Vec<Value<'vm>>) -> Value<'vm> {
         if (args.len() as u8) < self.arity {
             Value::PartialPrimOp(PartialPrimOp {
                 name: self.name,
@@ -39,7 +39,7 @@ pub struct PartialPrimOp<'vm> {
     pub name: &'static str,
     arity: u8,
     args: Vec<Value<'vm>>,
-    func: fn(&VM<'vm, '_>, Vec<Value<'vm>>) -> Value<'vm>,
+    func: fn(&'vm VM<'_>, Vec<Value<'vm>>) -> Value<'vm>,
 }
 
 impl PartialEq for PartialPrimOp<'_> {
@@ -49,7 +49,7 @@ impl PartialEq for PartialPrimOp<'_> {
 }
 
 impl<'vm> PartialPrimOp<'vm> {
-    pub fn call(mut self, vm: &VM<'vm, '_>, args: Vec<Value<'vm>>) -> Value<'vm> {
+    pub fn call(mut self, vm: &'vm VM<'_>, args: Vec<Value<'vm>>) -> Value<'vm> {
         let len = args.len() as u8;
         self.args.extend(args);
         if len < self.arity {

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
+        prog.funcs,
-            .into_iter()
-            .map(|f| Func {
-                env: OnceCell::new(),
-                param: f.param,
-                opcodes: f.opcodes,
-                compiled: None,
-            })
-            .collect(),
         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> {
+pub struct VM<'jit> {
-    thunks: Box<[Thunk<'vm>]>,
+    thunks: Box<[OpCodes]>,
-    funcs: Box<[Func<'vm>]>,
+    funcs: Box<[F]>,
     jit: JITContext<'jit>,
 }
 
-impl<'vm, 'jit> VM<'vm, 'jit> {
+impl<'vm, 'jit: 'vm> VM<'jit> {
-    fn new(thunks: Box<[OpCodes]>, funcs: Box<[Func<'vm>]>, jit: JITContext<'jit>) -> Self {
+    fn new(thunks: Box<[OpCodes]>, funcs: Box<[F]>, jit: JITContext<'jit>) -> Self {
-        let thunks = thunks
-            .into_iter()
-            .map(|opcodes| Thunk::new(opcodes))
-            .collect();
         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);
+                stack.push(Value::Thunk(Thunk::new(self.get_thunk(idx))))?
-                thunk.capture(env.clone());
+            }
-                stack.push(Value::ThunkRef(self.get_thunk(idx)))?
+            OpCode::CaptureEnv => {
+                match stack.tos()? {
+                    Value::Thunk(thunk) => thunk.capture(env.clone()),
+                    _ => ()
                 }
-            OpCode::LoadValue { idx } => {
-                stack.push(self.get_thunk(idx).force(self)?)?;
             }
             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::new(func, env.clone(), None)))?;
-                stack.push(Value::Func(func))?;
             }
             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();