Skip to content

Arena Allocator

Region-based memory management.

ArenaElement

rust
struct ArenaElement {
    value: *mut u8,
    drop: unsafe fn(*mut u8),
}
  • value A raw pointer to the allocated element in the arena.
  • drop A pointer to an unsafe function that will be called when the element in the arena is dropped.

This will execute self.drop with the pointer to the element whenever ArenaElement goes out of scope or is explicitly dropped.

rust
impl Drop for ArenaElement {
    #[inline(always)]
    fn drop(&mut self) {
        unsafe {
            (self.drop)(self.value);
        }
    }
}

The Arena struct will hold the ArenaElements:

rust
pub struct Arena {
    start: *mut u8,
    end: *mut u8,
    offset: *mut u8,
    elements: Vec<ArenaElement>,
    valid: Rc<Cell<bool>>,
}
  • start: A pointer to the start of the allocated memory region.
  • end: A pointer to the end of the allocated memory region.
  • offset: A pointer indicating the current allocation offset within the memory region.
  • elements: A vector of allocated ArenaElements.
  • valid: A reference-counted (Rc) boolean indicating if the arena is valid (not dropped)

Arena

To create a new Arena struct the method new is used. new allocates a new memory region specified in bytes, elements is empty, start, end and offset are calculated from the size_in_bytes parameter, valid is set to true.

rust
impl Arena {
    pub fn new(size_in_bytes: usize) -> Self {
        unsafe {
            let layout = alloc::Layout::from_size_align(size_in_bytes, 1).unwrap();
            let start = alloc::alloc(layout);
            let end = start.add(size_in_bytes);
            Self {
                start,
                end,
                offset: start,
                elements: Vec::new(),
                valid: Rc::new(Cell::new(true)),
            }
        }
    }

	// ...

The len function returns the number of bytes allocated so far, is_empty will return true if no bytes have been allocated. The capacity function returns the total capacity of the arena (in bytes).

rust
pub fn len(&self) -> usize {
    self.offset as usize - self.start as usize
}

pub fn is_empty(&self) -> bool {
	self.len() == 0
}

pub fn capacity(&self) -> usize {
    self.end as usize - self.start as usize
}

This method will clear the arena by resetting the elements vector, invalidating the valid flag and resetting the offset back to start.

rust
pub fn clear(&mut self) {
    self.valid.set(false);
    self.valid = Rc::new(Cell::new(true));
    self.elements.clear();
    self.offset = self.start;
}

Allocates space for a new object of type T and calls a closure f to construct the object in place. It also records the object in elements to ensure it is properly dropped when the arena is cleared or dropped.

rust
#[inline(always)]
pub fn alloc<T>(&mut self, f: impl FnOnce() -> T) -> ArenaBox<T> {
    #[inline(always)]
    unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
    where
        F: FnOnce() -> T,
    {
        ptr::write(ptr, f());
    }

    unsafe fn drop<T>(ptr: *mut u8) {
        std::ptr::drop_in_place(ptr.cast::<T>());
    }

    unsafe {
        let layout = alloc::Layout::new::<T>();
        let offset = self.offset.add(self.offset.align_offset(layout.align()));
        let next_offset = offset.add(layout.size());
        assert!(next_offset <= self.end, "not enough space in Arena");

        let result = ArenaBox {
            ptr: offset.cast(),
            valid: self.valid.clone(),
        };

        inner_writer(result.ptr, f);
        self.elements.push(ArenaElement {
            value: offset,
            drop: drop::<T>,
        });
        self.offset = next_offset;

        result
    }
}

When the Arena is dropped, the clear method is called.

rust
impl Drop for Arena {
    fn drop(&mut self) {
        self.clear();
    }
}

ArenaBox

As you can see, alloc returns an ArenaBox<T> type that we didn't introduce yet. The purpose of this type is to manage the allocated elements of the Arena without using ArenaElement.

rust
pub struct ArenaBox<T: ?Sized> {
    ptr: *mut T,
    valid: Rc<Cell<bool>>,
}
  • ptr: A raw pointer to the allocated object.
  • valid: The validity of the Arena where this element was allocated.
rust
impl<T: ?Sized> ArenaBox<T> {
    #[inline(always)]
    pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaBox<U> {
        ArenaBox {
            ptr: f(&mut self),
            valid: self.valid,
        }
    }

    fn validate(&self) {
        assert!(
            self.valid.get(),
            "attempted to dereference an ArenaRef after its Arena was cleared"
        );
    }
}

The map method allows to transform the inner type T in ArenaBox<T> into another type U within a new ArenaBox<U>.

The validate method ensures the arena is still valid before accessing the contained object. If the arena is cleared, it will panic.

rust
impl<T: ?Sized> Deref for ArenaBox<T> {
    type Target = T;

    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        self.validate();
        unsafe { &*self.ptr }
    }
}

impl<T: ?Sized> DerefMut for ArenaBox<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.validate();
        unsafe { &mut *self.ptr }
    }
}

The Deref and DerefMut implementations provide convenient access to the contained object. They first validate the arena and then dereference the raw pointer.

ArenaRef

ArenaRef is a thin wrapper around ArenaBox. It provides a way to create shared references to the object managed by ArenaBox

rust
pub struct ArenaRef<T: ?Sized>(ArenaBox<T>);

The Clone implementation for ArenaRef demonstrates how the reference count is managed:

rust
impl<T: ?Sized> Clone for ArenaRef<T> {
    fn clone(&self) -> Self {
        Self(ArenaBox {
            ptr: self.0.ptr,
            valid: self.0.valid.clone(),
        })
    }
}

self.0.valid.clone(): Cloning the Rc increments the reference count, ensuring that the validity state is shared among all ArenaRef instances.

The Deref implementation allows ArenaRef to be used like a regular reference:

rust
impl<T: ?Sized> Deref for ArenaRef<T> {
    type Target = T;

    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        self.0.deref()
    }
}

This calls the deref method on the inner ArenaBox, which validates the arena and then dereferences the raw pointer to provide access to the contained object.