use super::{BStackAllocator, BStackSlice};
use crate::BStack;
use core::{cell::Cell, marker::PhantomData, num::NonZeroU64};
use std::{fmt, io};
#[cfg(feature = "set")]
const ALSL_MAGIC: [u8; 8] = *b"ALSL\x00\x01\x00\x00";
#[cfg(feature = "set")]
const ALSL_MAGIC_PREFIX: [u8; 6] = *b"ALSL\x00\x01";
#[cfg(feature = "set")]
pub struct SlabBStackAllocator {
stack: BStack,
block_size: u64,
_not_sync: PhantomData<Cell<()>>,
}
#[cfg(feature = "set")]
impl SlabBStackAllocator {
const OFFSET_SIZE: u64 = 24;
const HEADER_SIZE: u64 = 24;
const ARENA_START: u64 = Self::OFFSET_SIZE + Self::HEADER_SIZE;
const FREE_HEAD_OFFSET: u64 = Self::OFFSET_SIZE + 16;
const MIN_BLOCK_SIZE: u64 = 8;
const SENTINEL: u64 = 0;
pub fn new(stack: BStack, block_size: u64) -> io::Result<Self> {
if !stack.is_empty()? {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"stack is not empty; use SlabBStackAllocator::open to reopen an existing allocator",
));
}
if block_size < Self::MIN_BLOCK_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"block_size ({block_size}) must be >= {}",
Self::MIN_BLOCK_SIZE
),
));
}
if usize::try_from(block_size).is_err() {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"block_size is too large for this platform",
));
}
let mut hdr = [0u8; Self::ARENA_START as usize];
let off = Self::OFFSET_SIZE as usize;
hdr[off..off + 8].copy_from_slice(&ALSL_MAGIC);
hdr[off + 8..off + 16].copy_from_slice(&block_size.to_le_bytes());
stack.push(hdr)?;
Ok(Self {
stack,
block_size,
_not_sync: PhantomData,
})
}
pub fn open(stack: BStack) -> io::Result<Self> {
if stack.is_empty()? {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"stack is empty; use SlabBStackAllocator::new to create a new allocator",
));
}
let stack_len = stack.len()?;
if stack_len < Self::ARENA_START {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"stack too short to contain allocator header",
));
}
let mut header = [0u8; Self::HEADER_SIZE as usize];
stack.get_into(Self::OFFSET_SIZE, &mut header)?;
if header[..ALSL_MAGIC_PREFIX.len()] != ALSL_MAGIC_PREFIX {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"invalid magic: not a SlabBStackAllocator file",
));
}
let stored_block_size = u64::from_le_bytes(header[8..16].try_into().unwrap());
if stored_block_size < Self::MIN_BLOCK_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("stored block_size ({stored_block_size}) is invalid"),
));
}
if usize::try_from(stored_block_size).is_err() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"stored block_size is too large for this platform",
));
}
let stored_free_head = u64::from_le_bytes(header[16..24].try_into().unwrap());
if stored_free_head != Self::SENTINEL
&& (stored_free_head < Self::ARENA_START
|| (stored_free_head - Self::ARENA_START) % stored_block_size != 0
|| stored_free_head >= stack_len)
{
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("stored free_head ({stored_free_head}) is not a valid block offset"),
));
}
let arena_bytes = stack_len.checked_sub(Self::ARENA_START).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidData,
"stack too short to contain allocator header",
)
})?;
if arena_bytes % stored_block_size != 0 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"stack tail is not aligned to block_size",
));
}
Ok(Self {
stack,
block_size: stored_block_size,
_not_sync: PhantomData,
})
}
pub fn block_size(&self) -> u64 {
self.block_size
}
fn pop_free_block(&self) -> io::Result<Option<NonZeroU64>> {
let head = u64::from_le_bytes(read_bstack!(self.stack, Self::FREE_HEAD_OFFSET => u64));
if head == Self::SENTINEL {
return Ok(None);
}
self.stack.set(
Self::FREE_HEAD_OFFSET,
read_bstack!(self.stack, head => u64),
)?;
self.stack.zero(head, self.block_size)?;
Ok(Some(head.try_into().unwrap()))
}
fn push_free_block(&self, block_start: u64) -> io::Result<()> {
self.stack.set(
block_start,
read_bstack!(self.stack, Self::FREE_HEAD_OFFSET => u64),
)?;
self.stack
.set(Self::FREE_HEAD_OFFSET, block_start.to_le_bytes())
}
fn push_free_blocks(&self, first_block: u64, count: u64) -> io::Result<()> {
if count == 0 {
return Ok(());
}
if count == 1 {
return self.push_free_block(first_block);
}
let old_head = read_bstack!(self.stack, Self::FREE_HEAD_OFFSET => u64);
let total_bytes = count.checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"freed region size overflows u64",
)
})?;
first_block.checked_add(total_bytes).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"freed region end offset overflows u64",
)
})?;
let buf_size = usize::try_from(total_bytes).map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
"freed region exceeds platform pointer size",
)
})?;
let mut buf = vec![0u8; buf_size];
for i in 0..count - 1 {
let next = first_block
.checked_add((i + 1).checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"next block index multiplication overflows u64",
)
})?)
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"next block offset overflows u64",
)
})?;
let off = usize::try_from(i.checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"free-list offset overflows u64",
)
})?)
.map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
"free-list offset overflows usize",
)
})?;
buf[off..off + 8].copy_from_slice(&next.to_le_bytes());
}
let last_off =
usize::try_from((count - 1).checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"last free-list offset overflows u64",
)
})?)
.map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
"last free-list offset overflows usize",
)
})?;
buf[last_off..last_off + 8].copy_from_slice(&old_head);
self.stack.set(first_block, buf)?;
self.stack
.set(Self::FREE_HEAD_OFFSET, first_block.to_le_bytes())
}
fn blocks_needed(&self, len: u64) -> u64 {
if len == 0 {
0
} else if len <= self.block_size {
1
} else {
len.div_ceil(self.block_size)
}
}
}
#[cfg(feature = "set")]
impl fmt::Debug for SlabBStackAllocator {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SlabBStackAllocator")
.field("block_size", &self.block_size)
.finish_non_exhaustive()
}
}
#[cfg(feature = "set")]
impl BStackAllocator for SlabBStackAllocator {
type Error = io::Error;
type Allocated<'a> = BStackSlice<'a, Self>;
fn stack(&self) -> &BStack {
&self.stack
}
fn into_stack(self) -> BStack {
self.stack
}
fn alloc(&self, len: u64) -> io::Result<BStackSlice<'_, Self>> {
if len == 0 {
return Ok(BStackSlice::empty(self));
}
if len <= self.block_size {
if let Some(block) = self.pop_free_block()? {
return Ok(unsafe { BStackSlice::from_raw_parts(self, block.into(), len) });
}
let offset = self.stack.extend(self.block_size)?;
return Ok(unsafe { BStackSlice::from_raw_parts(self, offset, len) });
}
let n = len.div_ceil(self.block_size);
let total = n.checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "allocation size overflows u64")
})?;
let offset = self.stack.extend(total)?;
Ok(unsafe { BStackSlice::from_raw_parts(self, offset, len) })
}
fn dealloc(&self, slice: BStackSlice<'_, Self>) -> io::Result<()> {
if slice.is_empty() && slice.start() == Self::SENTINEL {
return Ok(());
}
let n_blocks = self.blocks_needed(slice.len());
let backing_size = n_blocks.checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"deallocation size overflows u64",
)
})?;
let current_tail = self.stack.len()?;
let slice_end = slice.start().checked_add(backing_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"deallocation end offset overflows u64",
)
})?;
if slice.len() > self.block_size && slice_end == current_tail {
return self.stack.discard(backing_size);
}
self.push_free_blocks(slice.start(), n_blocks)
}
fn realloc<'a>(
&'a self,
slice: BStackSlice<'a, Self>,
new_len: u64,
) -> io::Result<BStackSlice<'a, Self>> {
if slice.is_empty() && slice.start() == Self::SENTINEL {
return self.alloc(new_len);
}
if new_len == 0 {
self.dealloc(slice)?;
return Ok(BStackSlice::empty(self));
}
if new_len == slice.len() {
return Ok(slice);
}
let old_n = self.blocks_needed(slice.len());
let new_n = self.blocks_needed(new_len);
if old_n == new_n {
if new_len > slice.len() {
self.stack.zero(slice.end(), new_len - slice.len())?;
}
return Ok(unsafe { BStackSlice::from_raw_parts(self, slice.start(), new_len) });
}
let old_backing = old_n.checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"old allocation size overflows u64",
)
})?;
let new_backing = new_n.checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"new allocation size overflows u64",
)
})?;
let current_tail = self.stack.len()?;
let is_tail = slice.start().checked_add(old_backing).ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "tail check overflows u64")
})? == current_tail;
if is_tail {
if new_n > old_n {
self.stack.extend(new_backing - old_backing)?;
if new_len > slice.len() {
self.stack.zero(slice.end(), new_len - slice.len())?;
}
} else {
self.stack.discard(old_backing - new_backing)?;
}
return Ok(unsafe { BStackSlice::from_raw_parts(self, slice.start(), new_len) });
}
if new_n < old_n {
let free_start = slice
.start()
.checked_add(new_n.checked_mul(self.block_size).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"free start multiplication overflows u64",
)
})?)
.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "free start overflows u64")
})?;
self.push_free_blocks(free_start, old_n - new_n)?;
return Ok(unsafe { BStackSlice::from_raw_parts(self, slice.start(), new_len) });
}
let buf_len = usize::try_from(new_backing).map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
"reallocation too large for this platform",
)
})?;
let mut data_buf = vec![0u8; buf_len];
let old_visible_len = usize::try_from(slice.len()).map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
"existing allocation too large for this platform",
)
})?;
self.stack
.get_into(slice.start(), &mut data_buf[..old_visible_len])?;
let new_ptr = self.stack.push(data_buf)?;
self.push_free_blocks(slice.start(), old_n)?;
Ok(unsafe { BStackSlice::from_raw_parts(self, new_ptr, new_len) })
}
}
#[cfg(all(test, feature = "set"))]
mod tests {
use super::SlabBStackAllocator;
use crate::BStack;
use crate::alloc::BStackAllocator;
use std::io::ErrorKind;
use std::sync::atomic::{AtomicU64, Ordering};
struct Guard(std::path::PathBuf);
impl Drop for Guard {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.0);
}
}
fn temp_path() -> std::path::PathBuf {
static COUNTER: AtomicU64 = AtomicU64::new(0);
let id = COUNTER.fetch_add(1, Ordering::Relaxed);
let pid = std::process::id();
std::env::temp_dir().join(format!("bstack_slab_{pid}_{id}.bin"))
}
fn empty_stack() -> (BStack, std::path::PathBuf) {
let path = temp_path();
(BStack::open(&path).unwrap(), path)
}
#[test]
fn new_initialises_header_and_reports_block_size() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let alloc = SlabBStackAllocator::new(stack, 16).unwrap();
assert_eq!(alloc.block_size(), 16);
assert_eq!(alloc.stack().len().unwrap(), 48);
}
#[test]
fn new_rejects_block_size_below_minimum() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let err = SlabBStackAllocator::new(stack, 7).unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvalidInput);
}
#[test]
fn new_rejects_nonempty_stack() {
let (stack, path) = empty_stack();
let _g = Guard(path);
stack.push(b"data").unwrap();
let err = SlabBStackAllocator::new(stack, 8).unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvalidInput);
}
#[test]
fn open_rejects_empty_stack() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let err = SlabBStackAllocator::open(stack).unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvalidInput);
}
#[test]
fn open_rejects_stack_too_short() {
let (stack, path) = empty_stack();
let _g = Guard(path.clone());
stack.push([0u8; 24]).unwrap(); drop(stack);
let err = SlabBStackAllocator::open(BStack::open(&path).unwrap()).unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvalidData);
}
#[test]
fn open_rejects_bad_magic() {
let (stack, path) = empty_stack();
let _g = Guard(path.clone());
stack.push([0u8; 48]).unwrap(); drop(stack);
let err = SlabBStackAllocator::open(BStack::open(&path).unwrap()).unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvalidData);
}
#[test]
fn open_rejects_invalid_stored_block_size() {
let (stack, path) = empty_stack();
let _g = Guard(path.clone());
let mut hdr = [0u8; 48];
hdr[24..32].copy_from_slice(b"ALSL\x00\x01\x00\x00");
hdr[32..40].copy_from_slice(&1u64.to_le_bytes());
stack.push(hdr).unwrap();
drop(stack);
let err = SlabBStackAllocator::open(BStack::open(&path).unwrap()).unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvalidData);
}
#[test]
fn open_rejects_misaligned_tail() {
let (stack, path) = empty_stack();
let _g = Guard(path.clone());
SlabBStackAllocator::new(stack, 8).unwrap();
let reopen = BStack::open(&path).unwrap();
reopen.extend(1).unwrap();
drop(reopen);
let err = SlabBStackAllocator::open(BStack::open(&path).unwrap()).unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvalidData);
}
#[test]
fn open_succeeds_and_restores_block_size() {
let (stack, path) = empty_stack();
let _g = Guard(path.clone());
SlabBStackAllocator::new(stack, 32).unwrap();
let alloc = SlabBStackAllocator::open(BStack::open(&path).unwrap()).unwrap();
assert_eq!(alloc.block_size(), 32);
}
#[test]
fn zero_alloc_returns_empty_slice() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let alloc = SlabBStackAllocator::new(stack, 8).unwrap();
let s = alloc.alloc(0).unwrap();
assert!(s.is_empty());
}
#[test]
fn dealloc_pushes_to_free_list_and_next_alloc_reuses_block() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let alloc = SlabBStackAllocator::new(stack, 8).unwrap();
let s1 = alloc.alloc(8).unwrap();
let offset1 = s1.start();
alloc.dealloc(s1).unwrap();
let s2 = alloc.alloc(8).unwrap();
assert_eq!(s2.start(), offset1);
}
#[test]
fn free_list_recycles_all_dealloc_d_blocks() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let alloc = SlabBStackAllocator::new(stack, 8).unwrap();
let a = alloc.alloc(8).unwrap();
let b = alloc.alloc(8).unwrap();
let c = alloc.alloc(8).unwrap();
let mut original = [a.start(), b.start(), c.start()];
alloc.dealloc(a).unwrap();
alloc.dealloc(b).unwrap();
alloc.dealloc(c).unwrap();
let r1 = alloc.alloc(8).unwrap();
let r2 = alloc.alloc(8).unwrap();
let r3 = alloc.alloc(8).unwrap();
let mut reused = [r1.start(), r2.start(), r3.start()];
original.sort();
reused.sort();
assert_eq!(reused, original);
}
#[test]
fn oversized_tail_dealloc_shrinks_stack() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let alloc = SlabBStackAllocator::new(stack, 8).unwrap();
let s = alloc.alloc(17).unwrap();
let tail_before = alloc.stack().len().unwrap();
assert_eq!(
s.start() + 24,
tail_before,
"allocation must be at the tail"
);
alloc.dealloc(s).unwrap();
assert_eq!(alloc.stack().len().unwrap(), tail_before - 24);
}
#[test]
fn write_and_read_round_trip() {
let (stack, path) = empty_stack();
let _g = Guard(path);
let alloc = SlabBStackAllocator::new(stack, 16).unwrap();
let s = alloc.alloc(12).unwrap();
s.write(b"hello world!").unwrap();
assert_eq!(s.read().unwrap(), b"hello world!");
}
#[test]
fn data_survives_reopen() {
let (stack, path) = empty_stack();
let _g = Guard(path.clone());
let alloc = SlabBStackAllocator::new(stack, 16).unwrap();
let s = alloc.alloc(5).unwrap();
let offset = s.start();
s.write(b"hello").unwrap();
drop(alloc);
let alloc2 = SlabBStackAllocator::open(BStack::open(&path).unwrap()).unwrap();
let s2 = unsafe { crate::alloc::BStackSlice::from_raw_parts(&alloc2, offset, 5) };
assert_eq!(s2.read().unwrap(), b"hello");
}
}