use super::memory::MemoryBudget;
use rquickjs::allocator::{Allocator, RustAllocator};
use std::{mem, ptr, sync::Arc};
const ALIGN: usize = mem::align_of::<u64>();
const HEADER: usize = if mem::size_of::<usize>() > ALIGN {
mem::size_of::<usize>()
} else {
ALIGN
};
pub(super) struct BoundedAllocator(pub(super) Arc<MemoryBudget>);
impl BoundedAllocator {
fn charge_size(size: usize) -> Option<usize> {
size.checked_add(ALIGN - 1)
.map(|n| n & !(ALIGN - 1))?
.checked_add(HEADER)
}
}
unsafe impl Allocator for BoundedAllocator {
fn alloc(&mut self, size: usize) -> *mut u8 {
let charge = Self::charge_size(size).unwrap_or(usize::MAX);
if !self.0.charge(charge) {
return ptr::null_mut();
}
let pointer = RustAllocator.alloc(size);
if pointer.is_null() {
self.0.release(charge);
self.0
.stopped
.store(true, std::sync::atomic::Ordering::SeqCst);
}
pointer
}
fn calloc(&mut self, count: usize, size: usize) -> *mut u8 {
let Some(size) = count.checked_mul(size) else {
self.0
.stopped
.store(true, std::sync::atomic::Ordering::SeqCst);
return ptr::null_mut();
};
if size == 0 {
return ptr::null_mut();
}
let pointer = self.alloc(size);
if !pointer.is_null() {
unsafe {
pointer.write_bytes(0, size);
}
}
pointer
}
unsafe fn dealloc(&mut self, pointer: *mut u8) {
if pointer.is_null() {
return;
}
unsafe {
self.0.release(RustAllocator::usable_size(pointer) + HEADER);
RustAllocator.dealloc(pointer);
}
}
unsafe fn realloc(&mut self, pointer: *mut u8, size: usize) -> *mut u8 {
if pointer.is_null() {
return self.alloc(size);
}
if size == 0 {
unsafe {
self.dealloc(pointer);
}
return ptr::null_mut();
}
let previous = unsafe { RustAllocator::usable_size(pointer) } + HEADER;
let next_charge = Self::charge_size(size).unwrap_or(usize::MAX);
let increase = next_charge.saturating_sub(previous);
if !self.0.charge(increase) {
return ptr::null_mut();
}
let next = unsafe { RustAllocator.realloc(pointer, size) };
if next.is_null() {
self.0.release(increase);
self.0
.stopped
.store(true, std::sync::atomic::Ordering::SeqCst);
} else {
self.0.release(previous.saturating_sub(next_charge));
}
next
}
unsafe fn usable_size(pointer: *mut u8) -> usize {
unsafe { RustAllocator::usable_size(pointer) }
}
}