feat: get rid of gc and cyclic thunk

src/vm/mod.rs

@@ -1,8 +1,6 @@
 use std::cell::RefCell;
 use std::rc::Rc;
 
-use gc_arena::arena::CollectionPhase;
-use gc_arena::{Arena, Collect, Gc, Mutation, Rootable};
 use hashbrown::{HashMap, HashSet};
 use inkwell::context::Context;
 
@@ -25,110 +23,65 @@ mod test;
 const STACK_SIZE: usize = 8 * 1024 / size_of::<Value>();
 const COLLECTOR_GRANULARITY: f64 = 1024.0;
 
-#[derive(Collect)]
-#[collect(require_static)]
 struct ContextWrapper(Context);
 
 pub fn run(mut prog: Program) -> Result<p::Value> {
-    let mut arena: Arena<Rootable![GcRoot<'_>]> = Arena::new(|mc| {
-        let jit = Gc::new(mc, ContextWrapper(Context::create()));
-        let mut thunks = std::mem::take(&mut prog.thunks);
-        thunks.iter_mut().for_each(|code| code.reverse());
-        let mut funcs = std::mem::take(&mut prog.funcs);
-        funcs
-            .iter_mut()
-            .for_each(|F { opcodes, .. }| opcodes.reverse());
-        let symbols = std::mem::take(&mut prog.symbols);
-        let symmap = std::mem::take(&mut prog.symmap);
-        let consts = std::mem::take(&mut prog.consts);
-        let vm = VM {
-            thunks,
-            funcs,
-            consts,
-            symbols: RefCell::new(symbols),
-            symmap: RefCell::new(symmap),
-            jit: JITContext::new(&jit.as_ref().0),
-        };
-        let vm = Gc::new(mc, vm);
-        GcRoot {
-            vm,
-            jit,
-            stack: Stack::new(),
-            envs: vec![vm_env(&vm, mc)],
-        }
-    });
+    let jit = Context::create();
+    prog.thunks.iter_mut().for_each(|code| code.reverse());
+    prog.funcs
+        .iter_mut()
+        .for_each(|F { opcodes, .. }| opcodes.reverse());
+    let vm = VM {
+        thunks: prog.thunks,
+        funcs: prog.funcs,
+        consts: prog.consts,
+        symbols: RefCell::new(prog.symbols),
+        symmap: RefCell::new(prog.symmap),
+        jit: JITContext::new(&jit),
+    };
     prog.top_level.reverse();
-    eval(prog.top_level, &mut arena, |val, root, _| {
-        Ok(val.to_public(&root.vm, &mut HashSet::new()))
-    })
+    Ok(eval(prog.top_level, &vm, vm_env(&vm))?.to_public(&vm, &mut HashSet::new()))
 }
 
-pub fn eval<T, F: for<'gc> FnOnce(Value<'gc>, &mut GcRoot<'gc>, &Mutation<'gc>) -> Result<T>>(
-    opcodes: Box<[OpCode]>,
-    arena: &mut Arena<impl for<'gc> Rootable<'gc, Root = GcRoot<'gc>>>,
-    f: F,
-) -> Result<T> {
+pub fn eval<'gc>(opcodes: Box<[OpCode]>, vm: &VM<'gc>, env: Rc<VmEnv<'gc>>) -> Result<Value<'gc>> {
     let mut opcodes = opcodes.into_vec();
+    let mut envs = vec![env];
+    let mut stack = Stack::<_, STACK_SIZE>::new();
     while let Some(opcode) = opcodes.pop() {
-        arena.mutate_root(|mc, root| {
-            let consq = single_op(
-                &root.vm,
-                opcode,
-                &mut root.stack,
-                root.envs.last_mut().unwrap(),
-                mc,
-            )?;
-            match consq {
-                Consq::NoOp => (),
-                Consq::Jmp(step) => opcodes.resize_with(opcodes.len() - step, || unreachable!()),
-                Consq::Force => {
-                    let thunk = root.stack.tos().as_ref().unwrap_thunk();
-                    let (code, env) = thunk.suspend(mc)?;
-                    opcodes.push(OpCode::InsertValue);
-                    opcodes.push(OpCode::PopEnv);
-                    opcodes.extend(code);
-                    root.envs.push(env);
-                }
-                Consq::PopEnv => {
-                    let _ = root.envs.pop().unwrap();
-                }
-                Consq::Call => {
-                    let arg = root.stack.pop();
-                    let func = root.stack.pop().unwrap_func();
-                    let env = func.env.enter_arg(arg, mc);
-                    let count = func.count.get();
-                    func.count.set(count + 1);
-                    if count > 1 {
-                        let compiled = func
-                            .compiled
-                            .borrow_mut(mc)
-                            .get_or_insert_with(|| root.vm.compile_func(func.func))
-                            .clone();
-                        let ret =
-                            unsafe { compiled(env.as_ref() as *const VmEnv, mc as *const _) };
-                        root.stack.push(ret.into())?;
-                    } else {
-                        root.envs.push(env);
-                        opcodes.push(OpCode::PopEnv);
-                        opcodes.extend(&func.func.opcodes);
-                    }
-                }
+        let consq = single_op(vm, opcode, &mut stack, envs.last_mut().unwrap())?;
+        match consq {
+            Consq::NoOp => (),
+            Consq::Jmp(step) => opcodes.resize_with(opcodes.len() - step, || unreachable!()),
+            Consq::Force => {
+                let thunk = stack.tos().as_ref().unwrap_thunk();
+                let (code, env) = thunk.suspend()?;
+                opcodes.push(OpCode::InsertValue);
+                opcodes.push(OpCode::PopEnv);
+                opcodes.extend(code);
+                envs.push(env);
             }
-            Result::Ok(())
-        })?;
-        if arena.metrics().allocation_debt() > COLLECTOR_GRANULARITY {
-            if arena.collection_phase() == CollectionPhase::Sweeping {
-                arena.collect_debt();
-            } else if let Some(marked) = arena.mark_debt() {
-                marked.start_sweeping();
+            Consq::PopEnv => {
+                let _ = envs.pop().unwrap();
+            }
+            Consq::Call => {
+                let arg = stack.pop();
+                let func = stack.pop().unwrap_func();
+                let env = func.env.clone().enter_arg(arg);
+                let count = func.count.get();
+                func.count.set(count + 1);
+                if count > 1 {
+                    let compiled = func.compiled.get_or_init(|| vm.compile_func(func.func));
+                    let ret = unsafe { compiled(env.as_ref() as *const VmEnv) };
+                    stack.push(ret.into())?;
+                } else {
+                    envs.push(env);
+                    opcodes.push(OpCode::PopEnv);
+                    opcodes.extend(&func.func.opcodes);
+                }
             }
         }
     }
-    arena.mutate_root(|mc, root| {
-        assert_eq!(root.stack.len(), 1);
-        let ret = root.stack.pop();
-        f(ret, root, mc)
-    })
+    stack.pop().ok()
 }
 
 enum Consq {
@@ -139,18 +92,12 @@ enum Consq {
     NoOp,
 }
 
-enum CycleAction {
-    Collect,
-    NoOp
-}
-
 #[inline(always)]
 fn single_op<'gc, const CAP: usize>(
-    vm: &'gc VM<'gc>,
+    vm: &VM<'gc>,
     opcode: OpCode,
     stack: &mut Stack<Value<'gc>, CAP>,
-    env: &mut Gc<'gc, VmEnv<'gc>>,
-    mc: &'gc Mutation<'gc>,
+    env: &mut Rc<VmEnv<'gc>>,
 ) -> Result<Consq> {
     match opcode {
         OpCode::Illegal => panic!("illegal opcode"),
@@ -161,13 +108,13 @@ fn single_op<'gc, const CAP: usize>(
             Const::String(x) => Value::String(Rc::new(x.into())),
             Const::Null => Value::Null,
         })?,
-        OpCode::LoadThunk { idx } => stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx), mc)))?,
+        OpCode::LoadThunk { idx } => stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx))))?,
         OpCode::LoadValue { idx } => {
-            stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx), mc)))?;
-            stack.tos().as_ref().unwrap_thunk().capture_env(*env, mc);
+            stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx))))?;
+            stack.tos().as_ref().unwrap_thunk().capture_env(env.clone());
             return Ok(Consq::Force);
         }
-        OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture_env(*env, mc),
+        OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture_env(env.clone()),
         OpCode::ForceValue => {
             if !stack.tos().is_thunk() {
                 return Ok(Consq::NoOp);
@@ -179,7 +126,11 @@ fn single_op<'gc, const CAP: usize>(
         }
         OpCode::InsertValue => {
             let val = stack.pop();
-            stack.tos().as_ref().unwrap_thunk().insert_value(val.clone(), mc);
+            stack
+                .tos()
+                .as_ref()
+                .unwrap_thunk()
+                .insert_value(val.clone());
             *stack.tos_mut() = val;
         }
         OpCode::Jmp { step } => return Ok(Consq::Jmp(step)),
@@ -195,11 +146,11 @@ fn single_op<'gc, const CAP: usize>(
                 let _ = stack.push(arg);
                 return Ok(Consq::Call);
             }
-            func.call(arg, vm, mc)?;
+            func.call(arg, vm)?;
         }
         OpCode::Func { idx } => {
             let func = vm.get_func(idx);
-            stack.push(Value::Func(Rc::new(Func::new(func, *env, mc))))?;
+            stack.push(Value::Func(Rc::new(Func::new(func, env.clone()))))?;
         }
         OpCode::Arg { level } => {
             stack.push(env.lookup_arg(level).clone())?;
@@ -251,12 +202,9 @@ fn single_op<'gc, const CAP: usize>(
         }
         OpCode::FinalizeLet => {
             let mut list = stack.pop().unwrap_list();
-            let map = list
-                .as_ref()
-                .clone()
-                .into_inner();
-            *env = env.enter_let(map, mc);
-            Rc::make_mut(&mut list).capture(*env, mc);
+            let map = list.as_ref().clone().into_inner();
+            *env = env.clone().enter_let(map);
+            Rc::make_mut(&mut list).capture(&Rc::downgrade(env));
         }
         OpCode::PushStaticAttr { name } => {
             let val = stack.pop();
@@ -305,7 +253,11 @@ fn single_op<'gc, const CAP: usize>(
             )?;
         }
         OpCode::LookUpLet { level, idx } => {
-            stack.push(env.lookup_let(level, idx).clone())?;
+            let val = env.lookup_let(level, idx);
+            if let Value::Thunk(thunk) = val {
+                thunk.upgrade();
+            }
+            stack.push(val.clone())?;
         }
         OpCode::LeaveEnv => *env = env.leave(),
         OpCode::EnterWithEnv => {
@@ -317,7 +269,7 @@ fn single_op<'gc, const CAP: usize>(
                 .iter()
                 .map(|(&k, v)| (k, v.clone()))
                 .collect_into(&mut new);
-            *env = env.enter_with(Gc::new(mc, new), mc);
+            *env = env.clone().enter_with(new.into());
         }
         OpCode::PopEnv => return Ok(Consq::PopEnv),
         OpCode::Assert => {
@@ -329,18 +281,7 @@ fn single_op<'gc, const CAP: usize>(
     Ok(Consq::NoOp)
 }
 
-#[derive(Collect)]
-#[collect(no_drop)]
-pub struct GcRoot<'gc, const CAP: usize = STACK_SIZE> {
-    vm: Gc<'gc, VM<'gc>>,
-    jit: Gc<'gc, ContextWrapper>,
-
-    stack: Stack<Value<'gc>, CAP>,
-    envs: Vec<Gc<'gc, VmEnv<'gc>>>,
-}
-
-#[derive(Constructor, Collect)]
-#[collect(no_drop)]
+#[derive(Constructor)]
 pub struct VM<'gc> {
     thunks: Box<[OpCodes]>,
     funcs: Box<[F]>,
@@ -351,12 +292,12 @@ pub struct VM<'gc> {
 }
 
 impl<'gc> VM<'gc> {
-    pub fn get_thunk(&self, idx: usize) -> &OpCodes {
-        &self.thunks[idx]
+    pub fn get_thunk(&self, idx: usize) -> &'gc OpCodes {
+        unsafe { &*(&self.thunks[idx] as *const _) }
     }
 
-    pub fn get_func(&self, idx: usize) -> &F {
-        &self.funcs[idx]
+    pub fn get_func(&self, idx: usize) -> &'gc F {
+        unsafe { &*(&self.funcs[idx] as *const _) }
     }
 
     pub fn get_sym(&self, idx: usize) -> Symbol {
@@ -380,7 +321,7 @@ impl<'gc> VM<'gc> {
         self.consts[idx].clone()
     }
 
-    pub fn compile_func(&'gc self, func: &'gc F) -> JITFunc<'gc> {
+    pub fn compile_func(&self, func: &'gc F) -> JITFunc<'gc> {
         self.jit
             .compile_seq(func.opcodes.iter().copied().rev(), self)
             .unwrap()