I need code advises for arena allocator

⚓ Rust    📅 2025-10-22    👤 surdeus    👁️ 4      

surdeus

Here are the arena specs I want to achieve :

  • single threaded usage
  • no individual item deallocation support (all items are dropped at once by dropping the arena container)
  • support reuse arena via reset

So far, my code is like this

#![feature(allocator_api)]

pub struct Arena {
    val: std::ptr::NonNull<u8>,
    handle: std::cell::UnsafeCell<*mut u8>,
    end: *mut u8,
    layout: std::alloc::Layout
}

impl Arena {
    #[inline(always)]
    pub fn new(size: usize) -> Self {
        unsafe {
            let layout = std::alloc::Layout::from_size_align_unchecked(size, std::mem::align_of::<usize>());

            let start = std::alloc::alloc(layout);
            if start.is_null() {
                std::alloc::handle_alloc_error(layout);
            }

            Self {
                val: std::ptr::NonNull::new_unchecked(start),
                handle: start.into(),
                end: start.add(size),
                layout: layout
            }
        }
    }

    #[inline(always)]
    fn inner_alloc<T>(&mut self, val: T) -> *mut T {
        unsafe {
            let handle = self.handle.get();
            let align = std::mem::align_of::<T>();
            //let size = std::mem::size_of::<T>();

            let handle_addr = (*handle) as usize;
            let aligned_addr = (handle_addr + (align - 1)) & !(align - 1);
            let ptr = aligned_addr as *mut T;
            let next = ptr.add(1) as *mut u8;
            debug_assert!(next <= self.end, "arena is out of memory");

            ptr.write(val);
            *handle = next;
            ptr
        }
    }

    #[inline(always)]
    pub fn alloc<T>(&mut self, val: T) -> &T {
      let ptr = self.inner_alloc(val);
      unsafe { & *ptr }
    }
    
    #[inline(always)]
    pub fn alloc_mut<T>(&mut self, val: T) -> &mut T {
      let ptr = self.inner_alloc(val);
      unsafe { &mut *ptr }
    }

    #[inline(always)]
    pub fn reset(&self) {
        let handle = self.handle.get();
        unsafe { *handle = self.val.as_ptr(); }
    }
}

impl Drop for Arena {
    #[inline(always)]
    fn drop(&mut self) {
        unsafe {
            std::alloc::dealloc(self.val.as_ptr(), self.layout);
        }
    }
}

unsafe impl std::alloc::Allocator for &Arena {
    #[inline(always)]
    fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {

      unsafe {
        let handle = self.handle.get();
        let size = layout.size();
        let align = layout.align();
        
        let handle_addr = (*handle) as usize;
        let aligned_addr = (handle_addr + (align - 1)) & !(align - 1);
        let ptr = aligned_addr as *mut u8;
        let next = ptr.add(size);
        debug_assert!(next <= self.end, "arena is out of memory");
        *handle = next;
        
        let offset = aligned_addr - self.val.as_ptr() as usize;
        let ptr = self.val.as_ptr().add(offset);

        return Ok(std::ptr::NonNull::new_unchecked(std::slice::from_raw_parts_mut(ptr, size)));
      }
    }
    
    #[inline(always)]
    unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: std::alloc::Layout) {

    }
}

fn main() {
  let mut arena = Arena::new(40);
        
        let a = arena.alloc("a");
        println!("{}", *a);
        
        let mut v = Vec::with_capacity_in(1, &arena);
        v.push(100);
        v.push(100);
        println!("{}, {}", v[0], v[1]);
        
        let mut w = Vec::with_capacity_in(1, &arena);
        w.push(10);
        println!("{}", w[0]);

        arena.reset();
        
        let mut w = Vec::with_capacity_in(1, &arena);
        w.push(10);
        println!("{}", w[0]);

       /* allocate more than the capacity will panic in debug build
       let mut w = Vec::with_capacity_in(100, &arena);
        w.push(10);
        println!("{}", w[0]);
       */

}

I'm not knowledgeable with raw pointer usage, this is the project I write to learn about it

I need code suggestions whether it is security or performance thing

1 post - 1 participant

Read full topic

🏷️ Rust_feed