use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
#[derive(Debug)]
pub struct CountingAllocator {
counter: AtomicUsize,
}
impl CountingAllocator {
pub const fn new() -> Self {
Self {
counter: AtomicUsize::new(0),
}
}
pub fn current_usage(&self) -> usize {
self.counter.load(SeqCst)
}
}
unsafe impl GlobalAlloc for CountingAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let ret = System.alloc(layout);
if !ret.is_null() {
self.counter.fetch_add(layout.size(), SeqCst);
}
return ret;
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
System.dealloc(ptr, layout);
self.counter.fetch_sub(layout.size(), SeqCst);
}
}