bump_scope/
bump_pool.rs

1use std::{
2    alloc::Layout,
3    mem::{self, ManuallyDrop},
4    ops::{Deref, DerefMut},
5    sync::{Mutex, MutexGuard, PoisonError},
6    vec::Vec,
7};
8
9use crate::{
10    Bump, BumpScope, ErrorBehavior, MinimumAlignment, SupportedMinimumAlignment,
11    alloc::{AllocError, Allocator},
12    maybe_default_allocator,
13};
14
15#[cfg(feature = "panic-on-alloc")]
16use crate::panic_on_error;
17
18macro_rules! make_pool {
19    ($($allocator_parameter:tt)*) => {
20        /// A pool of bump allocators.
21        ///
22        /// This type allows bump allocations in parallel, with the allocations' lifetimes tied to the pool.
23        ///
24        /// # Examples
25        ///
26        /// Using `BumpPool` with parallel iterators from [`rayon`](https://docs.rs/rayon):
27        /// ```
28        /// # use bump_scope::BumpPool;
29        /// # use rayon::prelude::{ParallelIterator, IntoParallelIterator};
30        /// # if cfg!(miri) { return } // rayon violates strict-provenance :(
31        /// #
32        /// let mut pool: BumpPool = BumpPool::new();
33        ///
34        /// let strings: Vec<&str> = (0..1000)
35        ///     .into_par_iter()
36        ///     .map_init(|| pool.get(), |bump, i| {
37        ///         // do some expensive work
38        ///         bump.alloc_fmt(format_args!("{i}")).into_ref()
39        ///     })
40        ///     .collect();
41        ///
42        /// dbg!(&strings);
43        ///
44        /// pool.reset();
45        ///
46        /// // memory of the strings is freed, trying to access `strings` will result in a lifetime error
47        /// // dbg!(&strings);
48        /// ```
49        ///
50        /// Using `BumpPool` with [`std::thread::scope`]:
51        /// ```
52        /// # use bump_scope::BumpPool;
53        /// let pool: BumpPool = BumpPool::new();
54        /// let (sender, receiver) = std::sync::mpsc::sync_channel(10);
55        ///
56        /// std::thread::scope(|s| {
57        ///     s.spawn(|| {
58        ///         let bump = pool.get();
59        ///         let string = bump.alloc_str("Hello");
60        ///         sender.send(string).unwrap();
61        ///         drop(sender);
62        ///     });
63        ///
64        ///     s.spawn(|| {
65        ///         for string in receiver {
66        ///             assert_eq!(string, "Hello");
67        ///         }
68        ///     });
69        /// });
70        /// ```
71        ///
72        #[doc(alias = "Herd")]
73        #[derive(Debug)]
74        pub struct BumpPool<
75            $($allocator_parameter)*,
76            const MIN_ALIGN: usize = 1,
77            const UP: bool = true,
78        > where
79            MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
80            A: Allocator,
81        {
82            bumps: Mutex<Vec<Bump<A, MIN_ALIGN, UP, true, true>>>,
83            allocator: A,
84        }
85    };
86}
87
88maybe_default_allocator!(make_pool);
89
90impl<A, const MIN_ALIGN: usize, const UP: bool> Default for BumpPool<A, MIN_ALIGN, UP>
91where
92    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
93    A: Allocator + Default,
94{
95    fn default() -> Self {
96        Self {
97            bumps: Mutex::default(),
98            allocator: Default::default(),
99        }
100    }
101}
102
103impl<A, const MIN_ALIGN: usize, const UP: bool> BumpPool<A, MIN_ALIGN, UP>
104where
105    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
106    A: Allocator + Default,
107{
108    /// Constructs a new `BumpPool`.
109    #[inline]
110    #[must_use]
111    pub fn new() -> Self {
112        Self::default()
113    }
114}
115
116impl<A, const MIN_ALIGN: usize, const UP: bool> BumpPool<A, MIN_ALIGN, UP>
117where
118    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
119    A: Allocator,
120{
121    /// Constructs a new `BumpPool` with the provided allocator.
122    #[inline]
123    #[must_use]
124    pub const fn new_in(allocator: A) -> Self {
125        Self {
126            bumps: Mutex::new(Vec::new()),
127            allocator,
128        }
129    }
130
131    /// [Resets](Bump::reset) all `Bump`s in this pool.
132    pub fn reset(&mut self) {
133        for bump in self.bumps() {
134            bump.reset();
135        }
136    }
137
138    /// Returns the vector of `Bump`s.
139    pub fn bumps(&mut self) -> &mut Vec<Bump<A, MIN_ALIGN, UP, true, true>> {
140        self.bumps.get_mut().unwrap_or_else(PoisonError::into_inner)
141    }
142
143    fn lock(&self) -> MutexGuard<'_, Vec<Bump<A, MIN_ALIGN, UP, true, true>>> {
144        self.bumps.lock().unwrap_or_else(PoisonError::into_inner)
145    }
146}
147
148impl<A, const MIN_ALIGN: usize, const UP: bool> BumpPool<A, MIN_ALIGN, UP>
149where
150    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
151    A: Allocator + Clone,
152{
153    /// Borrows a bump allocator from the pool.
154    /// With this `BumpPoolGuard` you can make allocations that live for as long as the pool lives.
155    ///
156    /// If this needs to create a new `Bump`, it will be constructed by calling <code>Bump::[new]\()</code>.
157    ///
158    /// [new]: Bump::new
159    ///
160    /// # Panics
161    /// Panics if the allocation fails.
162    #[must_use]
163    #[inline(always)]
164    #[cfg(feature = "panic-on-alloc")]
165    pub fn get(&self) -> BumpPoolGuard<'_, A, MIN_ALIGN, UP> {
166        panic_on_error(self.generic_get())
167    }
168
169    /// Borrows a bump allocator from the pool.
170    /// With this `BumpPoolGuard` you can make allocations that live for as long as the pool lives.
171    ///
172    /// If this needs to create a new `Bump`, it will be constructed by calling <code>Bump::[try_new]\()</code>.
173    ///
174    /// [try_new]: Bump::try_new
175    ///
176    /// # Errors
177    /// Errors if the allocation fails.
178    #[inline(always)]
179    pub fn try_get(&self) -> Result<BumpPoolGuard<'_, A, MIN_ALIGN, UP>, AllocError> {
180        self.generic_get()
181    }
182
183    pub(crate) fn generic_get<E: ErrorBehavior>(&self) -> Result<BumpPoolGuard<'_, A, MIN_ALIGN, UP>, E> {
184        let bump = match self.lock().pop() {
185            Some(bump) => bump,
186            None => Bump::generic_new_in(self.allocator.clone())?,
187        };
188
189        #[expect(deprecated)]
190        Ok(BumpPoolGuard {
191            pool: self,
192            bump: ManuallyDrop::new(bump),
193        })
194    }
195
196    /// Borrows a bump allocator from the pool.
197    /// With this `BumpPoolGuard` you can make allocations that live for as long as the pool lives.
198    ///
199    ///  If this needs to create a new `Bump`, it will be constructed by calling <code>Bump::[with_size]\(size)</code>.
200    ///
201    /// [with_size]: Bump::with_size
202    ///
203    /// # Panics
204    /// Panics if the allocation fails.
205    #[must_use]
206    #[inline(always)]
207    #[cfg(feature = "panic-on-alloc")]
208    pub fn get_with_size(&self, size: usize) -> BumpPoolGuard<'_, A, MIN_ALIGN, UP> {
209        panic_on_error(self.generic_get_with_size(size))
210    }
211
212    /// Borrows a bump allocator from the pool.
213    /// With this `BumpPoolGuard` you can make allocations that live for as long as the pool lives.
214    ///
215    ///  If this needs to create a new `Bump`, it will be constructed by calling <code>Bump::[try_with_size]\(size)</code>.
216    ///
217    /// [try_with_size]: Bump::try_with_size
218    ///
219    /// # Errors
220    /// Errors if the allocation fails.
221    #[inline(always)]
222    pub fn try_get_with_size(&self, size: usize) -> Result<BumpPoolGuard<'_, A, MIN_ALIGN, UP>, AllocError> {
223        self.generic_get_with_size(size)
224    }
225
226    pub(crate) fn generic_get_with_size<E: ErrorBehavior>(
227        &self,
228        size: usize,
229    ) -> Result<BumpPoolGuard<'_, A, MIN_ALIGN, UP>, E> {
230        let bump = match self.lock().pop() {
231            Some(bump) => bump,
232            None => Bump::generic_with_size_in(size, self.allocator.clone())?,
233        };
234
235        #[expect(deprecated)]
236        Ok(BumpPoolGuard {
237            pool: self,
238            bump: ManuallyDrop::new(bump),
239        })
240    }
241
242    /// Borrows a bump allocator from the pool.
243    /// With this `BumpPoolGuard` you can make allocations that live for as long as the pool lives.
244    ///
245    /// If this needs to create a new `Bump`, it will be constructed by calling <code>Bump::[with_capacity]\(layout)</code>.
246    ///
247    /// [with_capacity]: Bump::with_capacity
248    ///
249    /// # Panics
250    /// Panics if the allocation fails.
251    #[must_use]
252    #[inline(always)]
253    #[cfg(feature = "panic-on-alloc")]
254    pub fn get_with_capacity(&self, layout: Layout) -> BumpPoolGuard<'_, A, MIN_ALIGN, UP> {
255        panic_on_error(self.generic_get_with_capacity(layout))
256    }
257
258    /// Borrows a bump allocator from the pool.
259    /// With this `BumpPoolGuard` you can make allocations that live for as long as the pool lives.
260    ///
261    ///  If this needs to create a new `Bump`, it will be constructed by calling <code>Bump::[try_with_capacity]\(layout)</code>.
262    ///
263    /// [try_with_capacity]: Bump::try_with_capacity
264    ///
265    /// # Errors
266    /// Errors if the allocation fails.
267    #[inline(always)]
268    pub fn try_get_with_capacity(&self, layout: Layout) -> Result<BumpPoolGuard<'_, A, MIN_ALIGN, UP>, AllocError> {
269        self.generic_get_with_capacity(layout)
270    }
271
272    pub(crate) fn generic_get_with_capacity<E: ErrorBehavior>(
273        &self,
274        layout: Layout,
275    ) -> Result<BumpPoolGuard<'_, A, MIN_ALIGN, UP>, E> {
276        let bump = match self.lock().pop() {
277            Some(bump) => bump,
278            None => Bump::generic_with_capacity_in(layout, self.allocator.clone())?,
279        };
280
281        #[expect(deprecated)]
282        Ok(BumpPoolGuard {
283            pool: self,
284            bump: ManuallyDrop::new(bump),
285        })
286    }
287}
288
289macro_rules! make_pool_guard {
290    ($($allocator_parameter:tt)*) => {
291
292        /// This is a wrapper around [`Bump`] that mutably derefs to a [`BumpScope`] and returns its [`Bump`] back to the [`BumpPool`] on drop.
293        #[derive(Debug)]
294        pub struct BumpPoolGuard<
295            'a,
296            $($allocator_parameter)*,
297            const MIN_ALIGN: usize = 1,
298            const UP: bool = true,
299        > where
300            MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
301            A: Allocator,
302        {
303            bump: ManuallyDrop<Bump<A, MIN_ALIGN, UP, true, true>>,
304
305            #[doc(hidden)]
306            #[deprecated = "Swapping the pool can lead to Undefined Behavior due to dangling pointers, use `pool()` instead!"]
307            pub pool: &'a BumpPool<A, MIN_ALIGN, UP>,
308        }
309    };
310}
311
312maybe_default_allocator!(make_pool_guard);
313
314impl<'a, A, const MIN_ALIGN: usize, const UP: bool> BumpPoolGuard<'a, A, MIN_ALIGN, UP>
315where
316    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
317    A: Allocator,
318{
319    /// The [`BumpPool`], this [`BumpPoolGuard`] was created from.
320    pub fn pool(&self) -> &'a BumpPool<A, MIN_ALIGN, UP> {
321        #[expect(deprecated)]
322        self.pool
323    }
324}
325
326impl<'a, A, const MIN_ALIGN: usize, const UP: bool> Deref for BumpPoolGuard<'a, A, MIN_ALIGN, UP>
327where
328    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
329    A: Allocator,
330{
331    type Target = BumpScope<'a, A, MIN_ALIGN, UP, true, true>;
332
333    #[inline(always)]
334    fn deref(&self) -> &Self::Target {
335        unsafe { transmute_lifetime(self.bump.as_scope()) }
336    }
337}
338
339impl<A, const MIN_ALIGN: usize, const UP: bool> DerefMut for BumpPoolGuard<'_, A, MIN_ALIGN, UP>
340where
341    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
342    A: Allocator,
343{
344    #[inline(always)]
345    fn deref_mut(&mut self) -> &mut Self::Target {
346        unsafe { transmute_lifetime_mut(self.bump.as_mut_scope()) }
347    }
348}
349
350impl<A, const MIN_ALIGN: usize, const UP: bool> Drop for BumpPoolGuard<'_, A, MIN_ALIGN, UP>
351where
352    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
353    A: Allocator,
354{
355    fn drop(&mut self) {
356        let bump = unsafe { ManuallyDrop::take(&mut self.bump) };
357        #[expect(deprecated)]
358        self.pool.lock().push(bump);
359    }
360}
361
362// This exists as a "safer" transmute that only transmutes the `'a` lifetime parameter.
363#[expect(clippy::elidable_lifetime_names)]
364unsafe fn transmute_lifetime<'from, 'to, 'b, A, const MIN_ALIGN: usize, const UP: bool>(
365    scope: &'b BumpScope<'from, A, MIN_ALIGN, UP, true, true>,
366) -> &'b BumpScope<'to, A, MIN_ALIGN, UP, true, true> {
367    unsafe { mem::transmute(scope) }
368}
369
370// This exists as a "safer" transmute that only transmutes the `'a` lifetime parameter.
371#[expect(clippy::elidable_lifetime_names)]
372unsafe fn transmute_lifetime_mut<'from, 'to, 'b, A, const MIN_ALIGN: usize, const UP: bool>(
373    scope: &'b mut BumpScope<'from, A, MIN_ALIGN, UP, true, true>,
374) -> &'b mut BumpScope<'to, A, MIN_ALIGN, UP, true, true> {
375    unsafe { mem::transmute(scope) }
376}
bump_scope/bump_pool.rs

bump_scope/
bump_pool.rs
