rialo-s-program-entrypoint 0.4.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Custom heap allocator for Rialo programs that supports dynamic heap sizes.
//!
//! # Overview
//!
//! This module provides a bump allocator optimized for Rialo's execution model.
//! Unlike the default allocator which assumes a fixed 32KB heap, this implementation
//! automatically utilizes whatever heap size is allocated by the runtime, including
//! custom sizes requested via Compute Budget instructions.
//!
//! # How It Works
//!
//! The allocator stores a small header (4-8 bytes) at the start of the heap containing:
//! - Current allocation offset
//! - Optional global state (via generic parameter `G`)
//!
//! Memory is allocated by growing upward from `HEAP_START_ADDRESS`. When an allocation
//! would exceed available heap space, accessing that memory triggers a segfault, which
//! is how Rialo signals out-of-memory conditions. This design eliminates the need for
//! the allocator to know the heap size at compile time.
//!
//! # Key Assumptions
//!
//! This allocator relies on guarantees provided by the Rialo runtime:
//!
//! 1. **Heap location**: The heap always starts at address `0x300000000`
//! 2. **Zero-initialization**: The Rialo runtime zero-initializes the heap region
//!    before program execution begins
//! 3. **Segfault on overflow**: Accessing memory beyond the allocated heap causes
//!    a segfault that terminates the transaction
//!
//! These assumptions are part of Rialo's documented runtime behavior and are validated
//! in tests using a simulated heap environment.
//!
//! # Deallocation Behavior
//!
//! As a bump allocator, this implementation:
//! - Can reclaim space from the most recent allocation if deallocated
//! - Intentionally leaks memory for all other deallocations (by design)
//! - Is optimized for Rialo's short-lived transaction model where all memory
//!   is reclaimed when the transaction completes
//!
//! # Usage
//!
//! The allocator is typically set up via the `custom_heap_default!` macro in the
//! entrypoint crate. Programs don't interact with it directly - it's used automatically
//! by Rust's allocation APIs (`Vec`, `Box`, etc.).

use core::{
    alloc::{GlobalAlloc, Layout},
    cell::Cell,
};

/// Minimum guaranteed heap size.
///
/// NOTE: Actual heap size may be larger if requested via Compute Budget.
/// The allocator automatically uses all available heap space.
pub const MIN_HEAP_LENGTH: usize = 32 * 1024;

/// Bump allocator that grows upward from HEAP_START.
///
/// Generic parameter `G` allows for optional global state storage at the heap start.
/// Use `G = ()` (the default) for no global state.
///
/// # Safety
///
/// Only one instance should exist per program, and it must be set as the global allocator.
/// Creating multiple instances or using alongside another allocator is undefined behavior.
pub struct BumpAllocator<G = ()> {
    #[cfg(test)]
    ptr: core::ptr::NonNull<u8>,
    #[cfg(test)]
    layout: Layout,

    _phantom: core::marker::PhantomData<G>,
}

/// Header stored at the start of the heap containing allocator metadata.
///
/// The header is zero-initialized by the Rialo runtime (or explicitly by tests).
/// Using `Cell<u32>` provides interior mutability for updating the allocation offset.
#[repr(C)]
struct Header<G> {
    /// Offset from end of header to first free byte (not from heap start)
    used: Cell<u32>,
    /// Optional global state (zero-sized if G = ())
    global: G,
}

impl<G> Header<G> {
    /// Size of the header including alignment padding
    const SIZE: u32 = {
        let size = core::mem::size_of::<Header<G>>();
        // Size validation happens in header() where we check against MIN_HEAP_LENGTH
        size as u32
    };

    /// Get the offset from heap start to the first free byte
    #[inline(always)]
    fn get_end_offset(&self) -> u32 {
        self.used.get().wrapping_add(Self::SIZE)
    }

    /// Set the offset from heap start to the first free byte
    #[inline(always)]
    fn set_end_offset(&self, offset: u32) {
        self.used.set(offset.wrapping_sub(Self::SIZE));
    }
}

// Non-test (Rialo target) implementation
#[cfg(not(test))]
impl<G> BumpAllocator<G> {
    /// Start address of the memory region used for program heap.
    #[cfg(not(target_arch = "riscv64"))]
    const HEAP_START_ADDRESS: u64 = 0x300000000;
    #[cfg(target_arch = "riscv64")]
    pub const HEAP_START_ADDRESS: u64 = 0x100000;

    /// Creates a new allocator.
    ///
    /// # Safety
    ///
    /// - Only one BumpAllocator instance should exist per program
    /// - It must be set as the global allocator
    /// - Multiple instances or using alongside another allocator leads to undefined behavior
    /// - The Rialo runtime must have zero-initialized the heap region (guaranteed by spec)
    pub const unsafe fn new() -> Self {
        // SAFETY: Caller must ensure this is only called once and set as global allocator.
        // The Rialo runtime guarantees the heap region starting at HEAP_START_ADDRESS
        // is zero-initialized before program execution.
        Self {
            _phantom: core::marker::PhantomData,
        }
    }

    #[inline(always)]
    const fn heap_start(&self) -> *mut u8 {
        Self::HEAP_START_ADDRESS as *mut u8
    }

    #[inline(always)]
    fn to_offset(&self, ptr: *mut u8) -> u32 {
        let addr = ptr as u64;
        debug_assert!(
            addr >= Self::HEAP_START_ADDRESS && addr < Self::HEAP_START_ADDRESS + u32::MAX as u64,
            "Pointer outside valid heap range"
        );
        (addr - Self::HEAP_START_ADDRESS) as u32
    }

    #[allow(clippy::wrong_self_convention)]
    #[inline(always)]
    fn from_offset(&self, offset: u32) -> *mut u8 {
        (Self::HEAP_START_ADDRESS + offset as u64) as *mut u8
    }
}

// Test implementation with actual allocation
#[cfg(test)]
impl<G: bytemuck::Zeroable> BumpAllocator<G> {
    /// Creates a test allocator with specified heap size
    fn new_test(size: usize) -> Self {
        let size = size.min(u32::MAX as usize);
        assert!(
            size >= core::mem::size_of::<Header<G>>(),
            "Heap too small for header"
        );

        let align = core::mem::align_of::<Header<G>>().max(16);
        let layout = Layout::from_size_align(size, align).unwrap();

        // SAFETY: We're allocating with proper layout
        let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
        let ptr = core::ptr::NonNull::new(ptr).expect("Failed to allocate test heap");

        Self {
            ptr,
            layout,
            _phantom: core::marker::PhantomData,
        }
    }

    #[inline(always)]
    fn heap_start(&self) -> *mut u8 {
        self.ptr.as_ptr()
    }

    #[inline(always)]
    fn to_offset(&self, ptr: *mut u8) -> u32 {
        (ptr as usize - self.heap_start() as usize) as u32
    }

    #[allow(clippy::wrong_self_convention)]
    #[inline(always)]
    fn from_offset(&self, offset: u32) -> *mut u8 {
        self.heap_start().wrapping_add(offset as usize)
    }
}

#[cfg(test)]
impl<G> Drop for BumpAllocator<G> {
    fn drop(&mut self) {
        // SAFETY: ptr and layout match the allocation
        unsafe {
            std::alloc::dealloc(self.ptr.as_ptr(), self.layout);
        }
    }
}

impl<G: bytemuck::Zeroable> BumpAllocator<G> {
    /// Returns reference to the header at the start of the heap
    #[inline(always)]
    fn header(&self) -> &Header<G> {
        // Compile-time check: header must fit in minimum guaranteed heap
        const {
            assert!(
                core::mem::size_of::<Header<G>>() <= MIN_HEAP_LENGTH,
                "Header too large for minimum heap size"
            );
        }

        // SAFETY:
        // 1. On Rialo: HEAP_START_ADDRESS (0x300000000) is properly aligned by runtime
        // 2. In tests: Test allocator ensures proper alignment via Layout
        // 3. Header fits in heap (compile-time check above)
        // 4. Heap memory is zero-initialized (by Rialo runtime or test allocator)
        // 5. Header<G> is Zeroable, so zero-initialization is valid
        unsafe { &*self.heap_start().cast::<Header<G>>() }
    }

    /// Fast path allocation - assumes success is common case
    #[inline(always)]
    fn try_alloc_fast(&self, layout: Layout) -> Option<*mut u8> {
        let header = self.header();
        let current_offset = header.get_end_offset();

        let size = match u32::try_from(layout.size()) {
            Ok(s) => s,
            Err(_) => return None,
        };

        debug_assert!(layout.align().is_power_of_two());
        let align_mask = (layout.align() - 1) as u32;

        let aligned_offset = match current_offset.checked_add(align_mask) {
            Some(v) => v & !align_mask,
            None => return None,
        };

        #[allow(clippy::question_mark)]
        let end_offset = match aligned_offset.checked_add(size) {
            Some(end) => end,
            None => return None,
        };

        #[cfg(test)]
        if end_offset as usize > self.layout.size() {
            return None;
        }

        header.set_end_offset(end_offset);
        Some(self.from_offset(aligned_offset))
    }

    #[allow(clippy::question_mark)]
    /// Try to allocate at a specific pointer (used for in-place realloc)
    #[inline]
    fn try_alloc_at(&self, ptr: *mut u8, layout: Layout) -> Option<*mut u8> {
        let offset = self.to_offset(ptr);

        let size = match u32::try_from(layout.size()) {
            Ok(s) => s,
            Err(_) => return None,
        };

        let end_offset = match offset.checked_add(size) {
            Some(end) => end,
            None => return None,
        };

        #[cfg(test)]
        if end_offset as usize > self.layout.size() {
            return None;
        }

        self.header().set_end_offset(end_offset);
        Some(ptr)
    }

    /// Returns reference to global state reserved at heap start
    #[inline]
    pub fn global(&self) -> &G {
        &self.header().global
    }

    /// Returns amount of heap used (excluding header)
    #[cfg(test)]
    pub fn used(&self) -> usize {
        self.header().used.get() as usize
    }
}

// SAFETY: BumpAllocator correctly implements GlobalAlloc
unsafe impl<G: bytemuck::Zeroable> GlobalAlloc for BumpAllocator<G> {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // Fast path: assume allocation succeeds
        self.try_alloc_fast(layout).unwrap_or(core::ptr::null_mut())
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        // Only deallocate if this is the most recent allocation
        let header = self.header();
        let ptr_end = ptr.wrapping_add(layout.size());
        let end_offset = self.to_offset(ptr_end);

        if end_offset == header.get_end_offset() {
            // This was the last allocation, reclaim it
            header.set_end_offset(self.to_offset(ptr));
        }
        // Otherwise, bump allocator intentionally leaks (by design)
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let header = self.header();
        let ptr_end = ptr.wrapping_add(layout.size());
        let end_offset = self.to_offset(ptr_end);

        // Check if this is the last allocation
        if end_offset == header.get_end_offset() {
            // Last allocation - try to resize in place
            // SAFETY: Caller guarantees new layout is valid for the same alignment
            let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
            return self
                .try_alloc_at(ptr, new_layout)
                .unwrap_or(core::ptr::null_mut());
        }

        // Not the last allocation
        if new_size <= layout.size() {
            // Shrinking - return same pointer (leak extra space, this is bump allocator)
            return ptr;
        }

        // Growing non-last allocation - need new allocation and copy
        // SAFETY: Caller guarantees new layout is valid for the same alignment
        let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
        match self.try_alloc_fast(new_layout) {
            Some(new_ptr) => {
                // SAFETY:
                // - src is valid for reads of layout.size() bytes (caller guarantee)
                // - dst is valid for writes of new_size bytes (just allocated)
                // - Regions don't overlap (new allocation is after old in bump allocator)
                core::ptr::copy_nonoverlapping(ptr, new_ptr, layout.size());
                new_ptr
            }
            None => core::ptr::null_mut(),
        }
    }
}

#[cfg(test)]
mod unit_tests;