use r_efi::efi;
const POOL_ALIGNMENT: usize = 8usize;
#[repr(C)]
struct Marker(*mut u8);
fn align_request(size: usize, align: usize) -> usize {
if align > POOL_ALIGNMENT {
size + align
} else {
size
}
}
unsafe fn align_block(ptr: *mut u8, align: usize) -> *mut u8 {
if align > POOL_ALIGNMENT {
let offset = align - (ptr as usize & (align - 1));
assert!(offset >= POOL_ALIGNMENT);
assert!(POOL_ALIGNMENT >= core::mem::size_of::<Marker>());
assert!(POOL_ALIGNMENT >= core::mem::align_of::<Marker>());
let aligned = ptr.add(offset);
core::ptr::write((aligned as *mut Marker).offset(-1), Marker(ptr));
aligned
} else {
ptr
}
}
unsafe fn unalign_block(ptr: *mut u8, align: usize) -> *mut u8 {
if align > POOL_ALIGNMENT {
core::ptr::read((ptr as *mut Marker).offset(-1)).0
} else {
ptr
}
}
pub unsafe fn alloc(
system_table: *mut efi::SystemTable,
layout: core::alloc::Layout,
memory_type: efi::MemoryType,
) -> *mut u8 {
let align = layout.align();
let size = layout.size();
assert!(size > 0);
if size.checked_add(align).is_none() {
return core::ptr::null_mut();
}
let mut ptr: *mut core::ffi::c_void = core::ptr::null_mut();
let size_allocated = align_request(size, align);
let r = unsafe {
((*(*system_table).boot_services).allocate_pool)(
memory_type,
size_allocated,
&mut ptr,
)
};
if r.is_error() || ptr.is_null() {
core::ptr::null_mut()
} else {
unsafe { align_block(ptr as *mut u8, align) }
}
}
pub unsafe fn dealloc(
system_table: *mut efi::SystemTable,
ptr: *mut u8,
layout: core::alloc::Layout,
) {
assert!(!ptr.is_null());
let original = unalign_block(
ptr,
layout.align(),
) as *mut core::ffi::c_void;
let r = ((*(*system_table).boot_services).free_pool)(original);
assert!(!r.is_error());
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn align() {
let ptrsize = std::mem::size_of::<*mut ()>();
assert_eq!(POOL_ALIGNMENT, 8);
for i in 0..256 {
for j in &[1, 2, 4, 8, 16, 32, 64, 128] {
if *j <= 8 {
assert_eq!(align_request(i, *j), i);
} else {
assert!(align_request(i, *j) > i + ptrsize);
}
}
}
}
}