baryl 0.0.2

//! Claiming a named anchor, and the checkpoint heap as a Rust `Allocator`.

use core::ffi::{CStr, c_void};
#[cfg(feature = "component")]
use core::{
    alloc::{AllocError, Allocator, Layout},
    ptr::NonNull,
};

use crate::abi::{AnonPtr, SandboxSafe, SbxPtr};
#[cfg(feature = "component")]
use crate::abi::{BMap, BVec};

use super::{AllocRef, AllocRegion, baryl_leak_and_initialize};
#[cfg(feature = "component")]
use super::{baryl_ck_heap_alloc, baryl_ck_heap_free};

/// The `init` slot the ABI takes, over whatever closure the caller had.
unsafe extern "C" fn call_init(addr: u64, userdata: *mut c_void) {
    let f = unsafe { &mut *userdata.cast::<&mut dyn FnMut(u64)>() };
    f(addr);
}

impl AllocRef {
    /// Claim one `T` under `name`, in memory a checkpoint captures.
    ///
    /// The address is the same every time this name is claimed — this run, the
    /// next run, and after any restore — so this is how a component finds the
    /// state it left behind. `init` runs on the very first claim ever made
    /// under this name and never again: a restored run reads the bytes the
    /// checkpoint holds rather than re-initializing over them.
    ///
    /// Put nothing in `init` that this run needs done. Anything a restore has
    /// to redo belongs in `#[core(checkpoint_load)]`.
    ///
    /// A name claimed once is bound to `T`'s size and alignment. Claiming it
    /// again with a differently shaped `T` — or from the anonymous region —
    /// aborts the process, so change the layout of a claimed struct and every
    /// existing checkpoint stops restoring.
    ///
    /// # Panics
    ///
    /// If `init` is somehow invoked more than once for a single claim.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// #[repr(C)]
    /// struct Tally { blocks: u64, resets: u64 }
    /// impl_sandbox_safe!(Tally);
    ///
    /// #[core(init)]
    /// fn start(&mut self, t: &mut Control) {
    ///     // First run ever: zeroed and `init` runs. Every restore after:
    ///     // the same address, holding whatever the checkpoint saved.
    ///     let tally: SbxPtr<Tally> = t.subs.alloc.leak_and_initialize_sbx(
    ///         c"mycomponent.tally",
    ///         |addr| baryl::logging::info!("first claim at {addr:#x}"),
    ///     );
    ///     self.tally = tally;
    /// }
    /// ```
    pub fn leak_and_initialize_sbx<T: SandboxSafe>(
        &self,
        name: &CStr,
        init: impl FnOnce(u64),
    ) -> SbxPtr<T> {
        let mut once = Some(init);
        let mut call = |addr: u64| {
            let f = once.take().expect("an SbxRegion name inits at most once");
            f(addr);
        };
        let mut dyn_call: &mut dyn FnMut(u64) = &mut call;
        let addr = unsafe {
            baryl_leak_and_initialize(
                *self,
                name.as_ptr(),
                AllocRegion::SbxRegion,
                size_of::<T>() as u64,
                align_of::<T>() as u64,
                Some(call_init),
                (&raw mut dyn_call).cast(),
            )
        };
        SbxPtr::from_addr(addr)
    }

    /// Claim one `T` under `name`, in memory a checkpoint does **not** capture.
    ///
    /// The address is stable in the same way
    /// [`leak_and_initialize_sbx`](Self::leak_and_initialize_sbx)'s is, but the
    /// bytes are a fresh mapping each run: zero before `init`, and `init` runs
    /// on every claim rather than only the first. Use it for anything the run
    /// can rebuild from nothing — a cache, a lookup table, scratch space — and
    /// keep it out of the checkpoint's size.
    ///
    /// The same layout rule holds: a name is bound to the shape and the region
    /// it was first claimed under, and claiming it differently aborts.
    pub fn leak_and_initialize_anon<T: SandboxSafe>(
        &self,
        name: &CStr,
        mut init: impl FnMut(u64),
    ) -> AnonPtr<T> {
        let mut call = |addr: u64| init(addr);
        let mut dyn_call: &mut dyn FnMut(u64) = &mut call;
        let addr = unsafe {
            baryl_leak_and_initialize(
                *self,
                name.as_ptr(),
                AllocRegion::AnonRegion,
                size_of::<T>() as u64,
                align_of::<T>() as u64,
                Some(call_init),
                (&raw mut dyn_call).cast(),
            )
        };
        AnonPtr::from_addr(addr)
    }

    /// The checkpoint heap as an allocator, for `BVec` and `BMap`.
    #[cfg(feature = "component")]
    pub fn ck_heap(&self) -> SbxAlloc {
        SbxAlloc(*self)
    }
}

/// The checkpoint heap, as a Rust `Allocator`.
///
/// A zero-cost wrapper around the handle it came from — `Copy`, one pointer
/// wide, and itself storable in the pool, which is what lets a `BVec` or `BMap`
/// carry its own allocator inline. Blocks come back on restore along with the
/// rest of the checkpoint.
///
/// This is the only pool arena that can free, so it is where variable-sized
/// storage goes; fixed state belongs on a named anchor instead.
///
/// `allocate` answers `AllocError` when the heap is full — it never aborts, so
/// `BVec::try_new_in` and `BMap::try_with_capacity_in` give you back a `None`
/// you can act on.
///
/// # Examples
///
/// ```ignore
/// let heap: SbxAlloc = t.subs.alloc.ck_heap();
/// let seen: Option<BMap<u64, u64, SbxAlloc>> = BMap::try_with_capacity_in(heap, 4096);
/// ```
#[cfg(feature = "component")]
#[repr(transparent)]
#[derive(Clone, Copy, Debug)]
pub struct SbxAlloc(pub AllocRef);

// SAFETY: the handle is one pool VA and every block it hands back is checkpoint
// heap, so nothing reachable through it points outside the pool.
#[cfg(feature = "component")]
unsafe impl SandboxSafe for SbxAlloc {}

#[cfg(feature = "component")]
unsafe impl Allocator for SbxAlloc {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        let addr =
            unsafe { baryl_ck_heap_alloc(self.0, layout.size() as u64, layout.align() as u64) };
        NonNull::new(addr as *mut u8)
            .map(|p| NonNull::slice_from_raw_parts(p, layout.size()))
            .ok_or(AllocError)
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        unsafe {
            baryl_ck_heap_free(
                self.0,
                ptr.as_ptr() as u64,
                layout.size() as u64,
                layout.align() as u64,
            );
        }
    }
}

// The layouts C is compiled against, checked on the Rust side too.
const _: () = assert!(size_of::<AllocRef>() == 0x8, "the handle is one pointer");
#[cfg(feature = "component")]
const _: () = assert!(size_of::<BVec<u8, SbxAlloc>>() == 0x18, "handle, storage, count");
#[cfg(feature = "component")]
const _: () = assert!(
    size_of::<BMap<u64, u64, SbxAlloc>>() == 0x20,
    "the slot array and the live count"
);