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

Author: shruti2522

oscars/src/alloc/mempool3/alloc.rs

@@ -59,15 +59,24 @@ 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 (), *mut 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> {}
+
+impl<'pool, T: ?Sized> PoolPointer<'pool, T> {
     pub fn as_inner_ref(&self) -> &'pool T
     where
         T: 'pool,
@@ -80,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)
     }
 
@@ -95,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 {

oscars/src/collectors/mark_sweep_branded/ephemeron.rs

@@ -1,5 +1,5 @@
 use crate::{
-    alloc::mempool3::PoolItem,
+    alloc::mempool3::PoolPointer,
     collectors::mark_sweep_branded::{
         gc::Gc,
         gc_box::GcBox,
@@ -8,20 +8,21 @@ use crate::{
     },
 };
 use core::marker::PhantomData;
-use core::ptr::NonNull;
 
 pub struct Ephemeron<'id, K: Trace, V: Trace> {
-    pub(crate) key_ptr: NonNull<PoolItem<GcBox<K>>>,
+    pub(crate) key_ptr: Option<PoolPointer<'static, GcBox<K>>>,
     pub(crate) key_alloc_id: usize,
-    pub(crate) value_ptr: NonNull<PoolItem<GcBox<V>>>,
+    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 = unsafe { (*self.key_ptr.as_ptr()).0.alloc_id == self.key_alloc_id };
+        // 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,

oscars/src/collectors/mark_sweep_branded/gc.rs

@@ -1,7 +1,7 @@
 //! Core pointer types.
 
 use crate::{
-    alloc::mempool3::{PoolAllocator, PoolItem},
+    alloc::mempool3::{PoolAllocator, PoolPointer},
     collectors::mark_sweep_branded::{
         gc_box::GcBox,
         mutation_ctx::MutationContext,
@@ -19,7 +19,7 @@ pub(crate) type RootDropFn = unsafe fn(&mut PoolAllocator<'static>, NonNull<u8>)
 /// A transient pointer to a GC-managed value.
 #[derive(Debug)]
 pub struct Gc<'gc, T: Trace + ?Sized + 'gc> {
-    pub(crate) ptr: NonNull<PoolItem<GcBox<T>>>,
+    pub(crate) ptr: PoolPointer<'static, GcBox<T>>,
     pub(crate) _marker: PhantomData<(&'gc T, *const ())>,
 }
 
@@ -38,7 +38,7 @@ impl<'gc, T: Trace + 'gc> Gc<'gc, T> {
         // 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()).0.value }
+        unsafe { &(*self.ptr.as_ptr().as_ptr()).0.value }
     }
 }
 
@@ -61,14 +61,25 @@ pub(crate) struct RootNode<'id, T: Trace> {
     /// Intrusive list link
     pub(crate) link: RootLink,
     /// Pointer to the allocation
-    pub(crate) gc_ptr: NonNull<PoolItem<GcBox<T>>>,
+    pub(crate) gc_ptr: PoolPointer<'static, 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 ()>,
 }
 
+/// Type-erased version of [`RootNode`] for use during collection.
+///
+/// Since [`RootNode`] is `repr(C)` and `link` is always the first field,
+/// a `NonNull<RootLink>` from the sentinel iterator can be safely cast to
+/// `NonNull<ErasedRootNode>` to read `gc_ptr` without knowing `T`.
+#[repr(C)]
+pub(crate) struct ErasedRootNode {
+    pub(crate) link: RootLink,
+    pub(crate) gc_ptr: PoolPointer<'static, GcBox<()>>,
+}
+
 /// A handle that keeps a GC allocation live.
 #[must_use = "dropping a root unregisters it from the GC"]
 pub struct Root<'id, T: Trace> {

oscars/src/collectors/mark_sweep_branded/mod.rs

@@ -20,19 +20,18 @@ pub use mutation_ctx::MutationContext;
 pub use trace::{Finalize, Trace, Tracer};
 pub use weak::WeakGc;
 
-use crate::alloc::mempool3::PoolAllocator;
+use crate::alloc::mempool3::{PoolAllocError, PoolAllocator, PoolPointer};
 use core::cell::{Cell, RefCell};
 use core::marker::PhantomData;
 use core::ptr::NonNull;
 use gc_box::{DropFn, GcBox, GcColor};
-use root_link::{RootLink, RootSentinel};
+use root_link::RootSentinel;
 use rust_alloc::vec::Vec;
 
 /// Type-erased ephemeron registration.
 pub(crate) struct EphemeronEntry {
-    pub(crate) key_ptr: NonNull<u8>,
-    pub(crate) key_alloc_id: usize,
-    pub(crate) value_ptr: NonNull<u8>,
+    pub(crate) key_ptr: Option<PoolPointer<'static, GcBox<()>>>,
+    pub(crate) value_ptr: PoolPointer<'static, GcBox<()>>,
 }
 
 pub(crate) struct Collector {
@@ -61,21 +60,19 @@ impl Collector {
     /// Registers an ephemeron key/value pair for processing during collection.
     pub(crate) fn register_ephemeron(
         &self,
-        key_ptr: NonNull<u8>,
-        key_alloc_id: usize,
-        value_ptr: NonNull<u8>,
+        key_ptr: PoolPointer<'static, GcBox<()>>,
+        value_ptr: PoolPointer<'static, GcBox<()>>,
     ) {
         self.ephemerons.borrow_mut().push(EphemeronEntry {
-            key_ptr,
-            key_alloc_id,
+            key_ptr: Some(key_ptr),
             value_ptr,
         });
     }
 
     /// 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>>>,
+        gc_ptr: PoolPointer<'static, GcBox<T>>,
     ) -> NonNull<gc::RootNode<'id, T>> {
         unsafe fn drop_and_free<T: trace::Trace>(
             pool: &mut PoolAllocator<'static>,
@@ -121,7 +118,7 @@ impl Collector {
     pub(crate) fn alloc<'gc, T: trace::Trace + trace::Finalize + 'gc>(
         &'gc self,
         value: T,
-    ) -> Gc<'gc, T> {
+    ) -> Result<Gc<'gc, T>, PoolAllocError> {
         let alloc_id = self.generic_alloc_id.get();
 
         // Check for alloc_id wrap before incrementing.
@@ -151,51 +148,32 @@ impl Collector {
         }
 
         let mut pool = self.pool.borrow_mut();
-        let ptr = pool
-            .try_alloc(GcBox::new(
-                value,
-                gc_box::trace_value::<T>,
-                drop_and_free::<T>,
-                alloc_id,
-            ))
-            .expect("branded GC: pool allocation failed");
+        let ptr = pool.try_alloc(GcBox::new(
+            value,
+            gc_box::trace_value::<T>,
+            drop_and_free::<T>,
+            alloc_id,
+        ))?;
 
         drop(pool);
 
-        Gc {
-            ptr: ptr.as_ptr(),
+        Ok(Gc {
+            ptr: unsafe { ptr.extend_lifetime() },
             _marker: PhantomData,
-        }
+        })
     }
 
     /// Runs a collection cycle
     pub(crate) fn collect(&self) {
         let mut tracer = Tracer::new();
 
-        let gc_ptr_offset = core::mem::offset_of!(
-            crate::collectors::mark_sweep_branded::gc::RootNode<'static, i32>,
-            gc_ptr
-        );
-        debug_assert_eq!(
-            gc_ptr_offset,
-            core::mem::offset_of!(
-                crate::collectors::mark_sweep_branded::gc::RootNode<'static, u64>,
-                gc_ptr
-            ),
-            "gc_ptr offset must be stable across all T: Sized"
-        );
-
         for link_ptr in self.sentinel.iter() {
             unsafe {
-                // Read the `gc_ptr` field using the stable offset
-                let gc_ptr_ptr = link_ptr
-                    .as_ptr()
-                    .cast::<u8>()
-                    .add(gc_ptr_offset)
-                    .cast::<NonNull<u8>>();
-                let gc_ptr = gc_ptr_ptr.read();
-
-                tracer.mark_raw(gc_ptr.cast::<u8>());
+                // SAFETY: link_ptr points to the `link` field which is first in repr(C) RootNode.
+                // Casting to ErasedRootNode (also repr(C), same first two fields) lets us read
+                // gc_ptr without knowing T, avoiding manual offset arithmetic.
+                let erased = link_ptr.cast::<gc::ErasedRootNode>();
+                tracer.mark_raw((*erased.as_ptr()).gc_ptr.as_ptr().cast::<u8>());
             }
         }
 
@@ -207,15 +185,12 @@ impl Collector {
         loop {
             let mut any_newly_marked = false;
             for entry in self.ephemerons.borrow().iter() {
-                use crate::alloc::mempool3::PoolItem;
+                let Some(key_ptr) = entry.key_ptr else {
+                    continue;
+                };
                 unsafe {
-                    let key_box = entry.key_ptr.cast::<PoolItem<GcBox<()>>>();
-                    // Skip entries invalidated by a previous collection cycle.
-                    if (*key_box.as_ptr()).0.alloc_id != entry.key_alloc_id {
-                        continue;
-                    }
-                    if (*key_box.as_ptr()).0.color.get() != GcColor::White {
-                        any_newly_marked |= tracer.mark_raw(entry.value_ptr);
+                    if (*key_ptr.as_ptr().as_ptr()).0.color.get() != GcColor::White {
+                        any_newly_marked |= tracer.mark_raw(entry.value_ptr.as_ptr().cast::<u8>());
                     }
                 }
             }
@@ -253,12 +228,12 @@ impl Collector {
         }
 
         // Phase 4: remove ephemeron entries whose key was swept this cycle.
+        // A swept key has alloc_id set to FREED_ALLOC_ID by the sweep above.
+        // Using the Option lets us express the invalid state without a stored alloc_id.
         self.ephemerons.borrow_mut().retain(|entry| {
-            use crate::alloc::mempool3::PoolItem;
-            unsafe {
-                let key_box = entry.key_ptr.cast::<PoolItem<GcBox<()>>>();
-                (*key_box.as_ptr()).0.alloc_id == entry.key_alloc_id
-            }
+            entry.key_ptr.is_some_and(|key_ptr| unsafe {
+                (*key_ptr.as_ptr().as_ptr()).0.alloc_id != GcBox::<()>::FREED_ALLOC_ID
+            })
         });
     }
 }

oscars/src/collectors/mark_sweep_branded/mutation_ctx.rs

@@ -1,12 +1,16 @@
 //! `MutationContext<'id, 'gc>` handle.
 
-use crate::collectors::mark_sweep_branded::{
-    Collector,
-    ephemeron::Ephemeron,
-    gc::{Gc, Root},
-    root_link::RootLink,
-    trace::{Finalize, Trace},
-    weak::WeakGc,
+use crate::{
+    alloc::mempool3::{PoolAllocError, PoolPointer},
+    collectors::mark_sweep_branded::{
+        Collector,
+        ephemeron::Ephemeron,
+        gc::{Gc, Root},
+        gc_box::GcBox,
+        root_link::RootLink,
+        trace::{Finalize, Trace},
+        weak::WeakGc,
+    },
 };
 use core::marker::PhantomData;
 
@@ -18,17 +22,13 @@ pub struct MutationContext<'id, 'gc> {
 
 impl<'id, 'gc> MutationContext<'id, 'gc> {
     /// Allocates a value on the GC heap.
-    ///
-    /// # Panics
-    ///
-    /// Panics if the pool allocator fails to allocate.
-    pub fn alloc<T: Trace + Finalize + 'gc>(&self, value: T) -> Gc<'gc, T> {
+    pub fn alloc<T: Trace + Finalize + 'gc>(&self, value: T) -> Result<Gc<'gc, T>, PoolAllocError> {
         self.collector.alloc(value)
     }
 
     /// Downgrades a `Gc` into a weak reference
     pub fn alloc_weak<T: Trace + Finalize + 'gc>(&self, gc: Gc<'gc, T>) -> WeakGc<'id, T> {
-        let alloc_id = unsafe { (*gc.ptr.as_ptr()).0.alloc_id };
+        let alloc_id = unsafe { (*gc.ptr.as_ptr().as_ptr()).0.alloc_id };
         WeakGc {
             ptr: gc.ptr,
             alloc_id,
@@ -60,14 +60,16 @@ impl<'id, 'gc> MutationContext<'id, 'gc> {
         key: Gc<'gc, K>,
         value: Gc<'gc, V>,
     ) -> Ephemeron<'id, K, V> {
-        let key_alloc_id = unsafe { (*key.ptr.as_ptr()).0.alloc_id };
-        self.collector.register_ephemeron(
-            key.ptr.cast::<u8>(),
-            key_alloc_id,
-            value.ptr.cast::<u8>(),
-        );
+        let key_alloc_id = unsafe { (*key.ptr.as_ptr().as_ptr()).0.alloc_id };
+        // SAFETY: GcBox<K> and GcBox<V> are erased to GcBox<()>, the collector
+        // only reads the fixed size prefix fields via this pointer
+        let erased_key: PoolPointer<'static, GcBox<()>> =
+            unsafe { key.ptr.to_erased().to_typed_pool_pointer::<GcBox<()>>() };
+        let erased_value: PoolPointer<'static, GcBox<()>> =
+            unsafe { value.ptr.to_erased().to_typed_pool_pointer::<GcBox<()>>() };
+        self.collector.register_ephemeron(erased_key, erased_value);
         Ephemeron {
-            key_ptr: key.ptr,
+            key_ptr: Some(key.ptr),
             key_alloc_id,
             value_ptr: value.ptr,
             _marker: core::marker::PhantomData,

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

@@ -25,7 +25,9 @@ mod tests {
         // Assert offset is correct via offset_of-like macro conceptually
         let make_dummy = || RootNode::<i32> {
             link: RootLink::new(),
-            gc_ptr: core::ptr::NonNull::dangling(),
+            gc_ptr: unsafe {
+                crate::alloc::mempool3::PoolPointer::from_raw(core::ptr::NonNull::dangling())
+            },
             drop_fn: |_, _| {},
             collector_ptr: core::ptr::null(),
             _marker: core::marker::PhantomData,