feat: implement mark_sweep _branded based on approved API redesign (#83)

* feat: implement mark_sweep_branded based on approved API redesign

* fix

* add ephemeron support

* replace Tracer with TraceColor

* remove pool_entries Vec and allocation_count

* implement tri-color marking

* refactor: use PoolPointer, make Ephemeron::key_ptr Option<PoolPointer>

* refactor: move RootNode, Root and RootLink into separate root module

Author: shruti2522

oscars/Cargo.toml

@@ -33,4 +33,5 @@ default = ["mark_sweep"]
 std = []
 mark_sweep = []
 mark_sweep2 = ["mark_sweep"]
+mark_sweep_branded = ["mark_sweep"]
 thin-vec = ["dep:thin-vec", "mark_sweep"]

oscars/src/alloc/mempool3/alloc.rs

@@ -59,16 +59,28 @@ 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: ?Sized>(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)
+impl<'pool, T: ?Sized> core::fmt::Debug for PoolPointer<'pool, T> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "PoolPointer({:p})", self.0.as_ptr())
     }
+}
+
+impl<'pool, T: ?Sized> Clone for PoolPointer<'pool, T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<'pool, T: ?Sized> Copy for PoolPointer<'pool, T> {}
 
-    pub fn as_inner_ref(&self) -> &'pool T {
+impl<'pool, T: ?Sized> PoolPointer<'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 }
     }
@@ -77,7 +89,10 @@ impl<'pool, T> PoolPointer<'pool, T> {
         self.0
     }
 
-    pub fn to_erased(self) -> ErasedPoolPointer<'pool> {
+    pub fn to_erased(self) -> ErasedPoolPointer<'pool>
+    where
+        T: Sized,
+    {
         ErasedPoolPointer(self.0.cast::<u8>(), PhantomData)
     }
 
@@ -92,6 +107,12 @@ impl<'pool, T> PoolPointer<'pool, T> {
     }
 }
 
+impl<'pool, T> PoolPointer<'pool, T> {
+    pub(crate) unsafe fn from_raw(raw: NonNull<PoolItem<T>>) -> Self {
+        Self(raw, PhantomData)
+    }
+}
+
 // ==== SlotPool ==== //
 impl core::fmt::Debug for SlotPool {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
@@ -310,6 +331,18 @@ impl SlotPool {
         self.live.set(self.live.get().saturating_sub(1));
     }
 
+    /// Iterates over all live (allocated) slot pointers in this pool.
+    pub(crate) fn iter_live(&self) -> impl Iterator<Item = NonNull<u8>> + '_ {
+        (0..self.slot_count).filter_map(move |i| {
+            let chunk = self.bitmap_chunk(i);
+            if chunk.get() & (1u64 << (i % 64)) != 0 {
+                Some(self.slot_ptr(i))
+            } else {
+                None
+            }
+        })
+    }
+
     /// returns true when the pool is empty and safe to drop
     /// `live` tracks the count, so no bitmap scan is needed
     pub fn run_drop_check(&self) -> bool {

oscars/src/alloc/mempool3/mod.rs

@@ -19,6 +19,7 @@ pub enum PoolAllocError {
     LayoutError(LayoutError),
     OutOfMemory,
     AlignmentNotPossible,
+    AllocIdExhausted,
 }
 
 impl From<LayoutError> for PoolAllocError {
@@ -121,6 +122,13 @@ impl<'alloc> PoolAllocator<'alloc> {
         self.current_heap_size
     }
 
+    /// Iterates over every live slot pointer across all slot pools.
+    ///
+    /// Yields one `NonNull<u8>` per allocated (not yet freed) slot.
+    pub fn iter_live_slots(&self) -> impl Iterator<Item = core::ptr::NonNull<u8>> + '_ {
+        self.slot_pools.iter().flat_map(|pool| pool.iter_live())
+    }
+
     pub fn is_below_threshold(&self) -> bool {
         // keep 25% headroom so collection fires before the last page fills
         let margin = self.heap_threshold / 4;

oscars/src/collectors/mark_sweep_branded/cell.rs

@@ -0,0 +1,71 @@
+//! Interior mutability for GC-managed values.
+
+use crate::collectors::mark_sweep_branded::trace::{Finalize, Trace, Tracer};
+use core::cell::{Ref, RefCell, RefMut};
+use core::ops::{Deref, DerefMut};
+
+/// A GC-aware wrapper around [`RefCell<T>`].
+pub struct GcRefCell<T: Trace> {
+    inner: RefCell<T>,
+}
+
+impl<T: Trace> GcRefCell<T> {
+    /// Wraps `value` in a new `GcRefCell`.
+    pub fn new(value: T) -> Self {
+        Self {
+            inner: RefCell::new(value),
+        }
+    }
+
+    /// Acquires a shared borrow of the inner value.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the value is currently mutably borrowed.
+    pub fn borrow(&self) -> GcRef<'_, T> {
+        GcRef(self.inner.borrow())
+    }
+
+    /// Acquires a mutable borrow of the inner value.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the value is currently borrowed.
+    pub fn borrow_mut(&self) -> GcRefMut<'_, T> {
+        GcRefMut(self.inner.borrow_mut())
+    }
+}
+
+/// A shared borrow guard returned by [`GcRefCell::borrow`].
+pub struct GcRef<'a, T: Trace>(Ref<'a, T>);
+
+impl<T: Trace> Deref for GcRef<'_, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        &self.0
+    }
+}
+
+/// A mutable borrow guard returned by [`GcRefCell::borrow_mut`].
+pub struct GcRefMut<'a, T: Trace>(RefMut<'a, T>);
+
+impl<T: Trace> Deref for GcRefMut<'_, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        &self.0
+    }
+}
+
+impl<T: Trace> DerefMut for GcRefMut<'_, T> {
+    fn deref_mut(&mut self) -> &mut T {
+        &mut self.0
+    }
+}
+
+impl<T: Trace> Finalize for GcRefCell<T> {}
+
+impl<T: Trace> Trace for GcRefCell<T> {
+    fn trace(&mut self, tracer: &mut Tracer) {
+        self.inner.get_mut().trace(tracer);
+    }
+}

oscars/src/collectors/mark_sweep_branded/ephemeron.rs

@@ -0,0 +1,49 @@
+use crate::{
+    alloc::mempool3::PoolPointer,
+    collectors::mark_sweep_branded::{
+        gc::Gc,
+        gc_box::GcBox,
+        mutation_ctx::MutationContext,
+        trace::{Finalize, Trace, Tracer},
+    },
+};
+use core::marker::PhantomData;
+
+pub struct Ephemeron<'id, K: Trace, V: Trace> {
+    pub(crate) key_ptr: Option<PoolPointer<'static, GcBox<K>>>,
+    pub(crate) key_alloc_id: usize,
+    pub(crate) value_ptr: PoolPointer<'static, GcBox<V>>,
+    pub(crate) _marker: PhantomData<*mut &'id ()>,
+}
+
+impl<'id, K: Trace, V: Trace> Ephemeron<'id, K, V> {
+    /// Returns the value if the key is alive.
+    pub fn get_value<'gc>(&self, _cx: &MutationContext<'id, 'gc>) -> Option<Gc<'gc, V>> {
+        // SAFETY: `_cx` proves the collector is alive, alloc_id guards ABA
+        let key_alive = self
+            .key_ptr
+            .is_some_and(|p| unsafe { (*p.as_ptr().as_ptr()).0.alloc_id == self.key_alloc_id });
+        if key_alive {
+            Some(Gc {
+                ptr: self.value_ptr,
+                _marker: PhantomData,
+            })
+        } else {
+            None
+        }
+    }
+}
+
+impl<'id, K: Trace, V: Trace> Clone for Ephemeron<'id, K, V> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+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(&mut self, _tracer: &mut Tracer) {}
+}

oscars/src/collectors/mark_sweep_branded/gc.rs

@@ -0,0 +1,58 @@
+//! Core pointer types.
+
+use crate::{
+    alloc::mempool3::PoolPointer,
+    collectors::mark_sweep_branded::{
+        gc_box::GcBox,
+        trace::{Finalize, Trace},
+    },
+};
+use core::fmt;
+use core::marker::PhantomData;
+use core::ops::Deref;
+
+/// A transient pointer to a GC-managed value.
+#[derive(Debug)]
+pub struct Gc<'gc, T: Trace + ?Sized + 'gc> {
+    pub(crate) ptr: PoolPointer<'static, GcBox<T>>,
+    pub(crate) _marker: PhantomData<(&'gc T, *const ())>,
+}
+
+impl<'gc, T: Trace + ?Sized + 'gc> Copy for Gc<'gc, T> {}
+impl<'gc, T: Trace + ?Sized + 'gc> Clone for Gc<'gc, T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<'gc, T: Trace + 'gc> Gc<'gc, T> {
+    /// Returns a shared reference to the value.
+    #[inline]
+    pub fn get(&self) -> &T {
+        // SAFETY: `ptr` is non-null and valid for `'gc` by construction.
+        // The `'gc` lifetime is scoped to a `mutate()` closure, collection only occurs
+        // via `cx.collect()` within that same closure and `Gc<'gc, T>` can't
+        // escape the closure.
+        unsafe { &(*self.ptr.as_ptr().as_ptr()).0.value }
+    }
+}
+
+impl<'gc, T: Trace + fmt::Display + 'gc> fmt::Display for Gc<'gc, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Display::fmt(self.get(), f)
+    }
+}
+
+impl<'gc, T: Trace + 'gc> Deref for Gc<'gc, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        self.get()
+    }
+}
+
+impl<T: Trace> Finalize for Gc<'_, T> {}
+impl<T: Trace> Trace for Gc<'_, T> {
+    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

@@ -0,0 +1,67 @@
+//! The heap header wrapping every GC-managed value.
+
+use core::cell::Cell;
+use core::ptr::NonNull;
+
+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> {
+    /// 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
+    pub(crate) drop_fn: DropFn,
+    /// Allocation ID used to validate weak pointers.
+    pub(crate) alloc_id: usize,
+    /// The user value.
+    pub(crate) value: T,
+}
+
+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);
+    }
+}