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

src/ty/internal/mod.rs

@@ -1,10 +1,13 @@
-use hashbrown::HashSet;
-use std::cell::OnceCell;
-use std::cell::RefCell;
+use std::cell::Cell;
 use std::hash::Hash;
-use std::rc::{Rc, Weak};
+use std::mem::MaybeUninit;
+use std::ops::Deref;
 
 use derive_more::{IsVariant, Unwrap};
+use gc_arena::Mutation;
+use gc_arena::lock::RefLock;
+use gc_arena::{Collect, Gc, lock::GcRefLock};
+use hashbrown::HashSet;
 
 use super::common::*;
 use super::public as p;
@@ -25,38 +28,125 @@ pub use func::*;
 pub use list::List;
 pub use primop::*;
 
-#[derive(Debug, IsVariant, Unwrap, Clone)]
-pub enum Value<'jit: 'vm, 'vm> {
-    Const(Const),
-    Thunk(Rc<Thunk<'jit, 'vm>>),
-    AttrSet(Rc<AttrSet<'jit, 'vm>>),
-    List(Rc<List<'jit, 'vm>>),
-    Catchable(Catchable),
-    PrimOp(Rc<PrimOp<'jit, 'vm>>),
-    PartialPrimOp(Rc<PartialPrimOp<'jit, 'vm>>),
-    Func(Rc<Func<'jit, 'vm>>),
-    Builtins(Weak<AttrSet<'jit, 'vm>>),
+#[derive(Collect)]
+#[collect(unsafe_drop)]
+pub struct CoW<'gc, T: Clone + Collect> {
+    inner: Gc<'gc, CoWInner<T>>,
 }
 
-impl Hash for Value<'_, '_> {
+struct CoWInner<T: Clone + Collect> {
+    ref_count: Cell<usize>,
+    data: MaybeUninit<T>,
+}
+
+unsafe impl<T: Clone + Collect> Collect for CoWInner<T> {
+    fn trace(&self, cc: &gc_arena::Collection) {
+        unsafe { self.data.assume_init_ref() }.trace(cc);
+    }
+}
+
+impl<'gc, T: Clone + Collect> CoW<'gc, T> {
+    pub fn new(data: T, mc: &Mutation<'gc>) -> Self {
+        Self {
+            inner: Gc::new(mc, CoWInner::new(data)),
+        }
+    }
+
+    pub fn new_cyclic(datafn: impl FnOnce(Self) -> T, mc: &Mutation<'gc>) -> Self {
+        let inner = Gc::new(
+            mc,
+            CoWInner {
+                ref_count: Cell::new(1),
+                data: MaybeUninit::uninit(),
+            },
+        );
+        let data = datafn(CoW { inner });
+        unsafe { std::mem::transmute::<_, &mut MaybeUninit<_>>(&inner.data) }.write(data);
+        Self { inner }
+    }
+
+    pub fn make_mut(&mut self, mc: &Mutation<'gc>) -> &mut T {
+        if self.inner.ref_count.get() == 1 {
+            unsafe { std::mem::transmute(self.inner.data.assume_init_ref()) }
+        } else {
+            unsafe {
+                *self = CoW::new(self.inner.data.assume_init_ref().clone(), mc);
+                std::mem::transmute(self.inner.data.assume_init_ref())
+            }
+        }
+    }
+
+    pub fn as_ptr(&self) -> *const T {
+        self.as_ref() as *const T
+    }
+}
+
+impl<'gc, T: Clone + Collect> AsRef<T> for CoW<'gc, T> {
+    fn as_ref(&self) -> &T {
+        unsafe { self.inner.data.assume_init_ref() }
+    }
+}
+
+impl<'gc, T: Clone + Collect> Deref for CoW<'gc, T> {
+    type Target = T;
+    fn deref(&self) -> &Self::Target {
+        self.as_ref()
+    }
+}
+
+impl<'gc, T: Clone + Collect> Clone for CoW<'gc, T> {
+    fn clone(&self) -> Self {
+        self.inner.ref_count.set(self.inner.ref_count.get() + 1);
+        Self { inner: self.inner }
+    }
+}
+
+impl<'gc, T: Clone + Collect> Drop for CoW<'gc, T> {
+    fn drop(&mut self) {
+        self.inner.ref_count.set(self.inner.ref_count.get() - 1);
+    }
+}
+
+impl<T: Clone + Collect> CoWInner<T> {
+    fn new(data: T) -> Self {
+        Self {
+            ref_count: Cell::new(1),
+            data: MaybeUninit::new(data),
+        }
+    }
+}
+
+#[derive(IsVariant, Unwrap, Clone, Collect)]
+#[collect(no_drop)]
+pub enum Value<'gc> {
+    Const(Const),
+    Thunk(Thunk<'gc>),
+    AttrSet(CoW<'gc, AttrSet<'gc>>),
+    List(CoW<'gc, List<'gc>>),
+    Catchable(Catchable),
+    PrimOp(Gc<'gc, PrimOp<'gc>>),
+    PartialPrimOp(CoW<'gc, PartialPrimOp<'gc>>),
+    Func(Gc<'gc, Func<'gc>>),
+}
+
+impl Hash for Value<'_> {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         use Value::*;
         std::mem::discriminant(self).hash(state);
         match self {
             Const(x) => x.hash(state),
-            Thunk(x) => (x.as_ref() as *const self::Thunk).hash(state),
-            AttrSet(x) => (x.as_ref() as *const self::AttrSet).hash(state),
-            List(x) => (x.as_ref() as *const self::List).hash(state),
+            Thunk(x) => (x as *const self::Thunk).hash(state),
+            AttrSet(x) => x.as_ptr().hash(state),
+            List(x) => x.as_ptr().hash(state),
             Catchable(x) => x.hash(state),
             PrimOp(x) => (x.as_ref() as *const self::PrimOp).hash(state),
             PartialPrimOp(x) => (x.as_ref() as *const self::PartialPrimOp).hash(state),
             Func(x) => (x.as_ref() as *const self::Func).hash(state),
-            Builtins(x) => (x.as_ptr()).hash(state),
         }
     }
 }
 
-impl<'jit: 'vm, 'vm> Value<'jit, 'vm> {
+impl<'gc> Value<'gc> {
     fn eq_impl(&self, other: &Self) -> bool {
         use Value::*;
         match (self, other) {
@@ -68,49 +158,46 @@ impl<'jit: 'vm, 'vm> Value<'jit, 'vm> {
     }
 }
 
-impl<'jit: 'vm, 'vm> PartialEq for Value<'jit, 'vm> {
+impl<'jit: 'vm, 'vm, 'gc> PartialEq for Value<'gc> {
     fn eq(&self, other: &Self) -> bool {
         use Value::*;
         match (self, other) {
             (Const(a), Const(b)) => a.eq(b),
-            (AttrSet(a), AttrSet(b)) => {
-                (a.as_ref() as *const self::AttrSet).eq(&(b.as_ref() as *const _))
-            }
-            (List(a), List(b)) => (a.as_ref() as *const self::List).eq(&(b.as_ref() as *const _)),
-            (Builtins(a), Builtins(b)) => a.ptr_eq(b),
+            (AttrSet(a), AttrSet(b)) => a.as_ptr().eq(&b.as_ptr()),
+            (List(a), List(b)) => a.as_ptr().eq(&b.as_ptr()),
             _ => false,
         }
     }
 }
 
-impl Eq for Value<'_, '_> {}
+impl Eq for Value<'_> {}
 
-#[derive(Debug, IsVariant, Unwrap, Clone, PartialEq)]
-pub enum ValueAsRef<'v, 'vm: 'v, 'jit: 'vm> {
+#[derive(IsVariant, Unwrap, Clone)]
+pub enum ValueAsRef<'v, 'gc> {
     Const(&'v Const),
-    Thunk(&'v Thunk<'jit, 'vm>),
-    AttrSet(&'v AttrSet<'jit, 'vm>),
-    List(&'v List<'jit, 'vm>),
+    Thunk(&'v Thunk<'gc>),
+    AttrSet(&'v AttrSet<'gc>),
+    List(&'v List<'gc>),
     Catchable(&'v Catchable),
-    PrimOp(&'v PrimOp<'jit, 'vm>),
-    PartialPrimOp(&'v PartialPrimOp<'jit, 'vm>),
-    Func(&'v Func<'jit, 'vm>),
+    PrimOp(&'v PrimOp<'gc>),
+    PartialPrimOp(&'v PartialPrimOp<'gc>),
+    Func(&'v Func<'gc>),
 }
 
-#[derive(Debug, IsVariant, Unwrap, PartialEq)]
-pub enum ValueAsMut<'v, 'vm: 'v, 'jit: 'vm> {
+#[derive(IsVariant, Unwrap)]
+pub enum ValueAsMut<'v, 'gc> {
     Const(&'v mut Const),
-    Thunk(&'v Thunk<'jit, 'vm>),
-    AttrSet(&'v mut AttrSet<'jit, 'vm>),
-    List(&'v mut List<'jit, 'vm>),
+    Thunk(&'v Thunk<'gc>),
+    AttrSet(&'v mut AttrSet<'gc>),
+    List(&'v mut List<'gc>),
     Catchable(&'v mut Catchable),
-    PrimOp(&'v mut PrimOp<'jit, 'vm>),
-    PartialPrimOp(&'v mut PartialPrimOp<'jit, 'vm>),
-    Func(&'v Func<'jit, 'vm>),
+    PrimOp(&'v PrimOp<'gc>),
+    PartialPrimOp(&'v mut PartialPrimOp<'gc>),
+    Func(&'v Func<'gc>),
 }
 
-impl<'v, 'vm: 'v, 'jit: 'vm> Value<'jit, 'vm> {
-    pub fn as_ref(&'v self) -> ValueAsRef<'v, 'vm, 'jit> {
+impl<'gc, 'v> Value<'gc> {
+    pub fn as_ref(&'v self) -> ValueAsRef<'v, 'gc> {
         use Value::*;
         use ValueAsRef as R;
         match self {
@@ -122,29 +209,27 @@ impl<'v, 'vm: 'v, 'jit: 'vm> Value<'jit, 'vm> {
             PrimOp(x) => R::PrimOp(x),
             PartialPrimOp(x) => R::PartialPrimOp(x),
             Func(x) => R::Func(x),
-            Builtins(_) => unreachable!(),
         }
     }
 
-    pub fn as_mut(&'v mut self) -> ValueAsMut<'v, 'vm, 'jit> {
+    pub fn as_mut(&'v mut self, mc: &Mutation<'gc>) -> ValueAsMut<'v, 'gc> {
         use Value::*;
         use ValueAsMut as M;
         match self {
             Const(x) => M::Const(x),
             Thunk(x) => M::Thunk(x),
-            AttrSet(x) => M::AttrSet(Rc::make_mut(x)),
-            List(x) => M::List(Rc::make_mut(x)),
+            AttrSet(x) => M::AttrSet(x.make_mut(mc)),
+            List(x) => M::List(x.make_mut(mc)),
             Catchable(x) => M::Catchable(x),
-            PrimOp(x) => M::PrimOp(Rc::make_mut(x)),
-            PartialPrimOp(x) => M::PartialPrimOp(Rc::make_mut(x)),
+            PrimOp(x) => M::PrimOp(x),
+            PartialPrimOp(x) => M::PartialPrimOp(x.make_mut(mc)),
             Func(x) => M::Func(x),
-            Builtins(_) => unreachable!(),
         }
     }
 }
 
 use Value::Const as VmConst;
-impl<'jit, 'vm> Value<'jit, 'vm> {
+impl<'gc> Value<'gc> {
     pub fn ok(self) -> Result<Self> {
         Ok(self)
     }
@@ -164,7 +249,6 @@ impl<'jit, 'vm> Value<'jit, 'vm> {
             PrimOp(_) => "lambda",
             PartialPrimOp(_) => "lambda",
             Func(_) => "lambda",
-            Builtins(_) => "set",
         }
     }
 
@@ -176,16 +260,16 @@ impl<'jit, 'vm> Value<'jit, 'vm> {
         }
     }
 
-    pub fn call(&mut self, vm: &'vm VM<'jit>, arg: Self) -> Result<()> {
+    pub fn call(&mut self, arg: Self, vm: &VM, mc: &Mutation<'gc>) -> Result<()> {
         use Value::*;
         if matches!(arg, Value::Catchable(_)) {
             *self = arg;
             return Ok(());
         }
         *self = match self {
-            PrimOp(func) => func.call(vm, arg),
-            PartialPrimOp(func) => func.call(vm, arg),
-            Value::Func(func) => func.call(vm, arg),
+            PrimOp(func) => func.call(arg, vm, mc),
+            PartialPrimOp(func) => func.call(arg, vm, mc),
+            Value::Func(func) => func.call(arg, vm, mc),
             Catchable(_) => return Ok(()),
             _ => todo!(),
         }?;
@@ -316,9 +400,9 @@ impl<'jit, 'vm> Value<'jit, 'vm> {
         self
     }
 
-    pub fn push(&mut self, elem: Self) -> &mut Self {
+    pub fn push(&mut self, elem: Self, mc: &Mutation<'gc>) -> &mut Self {
         if let Value::List(list) = self {
-            Rc::make_mut(list).push(elem);
+            list.make_mut(mc).push(elem);
         } else if let Value::Catchable(_) = self {
         } else if let Value::Catchable(_) = elem {
             *self = elem;
@@ -328,23 +412,23 @@ impl<'jit, 'vm> Value<'jit, 'vm> {
         self
     }
 
-    pub fn concat(&mut self, other: Self) {
+    pub fn concat(&mut self, other: Self, mc: &Mutation<'gc>) {
         if let x @ Value::Catchable(_) = other {
             *self = x;
             return;
         }
         match (self, other) {
             (Value::List(a), Value::List(b)) => {
-                Rc::make_mut(a).concat(b.as_ref());
+                a.make_mut(mc).concat(b.as_ref());
             }
             (Value::Catchable(_), _) => (),
             _ => todo!(),
         }
     }
 
-    pub fn push_attr(&mut self, sym: usize, val: Self) -> &mut Self {
+    pub fn push_attr(&mut self, sym: usize, val: Self, mc: &Mutation<'gc>) -> &mut Self {
         if let Value::AttrSet(attrs) = self {
-            Rc::make_mut(attrs).push_attr(sym, val)
+            attrs.make_mut(mc).push_attr(sym, val);
         } else if let Value::Catchable(_) = self {
         } else if let Value::Catchable(_) = val {
             *self = val
@@ -354,21 +438,21 @@ impl<'jit, 'vm> Value<'jit, 'vm> {
         self
     }
 
-    pub fn update(&mut self, other: Self) {
+    pub fn update(&mut self, other: Self, mc: &Mutation<'gc>) {
         if let x @ Value::Catchable(_) = other {
             *self = x;
             return;
         }
         match (self, other) {
             (Value::AttrSet(a), Value::AttrSet(b)) => {
-                Rc::make_mut(a).update(b.as_ref());
+                a.make_mut(mc).update(b.as_ref());
             }
             (Value::Catchable(_), _) => (),
             _ => todo!(),
         }
     }
 
-    pub fn select(&mut self, sym: usize, vm: &'vm VM<'jit>) -> Result<&mut Self> {
+    pub fn select(&mut self, sym: usize, vm: &VM<'gc>) -> Result<&mut Self> {
         let val = match self {
             Value::AttrSet(attrs) => attrs
                 .select(sym)
@@ -418,106 +502,127 @@ impl<'jit, 'vm> Value<'jit, 'vm> {
         self
     }
 
-    pub fn force(&mut self, vm: &'vm VM<'jit>) -> Result<&mut Self> {
+    pub fn force(&mut self, vm: &VM<'gc>, mc: &Mutation<'gc>) -> Result<&mut Self> {
         if let Value::Thunk(thunk) = self {
-            let value = thunk.force(vm)?.clone();
+            let value = thunk.force(vm, mc)?.clone();
             *self = value
         }
         Ok(self)
     }
 
-    pub fn force_deep(&mut self, vm: &'vm VM<'jit>) -> Result<&mut Self> {
+    pub fn force_deep(&mut self, vm: &VM<'gc>, mc: &Mutation<'gc>) -> Result<&mut Self> {
         match self {
             Value::Thunk(thunk) => {
-                let mut value = thunk.force(vm)?.clone();
-                let _ = value.force_deep(vm)?;
+                let mut value = thunk.force(vm, mc)?.clone();
+                let _ = value.force_deep(vm, mc)?;
                 *self = value;
             }
-            Value::List(list) => Rc::make_mut(list).force_deep(vm)?,
-            Value::AttrSet(attrs) => Rc::make_mut(attrs).force_deep(vm)?,
+            Value::List(list) => list.make_mut(mc).force_deep(vm, mc)?,
+            Value::AttrSet(attrs) => attrs.make_mut(mc).force_deep(vm, mc)?,
             _ => (),
         }
         Ok(self)
     }
 
-    pub fn to_public(&self, vm: &'vm VM, seen: &mut HashSet<Value<'jit, 'vm>>) -> p::Value {
+    pub fn to_public(
+        &self,
+        vm: &VM<'gc>,
+        seen: &mut HashSet<Value<'gc>>,
+    ) -> p::Value {
         use self::Value::*;
         use p::Value;
         if seen.contains(self) {
             return Value::Repeated;
         }
-        seen.insert(self.clone());
         match self {
-            AttrSet(attrs) => attrs.to_public(vm, seen),
-            List(list) => list.to_public(vm, seen),
+            AttrSet(attrs) => {
+                seen.insert(self.clone());
+                attrs.to_public(vm, seen)
+            }
+            List(list) => {
+                seen.insert(self.clone());
+                list.to_public(vm, seen)
+            }
             Catchable(catchable) => Value::Catchable(catchable.clone()),
             Const(cnst) => Value::Const(cnst.clone()),
             Thunk(_) => Value::Thunk,
             PrimOp(primop) => Value::PrimOp(primop.name),
             PartialPrimOp(primop) => Value::PartialPrimOp(primop.name),
             Func(_) => Value::Func,
-            Builtins(x) => x.upgrade().unwrap().to_public(vm, seen),
         }
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct Thunk<'jit, 'vm> {
-    pub thunk: RefCell<_Thunk<'jit, 'vm>>,
+#[repr(transparent)]
+#[derive(Clone, Collect)]
+#[collect(no_drop)]
+pub struct Thunk<'gc> {
+    pub thunk: GcRefLock<'gc, _Thunk<'gc>>,
 }
 
-#[derive(Debug, IsVariant, Unwrap, Clone)]
-pub enum _Thunk<'jit, 'vm> {
-    Code(&'vm OpCodes, OnceCell<&'vm VmEnv<'jit, 'vm>>),
-    SuspendedFrom(*const Thunk<'jit, 'vm>),
-    Value(Value<'jit, 'vm>),
+#[derive(IsVariant, Unwrap)]
+pub enum _Thunk<'gc> {
+    Code(
+        &'gc OpCodes,
+        GcRefLock<'gc, Option<Gc<'gc, VmEnv<'gc>>>>,
+    ),
+    Suspended,
+    Value(Value<'gc>),
 }
 
-impl<'jit, 'vm> Thunk<'jit, 'vm> {
-    pub fn new(opcodes: &'vm OpCodes) -> Self {
+unsafe impl<'gc> Collect for _Thunk<'gc> {
+    fn trace(&self, cc: &gc_arena::Collection) {
+        use _Thunk::*;
+        match self {
+            Code(_, env) => env.trace(cc),
+            Value(val) => val.trace(cc),
+            _ => (),
+        }
+    }
+}
+
+impl<'gc> Thunk<'gc> {
+    pub fn new(opcodes: &'gc OpCodes, mc: &Mutation<'gc>) -> Self {
         Thunk {
-            thunk: RefCell::new(_Thunk::Code(opcodes, OnceCell::new())),
+            thunk: Gc::new(
+                mc,
+                RefLock::new(_Thunk::Code(opcodes, Gc::new(mc, RefLock::new(None)))),
+            ),
         }
     }
 
-    pub fn capture(&self, env: &'vm VmEnv<'jit, 'vm>) {
+    pub fn capture(&self, env: Gc<'gc, VmEnv<'gc>>, mc: &Mutation<'gc>) {
         if let _Thunk::Code(_, envcell) = &*self.thunk.borrow() {
-            envcell.get_or_init(|| env);
+            envcell.borrow_mut(mc).insert(env);
         }
     }
 
-    pub fn force(&self, vm: &'vm VM<'jit>) -> Result<&Value<'jit, 'vm>> {
+    pub fn force(&self, vm: &VM<'gc>, mc: &Mutation<'gc>) -> Result<Value<'gc>> {
         use _Thunk::*;
         match &*self.thunk.borrow() {
             Value(_) => {
                 return Ok(match unsafe { &*(&*self.thunk.borrow() as *const _) } {
-                    Value(value) => value,
+                    Value(value) => value.clone(),
                     _ => unreachable!(),
                 });
             }
-            SuspendedFrom(from) => {
-                return Err(Error::EvalError(format!(
-                    "thunk {:p} already suspended from {from:p} (infinite recursion encountered)",
-                    self as *const Thunk
-                )));
+            Suspended => {
+                return Err(Error::EvalError("infinite recursion encountered".into()));
             }
             Code(..) => (),
         }
-        let (opcodes, env) = std::mem::replace(
-            &mut *self.thunk.borrow_mut(),
-            _Thunk::SuspendedFrom(self as *const Thunk),
-        )
-        .unwrap_code();
-        let env = env.get().unwrap();
+        let (opcodes, env) =
+            std::mem::replace(&mut *self.thunk.borrow_mut(mc), _Thunk::Suspended).unwrap_code();
+        let env = env.as_ref().borrow().unwrap();
         let value = vm.eval(opcodes.iter().copied(), env)?;
-        let _ = std::mem::replace(&mut *self.thunk.borrow_mut(), _Thunk::Value(value));
+        let _ = std::mem::replace(&mut *self.thunk.borrow_mut(mc), _Thunk::Value(value));
         Ok(match unsafe { &*(&*self.thunk.borrow() as *const _) } {
-            Value(value) => value,
+            Value(value) => value.clone(),
             _ => unreachable!(),
         })
     }
 
-    pub fn value(&'vm self) -> Option<Value<'jit, 'vm>> {
+    pub fn value(&self) -> Option<Value<'gc>> {
         match &*self.thunk.borrow() {
             _Thunk::Value(value) => Some(value.clone()),
             _ => None,
@@ -525,7 +630,7 @@ impl<'jit, 'vm> Thunk<'jit, 'vm> {
     }
 }
 
-impl PartialEq for Thunk<'_, '_> {
+impl PartialEq for Thunk<'_> {
     fn eq(&self, _: &Self) -> bool {
         false
     }