fix(mempool3): increase SIZE_CLASSES to avoid buffer overruns on large objects

oscars/src/alloc/mempool3/mod.rs

@@ -28,19 +28,21 @@ impl From<LayoutError> for PoolAllocError {
     }
 }
 
-const SIZE_CLASSES: &[usize] = &[16, 24, 32, 48, 64, 96, 128, 192, 256, 512, 1024, 2048, 4096];
+const SIZE_CLASSES: &[usize] = &[
+    16, 24, 32, 48, 64, 96, 128, 192, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536,
+];
 
 #[inline(always)]
 fn size_class_index_for(size: usize) -> usize {
     // binary search over size classes
     let idx = SIZE_CLASSES.partition_point(|&sc| sc < size);
-    debug_assert!(
+    assert!(
         idx < SIZE_CLASSES.len(),
         "object size {size}B exceeds the largest size class ({}B); \
          consider adding a larger class",
         SIZE_CLASSES.last().unwrap()
     );
-    idx.min(SIZE_CLASSES.len() - 1)
+    idx
 }
 
 const DEFAULT_PAGE_SIZE: usize = 262_144;
@@ -58,7 +60,7 @@ pub struct PoolAllocator<'alloc> {
     // cached index of the last pool used by free_slot
     pub(crate) free_cache: Cell<usize>,
     // per size class cached index of the last pool used by alloc_slot
-    pub(crate) alloc_cache: [Cell<usize>; 12],
+    pub(crate) alloc_cache: [Cell<usize>; 17],
     // empty slot pools kept alive to avoid OS reallocation on the next cycle
     pub(crate) recycled_pools: Vec<SlotPool>,
     // maximum number of idle pages held across all size classes
@@ -78,20 +80,7 @@ impl<'alloc> Default for PoolAllocator<'alloc> {
             slot_pools: Vec::new(),
             bump_pages: Vec::new(),
             free_cache: Cell::new(usize::MAX),
-            alloc_cache: [
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-                Cell::new(usize::MAX),
-            ],
+            alloc_cache: core::array::from_fn(|_| Cell::new(usize::MAX)),
             recycled_pools: Vec::new(),
             // keep two empty pages per size class to reduce OS overhead
             max_recycled: SIZE_CLASSES.len() * 2,

oscars/src/alloc/mempool3/tests.rs

@@ -234,3 +234,33 @@ fn max_recycled_cap_respected() {
     assert_eq!(allocator.recycled_pools.len(), 1);
     assert!(allocator.current_heap_size < heap_before);
 }
+
+#[test]
+fn alloc_large_object_success() {
+    let mut allocator = PoolAllocator::default().with_page_size(16384);
+    // allocate a size that maps to a class > 2048
+    // e.g. 4096 size class => size 3000
+    struct LargeObject {
+        _data: [u8; 3000],
+    }
+    let _a = allocator
+        .try_alloc(LargeObject { _data: [0; 3000] })
+        .unwrap();
+    assert_eq!(allocator.pools_len(), 1);
+
+    // Test that the caching works for the new size classes
+    let _b = allocator
+        .try_alloc(LargeObject { _data: [0; 3000] })
+        .unwrap();
+    assert_eq!(allocator.pools_len(), 1); // should still fit in the same page
+}
+
+#[test]
+#[should_panic(expected = "object size")]
+fn alloc_oversized_object_panics() {
+    let mut allocator = PoolAllocator::default().with_page_size(128 * 1024);
+    struct OversizedObject {
+        _data: [u8; 70000],
+    }
+    let _ = allocator.try_alloc(OversizedObject { _data: [0; 70000] });
+}

oscars/src/collectors/mark_sweep_branded/gc.rs

@@ -19,16 +19,31 @@ pub struct Gc<'gc, T: Trace + ?Sized + 'gc> {
     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
+        let count = unsafe { &(*self.ptr.as_ptr().as_ptr()).0.root_count };
+        count.set(count.get() + 1);
+        Self {
+            ptr: self.ptr,
+            _marker: core::marker::PhantomData,
+        }
+    }
+}
+
+impl<'gc, T: Trace + ?Sized + 'gc> Drop for Gc<'gc, T> {
+    fn drop(&mut self) {
+        let count = unsafe { &(*self.ptr.as_ptr().as_ptr()).0.root_count };
+        if count.get() > 0 {
+            count.set(count.get() - 1);
+        }
     }
 }
 
 impl<'gc, T: Trace + ?Sized + 'gc> Gc<'gc, T> {
     #[inline]
     pub(crate) fn with_pointer(ptr: PoolPointer<'static, GcBox<T>>) -> Self {
+        let count = unsafe { &(*ptr.as_ptr().as_ptr()).0.root_count };
+        count.set(count.get() + 1);
         Self {
             ptr,
             _marker: PhantomData,
@@ -84,10 +99,12 @@ impl<'gc, T: Trace + ?Sized + 'gc> Gc<'gc, T> {
             .ptr
             .as_ptr()
             .cast::<crate::alloc::mempool3::PoolItem<GcBox<U>>>();
-        Gc {
+        let new_gc = Gc {
             ptr: unsafe { crate::alloc::mempool3::PoolPointer::from_raw(raw) },
             _marker: PhantomData,
-        }
+        };
+        core::mem::forget(self);
+        new_gc
     }
 
     /// Returns `true` if the inner value is of type `U`.
@@ -105,7 +122,7 @@ impl<'gc, T: Trace + ?Sized + 'gc> Gc<'gc, T> {
     #[allow(private_interfaces)]
     pub fn into_raw(self) -> core::ptr::NonNull<crate::alloc::mempool3::PoolItem<GcBox<T>>> {
         let ptr = self.ptr.as_ptr();
-        let _ = self;
+        core::mem::forget(self);
         ptr
     }
 

oscars/src/collectors/mark_sweep_branded/gc_box.rs

@@ -24,11 +24,14 @@ pub(crate) enum GcColor {
 /// Heap wrapper for a garbage-collected value.
 ///
 /// Allocated via [`PoolAllocator`].
+#[repr(C)]
 pub 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,
+    pub(crate) finalize_fn: unsafe fn(NonNull<u8>),
+    pub(crate) root_count: Cell<usize>,
     /// Type-erased finalize and free fn
     pub(crate) drop_fn: DropFn,
     /// Allocation ID used to validate weak pointers.
@@ -52,9 +55,17 @@ impl<T: Trace> GcBox<T> {
     ///
     /// Requires `T: Trace` for the `TypeId`.
     pub(crate) fn new(value: T, trace_fn: TraceFn, drop_fn: DropFn, alloc_id: usize) -> Self {
+        unsafe fn finalize_node<T: Trace>(ptr: NonNull<u8>) {
+            let item_ptr = ptr.cast::<PoolItem<GcBox<T>>>();
+            unsafe {
+                (*item_ptr.as_ptr()).0.value.run_finalizer();
+            }
+        }
         Self {
             color: Cell::new(GcColor::White),
             trace_fn,
+            finalize_fn: finalize_node::<T>,
+            root_count: Cell::new(0),
             drop_fn,
             alloc_id,
             type_id: typeid::of::<T>(),

oscars/src/collectors/mark_sweep_branded/mod.rs

@@ -45,6 +45,7 @@ pub struct Collector {
     pub(crate) sentinel: RootSentinel,
     pub(crate) generic_alloc_id: Cell<usize>,
     pub(crate) ephemerons: RefCell<Vec<EphemeronEntry>>,
+    pub(crate) is_sweeping: Cell<bool>,
 }
 
 impl Default for Collector {
@@ -61,6 +62,7 @@ impl Collector {
             sentinel: RootSentinel::new(),
             generic_alloc_id: Cell::new(0),
             ephemerons: RefCell::new(Vec::new()),
+            is_sweeping: Cell::new(false),
         }
     }
 
@@ -149,6 +151,10 @@ impl Collector {
     }
 
     /// Runs a collection cycle
+    pub fn is_sweeping(&self) -> bool {
+        self.is_sweeping.get()
+    }
+
     pub fn collect(&self) {
         self.collect_with_roots(|_| {})
     }
@@ -162,6 +168,26 @@ impl Collector {
 
         trace_external(&mut tracer);
 
+        for ptr in self.pool.borrow().iter_live_slots() {
+            unsafe {
+                let gc_box = &(*ptr
+                    .cast::<crate::alloc::mempool3::PoolItem<GcBox<()>>>()
+                    .as_ptr())
+                .0;
+                if gc_box.root_count.get() > 0 {
+                    let trace_fn_ptr = gc_box.trace_fn as *const ();
+                    if trace_fn_ptr.is_null() {
+                        panic!(
+                            "Collector::sweep: trace_fn is NULL for rooted object! alloc_id: {}, root_count: {}",
+                            gc_box.alloc_id,
+                            gc_box.root_count.get()
+                        );
+                    }
+                    (gc_box.trace_fn)(ptr, &mut tracer);
+                }
+            }
+        }
+
         for link_ptr in self.sentinel.iter() {
             unsafe {
                 // SAFETY: link_ptr points to the `link` field which is first in repr(C) RootNode.
@@ -197,7 +223,7 @@ impl Collector {
 
         // Phase 3: sweep all slots. Collect unmarked ones, then invalidate and free them.
         use crate::alloc::mempool3::PoolItem;
-        let dead: Vec<(NonNull<u8>, DropFn)> = {
+        let dead: Vec<(NonNull<u8>, DropFn, unsafe fn(NonNull<u8>))> = {
             let pool = self.pool.borrow();
             pool.iter_live_slots()
                 .filter_map(|ptr| unsafe {
@@ -206,14 +232,21 @@ impl Collector {
                         gc_box.color.set(GcColor::White);
                         None
                     } else {
-                        Some((ptr, gc_box.drop_fn))
+                        Some((ptr, gc_box.drop_fn, gc_box.finalize_fn))
                     }
                 })
                 .collect()
         };
+
+        for (ptr, _, finalize_fn) in &dead {
+            unsafe {
+                (finalize_fn)(*ptr);
+            }
+        }
+
         {
             let mut pool = self.pool.borrow_mut();
-            for (ptr, drop_fn) in dead {
+            for (ptr, drop_fn, _) in dead {
                 unsafe {
                     (*ptr.cast::<PoolItem<GcBox<()>>>().as_ptr()).0.alloc_id =
                         GcBox::<()>::FREED_ALLOC_ID;

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

@@ -6,7 +6,7 @@ fn ephemeron_value_survives_when_key_is_rooted() {
         let (root_key, eph) = ctx.mutate(|cx| {
             let key = cx.try_alloc(1u32).unwrap();
             let value = cx.try_alloc(42u32).unwrap();
-            let root_key = cx.root(key).unwrap();
+            let root_key = cx.root(key.clone()).unwrap();
             let eph = cx.alloc_ephemeron(&key, value);
             (root_key, eph)
         });
@@ -54,8 +54,8 @@ fn ephemeron_chain_fixpoint() {
             let a = cx.try_alloc(1u32).unwrap();
             let b = cx.try_alloc(2u32).unwrap();
             let c = cx.try_alloc(3u32).unwrap();
-            let root_a = cx.root(a).unwrap();
-            let eph_ab = cx.alloc_ephemeron(&a, b);
+            let root_a = cx.root(a.clone()).unwrap();
+            let eph_ab = cx.alloc_ephemeron(&a, b.clone());
             let eph_bc = cx.alloc_ephemeron(&b, c);
             (root_a, eph_ab, eph_bc)
         });

oscars/src/collectors/mark_sweep_branded/tests/ui/gc_cannot_store_outer_scopr.stderr

@@ -8,4 +8,4 @@ error[E0521]: borrowed data escapes outside of closure
    |                     -- `cx` is a reference that is only valid in the closure body
 ...
 23 |             holder = Some(Holder { gc });
-   |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `cx` escapes the closure body here
+   |             ^^^^^^ `cx` escapes the closure body here

oscars/src/collectors/mark_sweep_branded/trace.rs

@@ -96,6 +96,13 @@ impl<'a> Tracer<'a> {
             let gc_box = &(*gc.ptr.as_ptr().as_ptr()).0;
             if gc_box.color.get() == GcColor::White {
                 gc_box.color.set(GcColor::Gray);
+                let trace_fn_ptr = gc_box.trace_fn as *const ();
+                if trace_fn_ptr.is_null() {
+                    panic!(
+                        "Tracer::mark: trace_fn is NULL! alloc_id: {}",
+                        gc_box.alloc_id
+                    );
+                }
                 self.worklist
                     .push((gc.ptr.as_ptr().cast::<u8>(), gc_box.trace_fn));
             }