solar_data_structures/
bump_ext.rs

1use crate::RawThinSlice;
2use bumpalo::Bump;
3use smallvec::SmallVec;
4
5/// Extension trait for [`Bump`].
6#[expect(clippy::mut_from_ref)] // Arena.
7pub trait BumpExt {
8    /// Returns the number of bytes currently in use.
9    fn used_bytes(&self) -> usize;
10
11    /// Allocates a value as a slice of length 1.
12    fn alloc_as_slice<T>(&self, value: T) -> &mut [T];
13
14    /// Allocates an iterator by first collecting it into a (possibly stack-allocated) vector.
15    ///
16    /// Does not collect if the iterator is exact size, meaning `size_hint` returns equal values.
17    fn alloc_from_iter<T>(&self, iter: impl Iterator<Item = T>) -> &mut [T];
18
19    /// Allocates a vector of items on the arena.
20    ///
21    /// NOTE: This method does not drop the values, so you likely want to wrap the result in a
22    /// [`bumpalo::boxed::Box`] if `T: Drop`.
23    fn alloc_vec<T>(&self, values: Vec<T>) -> &mut [T];
24
25    /// Allocates a `SmallVector` of items on the arena.
26    ///
27    /// NOTE: This method does not drop the values, so you likely want to wrap the result in a
28    /// [`bumpalo::boxed::Box`] if `T: Drop`.
29    fn alloc_smallvec<A: smallvec::Array>(&self, values: SmallVec<A>) -> &mut [A::Item];
30
31    /// Allocates an array of items on the arena.
32    ///
33    /// NOTE: This method does not drop the values, so you likely want to wrap the result in a
34    /// [`bumpalo::boxed::Box`] if `T: Drop`.
35    fn alloc_array<T, const N: usize>(&self, values: [T; N]) -> &mut [T];
36
37    /// Allocates a slice of items on the arena and copies them in.
38    ///
39    /// # Safety
40    ///
41    /// If `T: Drop`, the resulting slice must not be wrapped in [`bumpalo::boxed::Box`], unless
42    /// ownership is moved as well, such as through [`alloc_vec`](Self::alloc_vec) and the other
43    /// methods in this trait.
44    unsafe fn alloc_slice_unchecked<'a, T>(&'a self, slice: &[T]) -> &'a mut [T];
45
46    // `RawThinSlice` methods
47
48    /// See [`alloc_as_slice`](Self::alloc_as_slice).
49    fn alloc_as_thin_slice<H, T>(&self, header: H, value: T) -> &mut RawThinSlice<H, T>;
50
51    /// See [`alloc_from_iter`](Self::alloc_from_iter).
52    fn alloc_from_iter_thin<H, T>(
53        &self,
54        header: H,
55        iter: impl Iterator<Item = T>,
56    ) -> &mut RawThinSlice<H, T>;
57
58    /// See [`alloc_vec`](Self::alloc_vec).
59    fn alloc_vec_thin<H, T>(&self, header: H, values: Vec<T>) -> &mut RawThinSlice<H, T>;
60
61    /// See [`alloc_smallvec`](Self::alloc_smallvec).
62    fn alloc_smallvec_thin<H, A: smallvec::Array>(
63        &self,
64        header: H,
65        values: SmallVec<A>,
66    ) -> &mut RawThinSlice<H, A::Item>;
67
68    /// See [`alloc_array`](Self::alloc_array).
69    fn alloc_array_thin<H, T, const N: usize>(
70        &self,
71        header: H,
72        values: [T; N],
73    ) -> &mut RawThinSlice<H, T>;
74
75    /// See [`alloc_slice_copy`](Bump::alloc_slice_copy).
76    fn alloc_thin_slice_copy<'a, H, T: Copy>(
77        &'a self,
78        header: H,
79        src: &[T],
80    ) -> &'a mut RawThinSlice<H, T> {
81        // SAFETY: `T` is `Copy`.
82        unsafe { self.alloc_thin_slice_unchecked(header, src) }
83    }
84
85    /// See [`alloc_slice_unchecked`](Self::alloc_slice_unchecked).
86    #[expect(clippy::missing_safety_doc)]
87    unsafe fn alloc_thin_slice_unchecked<'a, H, T>(
88        &'a self,
89        header: H,
90        src: &[T],
91    ) -> &'a mut RawThinSlice<H, T>;
92}
93
94impl BumpExt for Bump {
95    fn used_bytes(&self) -> usize {
96        // SAFETY: The data is not read, and the arena is not used during the iteration.
97        unsafe { self.iter_allocated_chunks_raw().map(|(_ptr, len)| len).sum::<usize>() }
98    }
99
100    #[inline]
101    fn alloc_as_slice<T>(&self, value: T) -> &mut [T] {
102        std::slice::from_mut(self.alloc(value))
103    }
104
105    #[inline]
106    fn alloc_from_iter<T>(&self, mut iter: impl Iterator<Item = T>) -> &mut [T] {
107        match iter.size_hint() {
108            (min, Some(max)) if min == max => self.alloc_slice_fill_with(min, |_| {
109                iter.next().expect("Iterator supplied too few elements")
110            }),
111            _ => self.alloc_smallvec(SmallVec::<[T; 8]>::from_iter(iter)),
112        }
113    }
114
115    #[inline]
116    fn alloc_vec<T>(&self, mut values: Vec<T>) -> &mut [T] {
117        if values.is_empty() {
118            return &mut [];
119        }
120
121        // SAFETY: The `Vec` is deallocated, but the elements are not dropped.
122        unsafe {
123            let r = self.alloc_slice_unchecked(values.as_slice());
124            values.set_len(0);
125            r
126        }
127    }
128
129    #[inline]
130    fn alloc_smallvec<A: smallvec::Array>(&self, mut values: SmallVec<A>) -> &mut [A::Item] {
131        if values.is_empty() {
132            return &mut [];
133        }
134
135        // SAFETY: See `alloc_vec`.
136        unsafe {
137            let r = self.alloc_slice_unchecked(values.as_slice());
138            values.set_len(0);
139            r
140        }
141    }
142
143    #[inline]
144    fn alloc_array<T, const N: usize>(&self, values: [T; N]) -> &mut [T] {
145        if values.is_empty() {
146            return &mut [];
147        }
148
149        let values = std::mem::ManuallyDrop::new(values);
150        // SAFETY: The values are not dropped.
151        unsafe { self.alloc_slice_unchecked(values.as_slice()) }
152    }
153
154    #[inline]
155    unsafe fn alloc_slice_unchecked<'a, T>(&'a self, src: &[T]) -> &'a mut [T] {
156        // Copied from `alloc_slice_copy`.
157        let layout = std::alloc::Layout::for_value(src);
158        let dst = self.alloc_layout(layout).cast::<T>();
159        unsafe {
160            std::ptr::copy_nonoverlapping(src.as_ptr(), dst.as_ptr(), src.len());
161            std::slice::from_raw_parts_mut(dst.as_ptr(), src.len())
162        }
163    }
164
165    #[inline]
166    fn alloc_as_thin_slice<H, T>(&self, header: H, value: T) -> &mut RawThinSlice<H, T> {
167        let value = std::mem::ManuallyDrop::new(value);
168        unsafe { self.alloc_thin_slice_unchecked(header, std::slice::from_ref(&*value)) }
169    }
170
171    #[inline]
172    fn alloc_from_iter_thin<H, T>(
173        &self,
174        header: H,
175        mut iter: impl Iterator<Item = T>,
176    ) -> &mut RawThinSlice<H, T> {
177        match iter.size_hint() {
178            (min, Some(max)) if min == max => {
179                RawThinSlice::<H, T>::from_arena_with(self, header, min, |_| {
180                    iter.next().expect("Iterator supplied too few elements")
181                })
182            }
183            _ => self.alloc_smallvec_thin(header, SmallVec::<[T; 8]>::from_iter(iter)),
184        }
185    }
186
187    #[inline]
188    fn alloc_vec_thin<H, T>(&self, header: H, mut values: Vec<T>) -> &mut RawThinSlice<H, T> {
189        // SAFETY: The `Vec` is deallocated, but the elements are not dropped.
190        unsafe {
191            let r = self.alloc_thin_slice_unchecked(header, values.as_slice());
192            values.set_len(0);
193            r
194        }
195    }
196
197    #[inline]
198    fn alloc_smallvec_thin<H, A: smallvec::Array>(
199        &self,
200        header: H,
201        mut values: SmallVec<A>,
202    ) -> &mut RawThinSlice<H, A::Item> {
203        // SAFETY: See `alloc_vec_thin`.
204        unsafe {
205            let r = self.alloc_thin_slice_unchecked(header, values.as_slice());
206            values.set_len(0);
207            r
208        }
209    }
210
211    #[inline]
212    fn alloc_array_thin<H, T, const N: usize>(
213        &self,
214        header: H,
215        values: [T; N],
216    ) -> &mut RawThinSlice<H, T> {
217        let values = std::mem::ManuallyDrop::new(values);
218        // SAFETY: The values are not dropped.
219        unsafe { self.alloc_thin_slice_unchecked(header, values.as_slice()) }
220    }
221
222    #[inline]
223    unsafe fn alloc_thin_slice_unchecked<'a, H, T>(
224        &'a self,
225        header: H,
226        src: &[T],
227    ) -> &'a mut RawThinSlice<H, T> {
228        RawThinSlice::from_arena(self, header, src)
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[derive(Debug, PartialEq, Eq)]
237    struct DropBomb(i32, bool);
238    impl DropBomb {
239        fn new(i: i32) -> Self {
240            Self(i, true)
241        }
242        fn defuse(&mut self) {
243            self.1 = false;
244        }
245    }
246    impl Drop for DropBomb {
247        fn drop(&mut self) {
248            if self.1 && !std::thread::panicking() {
249                panic!("boom");
250            }
251        }
252    }
253
254    #[test]
255    fn test_alloc_vec() {
256        let bump = Bump::new();
257        let vec = vec![DropBomb::new(1), DropBomb::new(2), DropBomb::new(3)];
258        let other_vec = vec![DropBomb::new(1), DropBomb::new(2), DropBomb::new(3)];
259        let slice = bump.alloc_vec(vec);
260        assert_eq!(slice, &other_vec[..]);
261        for item in slice {
262            item.defuse();
263        }
264        for mut item in other_vec {
265            item.defuse();
266        }
267    }
268
269    #[test]
270    fn test_alloc_vec_thin() {
271        let bump = Bump::new();
272        let vec = vec![DropBomb::new(1), DropBomb::new(2), DropBomb::new(3)];
273        let other_vec = vec![DropBomb::new(1), DropBomb::new(2), DropBomb::new(3)];
274        let raw_slice = bump.alloc_vec_thin(69usize, vec);
275        assert_eq!(*raw_slice.header(), 69usize);
276        let slice = &mut **raw_slice;
277        assert_eq!(slice, &other_vec[..]);
278        for item in slice {
279            item.defuse();
280        }
281        for mut item in other_vec {
282            item.defuse();
283        }
284    }
285
286    #[test]
287    fn test_alloc_thin_empty() {
288        let bump = Bump::new();
289        let data = Vec::<&'static str>::new();
290        let raw_slice = bump.alloc_vec_thin(69usize, data);
291        assert_eq!(*raw_slice.header(), 69usize);
292        assert!(raw_slice.as_slice().is_empty());
293    }
294}
solar_data_structures/bump_ext.rs

solar_data_structures/
bump_ext.rs
