arena.rs

use std::mem;

pub struct Arena {
    data: Vec<Vec<u8>>,
    next_size: usize,
}

impl Arena {
    pub fn with_capacity(capacity: usize) -> Arena {
        Arena {
            data: vec![Vec::with_capacity(capacity)],
            next_size: capacity,
        }
    }

    pub fn alloc<'a, T: Copy>(&mut self, value: T) -> &'a mut T where Self: 'a {
        let bytes = mem::size_of::<T>();
        let align = mem::align_of::<T>();
        let (ptr, len, capacity) = {
            let last = &self.data.last().unwrap();
            let len = last.len();
            (last.as_ptr() as usize + len, len, last.capacity())
        };
        let offset = ((ptr + align - 1) & !(align - 1)) - ptr;

        let result_ptr = if len + offset + bytes > capacity {
            let mut exact: Vec<T> = Vec::with_capacity(1);
            let (exact_ptr, exact_len, exact_cap) = (exact.as_mut_ptr(), exact.len(), exact.capacity());
            mem::forget(exact);
            let mut exact_bytes: Vec<u8> = unsafe { Vec::from_raw_parts(exact_ptr as *mut u8, exact_len * bytes, exact_cap * bytes) };
            let result = exact_bytes.as_mut_ptr();
            self.data.push(exact_bytes);

            self.next_size = self.next_size.saturating_add(self.next_size);
            self.data.push(Vec::with_capacity(self.next_size));

            result
        } else {
            let last = &mut self.data.last_mut().unwrap();
            unsafe {
                last.set_len(len + offset + bytes);
                last.as_mut_ptr().offset((len + offset) as isize)
            }
        };

        let result: &'a mut T = unsafe { mem::transmute(result_ptr) };
        *result = value;
        result
    }
}