Hello, I just finished my fixed size multi pool allocator. Now I need code review :>

From the previous review, I got this following :

• small alignment can't be used to alloc bigger alignment
• the size must be power of the alignment size

I created new project to address that with the following solution :

• there are multiple pool
• each pool can be configured to have different size and alignment
• the pool with size and alignment bigger than the data's size and alignment is chosen

I already checked it with Miri and fixed the reported errors, is there still something wrong that I don't know?

#![feature(allocator_api)]

pub struct Block {
    next: Option<std::ptr::NonNull<Block>>
}

pub struct StaticMultiPoolAlloc<const N: usize> {
    start: [std::ptr::NonNull<u8>; N],
    head: [std::cell::UnsafeCell<Option<std::ptr::NonNull<Block>>>; N],
    block_size: [usize; N],
    aligns: [usize; N],
    layout: [std::alloc::Layout; N]
}

fn align_up(size: usize, align: usize) -> usize {
    (size + align - 1) & !(align - 1)
}

impl<const N: usize> StaticMultiPoolAlloc<N> {
    pub fn new(num_blocks: [usize; N], block_size: [usize; N], align_multiply: [usize; N]) -> Self {
        _ = align_multiply.map(|val| {
            if val == 0 {
                panic!("Align multiply must be > 0");
            }
        });
        let block_align = std::mem::align_of::<Block>();
        let aligns: [usize; N] = align_multiply.map(|val| val * block_align);
        
        let base_block_size = std::mem::size_of::<Block>();
        let block_size: [usize; N] = std::array::from_fn(|i| {
            let size = block_size[i].max(base_block_size);
            align_up(size, block_align) 
        });
        
        let layout: [std::alloc::Layout; N] = std::array::from_fn(|i| std::alloc::Layout::from_size_align(
            num_blocks[i] * block_size[i],
            aligns[i]
        ).expect("Error in creating layout in fn new -> PoolAllocator"));
        
        // SAFETY: the layout is valid, because if not it will trigger panic
        
        let starts: [std::ptr::NonNull<u8>; N] = std::array::from_fn(|i| {
        
        let start = unsafe { std::alloc::alloc(layout[i]) };
        let start_ptr = std::ptr::NonNull::new(start).expect("Error OOM in fn new -> PoolAllocator");
        
        unsafe {
            for j in 0..num_blocks[i] {
                let current_ptr = start.add(j * block_size[i]).cast::<Block>();
                let next_ptr = if j < num_blocks[i] - 1 {
                    Some(std::ptr::NonNull::new_unchecked(start.add((j + 1) * block_size[i]).cast::<Block>()))
                } else {
                    None
                };
                
                (*current_ptr).next = next_ptr;
            }
        }
        
        start_ptr
        
        });
        
        let head: [std::cell::UnsafeCell<Option<std::ptr::NonNull<Block>>>; N] = std::array::from_fn(|i| {
            std::cell::UnsafeCell::new(Some(std::ptr::NonNull::new(starts[i].as_ptr().cast::<Block>()).unwrap()))
        });
        
        Self {
            start: starts,
            head,
            block_size,
            aligns,
            layout,
        }
        
    }
}

impl<const N: usize> Drop for StaticMultiPoolAlloc<N> {
    fn drop(&mut self) {
        // SAFETY: deallocate using the pointer. The layout is saved, so it is the correct layout
        for (i, val) in self.start.iter().enumerate() {
            unsafe {
                std::alloc::dealloc(val.as_ptr(), self.layout[i]);
            }
        }
    }
}

unsafe impl<const N: usize> std::alloc::Allocator for StaticMultiPoolAlloc<N> {
    #[inline(always)]
    fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        let size = layout.size();
        let align = layout.align();
        //let base_align = std::mem::align_of::<Block>();
        
        let index = self.aligns.iter().enumerate().position(|(i, &val)| {
            size <= self.block_size[i] && align <= val
        });

        let i = index.ok_or(std::alloc::AllocError)?;
        
        let head_ptr = unsafe { &mut *self.head[i].get() };
        if let Some(block) = *head_ptr {
            let taken_block = block;
            unsafe { *head_ptr = block.as_ref().next };
            let ptr = taken_block.cast::<u8>();
            return Ok(std::ptr::NonNull::slice_from_raw_parts(
                ptr, self.block_size[i]
            ));
        }
        
        Err(std::alloc::AllocError)
    }
    
    #[inline(always)]
    unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
        let new_block = ptr.cast::<Block>();
        let align = layout.align();
        
        let index = self.aligns.iter().enumerate().position(|(_, &val)| {
            align <= val
        });

        let i = index.ok_or(std::alloc::AllocError).unwrap();
        
        // dereference raw pointer is unsafe
        let head_ptr = unsafe { &mut *self.head[i].get() };
        unsafe { (*new_block.as_ptr()).next = *head_ptr };
        
        *head_ptr = Some(new_block);
    }
}

#[derive(Debug)]
#[repr(align(64))]
struct Tes(i32);

fn main() {
    /*
        [pool 1, num block: 4, each block: 80 byte, total byte: 80*4, align: 8*1]
        [pool 2, num block: 2, each block: 1024 byte, total byte: 1024*2, align: 8*8]
        total prealloc in entire pool: (80*4) + (1024*2)
        fixed size, panic if no enough size available
        single threaded usage, each thread has their own pool so there is no locking, maximizing cpu cache hit
    */
    let pool: StaticMultiPoolAlloc<2> = StaticMultiPoolAlloc::new([4, 2], [80, 1024], [1, 8]);
    let mut vec = Vec::with_capacity_in(2, &pool);
    vec.push(Tes(10));
    println!("{:?}", vec[0].0);
}

1 post - 1 participant

Read full topic