mod error;
use std::alloc::{Layout, alloc, dealloc};
use std::cmp;
use std::fmt::Debug;
use std::fmt::Display;
use std::marker::PhantomPinned;
use std::ptr::NonNull;
use std::ptr::null_mut;
pub use error::NoMemory;
const MIN_STRUCT_WORDS: usize = 32;
const MIN_CDATA_BYTES: usize = 256;
#[cfg(test)]
use std::cell::RefCell;
#[cfg(test)]
thread_local! {
static IKSEMEL_ALLOCATED: RefCell<usize> = const { RefCell::new(0) };
}
#[cfg(test)]
fn test_allocated_add(bytes: usize) {
IKSEMEL_ALLOCATED.with_borrow_mut(|cell| *cell += bytes);
}
#[cfg(not(test))]
fn test_allocated_add(_: usize) {}
#[cfg(test)]
fn test_allocated_sub(bytes: usize) {
IKSEMEL_ALLOCATED.with_borrow_mut(|cell| *cell -= bytes);
}
#[cfg(not(test))]
fn test_allocated_sub(_: usize) {}
#[cfg(test)]
pub(self) fn test_allocated() -> usize {
IKSEMEL_ALLOCATED.with_borrow(|cell| *cell)
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct ArenaStats {
pub chunks: u32,
pub allocated_bytes: usize,
pub used_bytes: usize,
}
impl Display for ArenaStats {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{} chunks, {} bytes allocated, {} bytes used",
self.chunks, self.allocated_bytes, self.used_bytes
)
}
}
#[repr(transparent)]
pub struct Arena {
head_ptr: *mut Head,
}
struct Head {
struct_chunk: *mut Chunk,
cdata_chunk: *mut Chunk,
alloc_layout: Layout,
_pin: PhantomPinned,
}
struct Chunks {
next: *mut Chunk,
}
impl Iterator for Chunks {
type Item = *mut Chunk;
fn next(&mut self) -> Option<Self::Item> {
if self.next.is_null() {
None
} else {
let chunk = self.next;
self.next = unsafe { (*chunk).next };
Some(chunk)
}
}
}
impl Head {
fn struct_chunks(&mut self) -> Chunks {
Chunks {
next: self.struct_chunk,
}
}
fn extra_struct_chunks(&mut self) -> Chunks {
Chunks {
next: unsafe { (*self.struct_chunk).next },
}
}
fn cdata_chunks(&mut self) -> Chunks {
Chunks {
next: self.cdata_chunk,
}
}
fn extra_cdata_chunks(&mut self) -> Chunks {
Chunks {
next: unsafe { (*self.cdata_chunk).next },
}
}
}
struct Chunk {
next: *mut Chunk,
size: usize,
used: usize,
last: *mut u8,
mem: *mut u8,
alloc_layout: Layout,
_pin: PhantomPinned,
}
impl Chunk {
fn raw_init(self: &mut Chunk, ptr: *mut u8, size: usize, alloc_layout: Layout) {
self.next = null_mut();
self.size = size;
self.used = 0;
self.last = ptr;
self.mem = ptr;
self.alloc_layout = alloc_layout;
}
fn clear(&mut self) {
self.used = 0;
self.last = self.mem;
}
fn add_chunk(self: &mut Chunk, size: usize) -> Result<NonNull<Chunk>, NoMemory> {
let data_layout = Layout::array::<u8>(size)?;
let chunk_layout = Layout::new::<Chunk>();
let (chunk_layout, data_offset) = chunk_layout.extend(data_layout)?;
let chunk_layout = chunk_layout.pad_to_align();
unsafe {
let ptr = alloc(chunk_layout);
if ptr.is_null() {
return Err(NoMemory);
}
test_allocated_add(chunk_layout.size());
let chunk = ptr as *mut Chunk;
(*chunk).raw_init(ptr.byte_add(data_offset), size, chunk_layout);
self.next = chunk;
Ok(NonNull::new_unchecked(chunk))
}
}
fn has_space(self: &mut Chunk, size: usize) -> bool {
size <= self.size && self.used + size <= self.size
}
fn has_aligned_space(self: &mut Chunk, layout: Layout) -> bool {
let size = layout.size();
let used_layout = Layout::from_size_align(self.used, layout.align()).unwrap();
let used_layout = used_layout.pad_to_align();
size <= self.size && used_layout.size() + size <= self.size
}
fn make_aligned_space(self: &mut Chunk, layout: Layout) -> Result<NonNull<u8>, NoMemory> {
let mut expected_next_size = self.size;
let mut current: *mut Chunk = self;
unsafe {
while !(*current).has_aligned_space(layout) {
expected_next_size *= 2;
let mut next = (*current).next;
if next.is_null() {
let data_size = cmp::max(expected_next_size, layout.size());
next = (*current).add_chunk(data_size)?.as_ptr();
}
current = next;
}
let used_layout = Layout::from_size_align((*current).used, layout.align())?;
let used_layout = used_layout.pad_to_align();
let offset = used_layout.size() - (*current).used;
let ptr = (*current).mem.byte_add((*current).used + offset);
(*current).last = ptr;
(*current).used += layout.size() + offset;
Ok(NonNull::new_unchecked(ptr))
}
}
fn make_space(self: &mut Chunk, size: usize) -> Result<NonNull<u8>, NoMemory> {
let mut expected_next_size = self.size;
let mut current: *mut Chunk = self;
unsafe {
while !(*current).has_space(size) {
expected_next_size *= 2;
let mut next = (*current).next;
if next.is_null() {
let data_size = cmp::max(expected_next_size, size);
next = (*current).add_chunk(data_size)?.as_ptr();
}
current = next;
}
let ptr = (*current).mem.byte_add((*current).used);
(*current).last = ptr;
(*current).used += size;
Ok(NonNull::new_unchecked(ptr))
}
}
fn find_adjacent_space(
self: &mut Chunk,
old_p: *const u8,
old_size: usize,
size: usize,
) -> Option<*mut Chunk> {
let mut current: *mut Chunk = self;
unsafe {
loop {
if std::ptr::addr_eq(old_p, (*current).last) {
let chunk_end = (*current).mem.byte_add((*current).used);
let old_end = old_p.byte_add(old_size);
if std::ptr::addr_eq(chunk_end, old_end) && (*current).has_space(size) {
return Some(current);
}
return None;
}
if (*current).next.is_null() {
return None;
}
current = (*current).next;
}
}
}
}
impl Arena {
pub fn new() -> Result<Arena, NoMemory> {
Self::with_chunk_sizes(0, 0)
}
pub fn with_chunk_sizes(struct_words: usize, cdata_bytes: usize) -> Result<Arena, NoMemory> {
let struct_words = cmp::max(struct_words, MIN_STRUCT_WORDS);
let struct_buf_layout = Layout::array::<*const usize>(struct_words)?;
let cdata_bytes = cmp::max(cdata_bytes, MIN_CDATA_BYTES);
let cdata_buf_layout = Layout::array::<u8>(cdata_bytes)?;
let head_layout = Layout::new::<Head>();
let (head_layout, struct_offset) = head_layout.extend(Layout::new::<Chunk>())?;
let (head_layout, cdata_offset) = head_layout.extend(Layout::new::<Chunk>())?;
let (head_layout, struct_buf_offset) = head_layout.extend(struct_buf_layout)?;
let (head_layout, cdata_buf_offset) = head_layout.extend(cdata_buf_layout)?;
let head_layout = head_layout.pad_to_align();
let head_ptr;
unsafe {
let ptr = alloc(head_layout);
if ptr.is_null() {
return Err(NoMemory);
}
test_allocated_add(head_layout.size());
head_ptr = ptr as *mut Head;
(*head_ptr).alloc_layout = head_layout;
let struct_ptr = ptr.byte_add(struct_offset);
let struct_chunk = struct_ptr as *mut Chunk;
(*head_ptr).struct_chunk = struct_chunk;
let cdata_ptr = ptr.byte_add(cdata_offset);
let cdata_chunk = cdata_ptr as *mut Chunk;
(*head_ptr).cdata_chunk = cdata_chunk;
let struct_buf_ptr = ptr.byte_add(struct_buf_offset);
(*struct_chunk).raw_init(struct_buf_ptr, struct_buf_layout.size(), head_layout);
let cdata_buf_ptr = ptr.byte_add(cdata_buf_offset);
(*cdata_chunk).raw_init(cdata_buf_ptr, cdata_buf_layout.size(), head_layout);
}
Ok(Arena { head_ptr })
}
pub fn alloc_struct<T>(&self) -> Result<NonNull<T>, NoMemory> {
unsafe {
let head = &mut *self.head_ptr;
let layout = Layout::new::<T>();
let ptr = (*head.struct_chunk).make_aligned_space(layout)?;
Ok(NonNull::new_unchecked(ptr.as_ptr() as *mut T))
}
}
pub fn push_str<'a>(&'a self, s: &str) -> Result<&'a str, NoMemory> {
let size = s.len();
unsafe {
let head = &mut *self.head_ptr;
let ptr = (*head.cdata_chunk).make_space(size)?.as_ptr();
std::ptr::copy_nonoverlapping(s.as_ptr(), ptr, size);
let slice = std::slice::from_raw_parts(ptr, size);
Ok(std::str::from_utf8_unchecked(slice))
}
}
pub fn concat_str<'a>(&'a self, old_s: &str, s: &str) -> Result<&'a str, NoMemory> {
unsafe {
let head = &mut *self.head_ptr;
let data_chunk = head.cdata_chunk;
let slice;
if let Some(chunk) =
(*data_chunk).find_adjacent_space(old_s.as_ptr(), old_s.len(), s.len())
{
let p = (*chunk).mem.byte_add((*chunk).used);
(*chunk).used += s.len();
std::ptr::copy_nonoverlapping(s.as_ptr(), p, s.len());
slice = std::slice::from_raw_parts(p.byte_sub(old_s.len()), old_s.len() + s.len());
} else {
let ptr = (*data_chunk).make_space(old_s.len() + s.len())?.as_ptr();
std::ptr::copy_nonoverlapping(old_s.as_ptr(), ptr, old_s.len());
let ptr2 = ptr.byte_add(old_s.len());
std::ptr::copy_nonoverlapping(s.as_ptr(), ptr2, s.len());
slice = std::slice::from_raw_parts(ptr, old_s.len() + s.len());
}
Ok(std::str::from_utf8_unchecked(slice))
}
}
pub fn stats(&self) -> ArenaStats {
let mut stats = ArenaStats {
chunks: 1,
allocated_bytes: 0,
used_bytes: 0,
};
unsafe {
let head = &mut *self.head_ptr;
stats.allocated_bytes += head.alloc_layout.size();
stats.used_bytes += (*head.struct_chunk).used;
stats.used_bytes += (*head.cdata_chunk).used;
for chunk in head.extra_struct_chunks() {
stats.chunks += 1;
stats.allocated_bytes += (*chunk).alloc_layout.size();
stats.used_bytes += (*chunk).used;
}
for chunk in head.extra_cdata_chunks() {
stats.chunks += 1;
stats.allocated_bytes += (*chunk).alloc_layout.size();
stats.used_bytes += (*chunk).used;
}
}
stats
}
pub fn into_empty_arena(self) -> Arena {
unsafe {
let head = &mut *self.head_ptr;
for chunk in (*head).struct_chunks() {
(*chunk).clear();
}
for chunk in (*head).cdata_chunks() {
(*chunk).clear();
}
}
self
}
}
impl Drop for Arena {
fn drop(&mut self) {
unsafe {
let head = &mut *self.head_ptr;
for chunk in (*head).extra_struct_chunks() {
test_allocated_sub((*chunk).alloc_layout.size());
let layout = (*chunk).alloc_layout;
dealloc(chunk as *mut u8, layout);
}
for chunk in (*head).extra_cdata_chunks() {
test_allocated_sub((*chunk).alloc_layout.size());
let layout = (*chunk).alloc_layout;
dealloc(chunk as *mut u8, layout);
}
test_allocated_sub(head.alloc_layout.size());
let layout = head.alloc_layout;
dealloc(self.head_ptr as *mut u8, layout);
}
}
}
impl Display for Arena {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Arena ({})", self.stats())
}
}
impl Debug for Arena {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
unsafe {
let head = &mut *self.head_ptr;
write!(f, "Arena (head[alloc: {}]", head.alloc_layout.size())?;
for chunk in (*head).struct_chunks() {
write!(
f,
", struct[alloc: {}, used: {}, size: {}]",
(*chunk).alloc_layout.size(),
(*chunk).used,
(*chunk).size
)?;
}
for chunk in (*head).cdata_chunks() {
write!(
f,
", cdata[alloc: {}, used: {}, size: {}]",
(*chunk).alloc_layout.size(),
(*chunk).used,
(*chunk).size
)?;
}
write!(f, ")")
}
}
}
#[cfg(test)]
mod tests;
#[cfg(doctest)]
struct MustNotCompileTests;