Skip to main content

subms_arena_allocator/features/
aligned.rs

1//! `AlignedBump`: bump arena exposing an `alloc_aligned(size, align)`
2//! convenience over the base `Layout`-driven raw API.
3//!
4//! The base `Bump::alloc_raw` already accepts an arbitrary `Layout`
5//! and respects its alignment. This module exists as a focused
6//! shape for callers who think in (bytes, alignment) pairs:
7//! cache-line aligned scratch buffers for SIMD, page-aligned scratch
8//! for DMA, etc.
9//!
10//! Fixed-capacity, single chunk. Panics on OOM; use `try_alloc_aligned`
11//! for the fallible form.
12
13use std::alloc::{Layout, alloc, dealloc};
14use std::slice;
15
16use crate::align_up;
17
18/// Bump arena exposing explicit per-allocation alignment.
19pub struct AlignedBump {
20    ptr: *mut u8,
21    layout: Layout,
22    cursor: usize,
23}
24
25impl AlignedBump {
26    /// New arena with the given capacity. The backing buffer itself is
27    /// allocated 64-byte aligned so cache-line requests within
28    /// `capacity` always succeed.
29    pub fn with_capacity(capacity: usize) -> Self {
30        let capacity = capacity.max(64);
31        // 64-byte chunk alignment so the first cache-line request
32        // costs zero padding.
33        let layout = Layout::from_size_align(capacity, 64).expect("layout");
34        let ptr = unsafe { alloc(layout) };
35        assert!(!ptr.is_null(), "OOM allocating aligned arena chunk");
36        Self {
37            ptr,
38            layout,
39            cursor: 0,
40        }
41    }
42
43    /// Allocate `size` bytes aligned to `align` (must be a power of two).
44    /// Panics if the request doesn't fit.
45    pub fn alloc_aligned(&mut self, size: usize, align: usize) -> &mut [u8] {
46        let cursor = self.cursor;
47        let cap = self.layout.size();
48        match self.try_alloc_aligned(size, align) {
49            Some(s) => s,
50            None => panic!(
51                "AlignedBump out of capacity: cursor={cursor} cap={cap} size={size} align={align}",
52            ),
53        }
54    }
55
56    /// Fallible aligned alloc. Returns `None` if the request doesn't fit.
57    pub fn try_alloc_aligned(&mut self, size: usize, align: usize) -> Option<&mut [u8]> {
58        assert!(
59            align.is_power_of_two(),
60            "align must be power of two: {align}"
61        );
62        let base = self.ptr as usize;
63        let aligned = align_up(base + self.cursor, align) - base;
64        let end = aligned.checked_add(size)?;
65        if end > self.layout.size() {
66            return None;
67        }
68        self.cursor = end;
69        unsafe {
70            let p = self.ptr.add(aligned);
71            Some(slice::from_raw_parts_mut(p, size))
72        }
73    }
74
75    /// Rewind. Buffer retained for reuse.
76    pub fn reset(&mut self) {
77        self.cursor = 0;
78    }
79
80    /// Total capacity.
81    pub fn capacity(&self) -> usize {
82        self.layout.size()
83    }
84
85    /// Bytes used so far.
86    pub fn used(&self) -> usize {
87        self.cursor
88    }
89}
90
91impl Drop for AlignedBump {
92    fn drop(&mut self) {
93        unsafe { dealloc(self.ptr, self.layout) };
94    }
95}
96
97#[cfg(test)]
98#[path = "aligned_tests.rs"]
99mod tests;