nimix 0.2.0

An allocator designed to be use by a GC
Documentation
use super::error::AllocError;
use super::constants::{FREE_MARK, LARGE_OBJECT_MIN};

use alloc::alloc::{alloc, dealloc, Layout};
use core::num::NonZero;
use core::sync::atomic::{AtomicU8, Ordering};
use core::ptr::write;

pub struct LargeBlock {
    layout: Layout,
    ptr: *mut u8,
    mark: *const AtomicU8
}

impl Drop for LargeBlock {
    fn drop(&mut self) {
        unsafe { dealloc(self.ptr, self.layout) }
    }
}

impl LargeBlock {
    pub fn new(obj_layout: Layout) -> Result<Self, AllocError> {
        debug_assert!(obj_layout.size() >= LARGE_OBJECT_MIN);

        let mark_layout = Layout::new::<AtomicU8>();
        let (obj_mark_layout, mark_offset) = obj_layout.extend(mark_layout)?;

        let block_layout = obj_mark_layout.pad_to_align();

        unsafe {
            let ptr = alloc(block_layout);

            if ptr.is_null() {
                return Err(AllocError::OOM);
            }

            let mark = ptr.add(mark_offset) as *const AtomicU8;
            write(mark as *mut AtomicU8, AtomicU8::new(FREE_MARK));

            let large_block = Self {
                layout: block_layout,
                ptr,
                mark
            };

            Ok(large_block)
        }
    }

    pub unsafe fn mark(ptr: *const u8, obj_layout: Layout, mark: NonZero<u8>) -> Result<(), AllocError> {
        let mark_layout = Layout::new::<AtomicU8>();
        let (_, mark_offset) = obj_layout.extend(mark_layout)?;
        let block_mark: *const AtomicU8 = ptr.add(mark_offset) as *const AtomicU8;

        (&*block_mark).store(mark.into(), Ordering::Relaxed);

        Ok(())
    }

    pub fn is_marked(&self, mark: NonZero<u8>) -> bool {
        unsafe { (&*self.mark).load(Ordering::Relaxed) == u8::from(mark) }
    }

    pub fn get_size(&self) -> usize {
        self.layout.size()
    }

    pub fn as_ptr(&self) -> *const u8 {
        self.ptr
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use core::ptr::write;

    #[test]
    fn new_large_block() {
        let align = 8;
        let block =  LargeBlock::new(Layout::from_size_align(LARGE_OBJECT_MIN, align).unwrap()).unwrap();

        assert!(block.get_size() > LARGE_OBJECT_MIN);
        assert_eq!(block.as_ptr() as usize % align, 0);
        assert_eq!(block.is_marked(NonZero::new(1).unwrap()), false);
    }

    #[test]
    fn mark_large() {
        let data = [0u8; LARGE_OBJECT_MIN];
        let align = 8;
        let layout = Layout::from_size_align(LARGE_OBJECT_MIN, align).unwrap();
        let block =  LargeBlock::new(layout).unwrap();

        unsafe { 
            write(block.as_ptr() as *mut [u8; LARGE_OBJECT_MIN], data);
            LargeBlock::mark(block.as_ptr(), layout, NonZero::new(1).unwrap()).unwrap();

            assert_eq!(&*(block.as_ptr() as *mut [u8; LARGE_OBJECT_MIN]), &data);
        }

        assert!(block.is_marked(NonZero::new(1).unwrap()));
    }
}