feat: gc-arena (WIP, does not compile)

src/vm/mod.rs

@@ -1,7 +1,9 @@
-use bumpalo::Bump;
+use std::cell::{Cell, RefCell};
+
+use gc_arena::lock::{GcRefLock, RefLock};
+use gc_arena::{Arena, Collect, DynamicRootSet, Gc, Mutation, Rootable};
 use hashbrown::{HashMap, HashSet};
-use inkwell::execution_engine::JitFunction;
-use std::cell::{Cell, OnceCell, RefCell};
+use inkwell::context::Context;
 
 use crate::builtins::env;
 use crate::bytecode::{BinOp, Func as F, OpCode, OpCodes, Program, UnOp};
@@ -20,37 +22,301 @@ use ecow::EcoString;
 mod test;
 
 const STACK_SIZE: usize = 8 * 1024 / size_of::<Value>();
+type GcArena = Arena<Rootable!['gc => GcRoot<'gc>]>;
 
-pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
-    let vm = VM {
-        thunks: prog.thunks,
-        funcs: prog.funcs,
-        symbols: RefCell::new(prog.symbols),
-        symmap: RefCell::new(prog.symmap),
-        consts: prog.consts,
-        bump: Bump::new(),
-        jit,
-    };
-    let env = env(&vm);
-    let mut seen = HashSet::new();
-    let value = vm
-        .eval(prog.top_level.into_iter(), vm.bump.alloc(env))?
-        .to_public(&vm, &mut seen);
-    Ok(value)
+#[derive(Collect)]
+#[collect(require_static)]
+struct ContextWrapper(Context);
+
+
+pub fn run(mut prog: Program) -> Result<p::Value> {
+    fn new<'gc>(prog: &mut Program, mc: &'gc Mutation<'gc>) -> GcRoot<'gc> {
+        let jit = Gc::new(mc, ContextWrapper(Context::create()));
+        let thunks = std::mem::take(&mut prog.thunks);
+        let funcs = std::mem::take(&mut prog.funcs);
+        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,
+            ctxs: DynamicRootSet::new(mc)
+        }
+    }
+    let arena: Arena<Rootable![GcRoot<'_>]> = Arena::new(|mc| {
+        new(&mut prog, mc)
+    });
+    eval(prog.top_level.into_iter(), &arena, |val| {
+        Ok(arena.mutate(|_, root: &GcRoot| {
+            val.to_public(&root.vm, &mut HashSet::new())
+        }))
+    })
 }
 
-#[derive(Constructor)]
-pub struct VM<'jit> {
+fn eval<'gc, T, F: for<'a> FnOnce(Value<'a>) -> Result<T>>(
+    opcodes: impl Iterator<Item = OpCode>,
+    arena: &Arena<impl for<'a> Rootable<'a, Root = GcRoot<'a>>>,
+    f: F
+) -> Result<T> {
+    let mut iter = opcodes.into_iter();
+    let dynroot: gc_arena::DynamicRoot<Rootable!['a => GcRefLock<'a, EvalContext<'a>>]> = arena.mutate(|mc, root| {
+        root.ctxs.stash(mc, Gc::new(mc, RefLock::new( EvalContext::new(env(&root.vm, mc)))))
+    });
+    while let Some(opcode) = iter.next() {
+        arena.mutate(|mc, root| {
+            let mut ctx_mut = root.ctxs.fetch(&dynroot).borrow_mut(mc);
+            let ctx = &mut *ctx_mut;
+            let jmp = single_op(&root.vm, opcode, &mut ctx.stack, &mut ctx.env, mc)?;
+            for _ in 0..jmp {
+                iter.next().unwrap();
+            }
+            Result::Ok(())
+        })?;
+    }
+    arena.mutate(|mc, root| {
+        let mut ctx = root.ctxs.fetch(&dynroot).borrow_mut(mc);
+        assert_eq!(ctx.stack.len(), 1);
+        let mut ret = ctx.stack.pop();
+        ret.force(&root.vm, mc);
+        f(ret)
+    })
+}
+
+#[inline(always)]
+fn single_op<'gc, const CAP: usize>(
+    vm: &'gc VM<'gc>,
+    opcode: OpCode,
+    stack: &mut Stack<Value<'gc>, CAP>,
+    env: &mut Gc<'gc, VmEnv<'gc>>,
+    mc: &'gc Mutation<'gc>,
+) -> Result<usize> {
+    match opcode {
+        OpCode::Illegal => panic!("illegal opcode"),
+        OpCode::Const { idx } => stack.push(Value::Const(vm.get_const(idx)))?,
+        OpCode::LoadThunk { idx } => stack.push(Value::Thunk(Thunk::new(vm.get_thunk(idx), mc)))?,
+        OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture(*env, mc),
+        OpCode::ForceValue => {
+            stack.tos_mut().force(vm, mc)?;
+        }
+        OpCode::Jmp { step } => return Ok(step),
+        OpCode::JmpIfFalse { step } => {
+            if let Value::Const(Const::Bool(false)) = stack.pop() {
+                return Ok(step);
+            }
+        }
+        OpCode::Call => {
+            let arg = stack.pop();
+            let func = stack.tos_mut();
+            func.force(vm, mc)?;
+            func.call(arg, vm, mc)?;
+        }
+        OpCode::Func { idx } => {
+            let func = vm.get_func(idx);
+            let compiled: GcRefLock<'gc, Option<JITFunc<'gc>>> = Gc::new(mc, RefLock::new(None));
+            stack.push(Value::Func(Gc::new(
+                mc,
+                Func::new(func, *env, compiled, Cell::new(0)),
+            )))?;
+        }
+        OpCode::UnOp { op } => {
+            use UnOp::*;
+            let value = stack.tos_mut();
+            value.force(vm, mc)?;
+            match op {
+                Neg => value.neg(),
+                Not => value.not(),
+            }
+        }
+        OpCode::BinOp { op } => {
+            use BinOp::*;
+            let mut rhs = stack.pop();
+            let lhs = stack.tos_mut();
+            lhs.force(vm, mc)?;
+            rhs.force(vm, mc)?;
+            match op {
+                Add => lhs.add(rhs),
+                Sub => {
+                    rhs.neg();
+                    lhs.add(rhs);
+                }
+                Mul => lhs.mul(rhs),
+                Div => lhs.div(rhs)?,
+                And => lhs.and(rhs),
+                Or => lhs.or(rhs),
+                Eq => Value::eq(lhs, rhs),
+                Lt => lhs.lt(rhs),
+                Con => lhs.concat(rhs, mc),
+                Upd => lhs.update(rhs, mc),
+            }
+        }
+        OpCode::ConcatString => {
+            let mut rhs = stack.pop();
+            rhs.force(vm, mc)?;
+            stack.tos_mut().concat_string(rhs);
+        }
+        OpCode::Path => {
+            todo!()
+        }
+        OpCode::List => {
+            stack.push(Value::List(CoW::new(List::empty(), mc)))?;
+        }
+        OpCode::PushElem => {
+            let elem = stack.pop();
+            stack.tos_mut().push(elem, mc);
+        }
+        OpCode::AttrSet { cap } => {
+            stack.push(Value::AttrSet(CoW::new(AttrSet::with_capacity(cap), mc)))?;
+        }
+        OpCode::FinalizeRec => {
+            let mut new = HashMap::new();
+            stack
+                .tos()
+                .as_ref()
+                .unwrap_attr_set()
+                .as_inner()
+                .iter()
+                .map(|(&k, v)| (k, v.clone()))
+                .collect_into(&mut new);
+            *env = env.enter_let(Gc::new(mc, new.into()), mc);
+            stack
+                .tos_mut()
+                .as_mut(mc)
+                .unwrap_attr_set()
+                .capture(*env, mc);
+        }
+        OpCode::PushStaticAttr { name } => {
+            let val = stack.pop();
+            stack.tos_mut().push_attr(name, val, mc);
+        }
+        OpCode::PushDynamicAttr => {
+            let val = stack.pop();
+            let mut sym = stack.pop();
+            sym.force(vm, mc)?.coerce_to_string();
+            let sym = vm.new_sym(sym.unwrap_const().unwrap_string());
+            stack.tos_mut().push_attr(sym, val, mc);
+        }
+        OpCode::Select { sym } => {
+            stack.tos_mut().force(vm, mc)?.select(sym, vm)?;
+        }
+        OpCode::SelectOrDefault { sym } => {
+            let default = stack.pop();
+            stack
+                .tos_mut()
+                .force(vm, mc)?
+                .select_with_default(sym, default)?;
+        }
+        OpCode::SelectDynamic => {
+            let mut val = stack.pop();
+            val.force(vm, mc)?;
+            val.coerce_to_string();
+            let sym = vm.new_sym(val.unwrap_const().unwrap_string());
+            stack.tos_mut().force(vm, mc)?.select(sym, vm)?;
+        }
+        OpCode::SelectDynamicOrDefault => {
+            let default = stack.pop();
+            let mut val = stack.pop();
+            val.force(vm, mc)?;
+            val.coerce_to_string();
+            let sym = vm.new_sym(val.unwrap_const().unwrap_string());
+            stack
+                .tos_mut()
+                .force(vm, mc)?
+                .select_with_default(sym, default)?;
+        }
+        OpCode::HasAttr { sym } => {
+            stack.tos_mut().force(vm, mc)?.has_attr(sym);
+        }
+        OpCode::HasDynamicAttr => {
+            let mut val = stack.pop();
+            val.coerce_to_string();
+            let sym = vm.new_sym(val.unwrap_const().unwrap_string());
+            stack.tos_mut().force(vm, mc)?.has_attr(sym);
+        }
+        OpCode::LookUp { sym } => {
+            stack.push(
+                env.lookup(&sym)
+                    .ok_or_else(|| Error::EvalError(format!("{} not found", vm.get_sym(sym))))?
+                    .clone(),
+            )?;
+        }
+        OpCode::EnterLetEnv => {
+            let mut new = HashMap::new();
+            stack
+                .pop()
+                .unwrap_attr_set()
+                .as_inner()
+                .iter()
+                .map(|(&k, v)| (k, v.clone()))
+                .collect_into(&mut new);
+            *env = env.enter_let(Gc::new(mc, new.into()), mc);
+        }
+        OpCode::LeaveLetEnv => *env = env.leave(),
+        OpCode::EnterWithEnv => {
+            let mut new = HashMap::new();
+            stack
+                .pop()
+                .unwrap_attr_set()
+                .as_inner()
+                .iter()
+                .map(|(&k, v)| (k, v.clone()))
+                .collect_into(&mut new);
+            *env = env.enter_with(Gc::new(mc, new.into()), mc);
+        }
+        OpCode::LeaveWithEnv => *env = env.leave(),
+        OpCode::Assert => {
+            if !stack.pop().unwrap_const().unwrap_bool() {
+                todo!()
+            }
+        }
+    }
+    Ok(0)
+}
+
+#[derive(Collect)]
+#[collect(no_drop)]
+pub struct GcRoot<'gc> {
+    vm: Gc<'gc, VM<'gc>>,
+    jit: Gc<'gc, ContextWrapper>,
+    ctxs: DynamicRootSet<'gc>
+}
+
+#[derive(Collect)]
+#[collect(no_drop)]
+pub struct EvalContext<'gc, const CAP: usize = STACK_SIZE> {
+    stack: Stack<Value<'gc>, CAP>,
+    env: Gc<'gc, VmEnv<'gc>>
+}
+
+impl<'gc, const CAP: usize> EvalContext<'gc, CAP> {
+    pub fn new(env: Gc<'gc, VmEnv<'gc>>) -> Self {
+        Self {
+            stack: Stack::new(),
+            env
+        }
+    }
+}
+
+#[derive(Constructor, Collect)]
+#[collect(no_drop)]
+pub struct VM<'gc> {
     thunks: Box<[OpCodes]>,
     funcs: Box<[F]>,
     symbols: RefCell<Vec<EcoString>>,
     symmap: RefCell<HashMap<EcoString, usize>>,
     consts: Box<[Const]>,
-    pub bump: Bump,
-    jit: JITContext<'jit>,
+    jit: JITContext<'gc>,
 }
 
-impl<'vm, 'jit: 'vm> VM<'jit> {
+impl<'gc> VM<'gc> {
     pub fn get_thunk(&self, idx: usize) -> &OpCodes {
         &self.thunks[idx]
     }
@@ -80,219 +346,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
         self.consts[idx].clone()
     }
 
-    pub fn eval(
-        &'vm self,
-        opcodes: impl Iterator<Item = OpCode>,
-        mut env: &'vm VmEnv<'jit, 'vm>,
-    ) -> Result<Value<'jit, 'vm>> {
-        let mut stack = Stack::<_, STACK_SIZE>::new();
-        let mut iter = opcodes.into_iter();
-        while let Some(opcode) = iter.next() {
-            let jmp = self.single_op(opcode, &mut stack, &mut env)?;
-            for _ in 0..jmp {
-                iter.next().unwrap();
-            }
-        }
-        assert_eq!(stack.len(), 1);
-        let mut ret = stack.pop();
-        ret.force(self)?;
-        Ok(ret)
-    }
-
-    #[inline(always)]
-    fn single_op<'s, const CAP: usize>(
-        &'vm self,
-        opcode: OpCode,
-        stack: &'s mut Stack<Value<'jit, 'vm>, CAP>,
-        env: &mut &'vm VmEnv<'jit, 'vm>,
-    ) -> Result<usize> {
-        match opcode {
-            OpCode::Illegal => panic!("illegal opcode"),
-            OpCode::Const { idx } => stack.push(Value::Const(self.get_const(idx)))?,
-            OpCode::LoadThunk { idx } => {
-                stack.push(Value::Thunk(Thunk::new(self.get_thunk(idx)).into()))?
-            }
-            OpCode::CaptureEnv => stack.tos().as_ref().unwrap_thunk().capture(*env),
-            OpCode::ForceValue => {
-                stack.tos_mut().force(self)?;
-            }
-            OpCode::Jmp { step } => return Ok(step),
-            OpCode::JmpIfFalse { step } => {
-                if let Value::Const(Const::Bool(false)) = stack.pop() {
-                    return Ok(step);
-                }
-            }
-            OpCode::Call => {
-                let arg = stack.pop();
-                let func = stack.tos_mut();
-                func.force(self)?;
-                func.call(self, arg)?;
-            }
-            OpCode::Func { idx } => {
-                let func = self.get_func(idx);
-                stack.push(Value::Func(
-                    Func::new(func, env, OnceCell::new(), Cell::new(0)).into(),
-                ))?;
-            }
-            OpCode::UnOp { op } => {
-                use UnOp::*;
-                let value = stack.tos_mut();
-                value.force(self)?;
-                match op {
-                    Neg => value.neg(),
-                    Not => value.not(),
-                }
-            }
-            OpCode::BinOp { op } => {
-                use BinOp::*;
-                let mut rhs = stack.pop();
-                let lhs = stack.tos_mut();
-                lhs.force(self)?;
-                rhs.force(self)?;
-                match op {
-                    Add => lhs.add(rhs),
-                    Sub => {
-                        rhs.neg();
-                        lhs.add(rhs);
-                    }
-                    Mul => lhs.mul(rhs),
-                    Div => lhs.div(rhs)?,
-                    And => lhs.and(rhs),
-                    Or => lhs.or(rhs),
-                    Eq => Value::eq(lhs, rhs),
-                    Lt => lhs.lt(rhs),
-                    Con => lhs.concat(rhs),
-                    Upd => lhs.update(rhs),
-                }
-            }
-            OpCode::ConcatString => {
-                let mut rhs = stack.pop();
-                rhs.force(self)?;
-                stack.tos_mut().concat_string(rhs);
-            }
-            OpCode::Path => {
-                todo!()
-            }
-            OpCode::List => {
-                stack.push(Value::List(List::empty().into()))?;
-            }
-            OpCode::PushElem => {
-                let elem = stack.pop();
-                stack.tos_mut().push(elem);
-            }
-            OpCode::AttrSet { cap } => {
-                stack.push(Value::AttrSet(AttrSet::with_capacity(cap).into()))?;
-            }
-            OpCode::FinalizeRec => {
-                let new = self.bump.alloc(HashMap::new_in(&self.bump));
-                *env = env.enter_let(
-                    stack
-                        .tos()
-                        .as_ref()
-                        .unwrap_attr_set()
-                        .as_inner()
-                        .iter()
-                        .map(|(&k, v)| (k, v.clone()))
-                        .collect_into(new),
-                    &self.bump,
-                );
-                stack.tos_mut().as_mut().unwrap_attr_set().capture(env);
-            }
-            OpCode::PushStaticAttr { name } => {
-                let val = stack.pop();
-                stack.tos_mut().push_attr(name, val);
-            }
-            OpCode::PushDynamicAttr => {
-                let val = stack.pop();
-                let mut sym = stack.pop();
-                sym.force(self)?.coerce_to_string();
-                let sym = self.new_sym(sym.unwrap_const().unwrap_string());
-                stack.tos_mut().push_attr(sym, val);
-            }
-            OpCode::Select { sym } => {
-                stack.tos_mut().force(self)?.select(sym, self)?;
-            }
-            OpCode::SelectOrDefault { sym } => {
-                let default = stack.pop();
-                stack
-                    .tos_mut()
-                    .force(self)?
-                    .select_with_default(sym, default)?;
-            }
-            OpCode::SelectDynamic => {
-                let mut val = stack.pop();
-                val.force(self)?;
-                val.coerce_to_string();
-                let sym = self.new_sym(val.unwrap_const().unwrap_string());
-                stack.tos_mut().force(self)?.select(sym, self)?;
-            }
-            OpCode::SelectDynamicOrDefault => {
-                let default = stack.pop();
-                let mut val = stack.pop();
-                val.force(self)?;
-                val.coerce_to_string();
-                let sym = self.new_sym(val.unwrap_const().unwrap_string());
-                stack
-                    .tos_mut()
-                    .force(self)?
-                    .select_with_default(sym, default)?;
-            }
-            OpCode::HasAttr { sym } => {
-                stack.tos_mut().force(self)?.has_attr(sym);
-            }
-            OpCode::HasDynamicAttr => {
-                let mut val = stack.pop();
-                val.coerce_to_string();
-                let sym = self.new_sym(val.unwrap_const().unwrap_string());
-                stack.tos_mut().force(self)?.has_attr(sym);
-            }
-            OpCode::LookUp { sym } => {
-                stack.push(
-                    env.lookup(&sym)
-                        .ok_or_else(|| {
-                            Error::EvalError(format!("{} not found", self.get_sym(sym)))
-                        })?
-                        .clone(),
-                )?;
-            }
-            OpCode::EnterLetEnv => {
-                let new = self.bump.alloc(HashMap::new_in(&self.bump));
-                *env = env.enter_let(
-                    stack
-                        .pop()
-                        .unwrap_attr_set()
-                        .into_inner()
-                        .iter()
-                        .map(|(&k, v)| (k, v.clone()))
-                        .collect_into(new),
-                    &self.bump,
-                );
-            }
-            OpCode::LeaveLetEnv => *env = env.leave(),
-            OpCode::EnterWithEnv => {
-                let new = self.bump.alloc(HashMap::new_in(&self.bump));
-                *env = env.enter_with(
-                    stack
-                        .pop()
-                        .unwrap_attr_set()
-                        .into_inner()
-                        .iter()
-                        .map(|(&k, v)| (k, v.clone()))
-                        .collect_into(new),
-                    &self.bump,
-                );
-            }
-            OpCode::LeaveWithEnv => *env = env.leave(),
-            OpCode::Assert => {
-                if !stack.pop().unwrap_const().unwrap_bool() {
-                    todo!()
-                }
-            }
-        }
-        Ok(0)
-    }
-
-    pub fn compile_func(&'vm self, func: &'vm F) -> JitFunction<'jit, JITFunc<'jit, 'vm>> {
+    pub fn compile_func(&'gc self, func: &'gc F) -> JITFunc<'gc> {
         self.jit.compile_function(func, self).unwrap()
     }
 }