Fix incorrect layouts in small-object arena alloc

As of now, the 'array' tests pass.

Add cfg!(zerogc_simple_debug_alloc) which adds internal padding bytes.
This was very helpful in finding this bug ;)

Author: Techcable

libs/simple/src/alloc.rs

@@ -15,6 +15,88 @@ use std::cell::RefCell;
 
 use zerogc_context::utils::AtomicCell;
 
+const DEBUG_INTERNAL_ALLOCATOR: bool = cfg!(zerogc_simple_debug_alloc);
+mod debug {
+    pub const PADDING: u32 = 0xDEADBEAF;
+    pub const UNINIT: u32 = 0xCAFEBABE;
+    pub const PADDING_TIMES: usize = 16;
+    pub const PADDING_BYTES: usize = PADDING_TIMES * 4;
+    pub unsafe fn pad_memory_block(ptr: *mut u8, size: usize) {
+        assert!(super::DEBUG_INTERNAL_ALLOCATOR);
+        let start = ptr.sub(PADDING_BYTES);
+        for i in 0..PADDING_TIMES {
+            (start as *mut u32).add(i).write(PADDING);
+        }
+        let end = ptr.add(size);
+        for i in 0..PADDING_TIMES {
+            (end as *mut u32).add(i).write(PADDING);
+        }
+    }
+    pub unsafe fn mark_memory_uninit(ptr: *mut u8, size: usize) {
+        assert!(super::DEBUG_INTERNAL_ALLOCATOR);
+        let (blocks, leftover) = (size / 4, size % 4);
+        for i in 0..blocks {
+            (ptr as *mut u32).add(i).write(UNINIT);
+        }
+        let leftover_ptr = ptr.add(blocks * 4);
+        debug_assert_eq!(leftover_ptr.wrapping_add(leftover), ptr.add(size));
+        for i in 0..leftover {
+            leftover_ptr.add(i).write(0xF0);
+        }
+    }
+    pub unsafe fn assert_padded(ptr: *mut u8, size: usize) {
+        assert!(super::DEBUG_INTERNAL_ALLOCATOR);
+        let start = ptr.sub(PADDING_BYTES);
+        let end = ptr.add(size);
+        let start_padding = std::slice::from_raw_parts(
+            start as *const u8 as *const u32,
+            PADDING_TIMES
+        );
+        let region = std::slice::from_raw_parts(
+            ptr as *const u8,
+            size
+        );
+        let end_padding = std::slice::from_raw_parts(
+            end as *const u8 as *const u32,
+            PADDING_TIMES
+        );
+        let print_memory_region = || {
+            use std::fmt::Write;
+            let mut res = String::new();
+            for &val in start_padding {
+                write!(&mut res, "{:X}", val).unwrap();
+            }
+            res.push_str("||");
+            for &b in region {
+                write!(&mut res, "{:X}", b).unwrap();
+            }
+            res.push_str("||");
+            for &val in end_padding {
+                write!(&mut res, "{:X}", val).unwrap();
+            }
+            res
+        };
+        // Closest to farthest
+        for (idx, &block) in start_padding.iter().rev().enumerate() {
+            if block == PADDING { continue }
+            assert_eq!(
+                block, PADDING,
+                "Unexpected start padding (offset -{}) w/ {}",
+                idx * 4,
+                print_memory_region()
+            );
+        }
+        for (idx, &block) in end_padding.iter().enumerate() {
+            if block == PADDING { continue }
+            assert_eq!(
+                block, PADDING,
+                "Unexpected end padding (offset {}) w/ {}",
+                idx * 4,
+                print_memory_region()
+            )
+        }
+    }
+}
 /// The minimum size of supported memory (in words)
 ///
 /// Since the header takes at least one word,
@@ -73,10 +155,6 @@ impl Chunk {
         }
     }
     #[inline]
-    fn current(&self) -> *mut u8 {
-        self.current.load()
-    }
-    #[inline]
     fn capacity(&self) -> usize {
         self.end as usize - self.start as usize
     }
@@ -200,7 +278,11 @@ pub(crate) struct FreeList {
     next: AtomicCell<Option<NonNull<FreeSlot>>>
 }
 impl FreeList {
-    pub(crate) unsafe fn add_free(&self, free: *mut UnknownHeader) {
+    unsafe fn add_free(&self, free: *mut UnknownHeader, size: usize) {
+        if DEBUG_INTERNAL_ALLOCATOR {
+            debug::assert_padded(free as *mut u8, size);
+            debug::mark_memory_uninit(free as *mut u8, size);
+        }
         let new_slot = free as *mut FreeSlot;
         let mut next = self.next.load();
         loop {
@@ -214,14 +296,6 @@ impl FreeList {
         }
     }
     #[inline]
-    pub(crate) fn next_free(&self) -> Option<NonNull<FreeSlot>> {
-        self.next.load()
-    }
-    #[inline]
-    pub(crate) unsafe fn set_next_free(&self, next: Option<NonNull<FreeSlot>>) {
-        self.next.store(next)
-    }
-    #[inline]
     fn take_free(&self) -> Option<NonNull<u8>> {
         loop {
             let next_free = match self.next.load() {
@@ -248,7 +322,7 @@ pub struct SmallArena {
 
 impl SmallArena {
     pub(crate) unsafe fn add_free(&self, obj: *mut UnknownHeader) {
-        self.free.add_free(obj)
+        self.free.add_free(obj, self.element_size)
     }
     #[cold] // Initialization is the slow path
     fn with_words(num_words: usize) -> SmallArena {
@@ -266,6 +340,15 @@ impl SmallArena {
         // Check the free list
         if let Some(free) = self.free.take_free() {
             free.cast()
+        } else if DEBUG_INTERNAL_ALLOCATOR {
+            let mem = self.state.alloc(self.element_size + debug::PADDING_BYTES * 2)
+                .as_ptr() as *mut u8;
+            unsafe {
+                let mem = mem.add(debug::PADDING_BYTES);
+                debug::pad_memory_block(mem, self.element_size);
+                debug::mark_memory_uninit(mem, self.element_size);
+                NonNull::new_unchecked(mem).cast()
+            }
         } else {
             self.state.alloc(self.element_size)
         }
@@ -314,9 +397,6 @@ impl SmallArenaList {
             }
         }
     }
-    pub fn iter(&self) -> impl Iterator<Item=&SmallArena> + '_ {
-        self.arenas.iter().filter_map(OnceCell::get)
-    }
     #[inline] // This should hopefully be constant folded away (layout is const)
     pub fn find(&self, layout: Layout) -> Option<&SmallArena> {
         if !fits_small_object(layout) {

libs/simple/src/lib.rs

@@ -351,6 +351,13 @@ pub(crate) struct SimpleAlloc {
     mark_inverted: AtomicBool,
     allocated_size: AtomicUsize
 }
+#[derive(Debug)]
+struct TargetLayout<H> {
+    value_layout: Layout,
+    header_layout: HeaderLayout<H>,
+    value_offset: usize,
+    overall_layout: Layout
+}
 impl SimpleAlloc {
     fn new() -> SimpleAlloc {
         SimpleAlloc {
@@ -392,14 +399,25 @@ impl SimpleAlloc {
                     unreachable!("Invalid collector id")
                 }
                 #[cfg(not(debug_assertions))] {
-                    std::hint::unreachable_unchecked()
+                    unsafe { std::hint::unreachable_unchecked() }
                 }
             }
         };
-        let (header, value_ptr) = if let Some(arena) = self.small_arenas.find(value_layout) {
-            unsafe { self.alloc_layout_small(arena, header_layout, value_layout) }
+        let (mut overall_layout, value_offset) = header_layout.layout()
+            .extend(value_layout).unwrap();
+        debug_assert_eq!(header_layout.value_offset(value_layout.align()), value_offset);
+        overall_layout  = overall_layout.pad_to_align();
+        debug_assert_eq!(
+            value_offset,
+            header_layout.value_offset(value_layout.align())
+        );
+        let target_layout = TargetLayout {
+            header_layout, value_offset, value_layout, overall_layout
+        };
+        let (header, value_ptr) = if let Some(arena) = self.small_arenas.find(target_layout.overall_layout) {
+            unsafe { self.alloc_layout_small(arena, target_layout) }
         } else {
-            self.alloc_layout_big(header_layout, value_layout)
+            self.alloc_layout_big(target_layout)
         };
         unsafe {
             header_layout.common_header(header).write(GcHeader::new(
@@ -413,45 +431,34 @@ impl SimpleAlloc {
         (header, value_ptr)
     }
     #[inline]
-    unsafe fn alloc_layout_small<H>(&self, arena: &SmallArena, header_layout: HeaderLayout<H>, value_layout: Layout) -> (*mut H, *mut u8) {
-        let (mut overall_layout, value_offset) = header_layout.layout()
-            .extend(value_layout).unwrap();
-        overall_layout = overall_layout.pad_to_align();
-        debug_assert_eq!(
-            value_offset,
-            header_layout.value_offset(value_layout.align())
-        );
+    unsafe fn alloc_layout_small<H>(&self, arena: &SmallArena, target_layout: TargetLayout<H>) -> (*mut H, *mut u8) {
         let ptr = arena.alloc();
         debug_assert_eq!(
-            ptr.as_ptr() as usize % header_layout.layout().align(),
+            ptr.as_ptr() as usize % target_layout.header_layout.layout().align(),
             0
         );
-        self.add_allocated_size(overall_layout.size());
+        self.add_allocated_size(target_layout.overall_layout.size());
         {
             let mut lock = self.small_objects.lock();
-            lock.push((ptr.as_ptr() as *mut u8).add(header_layout.common_header_offset).cast());
+            lock.push((ptr.as_ptr() as *mut u8).add(target_layout.header_layout.common_header_offset).cast());
         }
-        (ptr.as_ptr().cast(), (ptr.as_ptr() as *mut u8).add(value_offset))
+        (ptr.as_ptr().cast(), (ptr.as_ptr() as *mut u8).add(target_layout.value_offset))
     }
-    fn alloc_layout_big<H>(&self, header_layout: HeaderLayout<H>, value_layout: Layout) -> (*mut H, *mut u8) {
-        let (mut overall_layout, value_offset) = header_layout.layout()
-            .extend(value_layout).unwrap();
-        debug_assert_eq!(header_layout.value_offset(value_layout.align()), value_offset);
-        overall_layout  = overall_layout.pad_to_align();
+    fn alloc_layout_big<H>(&self, target_layout: TargetLayout<H>) -> (*mut H, *mut u8) {
         let header: *mut H;
         let value_ptr = unsafe {
-            header = std::alloc::alloc(overall_layout).cast();
-            (header as *mut u8).add(value_offset)
+            header = std::alloc::alloc(target_layout.overall_layout).cast();
+            (header as *mut u8).add(target_layout.value_offset)
         };
         {
             unsafe {
                 let mut objs = self.big_objects.lock();
                 let common_header = (header as *mut u8)
-                    .add(header_layout.common_header_offset)
+                    .add(target_layout.header_layout.common_header_offset)
                     .cast();
                 objs.push(BigGcObject::from_ptr(common_header));
             }
-            self.add_allocated_size(overall_layout.size());
+            self.add_allocated_size(target_layout.overall_layout.size());
         }
         (header, value_ptr)
     }
@@ -486,7 +493,7 @@ impl SimpleAlloc {
             }
             let overall_layout = type_info.determine_total_layout(freed_common_header);
             let actual_start = type_info.header_layout()
-                .from_common_header(freed_common_header) ;
+                .from_common_header(freed_common_header);
             self.small_arenas.find(overall_layout).unwrap().add_free(actual_start)
         });
         // Clear large objects