use std::mem::size_of;
use std::ops::Deref;
use std::cell::Cell;
use std::borrow::Cow;
use std::fmt;
const ARENA_BLOCK: usize = 64 * 1024;
pub struct Arena {
store: Cell<Vec<Vec<u8>>>,
ptr: Cell<*mut u8>,
offset: Cell<usize>,
}
pub struct Uninitialized<'arena, T: Copy> {
pointer: &'arena mut MaybeUninit<T>,
}
union MaybeUninit<T: Copy> {
value: T,
_uninit: (),
}
impl<'arena, T: Copy> Uninitialized<'arena, T> {
#[inline]
pub fn init(self, value: T) -> &'arena mut T {
unsafe {
self.pointer.value = value;
&mut self.pointer.value
}
}
#[inline]
pub unsafe fn as_ref(&self) -> &'arena T {
&*(&self.pointer.value as *const T)
}
#[inline]
pub unsafe fn as_mut_ref(self) -> &'arena mut T {
&mut self.pointer.value
}
#[inline]
pub unsafe fn from_raw(pointer: *mut T) -> Self {
Uninitialized {
pointer: &mut *(pointer as *mut MaybeUninit<T>),
}
}
}
impl<'arena, T: Copy> From<&'arena mut T> for Uninitialized<'arena, T> {
#[inline]
fn from(pointer: &'arena mut T) -> Self {
unsafe { Self::from_raw(pointer) }
}
}
#[derive(Clone, Copy, PartialEq)]
pub struct NulTermStr<'arena>(&'arena str);
impl<'arena> fmt::Debug for NulTermStr<'arena> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
fmt::Debug::fmt(self.0, f)
}
}
impl<'arena> fmt::Display for NulTermStr<'arena> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
fmt::Display::fmt(self.0, f)
}
}
impl<'arena> NulTermStr<'arena> {
pub unsafe fn byte_unchecked(&self, index: usize) -> u8 {
*self.0.as_ptr().add(index)
}
}
impl<'arena> AsRef<str> for NulTermStr<'arena> {
fn as_ref(&self) -> &str {
self.0
}
}
impl<'arena> Deref for NulTermStr<'arena> {
type Target = &'arena str;
fn deref(&self) -> &&'arena str {
&self.0
}
}
impl<'arena> From<NulTermStr<'arena>> for &'arena str {
fn from(nts: NulTermStr<'arena>) -> &'arena str {
nts.0
}
}
impl Arena {
pub fn new() -> Self {
let mut store = vec![Vec::with_capacity(ARENA_BLOCK)];
let ptr = store[0].as_mut_ptr();
Arena {
store: Cell::new(store),
ptr: Cell::new(ptr),
offset: Cell::new(0),
}
}
#[inline]
pub fn alloc<'arena, T: Sized + Copy>(&'arena self, value: T) -> &'arena mut T {
self.alloc_uninitialized().init(value)
}
#[inline]
pub fn alloc_uninitialized<'arena, T: Sized + Copy>(&'arena self) -> Uninitialized<'arena, T> {
Uninitialized {
pointer: unsafe { &mut *(self.require(size_of::<T>()) as *mut MaybeUninit<T>) },
}
}
pub fn alloc_slice<'arena, T: Copy>(&'arena self, val: &[T]) -> &'arena [T] {
let ptr = self.require(val.len() * size_of::<T>()) as *mut T;
unsafe {
use std::ptr::copy_nonoverlapping;
use std::slice::from_raw_parts;
copy_nonoverlapping(val.as_ptr(), ptr, val.len());
from_raw_parts(ptr, val.len())
}
}
pub fn alloc_lazy_slice<'arena, T, I: Iterator<Item=T>>(&'arena self, vals: I, n: usize) -> &'arena [T] {
let ptr = self.require(n * size_of::<T>()) as *mut T;
let mut i: usize = 0;
unsafe {
use std::slice::from_raw_parts;
for val in vals.take(n) {
*ptr.offset(i as isize) = val;
i += 1;
}
let diff = n - i;
self.reset_to( self.offset() - diff * size_of::<T>() );
from_raw_parts(ptr, i)
}
}
pub fn alloc_vec<'arena, T: Copy>(&'arena self, mut val: Vec<T>) -> &'arena [T] {
use std::{mem, slice};
let ptr = val.as_mut_ptr();
let cap = val.capacity();
let len = val.len();
mem::forget(val);
let out = self.alloc_byte_vec(unsafe {
Vec::from_raw_parts(ptr as _, 0, cap * size_of::<T>())
});
unsafe { slice::from_raw_parts(out as _, len) }
}
#[inline]
pub fn alloc_cow<'input, 'arena, T>(&'arena self, vals: Cow<'input, [T]>) -> &'arena [T]
where
T: Sized + Copy + 'input,
{
match vals {
Cow::Owned(vec) => self.alloc_vec(vec),
Cow::Borrowed(slice) => self.alloc_slice(slice),
}
}
pub fn alloc_str<'arena>(&'arena self, val: &str) -> &'arena str {
unsafe {
use std::str::from_utf8_unchecked;
from_utf8_unchecked(self.alloc_slice(val.as_bytes()))
}
}
pub fn alloc_nul_term_str<'arena>(&'arena self, val: &str) -> NulTermStr {
let len_with_zero = val.len() + 1;
let ptr = self.require(len_with_zero);
unsafe {
use std::ptr::copy_nonoverlapping;
use std::slice::from_raw_parts;
use std::str::from_utf8_unchecked;
copy_nonoverlapping(val.as_ptr(), ptr, val.len());
*ptr.add(val.len()) = 0;
NulTermStr(from_utf8_unchecked(from_raw_parts(ptr, val.len())))
}
}
pub fn alloc_string<'arena>(&'arena self, val: String) -> &'arena str {
let len = val.len();
let ptr = self.alloc_byte_vec(val.into_bytes());
unsafe {
use std::str::from_utf8_unchecked;
use std::slice::from_raw_parts;
from_utf8_unchecked(from_raw_parts(ptr, len))
}
}
#[inline]
fn alloc_byte_vec(&self, mut val: Vec<u8>) -> *mut u8 {
let ptr = val.as_mut_ptr();
let mut temp = self.store.replace(Vec::new());
temp.push(val);
self.store.replace(temp);
ptr
}
fn alloc_bytes(&self, size: usize) -> *mut u8 {
self.alloc_byte_vec(Vec::with_capacity(size))
}
#[inline]
fn require(&self, size: usize) -> *mut u8 {
if size > ARENA_BLOCK {
return self.alloc_bytes(size);
}
let size = match size % size_of::<usize>() {
0 => size,
n => size + (size_of::<usize>() - n),
};
let offset = self.offset.get();
let cap = offset + size;
if cap > ARENA_BLOCK {
self.grow();
self.offset.set(size);
self.ptr.get()
} else {
self.offset.set(cap);
unsafe { self.ptr.get().add(offset) }
}
}
fn grow(&self) {
let ptr = self.alloc_byte_vec(Vec::with_capacity(ARENA_BLOCK));
self.ptr.set(ptr);
}
#[doc(hidden)]
#[inline]
pub unsafe fn clear(&self) {
self.reset_to(0)
}
#[doc(hidden)]
#[inline]
pub unsafe fn offset(&self) -> usize {
self.offset.get()
}
#[doc(hidden)]
#[inline]
pub unsafe fn reset_to(&self, offset: usize) {
self.offset.set(offset)
}
}
unsafe impl Send for Arena {}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn allocate_some_stuff() {
let arena = Arena::new();
assert_eq!(arena.alloc(0u64), &0);
assert_eq!(arena.alloc(42u64), &42);
assert_eq!(arena.alloc(0x8000000u64), &0x8000000u64);
assert_eq!(arena.offset.get(), 8 * 3);
let mut arena = arena;
assert_eq!(arena.store.get_mut().len(), 1);
}
#[test]
fn allocate_some_vecs() {
let arena = Arena::new();
let vecs = vec![vec![1u64, 2, 3, 4], vec![7; ARENA_BLOCK * 2], vec![]];
for vec in vecs {
assert_eq!(arena.alloc_vec(vec.clone()), &vec[..]);
}
}
#[test]
fn allocate_some_cows() {
let arena = Arena::new();
let vecs = vec![vec![1u64, 2, 3, 4], vec![7; ARENA_BLOCK * 2], vec![]];
for vec in vecs {
assert_eq!(arena.alloc_cow(vec.clone().into()), &vec[..]);
}
}
#[test]
fn allocate_huge_heap() {
let arena = Arena::new();
assert_eq!(arena.alloc(0u64), &0);
assert_eq!(arena.alloc(42u64), &42);
arena.alloc_uninitialized::<[usize; 1024 * 1024]>();
assert_eq!(arena.offset.get(), 8 * 2);
assert_eq!(arena.alloc(0x8000000u64), &0x8000000u64);
assert_eq!(arena.offset.get(), 8 * 3);
let mut arena = arena;
assert_eq!(arena.store.get_mut().len(), 2);
assert_eq!(
arena.store.get_mut()[1].capacity(),
size_of::<usize>() * 1024 * 1024
);
}
#[test]
fn alloc_slice() {
let arena = Arena::new();
assert_eq!(arena.alloc_slice(&[10u16, 20u16]), &[10u16, 20u16][..]);
assert_eq!(arena.offset.get(), 8);
}
#[test]
fn alloc_lazy_slices() {
let arena = Arena::new();
let nums: [u32; 6] = [1, 2, 3, 4, 5, 1000];
let big_nums: [u32; 6] = [100, 200, 300, 400, 500, 1050];
let all_nums = arena.alloc_lazy_slice(nums.iter().map(|x| *x), 6);
let trunc_nums = arena.alloc_lazy_slice(big_nums.iter().map(|x| *x), 3);
let half_nums = arena.alloc_lazy_slice(nums[0..3].iter().map(|x| *x), 6);
assert!(nums.iter().eq(all_nums.iter()));
assert!(nums[0..3].iter().eq(half_nums.iter()));
assert!(big_nums[0..3].iter().eq(trunc_nums.iter()));
}
#[test]
fn aligns_slice_allocs() {
let arena = Arena::new();
assert_eq!(arena.alloc_slice(b"foo"), b"foo");
assert_eq!(arena.offset.get(), 8);
assert_eq!(arena.alloc_slice(b"doge to the moon!"), b"doge to the moon!");
assert_eq!(arena.offset.get(), 32);
}
#[test]
fn aligns_str_allocs() {
let arena = Arena::new();
assert_eq!(arena.alloc_str("foo"), "foo");
assert_eq!(arena.offset.get(), 8);
assert_eq!(arena.alloc_str("doge to the moon!"), "doge to the moon!");
assert_eq!(arena.offset.get(), 32);
}
#[test]
fn alloc_nul_term_str() {
let arena = Arena::new();
let nts = arena.alloc_nul_term_str("abcdefghijk");
let allocated = unsafe { ::std::slice::from_raw_parts(nts.as_ptr(), 12) };
assert_eq!(arena.offset.get(), 16);
assert_eq!(
allocated,
"abcdefghijk\u{0}".as_bytes(),
);
assert_eq!(&**nts, "abcdefghijk");
}
}