use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicIsize, Ordering};
use wasm_bindgen::JsCast;
use wasm_bindgen::prelude::wasm_bindgen;
#[wasm_bindgen]
extern "C" {
type Memory;
#[wasm_bindgen(method, getter)]
fn buffer(this: &Memory) -> MaybeSharedArrayBuffer;
type MaybeSharedArrayBuffer;
#[wasm_bindgen(method, getter = byteLength)]
fn byte_length(this: &MaybeSharedArrayBuffer) -> f64;
}
struct CountingAllocator<A> {
inner: A,
allocated_now: AtomicIsize,
}
impl<A> CountingAllocator<A> {
const fn new(inner: A) -> Self {
Self {
inner,
allocated_now: AtomicIsize::new(0),
}
}
fn allocated_now(&self) -> usize {
self.allocated_now
.load(Ordering::Relaxed)
.try_into()
.unwrap_or(0)
}
}
unsafe impl<A: GlobalAlloc> GlobalAlloc for CountingAllocator<A> {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
self.allocated_now
.fetch_add(layout.size() as isize, Ordering::Relaxed);
unsafe { self.inner.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
self.allocated_now
.fetch_sub(layout.size() as isize, Ordering::Relaxed);
unsafe { self.inner.dealloc(ptr, layout) };
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
self.allocated_now
.fetch_add(layout.size() as isize, Ordering::Relaxed);
unsafe { self.inner.alloc_zeroed(layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
self.allocated_now.fetch_add(
new_size as isize - layout.size() as isize,
Ordering::Relaxed,
);
unsafe { self.inner.realloc(ptr, layout, new_size) }
}
}
#[global_allocator]
static ALLOCATOR: CountingAllocator<System> = CountingAllocator::new(System);
pub(crate) fn get_used() -> (usize, usize) {
let client_heap = wasm_bindgen::memory()
.unchecked_into::<Memory>()
.buffer()
.byte_length() as u64;
(ALLOCATOR.allocated_now(), client_heap as usize)
}
#[cfg(all(feature = "talc-allocator", feature = "trace-allocator"))]
compile_error!("Cannot enable both `talc-allocator` and `trace-allocator` features");