use std::alloc::{Layout, alloc, dealloc};
use std::slice;
use crate::align_up;
pub struct AlignedBump {
ptr: *mut u8,
layout: Layout,
cursor: usize,
}
impl AlignedBump {
pub fn with_capacity(capacity: usize) -> Self {
let capacity = capacity.max(64);
let layout = Layout::from_size_align(capacity, 64).expect("layout");
let ptr = unsafe { alloc(layout) };
assert!(!ptr.is_null(), "OOM allocating aligned arena chunk");
Self {
ptr,
layout,
cursor: 0,
}
}
pub fn alloc_aligned(&mut self, size: usize, align: usize) -> &mut [u8] {
let cursor = self.cursor;
let cap = self.layout.size();
match self.try_alloc_aligned(size, align) {
Some(s) => s,
None => panic!(
"AlignedBump out of capacity: cursor={cursor} cap={cap} size={size} align={align}",
),
}
}
pub fn try_alloc_aligned(&mut self, size: usize, align: usize) -> Option<&mut [u8]> {
assert!(
align.is_power_of_two(),
"align must be power of two: {align}"
);
let base = self.ptr as usize;
let aligned = align_up(base + self.cursor, align) - base;
let end = aligned.checked_add(size)?;
if end > self.layout.size() {
return None;
}
self.cursor = end;
unsafe {
let p = self.ptr.add(aligned);
Some(slice::from_raw_parts_mut(p, size))
}
}
pub fn reset(&mut self) {
self.cursor = 0;
}
pub fn capacity(&self) -> usize {
self.layout.size()
}
pub fn used(&self) -> usize {
self.cursor
}
}
impl Drop for AlignedBump {
fn drop(&mut self) {
unsafe { dealloc(self.ptr, self.layout) };
}
}
#[cfg(test)]
#[path = "aligned_tests.rs"]
mod tests;