use std::alloc::{alloc, dealloc, Layout};
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::sync::mutex::Mutex;
pub struct SystemHeap {
lock: Mutex<()>,
allocated_bytes: AtomicUsize,
allocation_count: AtomicUsize,
}
impl SystemHeap {
pub fn new() -> Self {
Self {
lock: Mutex::new(()),
allocated_bytes: AtomicUsize::new(0),
allocation_count: AtomicUsize::new(0),
}
}
pub fn alloc(&self, layout: Layout) -> *mut u8 {
let _guard = self.lock.lock();
let ptr = unsafe { alloc(layout) };
if !ptr.is_null() {
self.allocated_bytes.fetch_add(layout.size(), Ordering::Relaxed);
self.allocation_count.fetch_add(1, Ordering::Relaxed);
}
ptr
}
pub unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
let _guard = self.lock.lock();
#[cfg(feature = "debug")]
{
crate::debug::poison::poison_freed(ptr, layout.size());
}
dealloc(ptr, layout);
self.allocated_bytes.fetch_sub(layout.size(), Ordering::Relaxed);
}
pub fn alloc_typed<T>(&self) -> *mut T {
let layout = Layout::new::<T>();
self.alloc(layout) as *mut T
}
pub unsafe fn dealloc_typed<T>(&self, ptr: *mut T) {
let layout = Layout::new::<T>();
self.dealloc(ptr as *mut u8, layout);
}
pub fn allocated_bytes(&self) -> usize {
self.allocated_bytes.load(Ordering::Relaxed)
}
pub fn allocation_count(&self) -> usize {
self.allocation_count.load(Ordering::Relaxed)
}
}
impl Default for SystemHeap {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_heap_allocation() {
let heap = SystemHeap::new();
let ptr = heap.alloc_typed::<u64>();
assert!(!ptr.is_null());
assert_eq!(heap.allocated_bytes(), std::mem::size_of::<u64>());
unsafe {
heap.dealloc_typed(ptr);
}
assert_eq!(heap.allocated_bytes(), 0);
}
}