implement tri-color marking

Author: shruti2522

oscars/src/alloc/mempool3/alloc.rs

@@ -61,14 +61,17 @@ impl<'pool> ErasedPoolPointer<'pool> {
 /// typed pointer into a pool slot
 #[derive(Debug, Clone, Copy)]
 #[repr(transparent)]
-pub struct PoolPointer<'pool, T>(NonNull<PoolItem<T>>, PhantomData<&'pool T>);
+pub struct PoolPointer<'pool, T>(NonNull<PoolItem<T>>, PhantomData<(&'pool (), *mut T)>);
 
 impl<'pool, T> PoolPointer<'pool, T> {
     pub(crate) unsafe fn from_raw(raw: NonNull<PoolItem<T>>) -> Self {
         Self(raw, PhantomData)
     }
 
-    pub fn as_inner_ref(&self) -> &'pool T {
+    pub fn as_inner_ref(&self) -> &'pool T
+    where
+        T: 'pool,
+    {
         // SAFETY: pointer is valid and properly aligned
         unsafe { &(*self.0.as_ptr()).0 }
     }

oscars/src/alloc/mempool3/mod.rs

@@ -255,67 +255,6 @@ impl<'alloc> PoolAllocator<'alloc> {
         }
     }
 
-    /// Allocates a raw slot of the given size without type constraints.
-    /// Returns a pointer to uninitialized memory.
-    ///
-    /// # Safety
-    /// Caller must initialize the memory before use and ensure proper cleanup.
-    pub fn alloc_slot_raw(&mut self, size: usize) -> Result<NonNull<u8>, PoolAllocError> {
-        let sc_idx = size_class_index_for(size);
-        let slot_size = SIZE_CLASSES.get(sc_idx).copied().unwrap_or(size);
-
-        let cached_idx = self.alloc_cache[sc_idx].get();
-        if cached_idx < self.slot_pools.len() {
-            let pool = &self.slot_pools[cached_idx];
-            if pool.slot_size == slot_size
-                && let Some(slot_ptr) = pool.alloc_slot()
-            {
-                return Ok(slot_ptr);
-            }
-        }
-
-        // try existing pools with matching slot_size first
-        for (i, pool) in self.slot_pools.iter().enumerate().rev() {
-            if pool.slot_size == slot_size
-                && let Some(slot_ptr) = pool.alloc_slot()
-            {
-                self.alloc_cache[sc_idx].set(i);
-                return Ok(slot_ptr);
-            }
-        }
-
-        // need a new pool for this size class
-        // try the recycle list first
-        if let Some(pos) = self
-            .recycled_pools
-            .iter()
-            .rposition(|p| p.slot_size == slot_size)
-        {
-            let pool = self.recycled_pools.swap_remove(pos);
-            let slot_ptr = pool.alloc_slot().ok_or(PoolAllocError::OutOfMemory)?;
-            let insert_idx = self.slot_pools.len();
-            let (base, end) = pool.slot_range();
-            let spos = self.sorted_ranges.partition_point(|&(b, _, _)| b < base);
-            self.sorted_ranges.insert(spos, (base, end, insert_idx));
-            self.slot_pools.push(pool);
-            self.alloc_cache[sc_idx].set(insert_idx);
-            return Ok(slot_ptr);
-        }
-
-        // Allocate a fresh page from the OS
-        let total = self.page_size.max(slot_size * 4);
-        let new_pool = SlotPool::try_init(slot_size, total, 16)?;
-        self.current_heap_size += new_pool.layout.size();
-        let slot_ptr = new_pool.alloc_slot().ok_or(PoolAllocError::OutOfMemory)?;
-        let insert_idx = self.slot_pools.len();
-        let (base, end) = new_pool.slot_range();
-        let spos = self.sorted_ranges.partition_point(|&(b, _, _)| b < base);
-        self.sorted_ranges.insert(spos, (base, end, insert_idx));
-        self.slot_pools.push(new_pool);
-        self.alloc_cache[sc_idx].set(insert_idx);
-        Ok(slot_ptr)
-    }
-
     /// drops the value at `ptr` and returns the slot to the allocator
     ///
     /// # Safety

oscars/src/collectors/mark_sweep_branded/cell.rs

@@ -1,6 +1,6 @@
 //! Interior mutability for GC-managed values.
 
-use crate::collectors::mark_sweep_branded::trace::{Finalize, Trace, TraceColor};
+use crate::collectors::mark_sweep_branded::trace::{Finalize, Trace, Tracer};
 use core::cell::{Ref, RefCell, RefMut};
 use core::ops::{Deref, DerefMut};
 
@@ -65,7 +65,7 @@ impl<T: Trace> DerefMut for GcRefMut<'_, T> {
 impl<T: Trace> Finalize for GcRefCell<T> {}
 
 impl<T: Trace> Trace for GcRefCell<T> {
-    fn trace(&self, color: &TraceColor) {
-        self.inner.borrow().trace(color);
+    fn trace(&mut self, tracer: &mut Tracer) {
+        self.inner.get_mut().trace(tracer);
     }
 }

oscars/src/collectors/mark_sweep_branded/ephemeron.rs

@@ -4,7 +4,7 @@ use crate::{
         gc::Gc,
         gc_box::GcBox,
         mutation_ctx::MutationContext,
-        trace::{Finalize, Trace, TraceColor},
+        trace::{Finalize, Trace, Tracer},
     },
 };
 use core::marker::PhantomData;
@@ -44,5 +44,5 @@ impl<'id, K: Trace, V: Trace> Copy for Ephemeron<'id, K, V> {}
 impl<'id, K: Trace, V: Trace> Finalize for Ephemeron<'id, K, V> {}
 
 impl<'id, K: Trace, V: Trace> Trace for Ephemeron<'id, K, V> {
-    fn trace(&self, _color: &TraceColor) {}
+    fn trace(&mut self, _tracer: &mut Tracer) {}
 }

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,21 +88,21 @@ 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);
         }
     }
 }
 
 impl<T: Trace> Finalize for Gc<'_, T> {}
 impl<T: Trace> Trace for Gc<'_, T> {
-    fn trace(&self, color: &crate::collectors::mark_sweep_branded::trace::TraceColor) {
-        color.mark(self);
+    fn trace(&mut self, tracer: &mut crate::collectors::mark_sweep_branded::trace::Tracer) {
+        tracer.mark(self);
     }
 }

oscars/src/collectors/mark_sweep_branded/gc_box.rs

@@ -3,17 +3,29 @@
 use core::cell::Cell;
 use core::ptr::NonNull;
 
-use crate::alloc::mempool3::PoolAllocator;
-use crate::collectors::mark_sweep_branded::trace::TraceFn;
+use crate::alloc::mempool3::{PoolAllocator, PoolItem};
+use crate::collectors::mark_sweep_branded::trace::{Trace, TraceFn, Tracer};
 
 pub(crate) type DropFn = unsafe fn(&mut PoolAllocator<'static>, NonNull<u8>);
 
+/// The tri-color marking state of a [`GcBox`]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[repr(u8)]
+pub(crate) enum GcColor {
+    /// Not yet reached by mark phase
+    White = 0,
+    /// Reached and queued in the worklist, children not yet traced.
+    Gray = 1,
+    /// Reached and dequeued from the worklist, all children traced
+    Black = 2,
+}
+
 /// Heap wrapper for a garbage-collected value.
 ///
 /// Allocated via [`PoolAllocator`].
 pub(crate) struct GcBox<T: ?Sized> {
-    /// Reachability flag set by the mark phase.
-    pub(crate) marked: Cell<bool>,
+    /// tricolor marking state, updated by the mark phase
+    pub(crate) color: Cell<GcColor>,
     /// Type-erased trace function.
     pub(crate) trace_fn: TraceFn,
     /// Type-erased finalize and free fn
@@ -27,3 +39,29 @@ pub(crate) struct GcBox<T: ?Sized> {
 impl<T: ?Sized> GcBox<T> {
     pub(crate) const FREED_ALLOC_ID: usize = usize::MAX;
 }
+
+impl<T> GcBox<T> {
+    /// Create a [`GcBox`] for `value`, `color` starts as [`GcColor::White`]
+    pub(crate) fn new(value: T, trace_fn: TraceFn, drop_fn: DropFn, alloc_id: usize) -> Self {
+        Self {
+            color: Cell::new(GcColor::White),
+            trace_fn,
+            drop_fn,
+            alloc_id,
+            value,
+        }
+    }
+}
+
+/// type-erased trace function for a `GcBox<T>` slot.
+///
+/// # Safety
+///
+/// `ptr` must point to a live `PoolItem<GcBox<T>>` in the pool allocator
+pub(crate) unsafe fn trace_value<T: Trace>(ptr: NonNull<u8>, tracer: &mut Tracer<'_>) {
+    let pool_item_ptr = ptr.cast::<PoolItem<GcBox<T>>>();
+    unsafe {
+        (*pool_item_ptr.as_ptr()).0.color.set(GcColor::Black);
+        (*pool_item_ptr.as_ptr()).0.value.trace(tracer);
+    }
+}