ort_openrouter_cli/common/

alloc.rs

//! ort: Open Router CLI
//! https://github.com/grahamking/ort
//!
//! MIT License
//! Copyright (c) 2025 Graham King
//!
//! Arena allocator.
//! We map anonymous memory in arena-sized chunks and fill allocation
//! requests from the current chunk. This reduces syscalls to one per chunk,
//! often one per whole program. We never munmap old sections, memory stays
//! alive for duration of run.

#![allow(static_mut_refs)]

use core::alloc::Layout;
use core::ffi::c_void;
use core::ptr;

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;

/// Allocate a minimum of one MiB at a time.
const ARENA_SIZE: usize = 1024 * 1024;
const PAGE_SIZE: usize = 4096;

struct Arena {
    base: *mut u8,
    size: usize,
    offset: usize,
}

static mut ARENA: Arena = Arena {
    base: ptr::null_mut(),
    size: 0,
    offset: 0,
};

pub struct ArenaAlloc;

#[inline]
unsafe fn map_chunk(min_size: usize) -> Arena {
    let size = min_size.max(ARENA_SIZE).next_multiple_of(PAGE_SIZE);
    let base = syscall::mmap(
        ptr::null_mut(),
        size,
        syscall::PROT_READ | syscall::PROT_WRITE,
        syscall::MAP_PRIVATE | syscall::MAP_ANONYMOUS,
        -1,
        0,
    ) as *mut u8;

    if base.is_null() {
        let msg = c"mmap failed in common/alloc.rs\n";
        syscall::write(2, msg.as_ptr() as *const c_void, msg.count_bytes());
        syscall::exit(1);
    }

    Arena {
        base,
        size,
        offset: 0,
    }
}

// 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) };
        }

        let alloc_size = layout.size().next_multiple_of(ALIGN);

        unsafe {
            if ARENA.base.is_null() || ARENA.offset + alloc_size > ARENA.size {
                ARENA = map_chunk(alloc_size);
            }

            let ptr = ARENA.base.add(ARENA.offset);
            ARENA.offset += alloc_size;
            ptr
        }
    }

    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) };
        }

        // Anonymous mmap memory is zeroed by the kernel
        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()
*/