ort_openrouter_cli/common/

alloc.rs

//! ort: Open Router CLI
//! https://github.com/grahamking/ort
//!
//! MIT License
//! Copyright (c) 2025 Graham King
//!
//! Arena allocator. No heap allocations for entire program run.
//! We statically allocate a large chunk of memory (`.bss` segment),
//! and use that to fill allocation requests. This avoids doing
//! any syscalls to get memory.

#![allow(static_mut_refs)]

use core::alloc::Layout;
use core::ffi::c_void;
use core::sync::atomic::{AtomicUsize, Ordering};

use crate::syscall;

#[cfg(feature = "panic-on-realloc")]
static mut IS_FIRST_REALLOC: bool = true;

#[cfg(feature = "print-allocations")]
use crate::common::utils::to_ascii;

/// All allocated memory locations will have this alignment, max_align_t
const ALIGN: usize = 16;

// How much memory to allocate total. Don't exceed this!
//const MEM_SIZE: usize = 2 * 1024 * 1024;
// In debug mode to print panics, need > 8MiB
const MEM_SIZE: usize = 16 * 1024 * 1024;

#[repr(align(16))]
struct Heap(pub [u8; MEM_SIZE]);

static mut HEAP: Heap = Heap([0u8; MEM_SIZE]);
static mut OFFSET: AtomicUsize = AtomicUsize::new(0);

pub struct ArenaAlloc;

// In case you were wondering, yes all three methods get used. Rust does
// a bnuch of alloc_zeroed and realloc.
//
// Build with feature "print-allocations" to see memory being allocated:
// cargo build --features="print-allocations"
// cargo build --release --features="print-allocations" -Zbuild-std="core,alloc"
//
// There's a Python script at the end of this file to summarize the output.
//
// Normal / prompt usage seems to peak under 64 Kib of active memory.
//
// `ort list` peaks around 180 Kib because it has one large (~128 KiB)
// allocation which is a string holding the names of all models, so we can sort them.
unsafe impl core::alloc::GlobalAlloc for ArenaAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        #[cfg(feature = "print-allocations")]
        {
            let mut buf = [0u8; 16];
            buf[0] = b'+';
            let len = to_ascii(layout.size(), &mut buf[1..]);
            crate::syscall::write(2, buf.as_ptr().cast(), len);

            // Also print alignment
            //let mut buf = [0u8; 16];
            //buf[0] = b' ';
            //let len = to_ascii(layout.align(), &mut buf[1..]);
            //unsafe { crate::libc::write(2, buf.as_ptr().cast(), len) };
        }

        // Allocate in multiples of 16 bytes so that we stay aligned for next alloc.
        let alloc_size = layout.size().next_multiple_of(ALIGN);

        unsafe {
            // Read current offset, for this allocation, and move it by alloc_size for next
            // allocation. fetch_add returns the value before addition.
            let current_offset = OFFSET.fetch_add(alloc_size, Ordering::Relaxed);

            if OFFSET.load(Ordering::Relaxed) > MEM_SIZE {
                let msg = c"Out of memory, common/alloc.rs MEM_SIZE\n";
                syscall::write(2, msg.as_ptr() as *const c_void, msg.count_bytes());
                syscall::exit(1);
            }
            HEAP.0.as_mut_ptr().add(current_offset)
        }
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        #[cfg(feature = "print-allocations")]
        {
            let mut buf = [0u8; 16];
            buf[0] = b'+';
            let len = to_ascii(layout.size(), &mut buf[1..]);
            crate::syscall::write(2, buf.as_ptr().cast(), len);

            // Also print alignment
            //let mut buf = [0u8; 16];
            //buf[0] = b' ';
            //let len = to_ascii(layout.align(), &mut buf[1..]);
            //unsafe { crate::libc::write(2, buf.as_ptr().cast(), len) };
        }

        // .bss segment is already zeroed
        unsafe { self.alloc(layout) }
    }

    #[allow(unused_variables)]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        #[cfg(feature = "print-allocations")]
        {
            let mut buf = [0u8; 16];
            buf[0] = b'-';
            let len = to_ascii(layout.size(), &mut buf[1..]);
            crate::syscall::write(2, buf.as_ptr().cast(), len);
        }

        // we never free, program runtime is short
    }

    // ort is tuned to avoid reallocations, so we don't override it.
    // Parent implements realloc as alloc-and-copy, which is fine because it should not happen.
    //
    // Keep here to panic-on-realloc when code changes.
    /*
    #[allow(unused_variables)]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        #[cfg(feature = "print-allocations")]
        {
            let mut buf = [0u8; 16];
            buf[0] = b'\\';
            let len = to_ascii(layout.size(), &mut buf[1..]);
            crate::syscall::write(2, buf.as_ptr().cast(), len);

            buf[0] = b'/';
            let len = to_ascii(new_size, &mut buf[1..]);
            crate::syscall::write(2, buf.as_ptr().cast(), len);

            // Also print alignment
            //let mut buf = [0u8; 16];
            //buf[0] = b' ';
            //let len = to_ascii(layout.align(), &mut buf[1..]);
            //unsafe { crate::libc::write(2, buf.as_ptr().cast(), len) };
        }

        #[cfg(feature = "panic-on-realloc")]
        unsafe {
            // The panic machinery uses realloc, so we must immediately disable this or we would
            // panic during the panic, and get no useful stack trace.
            if IS_FIRST_REALLOC {
                IS_FIRST_REALLOC = false;
                panic!("realloc {} -> {}", layout.size(), new_size);
            }
        }
    }
    */
}

/*
"""Print running totals from allocs.txt and report the maximum cumulative value."""

from pathlib import Path


def main() -> None:
    path = Path("allocs.txt")
    if not path.is_file():
        raise SystemExit("allocs.txt not found in the current directory")

    total = 0
    max_total = None  # Highest cumulative total
    max_plus = None   # Largest individual + value
    max_minus = None  # Largest (most negative) individual - value

    with path.open() as fh:
        for raw_line in fh:
            line = raw_line.strip()
            if not line:
                continue  # Skip blank lines silently

            try:
                delta = int(line)
            except ValueError as exc:
                raise SystemExit(f"Invalid line in allocs.txt: {line!r}") from exc
            # realloc indicators
            line = line.replace('/', '+').replace('\\', '-')

            if delta > 0:
                max_plus = delta if max_plus is None else max(max_plus, delta)
            elif delta < 0:
                max_minus = delta if max_minus is None else min(max_minus, delta)

            total += delta
            max_total = total if max_total is None else max(max_total, total)
            print(f"{line} {total}")

    print()  # Blank line before the summary, matching the example
    print(f"Max: {max_total if max_total is not None else 0}")
    print(f"Largest +: {max_plus if max_plus is not None else 0}")
    print(f"Largest -: {max_minus if max_minus is not None else 0}")


if __name__ == "__main__":
    main()
*/