subms_arena_allocator/
lib.rs

1//! Bump-pointer arena. Allocate fixed-layout values into a growing byte buffer;
2//! `reset()` rewinds the bump cursor so the next request can reuse the entire
3//! arena.
4//!
5//! **Important: Drop is NOT run on items in the arena.** Anything with a
6//! non-trivial `Drop` (e.g. `String`, `Vec`) will leak its heap allocation
7//! if you store it in this arena. See [`Bump::alloc_copy`] which is
8//! constrained to `Copy` types as the safe public surface.
9//!
10//! ```
11//! use subms_arena_allocator::Bump;
12//! let mut a = Bump::with_capacity(1024);
13//! let x: &mut u32 = a.alloc_copy(42u32);
14//! assert_eq!(*x, 42);
15//! a.reset();
16//! ```
17
18use std::alloc::{Layout, alloc, dealloc};
19use std::cell::Cell;
20use std::mem;
21use std::ptr;
22
23/// Bump-pointer arena.
24pub struct Bump {
25    chunks: Vec<Chunk>,
26    /// Current chunk's bump cursor (byte offset into chunks.last()).
27    cursor: Cell<usize>,
28}
29
30struct Chunk {
31    ptr: *mut u8,
32    layout: Layout,
33}
34
35impl Drop for Chunk {
36    fn drop(&mut self) {
37        unsafe { dealloc(self.ptr, self.layout) };
38    }
39}
40
41impl Bump {
42    /// New empty arena. First chunk is allocated on first use.
43    pub fn new() -> Self {
44        Self::with_capacity(4096)
45    }
46
47    /// New arena, pre-allocating the first chunk at `initial_bytes` bytes.
48    pub fn with_capacity(initial_bytes: usize) -> Self {
49        let initial_bytes = initial_bytes.max(64);
50        let layout = Layout::from_size_align(initial_bytes, 16).expect("layout");
51        let ptr = unsafe { alloc(layout) };
52        assert!(!ptr.is_null(), "OOM allocating first arena chunk");
53        Self {
54            chunks: vec![Chunk { ptr, layout }],
55            cursor: Cell::new(0),
56        }
57    }
58
59    /// Allocate a `Copy` value. Returns an exclusive reference valid until
60    /// the next call to [`reset`](Bump::reset) or arena drop.
61    pub fn alloc_copy<T: Copy>(&mut self, value: T) -> &mut T {
62        let layout = Layout::new::<T>();
63        let p = self.alloc_raw(layout);
64        unsafe {
65            ptr::write(p as *mut T, value);
66            &mut *(p as *mut T)
67        }
68    }
69
70    /// Allocate `size` bytes, aligned to `align`. Returns a raw pointer; the
71    /// caller is responsible for writing into it. Used by `alloc_copy`.
72    pub fn alloc_raw(&mut self, layout: Layout) -> *mut u8 {
73        let size = layout.size();
74        let align = layout.align();
75        // align math: round cursor up to alignment boundary
76        let cur = self.cursor.get();
77        let aligned = align_up(cur, align);
78        let end = aligned + size;
79        let chunk = self.chunks.last().expect("at least one chunk");
80        if end <= chunk.layout.size() {
81            self.cursor.set(end);
82            return unsafe { chunk.ptr.add(aligned) };
83        }
84        // Grow: allocate a new chunk twice the previous size, at least `size + align`.
85        self.grow(size + align);
86        let chunk = self.chunks.last().unwrap();
87        let aligned = align_up(0, align);
88        self.cursor.set(aligned + size);
89        unsafe { chunk.ptr.add(aligned) }
90    }
91
92    fn grow(&mut self, min_bytes: usize) {
93        let last = self.chunks.last().expect("at least one chunk");
94        let new_size = (last.layout.size() * 2).max(min_bytes);
95        let layout = Layout::from_size_align(new_size, 16).expect("layout");
96        let ptr = unsafe { alloc(layout) };
97        assert!(!ptr.is_null(), "OOM growing arena");
98        self.chunks.push(Chunk { ptr, layout });
99        self.cursor.set(0);
100    }
101
102    /// Rewind to empty without freeing chunks. The next `alloc_*` call reuses
103    /// the same chunks.
104    pub fn reset(&mut self) {
105        // Drop all chunks except the largest, then keep that one and reset cursor.
106        if self.chunks.len() > 1 {
107            let largest = self
108                .chunks
109                .iter()
110                .enumerate()
111                .max_by_key(|(_, c)| c.layout.size())
112                .map(|(i, _)| i)
113                .unwrap();
114            let keeper = self.chunks.swap_remove(largest);
115            self.chunks.clear();
116            self.chunks.push(keeper);
117        }
118        self.cursor.set(0);
119    }
120
121    /// Total bytes currently allocated across all chunks.
122    pub fn total_capacity(&self) -> usize {
123        self.chunks.iter().map(|c| c.layout.size()).sum()
124    }
125}
126
127impl Default for Bump {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133/// `(p + align - 1) & !(align - 1)` - rounds up to the next aligned address.
134fn align_up(p: usize, align: usize) -> usize {
135    debug_assert!(align.is_power_of_two(), "alignment must be a power of two");
136    (p + align - 1) & !(align - 1)
137}
138
139#[cfg(feature = "harness")]
140pub mod recipe;
141
142// Silence the unused-mem warning for systems where this gets vendored.
143#[allow(dead_code)]
144fn _silence_unused() {
145    let _ = mem::size_of::<()>();
146}
subms_arena_allocator/lib.rs

subms_arena_allocator/
lib.rs
