Skip to main content

subms_arena_allocator/features/
growable.rs

1//! `GrowableBump`: a bump arena that allocates a fresh chunk (at twice
2//! the previous size, or the request size, whichever is larger) when
3//! the active chunk runs out of room. Previous chunks are retained
4//! until `reset()` or drop so any references handed out remain valid.
5//!
6//! Trade-off vs the fixed-capacity base: grow events are not free -
7//! they cost an allocator round-trip and break the steady-state p99.
8//! `reset()` keeps only the largest chunk, so steady-state workloads
9//! converge on a single chunk after the first round.
10
11use std::alloc::{Layout, alloc, dealloc};
12use std::ptr;
13
14use crate::align_up;
15
16/// Multi-chunk bump-pointer arena that auto-grows on exhaustion.
17pub struct GrowableBump {
18    chunks: Vec<Chunk>,
19    /// Cursor into `chunks.last()`.
20    cursor: usize,
21}
22
23struct Chunk {
24    ptr: *mut u8,
25    layout: Layout,
26}
27
28impl Drop for Chunk {
29    fn drop(&mut self) {
30        unsafe { dealloc(self.ptr, self.layout) };
31    }
32}
33
34impl GrowableBump {
35    /// New arena with a 4 KiB initial chunk.
36    pub fn new() -> Self {
37        Self::with_capacity(4096)
38    }
39
40    /// New arena with the requested initial chunk size (64-byte floor,
41    /// 16-byte aligned chunk allocation).
42    pub fn with_capacity(initial: usize) -> Self {
43        let initial = initial.max(64);
44        let layout = Layout::from_size_align(initial, 16).expect("layout");
45        let ptr = unsafe { alloc(layout) };
46        assert!(!ptr.is_null(), "OOM allocating first growable chunk");
47        Self {
48            chunks: vec![Chunk { ptr, layout }],
49            cursor: 0,
50        }
51    }
52
53    /// Allocate a `Copy` value. Grows on exhaustion.
54    pub fn alloc_copy<T: Copy>(&mut self, value: T) -> &mut T {
55        let layout = Layout::new::<T>();
56        let p = self.alloc_raw(layout);
57        unsafe {
58            ptr::write(p as *mut T, value);
59            &mut *(p as *mut T)
60        }
61    }
62
63    /// Allocate `layout.size()` bytes aligned to `layout.align()`.
64    /// Grows on exhaustion.
65    pub fn alloc_raw(&mut self, layout: Layout) -> *mut u8 {
66        let size = layout.size();
67        let align = layout.align();
68        let last_ptr;
69        let last_size;
70        {
71            let last = self.chunks.last().expect("at least one chunk");
72            last_ptr = last.ptr;
73            last_size = last.layout.size();
74        }
75        let base = last_ptr as usize;
76        let aligned = align_up(base + self.cursor, align) - base;
77        let end = aligned + size;
78        if end <= last_size {
79            self.cursor = end;
80            return unsafe { last_ptr.add(aligned) };
81        }
82        // Grow: 2x the previous chunk size, but at least `size + align`
83        // so the new chunk can definitely serve the pending request.
84        self.grow(size + align);
85        let last_ptr = self.chunks.last().unwrap().ptr;
86        let base = last_ptr as usize;
87        let aligned = align_up(base, align) - base;
88        self.cursor = aligned + size;
89        unsafe { last_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 = 0;
100    }
101
102    /// Rewind every chunk. Keeps only the largest chunk; smaller
103    /// chunks are dropped. Subsequent allocations reuse the kept chunk
104    /// without further grow events for workloads within its size.
105    pub fn reset(&mut self) {
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 = 0;
119    }
120
121    /// Bytes allocated across every retained chunk.
122    pub fn total_capacity(&self) -> usize {
123        self.chunks.iter().map(|c| c.layout.size()).sum()
124    }
125
126    /// Number of chunks currently retained.
127    pub fn chunk_count(&self) -> usize {
128        self.chunks.len()
129    }
130}
131
132impl Default for GrowableBump {
133    fn default() -> Self {
134        Self::new()
135    }
136}
137
138#[cfg(test)]
139#[path = "growable_tests.rs"]
140mod tests;