memlink_msdk/
arena.rs

1//! Bump-pointer arena allocator for temporary module allocations.
2//!
3//! Provides fast allocation for temporary data that lives only during
4//! module call execution. All memory is reclaimed at once via reset().
5
6use std::marker::PhantomData;
7use std::ptr::{self, NonNull};
8use std::sync::atomic::{AtomicUsize, Ordering};
9
10pub struct Arena {
11    base: NonNull<u8>,
12    capacity: usize,
13    used: AtomicUsize,
14    _marker: PhantomData<*mut u8>,
15}
16
17unsafe impl Send for Arena {}
18unsafe impl Sync for Arena {}
19
20impl Arena {
21    pub unsafe fn new(base: *mut u8, capacity: usize) -> Self {
22        debug_assert!(!base.is_null(), "Arena base pointer must not be null");
23        debug_assert!(capacity > 0, "Arena capacity must be > 0");
24
25        Arena {
26            base: NonNull::new_unchecked(base),
27            capacity,
28            used: AtomicUsize::new(0),
29            _marker: PhantomData,
30        }
31    }
32
33    #[allow(clippy::mut_from_ref)]
34    pub fn alloc<T>(&self) -> Option<&mut T> {
35        let size = std::mem::size_of::<T>();
36        let align = std::mem::align_of::<T>();
37
38        let offset = self.reserve_with_alignment(size, align)?;
39
40        Some(unsafe { &mut *(self.base.as_ptr().add(offset) as *mut T) })
41    }
42
43    #[allow(clippy::mut_from_ref)]
44    pub fn alloc_with<T>(&self, value: T) -> Option<&mut T> {
45        let slot = self.alloc::<T>()?;
46        unsafe { ptr::write(slot, value) };
47        Some(slot)
48    }
49
50    #[allow(clippy::mut_from_ref)]
51    pub fn alloc_bytes(&self, len: usize) -> Option<&mut [u8]> {
52        if len == 0 {
53            return Some(&mut []);
54        }
55
56        let offset = self.reserve(len)?;
57
58        Some(unsafe {
59            std::slice::from_raw_parts_mut(self.base.as_ptr().add(offset), len)
60        })
61    }
62
63    pub fn base_ptr(&self) -> *mut u8 {
64        self.base.as_ptr()
65    }
66
67    pub fn capacity(&self) -> usize {
68        self.capacity
69    }
70
71    pub fn used(&self) -> usize {
72        self.used.load(Ordering::Relaxed)
73    }
74
75    pub fn remaining(&self) -> usize {
76        self.capacity.saturating_sub(self.used())
77    }
78
79    pub fn usage(&self) -> f32 {
80        self.used() as f32 / self.capacity as f32
81    }
82
83    pub fn reset(&self) {
84        self.used.store(0, Ordering::Relaxed);
85    }
86
87    fn reserve(&self, size: usize) -> Option<usize> {
88        if size == 0 {
89            return Some(self.used.load(Ordering::Relaxed));
90        }
91
92        loop {
93            let current = self.used.load(Ordering::Relaxed);
94            let new_used = current.checked_add(size)?;
95
96            if new_used > self.capacity {
97                return None;
98            }
99
100            match self.used.compare_exchange(
101                current,
102                new_used,
103                Ordering::AcqRel,
104                Ordering::Relaxed,
105            ) {
106                Ok(_) => return Some(current),
107                Err(_) => continue,
108            }
109        }
110    }
111
112    fn reserve_with_alignment(&self, size: usize, align: usize) -> Option<usize> {
113        if size == 0 {
114            return Some(self.used.load(Ordering::Relaxed));
115        }
116
117        loop {
118            let current = self.used.load(Ordering::Relaxed);
119
120            let aligned = (current + align - 1) & !(align - 1);
121            let new_used = aligned.checked_add(size)?;
122
123            if new_used > self.capacity {
124                return None;
125            }
126
127            match self.used.compare_exchange(
128                current,
129                new_used,
130                Ordering::AcqRel,
131                Ordering::Relaxed,
132            ) {
133                Ok(_) => return Some(aligned),
134                Err(_) => continue,
135            }
136        }
137    }
138}
139
140impl Drop for Arena {
141    fn drop(&mut self) {
142        self.reset();
143    }
144}
memlink_msdk/arena.rs

memlink_msdk/
arena.rs
