#![allow(clippy::transmute_ptr_to_ptr)]
#![allow(clippy::mut_from_ref)]
mod utils;
use std::{
alloc::Layout,
cell::{Cell, RefCell},
collections::LinkedList,
fmt,
mem::{size_of, transmute, MaybeUninit},
slice,
};
use utils::{alignment_offset, min_alignment, repeat_layout};
#[derive(Default)]
pub struct Arena {
blocks: RefCell<LinkedList<Vec<MaybeUninit<u8>>>>,
min_block_size: usize,
growth_strategy: GrowthStrategy,
max_waste_percentage: usize,
stat_space_occupied: Cell<usize>,
stat_space_allocated: Cell<usize>,
}
impl fmt::Debug for Arena {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("Arena")
.field("blocks.len():", &self.blocks.borrow().len())
.field("min_block_size", &self.min_block_size)
.field("max_waste_percentage", &self.max_waste_percentage)
.field("stat_space_occupied", &self.stat_space_occupied)
.field("stat_space_allocated", &self.stat_space_allocated)
.finish()
}
}
impl Arena {
pub fn new() -> Arena {
Arena {
blocks: RefCell::new(LinkedList::new()),
min_block_size: 1 << 10, growth_strategy: GrowthStrategy::Constant,
max_waste_percentage: 20,
stat_space_occupied: Cell::new(0),
stat_space_allocated: Cell::new(0),
}
}
pub fn with_block_size(self, block_size: usize) -> Arena {
assert!(
block_size > 0,
"Initial block size must be greater \
than zero"
);
assert!(
self.blocks.borrow().is_empty(),
"Cannot change initial block size after \
blocks have already been allocated"
);
Arena {
min_block_size: block_size,
..self
}
}
pub fn with_max_waste_percentage(self, max_waste_percentage: usize) -> Arena {
assert!(
max_waste_percentage > 0 && max_waste_percentage <= 100,
"The max waste percentage must be between 1 and 100"
);
Arena {
max_waste_percentage,
..self
}
}
pub fn with_growth_strategy(self, growth_strategy: GrowthStrategy) -> Arena {
Arena {
growth_strategy,
..self
}
}
#[inline]
pub fn alloc<T: Copy>(&self, value: T) -> &mut T {
let memory = self.alloc_uninit();
unsafe {
*memory.as_mut_ptr() = value;
}
unsafe { transmute(memory) }
}
#[inline]
pub fn alloc_array<T: Copy>(&self, value: T, len: usize) -> &mut [T] {
let memory = self.alloc_array_uninit(len);
for v in memory.iter_mut() {
unsafe {
*v.as_mut_ptr() = value;
}
}
unsafe { transmute(memory) }
}
#[inline]
pub fn copy_slice<T: Copy>(&self, slice: &[T]) -> &mut [T] {
let memory = self.alloc_array_uninit(slice.len());
for (v, slice_item) in memory.iter_mut().zip(slice.iter()) {
unsafe {
*v.as_mut_ptr() = *slice_item;
}
}
unsafe { transmute(memory) }
}
#[inline]
pub fn copy_str(&self, text: &str) -> &mut str {
let memory = self.alloc_array_uninit::<u8>(text.len());
for (byte, text_byte) in memory.iter_mut().zip(text.as_bytes().iter()) {
unsafe {
*byte.as_mut_ptr() = *text_byte;
}
}
unsafe { std::str::from_utf8_unchecked_mut(transmute(memory)) }
}
#[inline]
pub fn alloc_align<T: Copy>(&self, value: T, align: usize) -> &mut T {
let memory = self.alloc_align_uninit(align);
unsafe {
*memory.as_mut_ptr() = value;
}
unsafe { transmute(memory) }
}
#[inline]
pub fn alloc_array_align<T: Copy>(&self, value: T, len: usize, align: usize) -> &mut [T] {
let memory = self.alloc_array_align_uninit(len, align);
for v in memory.iter_mut() {
unsafe {
*v.as_mut_ptr() = value;
}
}
unsafe { transmute(memory) }
}
#[inline]
pub fn copy_slice_align<T: Copy>(&self, slice: &[T], align: usize) -> &mut [T] {
let memory = self.alloc_array_align_uninit(slice.len(), align);
for (v, slice_item) in memory.iter_mut().zip(slice.iter()) {
unsafe {
*v.as_mut_ptr() = *slice_item;
}
}
unsafe { transmute(memory) }
}
#[inline]
pub fn alloc_uninit<T: Copy>(&self) -> &mut MaybeUninit<T> {
assert!(
size_of::<T>() > 0,
"`Arena` does not support zero-sized types."
);
let memory = self.alloc_raw(&Layout::new::<T>()) as *mut MaybeUninit<T>;
unsafe { memory.as_mut().unwrap() }
}
#[inline]
pub fn alloc_array_uninit<T: Copy>(&self, len: usize) -> &mut [MaybeUninit<T>] {
assert!(
size_of::<T>() > 0,
"`Arena` does not support zero-sized types."
);
let layout = &repeat_layout(&Layout::new::<T>(), len);
let memory = self.alloc_raw(&layout) as *mut MaybeUninit<T>;
unsafe { slice::from_raw_parts_mut(memory, len) }
}
#[inline]
pub fn alloc_align_uninit<T: Copy>(&self, align: usize) -> &mut MaybeUninit<T> {
assert!(
size_of::<T>() > 0,
"`Arena` does not support zero-sized types."
);
let layout = min_alignment(&Layout::new::<T>(), align);
let memory = self.alloc_raw(&layout) as *mut MaybeUninit<T>;
unsafe { memory.as_mut().unwrap() }
}
#[inline]
pub fn alloc_array_align_uninit<T: Copy>(
&self,
len: usize,
align: usize,
) -> &mut [MaybeUninit<T>] {
assert!(
size_of::<T>() > 0,
"`Arena` does not support zero-sized types."
);
let layout = min_alignment(&repeat_layout(&Layout::new::<T>(), len), align);
let memory = self.alloc_raw(&layout) as *mut MaybeUninit<T>;
unsafe { slice::from_raw_parts_mut(memory, len) }
}
pub fn alloc_raw(&self, layout: &Layout) -> *mut MaybeUninit<u8> {
let alignment = layout.align();
let size = layout.size();
let mut blocks = self.blocks.borrow_mut();
if blocks.is_empty() {
blocks.push_front(Vec::with_capacity(self.min_block_size));
self.stat_space_occupied
.set(self.stat_space_occupied.get() + self.min_block_size);
}
if size == 0 {
return blocks.front_mut().unwrap().as_mut_ptr();
}
let start_index_proposal = {
let cur_block = blocks.front().unwrap();
let block_addr = cur_block.as_ptr() as usize;
let block_filled = cur_block.len();
block_filled + alignment_offset(block_addr + block_filled, alignment)
};
if (start_index_proposal + size) <= blocks.front().unwrap().capacity() {
let cur_block = blocks.front_mut().unwrap();
let new_len = (start_index_proposal + size).max(cur_block.len());
unsafe { cur_block.set_len(new_len) };
self.stat_space_allocated
.set(self.stat_space_allocated.get() + size);
unsafe { cur_block.as_mut_ptr().add(start_index_proposal) }
}
else {
let next_shared_size = match self.growth_strategy {
GrowthStrategy::Constant => self.min_block_size,
GrowthStrategy::Percentage(perc) => {
let a = self.stat_space_occupied.get() / 100 * perc as usize;
let b = a % self.min_block_size;
self.min_block_size.max(a - b)
}
};
let waste_percentage = {
let block = blocks.front().unwrap();
let w1 = ((block.capacity() - block.len()) * 100) / block.capacity();
let w2 = ((self.stat_space_occupied.get() - self.stat_space_allocated.get()) * 100)
/ self.stat_space_occupied.get();
w1.min(w2)
};
let is_shared_block = (size + alignment) <= next_shared_size
&& waste_percentage <= self.max_waste_percentage;
let new_block_size = if is_shared_block {
next_shared_size
} else {
size + alignment - 1
};
self.stat_space_occupied
.set(self.stat_space_occupied.get() + new_block_size);
self.stat_space_allocated
.set(self.stat_space_allocated.get() + size);
let new_block = {
if is_shared_block {
blocks.push_front(Vec::with_capacity(new_block_size));
blocks.front_mut().unwrap()
} else {
blocks.push_back(Vec::with_capacity(new_block_size));
blocks.back_mut().unwrap()
}
};
let start_index = alignment_offset(new_block.as_ptr() as usize, alignment);
unsafe { new_block.set_len(start_index + size) };
unsafe { new_block.as_mut_ptr().add(start_index) }
}
}
pub fn clear(&mut self) {
unsafe { self.clear_unchecked() }
}
pub unsafe fn clear_unchecked(&self) {
let mut blocks = self.blocks.borrow_mut();
blocks.clear();
self.stat_space_occupied.set(0);
self.stat_space_allocated.set(0);
}
}
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum GrowthStrategy {
Constant,
Percentage(u8),
}
impl Default for GrowthStrategy {
fn default() -> GrowthStrategy {
GrowthStrategy::Constant
}
}