solar_data_structures/
bump_ext.rs

1use bumpalo::Bump;
2use smallvec::SmallVec;
3
4/// Extension trait for [`Bump`].
5#[allow(clippy::mut_from_ref)] // Arena.
6pub trait BumpExt {
7    /// Returns the number of bytes currently in use.
8    fn used_bytes(&self) -> usize;
9
10    /// Allocates a value as a slice of length 1.
11    fn alloc_as_slice<T>(&self, value: T) -> &mut [T];
12
13    /// Allocates an iterator by first collecting it into a (possibly stack-allocated) vector.
14    ///
15    /// Does not collect if the iterator is exact size, meaning `size_hint` returns equal values.
16    fn alloc_from_iter<T>(&self, iter: impl Iterator<Item = T>) -> &mut [T];
17
18    /// Allocates a vector of items on the arena.
19    ///
20    /// NOTE: This method does not drop the values, so you likely want to wrap the result in a
21    /// [`bumpalo::boxed::Box`] if `T: Drop`.
22    fn alloc_vec<T>(&self, values: Vec<T>) -> &mut [T];
23
24    /// Allocates a `SmallVector` of items on the arena.
25    ///
26    /// NOTE: This method does not drop the values, so you likely want to wrap the result in a
27    /// [`bumpalo::boxed::Box`] if `T: Drop`.
28    fn alloc_smallvec<A: smallvec::Array>(&self, values: SmallVec<A>) -> &mut [A::Item];
29
30    /// Allocates an array of items on the arena.
31    ///
32    /// NOTE: This method does not drop the values, so you likely want to wrap the result in a
33    /// [`bumpalo::boxed::Box`] if `T: Drop`.
34    fn alloc_array<T, const N: usize>(&self, values: [T; N]) -> &mut [T];
35
36    /// Allocates a slice of items on the arena and copies them in.
37    ///
38    /// # Safety
39    ///
40    /// If `T: Drop`, the resulting slice must not be wrapped in [`bumpalo::boxed::Box`], unless
41    /// ownership is moved as well, such as through [`alloc_vec`](Self::alloc_vec) and the other
42    /// methods in this trait.
43    unsafe fn alloc_slice_unchecked<'a, T>(&'a self, slice: &[T]) -> &'a mut [T];
44}
45
46impl BumpExt for Bump {
47    fn used_bytes(&self) -> usize {
48        // SAFETY: The data is not read, and the arena is not used during the iteration.
49        unsafe { self.iter_allocated_chunks_raw().map(|(_ptr, len)| len).sum::<usize>() }
50    }
51
52    #[inline]
53    fn alloc_as_slice<T>(&self, value: T) -> &mut [T] {
54        std::slice::from_mut(self.alloc(value))
55    }
56
57    #[inline]
58    fn alloc_from_iter<T>(&self, mut iter: impl Iterator<Item = T>) -> &mut [T] {
59        match iter.size_hint() {
60            (min, Some(max)) if min == max => self.alloc_slice_fill_with(min, |_| {
61                iter.next().expect("Iterator supplied too few elements")
62            }),
63            _ => self.alloc_smallvec(SmallVec::<[T; 8]>::from_iter(iter)),
64        }
65    }
66
67    #[inline]
68    fn alloc_vec<T>(&self, mut values: Vec<T>) -> &mut [T] {
69        if values.is_empty() {
70            return &mut [];
71        }
72
73        // SAFETY: The `Vec` is deallocated, but the elements are not dropped.
74        unsafe {
75            let r = self.alloc_slice_unchecked(values.as_slice());
76            values.set_len(0);
77            r
78        }
79    }
80
81    #[inline]
82    fn alloc_smallvec<A: smallvec::Array>(&self, mut values: SmallVec<A>) -> &mut [A::Item] {
83        if values.is_empty() {
84            return &mut [];
85        }
86
87        // SAFETY: See `alloc_vec`.
88        unsafe {
89            let r = self.alloc_slice_unchecked(values.as_slice());
90            values.set_len(0);
91            r
92        }
93    }
94
95    #[inline]
96    fn alloc_array<T, const N: usize>(&self, values: [T; N]) -> &mut [T] {
97        if values.is_empty() {
98            return &mut [];
99        }
100
101        let values = std::mem::ManuallyDrop::new(values);
102        // SAFETY: See `alloc_vec`.
103        unsafe { self.alloc_slice_unchecked(values.as_slice()) }
104    }
105
106    #[inline]
107    unsafe fn alloc_slice_unchecked<'a, T>(&'a self, src: &[T]) -> &'a mut [T] {
108        // Copied from `alloc_slice_copy`.
109        let layout = std::alloc::Layout::for_value(src);
110        let dst = self.alloc_layout(layout).cast::<T>();
111        unsafe {
112            std::ptr::copy_nonoverlapping(src.as_ptr(), dst.as_ptr(), src.len());
113            std::slice::from_raw_parts_mut(dst.as_ptr(), src.len())
114        }
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[derive(Debug, PartialEq, Eq)]
123    struct DropBomb(i32, bool);
124    impl DropBomb {
125        fn new(i: i32) -> Self {
126            Self(i, true)
127        }
128        fn defuse(&mut self) {
129            self.1 = false;
130        }
131    }
132    impl Drop for DropBomb {
133        fn drop(&mut self) {
134            if self.1 && !std::thread::panicking() {
135                panic!("boom");
136            }
137        }
138    }
139
140    #[test]
141    fn test_alloc_vec() {
142        let bump = Bump::new();
143        let vec = vec![DropBomb::new(1), DropBomb::new(2), DropBomb::new(3)];
144        let other_vec = vec![DropBomb::new(1), DropBomb::new(2), DropBomb::new(3)];
145        let slice = bump.alloc_vec(vec);
146        assert_eq!(slice, &other_vec[..]);
147        for item in slice {
148            item.defuse();
149        }
150        for mut item in other_vec {
151            item.defuse();
152        }
153    }
154}
solar_data_structures/bump_ext.rs

solar_data_structures/
bump_ext.rs
