feat: adds GCRootHandle which roots the contained handle in the GC

crates/mun_gc/src/handle.rs

@@ -1,3 +1,7 @@
+use crate::{GCRuntime, Type};
+use std::marker::PhantomData;
+use std::sync::{Arc, Weak};
+
 /// A `GCHandle` is what you interact with outside of the allocator. It is a pointer to a piece of
 /// memory that points to the actual data stored in memory.
 ///
@@ -16,14 +20,18 @@ unsafe impl Sync for GCHandle {}
 /// semantics used by the runtime.
 pub type RawGCHandle = *const *mut std::ffi::c_void;
 
-impl GCHandle {
+pub trait HasGCHandlePtr {
     /// Returns a pointer to the referenced memory.
     ///
     /// # Safety
     ///
     /// This method is unsafe because casting to a generic type T may be unsafe. We don't know the
     /// type of the stored data.
-    pub unsafe fn get_ptr<T: Sized>(self) -> *mut T {
+    unsafe fn get_ptr<T: Sized>(&self) -> *mut T;
+}
+
+impl HasGCHandlePtr for GCHandle {
+    unsafe fn get_ptr<T: Sized>(&self) -> *mut T {
         (*self.0).cast::<T>()
     }
 }
@@ -39,3 +47,64 @@ impl Into<GCHandle> for RawGCHandle {
         GCHandle(self)
     }
 }
+
+/// A `GCHandle` that automatically roots and unroots its internal `GCHandle`.
+pub struct GCRootHandle<T: Type, G: GCRuntime<T>> {
+    handle: GCHandle,
+    runtime: Weak<G>,
+    ty: PhantomData<T>,
+}
+
+impl<T: Type, G: GCRuntime<T>> Clone for GCRootHandle<T, G> {
+    fn clone(&self) -> Self {
+        if let Some(runtime) = self.runtime.upgrade() {
+            unsafe { runtime.root(self.handle) }
+        }
+        Self {
+            handle: self.handle,
+            runtime: self.runtime.clone(),
+            ty: Default::default(),
+        }
+    }
+}
+
+impl<T: Type, G: GCRuntime<T>> GCRootHandle<T, G> {
+    /// Constructs a new GCRootHandle from a runtime and a handle
+    ///
+    /// # Safety
+    ///
+    /// This method is unsafe because the passed GCHandle could point to random memory.
+    pub unsafe fn new(runtime: &Arc<G>, handle: GCHandle) -> Self {
+        runtime.root(handle);
+        Self {
+            handle,
+            runtime: Arc::downgrade(runtime),
+            ty: Default::default(),
+        }
+    }
+
+    /// Returns the handle of this instance
+    pub fn handle(&self) -> GCHandle {
+        self.handle
+    }
+}
+
+impl<T: Type, G: GCRuntime<T>> Into<GCHandle> for GCRootHandle<T, G> {
+    fn into(self) -> GCHandle {
+        self.handle
+    }
+}
+
+impl<T: Type, G: GCRuntime<T>> Drop for GCRootHandle<T, G> {
+    fn drop(&mut self) {
+        if let Some(runtime) = self.runtime.upgrade() {
+            unsafe { runtime.unroot(self.handle) }
+        }
+    }
+}
+
+impl<T: Type, G: GCRuntime<T>> HasGCHandlePtr for GCRootHandle<T, G> {
+    unsafe fn get_ptr<R: Sized>(&self) -> *mut R {
+        self.handle.get_ptr()
+    }
+}

crates/mun_gc/src/lib.rs

@@ -1,7 +1,7 @@
 mod handle;
 mod mark_sweep;
 
-pub use handle::{GCHandle, RawGCHandle};
+pub use handle::{GCHandle, GCRootHandle, HasGCHandlePtr, RawGCHandle};
 pub use mark_sweep::MarkSweep;
 
 /// A trait used by the GC to identify an object.
@@ -33,10 +33,21 @@ pub trait GCRuntime<T: Type>: Send + Sync {
     unsafe fn object_type(&self, obj: GCHandle) -> T;
 
     /// Tell the runtime that the specified object should be considered a root which keeps all other
-    /// objects it references alive.
+    /// objects it references alive. Objects marked as root, must also be unrooted before they can
+    /// be collected. Internally this increments a root refcount.
     ///
     /// # Safety
     ///
     /// This method is unsafe because the passed GCHandle could point to random memory.
-    unsafe fn set_root(&self, obj: GCHandle, is_root: bool);
+    unsafe fn root(&self, obj: GCHandle);
+
+    /// Tell the runtime that the specified object should unrooted which keeps all other
+    /// objects it references alive. Objects marked as root, must also be unrooted before they can
+    /// be collected. Internally this decrements a root refcount. When the refcount reaches 0, the
+    /// object is considered non-rooted.
+    ///
+    /// # Safety
+    ///
+    /// This method is unsafe because the passed GCHandle could point to random memory.
+    unsafe fn unroot(&self, obj: GCHandle);
 }

crates/mun_gc/src/mark_sweep.rs

@@ -80,7 +80,11 @@ impl<T: Type + Clone> GCRuntime<T> for MarkSweep<T> {
         src.ty.clone()
     }
 
-    unsafe fn set_root(&self, _obj: GCHandle, _is_root: bool) {
+    unsafe fn root(&self, _obj: GCHandle) {
+        // NOOP
+    }
+
+    unsafe fn unroot(&self, _obj: GCHandle) {
         // NOOP
     }
 }

crates/mun_runtime/src/allocator.rs

@@ -23,3 +23,7 @@ impl gc::Type for RawTypeInfo {
 
 /// Defines an allocator used by the `Runtime`
 pub type Allocator = gc::MarkSweep<RawTypeInfo>;
+
+pub use gc::GCHandle;
+
+pub type GCRootHandle = gc::GCRootHandle<RawTypeInfo, Allocator>;

crates/mun_runtime/src/lib.rs

@@ -245,8 +245,8 @@ impl Runtime {
     }
 
     /// Returns the runtime's allocator.
-    pub fn get_allocator(&self) -> Arc<Allocator> {
-        self.allocator.clone()
+    pub fn allocator(&self) -> &Arc<Allocator> {
+        &self.allocator
     }
 
     /// Updates the state of the runtime. This includes checking for file changes, and reloading

crates/mun_runtime/src/struct.rs

@@ -1,12 +1,13 @@
+use crate::allocator::{GCHandle, GCRootHandle};
 use crate::{
     marshal::Marshal,
     reflection::{
         equals_argument_type, equals_return_type, ArgumentReflection, ReturnTypeReflection,
     },
     Runtime,
 };
-use gc::GCHandle;
 use gc::GCRuntime;
+use gc::HasGCHandlePtr;
 use std::cell::RefCell;
 use std::ptr::{self, NonNull};
 use std::rc::Rc;
@@ -29,7 +30,7 @@ impl RawStruct {
 /// TODO: Handle destruction of `struct(value)`
 pub struct StructRef {
     runtime: Rc<RefCell<Runtime>>,
-    raw: RawStruct,
+    handle: GCRootHandle,
     info: abi::StructInfo,
 }
 
@@ -40,16 +41,20 @@ impl StructRef {
     fn new(runtime: Rc<RefCell<Runtime>>, type_info: &abi::TypeInfo, raw: RawStruct) -> StructRef {
         assert!(type_info.group.is_struct());
 
+        let handle = {
+            let runtime_ref = runtime.borrow();
+            unsafe { GCRootHandle::new(runtime_ref.allocator(), raw.0) }
+        };
         Self {
             runtime,
-            raw,
+            handle,
             info: type_info.as_struct().unwrap().clone(),
         }
     }
 
     /// Consumes the `Struct`, returning a raw Mun struct.
     pub fn into_raw(self) -> RawStruct {
-        self.raw
+        RawStruct(self.handle.handle())
     }
 
     /// Retrieves its struct information.
@@ -65,7 +70,7 @@ impl StructRef {
     unsafe fn offset_unchecked<T>(&self, field_idx: usize) -> NonNull<T> {
         let offset = *self.info.field_offsets().get_unchecked(field_idx);
         // self.raw is never null
-        NonNull::new_unchecked(self.raw.get_ptr().add(offset as usize).cast::<T>())
+        NonNull::new_unchecked(self.handle.get_ptr::<u8>().add(offset as usize).cast::<T>())
     }
 
     /// Retrieves the value of the field corresponding to the specified `field_name`.
@@ -144,7 +149,7 @@ impl ArgumentReflection for StructRef {
     }
 
     fn marshal(self) -> Self::Marshalled {
-        self.raw
+        self.into_raw()
     }
 }
 
@@ -184,7 +189,10 @@ impl Marshal<StructRef> for RawStruct {
             // struct value.
 
             // Create a new object using the runtime's intrinsic
-            let gc_handle = runtime.borrow().get_allocator().alloc_object(type_info_ptr);
+            let gc_handle = {
+                let runtime_ref = runtime.borrow();
+                runtime_ref.allocator().alloc_object(type_info_ptr)
+            };
 
             // Construct
             let src = ptr.cast::<u8>().as_ptr() as *const _;