Minimal arena allocators in Rust

!minimal_arenas.md

Minimal arena allocators in Rust

arena_refcell.rs

use std::cell::RefCell;

#[derive(Debug)]
struct Arena<T> {
    storage: RefCell<Vec<Vec<T>>>,
}

impl<T> Arena<T> {
    fn with_capacity(cap: usize) -> Self {
        Arena {
            storage: RefCell::new(vec![Vec::with_capacity(cap)]),
        }
    }

    fn alloc(&self, t: T) -> &mut T {
        let mut storage = self.storage.borrow_mut();
        let current_block = storage.last().unwrap();
        if current_block.len() < current_block.capacity() {
            let current_block = storage.last_mut().unwrap();
            current_block.push(t);
        } else {
            let mut next_block = Vec::with_capacity(current_block.capacity() * 2);
            next_block.push(t);
            storage.push(next_block);
        }
        let next_block = storage.last_mut().unwrap();
        let len = next_block.len();
        unsafe { &mut *next_block.as_mut_ptr().add(len - 1) }
    }
}

fn main() {
    let arena = Arena::with_capacity(1);
    dbg!((arena.alloc(1), arena.alloc(2)));
    dbg!(vec![arena.alloc(3), arena.alloc(4), arena.alloc(5)]);
    dbg!(arena);
}

arena_unsafecell.rs

use std::cell::UnsafeCell;

#[derive(Debug)]
struct Arena<T> {
    storage: UnsafeCell<Vec<Vec<T>>>,
}

impl<T> Arena<T> {
    fn with_capacity(cap: usize) -> Self {
        Arena {
            storage: UnsafeCell::new(vec![Vec::with_capacity(cap)]),
        }
    }

    fn alloc(&self, t: T) -> &mut T {
        let storage = unsafe { &mut *self.storage.get() };
        let current_block = storage.last().unwrap();
        // We need to manually ensure not to exceed capacity because that would invalidate previous
        // references.
        if current_block.len() < current_block.capacity() {
            let current_block = storage.last_mut().unwrap();
            current_block.push(t);
            return current_block.last_mut().unwrap();
        } else {
            let mut next_block = Vec::with_capacity(current_block.capacity() * 2);
            next_block.push(t);
            storage.push(next_block);
            return storage.last_mut().unwrap().last_mut().unwrap();
        }
    }
}

fn main() {
    let arena = Arena::with_capacity(1);
    dbg!((arena.alloc(1), arena.alloc(2)));
    dbg!(vec![arena.alloc(3), arena.alloc(4), arena.alloc(5)]);
    dbg!(arena.storage.into_inner());
}