use super::align::assert_aligned;
use alloc::alloc;
use core::alloc::Layout;
use core::marker::PhantomData;
pub struct Heap<AlignAs>(PhantomData<AlignAs>);
impl<AlignAs> Heap<AlignAs> {
pub const fn new() -> Self {
Self(PhantomData)
}
#[expect(
clippy::unused_self,
reason = "Using a static method is harder in presence of generic parameters"
)]
pub fn alloc(&self, n: usize) -> *mut u8 {
assert_aligned::<AlignAs>(n);
assert_ne!(n, 0, "Allocating 0 bytes is invalid");
isize::try_from(n.next_multiple_of(align_of::<AlignAs>())).expect("Too big allocation");
let layout = unsafe { Layout::from_size_align_unchecked(n, align_of::<AlignAs>()) };
unsafe { alloc::alloc(layout) }
}
#[expect(
clippy::unused_self,
reason = "Using a static method is harder in presence of generic parameters"
)]
pub unsafe fn dealloc(&self, ptr: *mut u8, n: usize) {
let layout = unsafe { Layout::from_size_align_unchecked(n, align_of::<AlignAs>()) };
unsafe { alloc::dealloc(ptr, layout) }
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
#[should_panic]
fn alloc_zero() {
Heap::<u8>::new().alloc(0);
}
#[test]
#[should_panic]
fn alloc_unaligned() {
Heap::<u16>::new().alloc(3);
}
#[test]
#[should_panic]
fn alloc_large() {
Heap::<u8>::new().alloc((isize::MAX as usize) + 1);
}
#[test]
fn overaligned() {
#[repr(align(256))]
struct Overaligned;
let heap = Heap::<Overaligned>::new();
let ptr = heap.alloc(256);
assert_eq!(ptr.addr() % 256, 0);
unsafe {
heap.dealloc(ptr, 256);
}
}
#[test]
fn unique() {
let heap = Heap::<u8>::new();
let ptr1 = unsafe { &mut *heap.alloc(1) };
let ptr2 = unsafe { &mut *heap.alloc(1) };
*ptr1 = 1;
*ptr2 = 2;
assert_eq!(*ptr1, 1);
assert_eq!(*ptr2, 2);
unsafe {
heap.dealloc(ptr1, 1);
}
unsafe {
heap.dealloc(ptr2, 1);
}
}
}