nub_rt/
alloc.rs

1//! A bump allocator with an explicit reset.
2//!
3//! Programs that need `alloc` on a freestanding target need a
4//! `#[global_allocator]`, and a bump arena is the cheapest thing that
5//! works: allocation is a pointer add, and a program that runs once and
6//! exits never needs to free.
7//!
8//! [`BumpAlloc::reset`] is what makes the arena re-usable. Without it a
9//! program is single-shot per instance — the second invocation walks
10//! off the end of the arena and the allocation fails, which surfaces as
11//! a guest panic rather than anything legible. That matters for any
12//! caller that invokes the same instance twice, benchmark harnesses
13//! very much included: measuring steady-state execution means running
14//! the same instance repeatedly.
15//!
16//! Resetting is not free of consequences: it invalidates every live
17//! allocation at once. It is only sound at a point where nothing from
18//! the previous invocation is still borrowed — i.e. at an entry point,
19//! before any work begins. [`bump_allocator!`](crate::bump_allocator)
20//! generates a `reset_heap()` for exactly that use.
21
22use core::alloc::{GlobalAlloc, Layout};
23use core::cell::UnsafeCell;
24
25/// A fixed-size bump arena of `N` bytes.
26///
27/// `dealloc` is a no-op; space is reclaimed only by [`reset`](Self::reset).
28pub struct BumpAlloc<const N: usize> {
29    heap: UnsafeCell<[u8; N]>,
30    pos: UnsafeCell<usize>,
31}
32
33// SAFETY: PVM2 programs are single-threaded — the engine runs one
34// instruction stream per instance and there is no way for a guest to
35// create a thread. Without that, the `UnsafeCell` accesses below would
36// need synchronization.
37unsafe impl<const N: usize> Sync for BumpAlloc<N> {}
38
39impl<const N: usize> BumpAlloc<N> {
40    pub const fn new() -> Self {
41        BumpAlloc {
42            heap: UnsafeCell::new([0; N]),
43            pos: UnsafeCell::new(0),
44        }
45    }
46
47    /// Free everything at once, making the arena usable again.
48    ///
49    /// # Safety
50    ///
51    /// Every pointer previously handed out becomes dangling. Call only
52    /// where no allocation from the previous run is still reachable —
53    /// at an entry point, before any work.
54    pub unsafe fn reset(&self) {
55        unsafe { *self.pos.get() = 0 };
56    }
57
58    /// High-water mark, in bytes. Useful for sizing an arena.
59    pub fn used(&self) -> usize {
60        unsafe { *self.pos.get() }
61    }
62}
63
64impl<const N: usize> Default for BumpAlloc<N> {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70unsafe impl<const N: usize> GlobalAlloc for BumpAlloc<N> {
71    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
72        let pos = unsafe { &mut *self.pos.get() };
73        let aligned = (*pos + layout.align() - 1) & !(layout.align() - 1);
74        let next = aligned + layout.size();
75        if next > N {
76            return core::ptr::null_mut();
77        }
78        *pos = next;
79        unsafe { (*self.heap.get()).as_mut_ptr().add(aligned) }
80    }
81
82    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
83}
84
85/// Install a [`BumpAlloc`] of `$bytes` as the program's global
86/// allocator, and define `reset_heap()` beside it.
87///
88/// A macro rather than a static in this crate because
89/// `#[global_allocator]` must be unique per binary: a static here would
90/// force the arena on every program that links `nub-rt`, including the
91/// ones that never allocate.
92///
93/// ```ignore
94/// nub_rt::bump_allocator!(64 * 1024);
95///
96/// #[nub_rt::endpoint(0)]
97/// fn run(_: u64) -> u64 {
98///     reset_heap();          // re-entrant: safe to invoke repeatedly
99///     my_kernel() as u64
100/// }
101/// ```
102///
103/// On host targets it expands to a `reset_heap()` that does nothing, so
104/// the same source builds both ways.
105#[macro_export]
106macro_rules! bump_allocator {
107    ($bytes:expr) => {
108        #[cfg(target_os = "none")]
109        #[global_allocator]
110        static __NUB_RT_HEAP: $crate::alloc::BumpAlloc<{ $bytes }> =
111            $crate::alloc::BumpAlloc::new();
112
113        /// Release everything allocated so far.
114        ///
115        /// Call at the top of an entry point, before any work — at that
116        /// moment nothing from a previous invocation is reachable, which
117        /// is what makes it sound.
118        #[cfg(target_os = "none")]
119        fn reset_heap() {
120            // SAFETY: called at entry, before any allocation of this
121            // invocation exists and after every allocation of the
122            // previous one has gone out of scope.
123            unsafe { __NUB_RT_HEAP.reset() };
124        }
125
126        /// No-op on host: the system allocator needs no reset.
127        #[cfg(not(target_os = "none"))]
128        #[allow(dead_code)]
129        fn reset_heap() {}
130    };
131}
nub_rt/alloc.rs

nub_rt/
alloc.rs
