allocate RootNode from dedicated pool

Author: shruti2522

oscars/src/collectors/mark_sweep_branded/gc.rs

@@ -1,7 +1,7 @@
 //! Core pointer types.
 
 use crate::{
-    alloc::mempool3::PoolItem,
+    alloc::mempool3::{PoolAllocator, PoolItem},
     collectors::mark_sweep_branded::{
         gc_box::GcBox,
         mutation_ctx::MutationContext,
@@ -13,7 +13,8 @@ use core::fmt;
 use core::marker::PhantomData;
 use core::ops::Deref;
 use core::ptr::NonNull;
-use rust_alloc::boxed::Box;
+
+pub(crate) type RootDropFn = unsafe fn(&mut PoolAllocator<'static>, NonNull<u8>);
 
 /// A transient pointer to a GC-managed value.
 #[derive(Debug)]
@@ -61,6 +62,10 @@ pub(crate) struct RootNode<'id, T: Trace> {
     pub(crate) link: RootLink,
     /// Pointer to the allocation
     pub(crate) gc_ptr: NonNull<PoolItem<GcBox<T>>>,
+    /// Type-erased drop function for freeing this RootNode
+    pub(crate) drop_fn: RootDropFn,
+    /// Raw pointer to the Collector for freeing this node
+    pub(crate) collector_ptr: *const crate::collectors::mark_sweep_branded::Collector,
     pub(crate) _marker: PhantomData<*mut &'id ()>,
 }
 
@@ -83,14 +88,14 @@ impl<'id, T: Trace> Root<'id, T> {
 
 impl<'id, T: Trace> Drop for Root<'id, T> {
     fn drop(&mut self) {
-        // SAFETY:
-        // * `self.raw` was created by `Box::into_raw`
-        // * The address is stable.
         unsafe {
-            let node = Box::from_raw(self.raw.as_ptr());
-            if node.link.is_linked() {
-                RootLink::unlink(NonNull::from(&node.link));
+            let node_ref = self.raw.as_ref();
+            if node_ref.link.is_linked() {
+                RootLink::unlink(NonNull::from(&node_ref.link));
             }
+            // SAFETY: collector_ptr is valid for the lifetime of the GcContext
+            let collector = &*node_ref.collector_ptr;
+            collector.free_root_node(self.raw.cast::<u8>(), node_ref.drop_fn);
         }
     }
 }

oscars/src/collectors/mark_sweep_branded/mod.rs

@@ -41,6 +41,8 @@ pub(crate) struct Collector {
     // and we ensure that `Gc` objects and pool allocations do not outlive
     // the `Collector` instance
     pub(crate) pool: RefCell<PoolAllocator<'static>>,
+    /// Dedicated pool for RootNode allocations
+    pub(crate) root_pool: RefCell<PoolAllocator<'static>>,
     pub(crate) sentinel: RootSentinel,
     pub(crate) generic_alloc_id: Cell<usize>,
     pub(crate) ephemerons: RefCell<Vec<EphemeronEntry>>,
@@ -50,6 +52,7 @@ impl Collector {
     fn new() -> Self {
         Self {
             pool: RefCell::new(PoolAllocator::default()),
+            root_pool: RefCell::new(PoolAllocator::default()),
             sentinel: RootSentinel::new(),
             generic_alloc_id: Cell::new(0),
             ephemerons: RefCell::new(Vec::new()),
@@ -70,6 +73,54 @@ impl Collector {
         });
     }
 
+    /// Allocates a RootNode from the dedicated root pool.
+    pub(crate) fn alloc_root_node<'id, T: trace::Trace>(
+        &self,
+        gc_ptr: NonNull<crate::alloc::mempool3::PoolItem<GcBox<T>>>,
+    ) -> NonNull<gc::RootNode<'id, T>> {
+        use crate::alloc::mempool3::PoolItem;
+
+        unsafe fn drop_and_free<T: trace::Trace>(
+            pool: &mut PoolAllocator<'static>,
+            ptr: NonNull<u8>,
+        ) {
+            unsafe {
+                let typed_ptr = ptr.cast::<PoolItem<gc::RootNode<'_, T>>>();
+                core::ptr::drop_in_place(typed_ptr.as_ptr());
+                pool.free_slot(ptr);
+            }
+        }
+
+        let size = mem::size_of::<PoolItem<gc::RootNode<'id, T>>>();
+        let mut pool = self.root_pool.borrow_mut();
+        let slot_ptr = pool
+            .alloc_slot_raw(size)
+            .expect("root pool allocation failed");
+
+        unsafe {
+            let typed_ptr = slot_ptr.cast::<PoolItem<gc::RootNode<'id, T>>>();
+            core::ptr::write(
+                typed_ptr.as_ptr(),
+                PoolItem(gc::RootNode {
+                    link: root_link::RootLink::new(),
+                    gc_ptr,
+                    drop_fn: drop_and_free::<T>,
+                    collector_ptr: self as *const Collector,
+                    _marker: PhantomData,
+                }),
+            );
+            NonNull::new_unchecked(&mut (*typed_ptr.as_ptr()).0)
+        }
+    }
+
+    /// Frees a RootNode back to the root pool.
+    pub(crate) fn free_root_node(&self, ptr: NonNull<u8>, drop_fn: gc::RootDropFn) {
+        let mut pool = self.root_pool.borrow_mut();
+        unsafe {
+            (drop_fn)(&mut pool, ptr);
+        }
+    }
+
     /// Allocates a value from the pool.
     ///
     /// # Panics
@@ -161,8 +212,6 @@ impl Collector {
 
     /// Runs a collection cycle
     pub(crate) fn collect(&self) {
-        let sentinel_ptr = self.sentinel.as_ptr();
-
         let color = TraceColor::new();
 
         let gc_ptr_offset = core::mem::offset_of!(
@@ -178,7 +227,7 @@ impl Collector {
             "gc_ptr offset must be stable across all T: Sized"
         );
 
-        for link_ptr in RootLink::iter_from_sentinel(sentinel_ptr) {
+        for link_ptr in self.sentinel.iter() {
             unsafe {
                 // Read the `gc_ptr` field using the stable offset
                 let gc_ptr_ptr = link_ptr
@@ -260,6 +309,28 @@ impl Drop for Collector {
     /// Frees all remaining allocations
     fn drop(&mut self) {
         use crate::alloc::mempool3::PoolItem;
+
+        // Free all root nodes first
+        let all_roots: Vec<(NonNull<u8>, gc::RootDropFn)> = self
+            .root_pool
+            .borrow()
+            .iter_live_slots()
+            .map(|ptr| unsafe {
+                let drop_fn = (*ptr.cast::<PoolItem<gc::RootNode<'_, ()>>>().as_ptr())
+                    .0
+                    .drop_fn;
+                (ptr, drop_fn)
+            })
+            .collect();
+        let mut root_pool = self.root_pool.borrow_mut();
+        for (ptr, drop_fn) in all_roots {
+            unsafe {
+                (drop_fn)(&mut root_pool, ptr);
+            }
+        }
+        drop(root_pool);
+
+        // Then free all GC allocations
         let all: Vec<(NonNull<u8>, DropFn)> = self
             .pool
             .borrow()

oscars/src/collectors/mark_sweep_branded/mutation_ctx.rs

@@ -3,14 +3,12 @@
 use crate::collectors::mark_sweep_branded::{
     Collector,
     ephemeron::Ephemeron,
-    gc::{Gc, Root, RootNode},
+    gc::{Gc, Root},
     root_link::RootLink,
     trace::{Finalize, Trace},
     weak::WeakGc,
 };
 use core::marker::PhantomData;
-use core::ptr::NonNull;
-use rust_alloc::boxed::Box;
 
 /// Handle for GC allocations
 pub struct MutationContext<'id, 'gc> {
@@ -40,15 +38,7 @@ impl<'id, 'gc> MutationContext<'id, 'gc> {
 
     /// Promotes a `Gc` pointer to a `Root`
     pub fn root<T: Trace + Finalize + 'gc>(&self, gc: Gc<'gc, T>) -> Root<'id, T> {
-        let gc_ptr = gc.ptr;
-
-        let node = Box::new(RootNode {
-            link: RootLink::new(),
-            gc_ptr,
-            _marker: PhantomData,
-        });
-
-        let raw = unsafe { NonNull::new_unchecked(Box::into_raw(node)) };
+        let raw = self.collector.alloc_root_node(gc.ptr);
 
         // SAFETY: `raw` points to a stable `RootNode`.
         unsafe {

oscars/src/collectors/mark_sweep_branded/root_link.rs

@@ -67,13 +67,6 @@ impl RootLink {
             node_ref.next.set(None);
         }
     }
-
-    /// Returns an iterator over all nodes after `sentinel`.
-    pub(crate) fn iter_from_sentinel(sentinel: NonNull<Self>) -> RootLinkIter {
-        // SAFETY: sentinel is pinned in Collector and outlives the iteration.
-        let first = unsafe { sentinel.as_ref().next.get() };
-        RootLinkIter { current: first }
-    }
 }
 
 /// Iterator over root links in the intrusive list.
@@ -109,4 +102,10 @@ impl RootSentinel {
             NonNull::new_unchecked(self.0.as_ref().get_ref() as *const RootLink as *mut RootLink)
         }
     }
+
+    /// Returns an iterator over all root nodes after this sentinel.
+    pub(crate) fn iter(&self) -> RootLinkIter {
+        let first = self.0.as_ref().next.get();
+        RootLinkIter { current: first }
+    }
 }

oscars/src/collectors/mark_sweep_branded/tests/api_compliance.rs

@@ -26,6 +26,8 @@ mod tests {
         let make_dummy = || RootNode::<i32> {
             link: RootLink::new(),
             gc_ptr: core::ptr::NonNull::dangling(),
+            drop_fn: |_, _| {},
+            collector_ptr: core::ptr::null(),
             _marker: core::marker::PhantomData,
         };