magi-code 0.96.1

Repository-aware CLI coding agent for terminal work
Documentation
use super::memory::MemoryBudget;
use rquickjs::allocator::{Allocator, RustAllocator};
use std::{mem, ptr, sync::Arc};

// Matches pinned rquickjs-core's RustAllocator layout on both 32- and 64-bit hosts.
const ALIGN: usize = mem::align_of::<u64>();
const HEADER: usize = if mem::size_of::<usize>() > ALIGN {
    mem::size_of::<usize>()
} else {
    ALIGN
};

/// Allocation refusal stays latched even if JavaScript catches its OOM exception.
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)
    }
}

// SAFETY: memory layout and ownership stay entirely with RustAllocator. Accounting
// never determines a pointer's layout; a refused realloc leaves its input intact.
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() {
            // SAFETY: allocation above contains at least size writable bytes.
            unsafe {
                pointer.write_bytes(0, size);
            }
        }
        pointer
    }

    unsafe fn dealloc(&mut self, pointer: *mut u8) {
        if pointer.is_null() {
            return;
        }
        // SAFETY: caller guarantees a live allocation from this allocator.
        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 {
            // SAFETY: same allocation ownership as this call.
            unsafe {
                self.dealloc(pointer);
            }
            return ptr::null_mut();
        }
        // SAFETY: caller guarantees a live allocation from this allocator.
        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();
        }
        // SAFETY: original pointer and requested size passed unchanged.
        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 {
        // SAFETY: caller guarantees a live allocation from this allocator.
        unsafe { RustAllocator::usable_size(pointer) }
    }
}