use std::alloc::{Layout, alloc as std_alloc, realloc as std_realloc};
use super::i_lua_allocator::ILuaAllocator;
#[repr(C)]
struct BlockHeader {
capacity: usize,
align: usize,
}
impl BlockHeader {
const LAYOUT: Layout = match Layout::from_size_align(size_of::<Self>(), align_of::<Self>()) {
Ok(l) => l,
Err(_) => panic!("块头布局非法"),
};
}
#[derive(Default)]
pub struct LuaManagedAllocator {
allocated_bytes: usize,
infallible_depth: u32,
}
impl LuaManagedAllocator {
pub fn allocated_bytes(&self) -> usize {
self.allocated_bytes
}
fn layout_for(size: usize, align: usize) -> Option<Layout> {
Layout::from_size_align(
size.max(1).checked_add(BlockHeader::LAYOUT.size())?,
align.max(BlockHeader::LAYOUT.align()),
)
.ok()
}
unsafe fn header_of(ptr: *mut u8) -> *mut BlockHeader {
unsafe { ptr.cast::<BlockHeader>().sub(1) }
}
}
impl ILuaAllocator for LuaManagedAllocator {
fn enter_infallible_allocation_region(&mut self) {
self.infallible_depth += 1;
}
fn try_exit_infallible_allocation_region(&mut self) -> bool {
self.infallible_depth = self.infallible_depth.saturating_sub(1);
true
}
fn allocate_new(&mut self, size: usize, align: usize) -> Option<*mut u8> {
let layout = Self::layout_for(size, align)?;
let ptr = unsafe { std_alloc(layout) };
if ptr.is_null() {
return None;
}
let header = ptr.cast::<BlockHeader>();
unsafe {
(*header) = BlockHeader {
capacity: layout.size(),
align: layout.align(),
};
}
self.allocated_bytes += layout.size();
Some(unsafe { ptr.add(BlockHeader::LAYOUT.size()) })
}
unsafe fn resize_allocation(&mut self, ptr: *mut u8, new_size: usize) -> Option<*mut u8> {
if ptr.is_null() {
return self.allocate_new(new_size, 1);
}
let header = unsafe { Self::header_of(ptr) };
let (capacity, align) = unsafe { ((*header).capacity, (*header).align) };
let new_layout = Self::layout_for(new_size, align)?;
let old_layout = Layout::from_size_align(capacity, align).ok()?;
let base = header.cast::<u8>();
let resized = unsafe { std_realloc(base, old_layout, new_layout.size()) };
if resized.is_null() {
return None;
}
self.allocated_bytes = self.allocated_bytes.saturating_sub(capacity) + new_layout.size();
unsafe {
(*resized.cast::<BlockHeader>()) = BlockHeader {
capacity: new_layout.size(),
align: new_layout.align(),
};
}
Some(unsafe { resized.add(BlockHeader::LAYOUT.size()) })
}
}
impl Drop for LuaManagedAllocator {
fn drop(&mut self) {
self.allocated_bytes = 0;
}
}