bump_scope/traits/
bump_allocator_ext.rs

1#![expect(clippy::missing_safety_doc)]
2
3use core::{alloc::Layout, num::NonZeroUsize, ptr::NonNull};
4
5use crate::{
6    BaseAllocator, Bump, BumpAllocator, BumpAllocatorScope, BumpScope, MinimumAlignment, MutBumpAllocator,
7    MutBumpAllocatorScope, SizedTypeProperties, SupportedMinimumAlignment, WithoutDealloc, WithoutShrink,
8    alloc::AllocError,
9    bump_down,
10    polyfill::non_null,
11    stats::{AnyStats, Stats},
12    traits::assert_implements,
13    up_align_usize_unchecked,
14};
15
16#[cfg(feature = "panic-on-alloc")]
17use crate::{handle_alloc_error, panic_on_error, private::capacity_overflow};
18
19/// An extension trait for [`BumpAllocator`]s.
20///
21/// Its main purpose is to provide methods that are optimized for a certain `T` and error behavior.
22///
23/// It also provides [`stats`] to get a `Bump` specific `Stats` object.
24///
25/// **Note:** This trait is not automatically implemented for all `BumpAllocator`s
26/// because it is meant to provide specialized methods and types for better performance.
27/// A blanket implementation for all `BumpAllocators` would defeat that purpose, at least
28/// until some form of specialization is stabilized.
29///
30/// [`stats`]: BumpAllocatorExt::stats
31pub unsafe trait BumpAllocatorExt: BumpAllocator {
32    /// The type returned by the [stats](BumpAllocatorExt::stats) method.
33    type Stats<'b>: Into<AnyStats<'b>>
34    where
35        Self: 'b;
36
37    /// Returns a type which provides statistics about the memory usage of the bump allocator.
38    fn stats(&self) -> Self::Stats<'_>;
39
40    /// A specialized version of [`allocate`](crate::alloc::Allocator::allocate).
41    ///
42    /// Behaves like the following code:
43    /// ```
44    /// # use core::{alloc::Layout, ptr::NonNull};
45    /// # #[expect(dead_code)]
46    /// # trait MyExt: bump_scope::BumpAllocator {
47    /// #     unsafe fn my_ext_fn(&self, layout: Layout) -> NonNull<u8> {
48    /// self.allocate(layout).unwrap().cast()
49    /// #     }
50    /// # }
51    /// ```
52    ///
53    /// # Panics
54    ///
55    /// Panics if the allocation fails.
56    #[cfg(feature = "panic-on-alloc")]
57    fn allocate_layout(&self, layout: Layout) -> NonNull<u8>;
58
59    /// A specialized version of [`allocate`](crate::alloc::Allocator::allocate).
60    ///
61    /// Behaves like the following code:
62    /// ```
63    /// # use core::{alloc::Layout, ptr::NonNull};
64    /// # use bump_scope::alloc::AllocError;
65    /// # #[expect(dead_code)]
66    /// # trait MyExt: bump_scope::BumpAllocator {
67    /// #     unsafe fn my_ext_fn(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
68    /// Ok(self.allocate(layout)?.cast())
69    /// #     }
70    /// # }
71    /// ```
72    ///
73    /// # Errors
74    ///
75    /// Errors if the allocation fails.
76    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError>;
77
78    /// A specialized version of [`allocate`](crate::alloc::Allocator::allocate).
79    ///
80    /// Behaves like the following code:
81    /// ```
82    /// # use core::{alloc::Layout, ptr::NonNull};
83    /// # type T = i32;
84    /// # #[expect(dead_code)]
85    /// # trait MyExt: bump_scope::BumpAllocator {
86    /// #     unsafe fn my_ext_fn(&self) -> NonNull<T> {
87    /// self.allocate(Layout::new::<T>()).unwrap().cast()
88    /// #     }
89    /// # }
90    /// ```
91    ///
92    /// # Panics
93    ///
94    /// Panics if the allocation fails.
95    #[cfg(feature = "panic-on-alloc")]
96    fn allocate_sized<T>(&self) -> NonNull<T>;
97
98    /// A specialized version of [`allocate`](crate::alloc::Allocator::allocate).
99    ///
100    /// Behaves like the following code:
101    /// ```
102    /// # use core::{alloc::Layout, ptr::NonNull};
103    /// # use bump_scope::alloc::AllocError;
104    /// # type T = i32;
105    /// # #[expect(dead_code)]
106    /// # trait MyExt: bump_scope::BumpAllocator {
107    /// #     unsafe fn my_ext_fn(&self) -> Result<NonNull<T>, AllocError> {
108    /// Ok(self.allocate(Layout::new::<T>())?.cast())
109    /// #     }
110    /// # }
111    /// ```
112    ///
113    /// # Errors
114    ///
115    /// Errors if the allocation fails.
116    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError>;
117
118    /// A specialized version of [`allocate`](crate::alloc::Allocator::allocate).
119    ///
120    /// Behaves like the following code:
121    /// ```
122    /// # use core::{alloc::Layout, ptr::NonNull};
123    /// # type T = i32;
124    /// # #[expect(dead_code)]
125    /// # trait MyExt: bump_scope::BumpAllocator {
126    /// #     unsafe fn my_ext_fn(&self, len: usize) -> NonNull<T> {
127    /// self.allocate(Layout::array::<T>(len).unwrap()).unwrap().cast()
128    /// #     }
129    /// # }
130    /// ```
131    ///
132    /// # Panics
133    ///
134    /// Panics if the allocation fails.
135    #[cfg(feature = "panic-on-alloc")]
136    fn allocate_slice<T>(&self, len: usize) -> NonNull<T>;
137
138    /// A specialized version of [`allocate`](crate::alloc::Allocator::allocate).
139    ///
140    /// Behaves like the following code:
141    /// ```
142    /// # use core::{alloc::Layout, ptr::NonNull};
143    /// # use bump_scope::alloc::AllocError;
144    /// # type T = i32;
145    /// # #[expect(dead_code)]
146    /// # trait MyExt: bump_scope::BumpAllocator {
147    /// #     unsafe fn my_ext_fn(&self, len: usize) -> Result<NonNull<T>, AllocError> {
148    /// Ok(self.allocate(Layout::array::<T>(len).map_err(|_| AllocError)?)?.cast())
149    /// #     }
150    /// # }
151    /// ```
152    ///
153    /// # Errors
154    ///
155    /// Errors if the allocation fails.
156    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError>;
157
158    /// A specialized version of [`shrink`](crate::alloc::Allocator::shrink).
159    ///
160    /// Behaves like the following code except that it returns `None`
161    /// when the allocation remains unchanged and the pointer stays valid.
162    /// ```
163    /// # use core::{alloc::Layout, ptr::NonNull};
164    /// # type T = i32;
165    /// # #[expect(dead_code)]
166    /// # trait MyExt: bump_scope::BumpAllocator {
167    /// #     unsafe fn my_ext_fn(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> NonNull<T> {
168    /// #         unsafe {
169    /// self.shrink(ptr.cast(),
170    ///     Layout::array::<T>(old_len).unwrap_unchecked(),
171    ///     Layout::array::<T>(new_len).unwrap_unchecked(),
172    /// ).unwrap_unchecked().cast()
173    /// #         }
174    /// #     }
175    /// # }
176    /// ```
177    ///
178    /// # Safety
179    ///
180    /// Same safety conditions as for the code above apply.
181    ///
182    /// [shrink]: crate::alloc::Allocator::shrink
183    /// [array]: Layout::array
184    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>>;
185
186    /// A specialized version of [`prepare_allocation`].
187    ///
188    /// Returns a `[T]` of free space in the bump allocator.
189    ///
190    /// [`prepare_allocation`]: crate::BumpAllocator::prepare_allocation
191    ///
192    /// # Panics
193    ///
194    /// Panics if the allocation fails.
195    #[cfg(feature = "panic-on-alloc")]
196    fn prepare_slice_allocation<T>(&self, cap: usize) -> NonNull<[T]>;
197
198    /// A specialized version of [`prepare_allocation`].
199    ///
200    /// Returns a `[T]` of free space in the bump allocator.
201    ///
202    /// [`prepare_allocation`]: crate::BumpAllocator::prepare_allocation
203    ///
204    /// # Errors
205    ///
206    /// Errors if the allocation fails.
207    fn try_prepare_slice_allocation<T>(&self, cap: usize) -> Result<NonNull<[T]>, AllocError>;
208
209    /// A specialized version of [`allocate_prepared`].
210    ///
211    /// Allocates part of the free space returned from a
212    /// <code>([try_](BumpAllocatorExt::try_prepare_slice_allocation))[prepare_slice_allocation](BumpAllocatorExt::prepare_slice_allocation)</code>
213    /// call.
214    ///
215    /// # Safety
216    /// - `ptr..ptr + cap` must be the pointer range returned from
217    ///   <code>([try_](BumpAllocatorExt::try_prepare_slice_allocation))[prepare_slice_allocation](BumpAllocatorExt::prepare_slice_allocation)</code>.
218    /// - no allocation, grow, shrink or deallocate must have taken place since then
219    /// - no resets must have taken place since then
220    /// - `len` must be less than or equal to `cap`
221    ///
222    /// [`allocate_prepared`]: BumpAllocator::allocate_prepared
223    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]>;
224
225    /// A specialized version of [`allocate_prepared_rev`].
226    ///
227    /// Allocates part of the free space returned from a
228    /// <code>([try_](BumpAllocatorExt::try_prepare_slice_allocation))[prepare_slice_allocation](BumpAllocatorExt::prepare_slice_allocation)</code>
229    /// call.
230    ///
231    /// # Safety
232    /// - `ptr - cap..ptr` must be the pointer range returned from
233    ///   <code>([try_](BumpAllocatorExt::try_prepare_slice_allocation))[prepare_slice_allocation](BumpAllocatorExt::prepare_slice_allocation)</code>.
234    /// - no allocation, grow, shrink or deallocate must have taken place since then
235    /// - no resets must have taken place since then
236    /// - `len` must be less than or equal to `cap`
237    ///
238    /// [`allocate_prepared_rev`]: crate::BumpAllocator::allocate_prepared_rev
239    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]>;
240}
241
242assert_implements! {
243    [BumpAllocatorExt + ?Sized]
244
245    dyn BumpAllocator
246    &dyn BumpAllocator
247    &mut dyn BumpAllocator
248
249    dyn BumpAllocatorScope
250    &dyn BumpAllocatorScope
251    &mut dyn BumpAllocatorScope
252
253    dyn MutBumpAllocator
254    &dyn MutBumpAllocator
255    &mut dyn MutBumpAllocator
256
257    dyn MutBumpAllocatorScope
258    &dyn MutBumpAllocatorScope
259    &mut dyn MutBumpAllocatorScope
260}
261
262unsafe impl BumpAllocatorExt for dyn BumpAllocator + '_ {
263    type Stats<'b>
264        = AnyStats<'b>
265    where
266        Self: 'b;
267
268    #[inline(always)]
269    fn stats(&self) -> AnyStats<'_> {
270        self.any_stats()
271    }
272
273    #[inline(always)]
274    #[cfg(feature = "panic-on-alloc")]
275    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
276        allocate_layout(self, layout)
277    }
278
279    #[inline(always)]
280    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
281        try_allocate_layout(self, layout)
282    }
283
284    #[inline(always)]
285    #[cfg(feature = "panic-on-alloc")]
286    fn allocate_sized<T>(&self) -> NonNull<T> {
287        allocate_sized(self)
288    }
289
290    #[inline(always)]
291    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
292        try_allocate_sized(self)
293    }
294
295    #[inline(always)]
296    #[cfg(feature = "panic-on-alloc")]
297    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
298        allocate_slice(self, len)
299    }
300
301    #[inline(always)]
302    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
303        try_allocate_slice(self, len)
304    }
305
306    #[inline(always)]
307    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
308        unsafe { shrink_slice(self, ptr, old_len, new_len) }
309    }
310
311    #[inline(always)]
312    #[cfg(feature = "panic-on-alloc")]
313    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
314        prepare_slice_allocation(self, len)
315    }
316
317    #[inline(always)]
318    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
319        try_prepare_slice_allocation(self, len)
320    }
321
322    #[inline(always)]
323    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
324        unsafe { allocate_prepared_slice(self, ptr, len, cap) }
325    }
326
327    #[inline(always)]
328    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
329        unsafe { allocate_prepared_slice_rev(self, ptr, len, cap) }
330    }
331}
332
333unsafe impl BumpAllocatorExt for dyn MutBumpAllocator + '_ {
334    type Stats<'b>
335        = AnyStats<'b>
336    where
337        Self: 'b;
338
339    #[inline(always)]
340    fn stats(&self) -> AnyStats<'_> {
341        self.any_stats()
342    }
343
344    #[inline(always)]
345    #[cfg(feature = "panic-on-alloc")]
346    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
347        allocate_layout(self, layout)
348    }
349
350    #[inline(always)]
351    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
352        try_allocate_layout(self, layout)
353    }
354
355    #[inline(always)]
356    #[cfg(feature = "panic-on-alloc")]
357    fn allocate_sized<T>(&self) -> NonNull<T> {
358        allocate_sized(self)
359    }
360
361    #[inline(always)]
362    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
363        try_allocate_sized(self)
364    }
365
366    #[inline(always)]
367    #[cfg(feature = "panic-on-alloc")]
368    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
369        allocate_slice(self, len)
370    }
371
372    #[inline(always)]
373    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
374        try_allocate_slice(self, len)
375    }
376
377    #[inline(always)]
378    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
379        unsafe { shrink_slice(self, ptr, old_len, new_len) }
380    }
381
382    #[inline(always)]
383    #[cfg(feature = "panic-on-alloc")]
384    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
385        prepare_slice_allocation(self, len)
386    }
387
388    #[inline(always)]
389    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
390        try_prepare_slice_allocation(self, len)
391    }
392
393    #[inline(always)]
394    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
395        unsafe { allocate_prepared_slice(self, ptr, len, cap) }
396    }
397
398    #[inline(always)]
399    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
400        unsafe { allocate_prepared_slice_rev(self, ptr, len, cap) }
401    }
402}
403
404unsafe impl BumpAllocatorExt for dyn BumpAllocatorScope<'_> + '_ {
405    type Stats<'b>
406        = AnyStats<'b>
407    where
408        Self: 'b;
409
410    #[inline(always)]
411    fn stats(&self) -> AnyStats<'_> {
412        self.any_stats()
413    }
414
415    #[inline(always)]
416    #[cfg(feature = "panic-on-alloc")]
417    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
418        allocate_layout(self, layout)
419    }
420
421    #[inline(always)]
422    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
423        try_allocate_layout(self, layout)
424    }
425
426    #[inline(always)]
427    #[cfg(feature = "panic-on-alloc")]
428    fn allocate_sized<T>(&self) -> NonNull<T> {
429        allocate_sized(self)
430    }
431
432    #[inline(always)]
433    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
434        try_allocate_sized(self)
435    }
436
437    #[inline(always)]
438    #[cfg(feature = "panic-on-alloc")]
439    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
440        allocate_slice(self, len)
441    }
442
443    #[inline(always)]
444    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
445        try_allocate_slice(self, len)
446    }
447
448    #[inline(always)]
449    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
450        unsafe { shrink_slice(self, ptr, old_len, new_len) }
451    }
452
453    #[inline(always)]
454    #[cfg(feature = "panic-on-alloc")]
455    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
456        prepare_slice_allocation(self, len)
457    }
458
459    #[inline(always)]
460    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
461        try_prepare_slice_allocation(self, len)
462    }
463
464    #[inline(always)]
465    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
466        unsafe { allocate_prepared_slice(self, ptr, len, cap) }
467    }
468
469    #[inline(always)]
470    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
471        unsafe { allocate_prepared_slice_rev(self, ptr, len, cap) }
472    }
473}
474
475unsafe impl BumpAllocatorExt for dyn MutBumpAllocatorScope<'_> + '_ {
476    type Stats<'b>
477        = AnyStats<'b>
478    where
479        Self: 'b;
480
481    #[inline(always)]
482    fn stats(&self) -> AnyStats<'_> {
483        self.any_stats()
484    }
485
486    #[inline(always)]
487    #[cfg(feature = "panic-on-alloc")]
488    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
489        allocate_layout(self, layout)
490    }
491
492    #[inline(always)]
493    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
494        try_allocate_layout(self, layout)
495    }
496
497    #[inline(always)]
498    #[cfg(feature = "panic-on-alloc")]
499    fn allocate_sized<T>(&self) -> NonNull<T> {
500        allocate_sized(self)
501    }
502
503    #[inline(always)]
504    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
505        try_allocate_sized(self)
506    }
507
508    #[inline(always)]
509    #[cfg(feature = "panic-on-alloc")]
510    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
511        allocate_slice(self, len)
512    }
513
514    #[inline(always)]
515    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
516        try_allocate_slice(self, len)
517    }
518
519    #[inline(always)]
520    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
521        unsafe { shrink_slice(self, ptr, old_len, new_len) }
522    }
523
524    #[inline(always)]
525    #[cfg(feature = "panic-on-alloc")]
526    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
527        prepare_slice_allocation(self, len)
528    }
529
530    #[inline(always)]
531    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
532        try_prepare_slice_allocation(self, len)
533    }
534
535    #[inline(always)]
536    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
537        unsafe { allocate_prepared_slice(self, ptr, len, cap) }
538    }
539
540    #[inline(always)]
541    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
542        unsafe { allocate_prepared_slice_rev(self, ptr, len, cap) }
543    }
544}
545
546#[inline]
547#[cfg(feature = "panic-on-alloc")]
548fn allocate_layout(bump: impl BumpAllocator, layout: Layout) -> NonNull<u8> {
549    match bump.allocate(layout) {
550        Ok(ptr) => ptr.cast(),
551        Err(AllocError) => handle_alloc_error(layout),
552    }
553}
554
555#[inline]
556fn try_allocate_layout(bump: impl BumpAllocator, layout: Layout) -> Result<NonNull<u8>, AllocError> {
557    match bump.allocate(layout) {
558        Ok(ptr) => Ok(ptr.cast()),
559        Err(err) => Err(err),
560    }
561}
562
563#[inline]
564#[cfg(feature = "panic-on-alloc")]
565fn allocate_sized<T>(bump: impl BumpAllocator) -> NonNull<T> {
566    let layout = Layout::new::<T>();
567
568    match bump.allocate(layout) {
569        Ok(ptr) => ptr.cast(),
570        Err(AllocError) => handle_alloc_error(Layout::new::<T>()),
571    }
572}
573
574#[inline]
575fn try_allocate_sized<T>(bump: impl BumpAllocator) -> Result<NonNull<T>, AllocError> {
576    match bump.allocate(Layout::new::<T>()) {
577        Ok(ptr) => Ok(ptr.cast()),
578        Err(err) => Err(err),
579    }
580}
581
582#[inline]
583#[cfg(feature = "panic-on-alloc")]
584fn allocate_slice<T>(bump: impl BumpAllocator, len: usize) -> NonNull<T> {
585    let Ok(layout) = Layout::array::<T>(len) else {
586        invalid_slice_layout()
587    };
588
589    match bump.allocate(layout) {
590        Ok(ptr) => ptr.cast(),
591        Err(AllocError) => handle_alloc_error(layout),
592    }
593}
594
595#[inline]
596fn try_allocate_slice<T>(bump: impl BumpAllocator, len: usize) -> Result<NonNull<T>, AllocError> {
597    let Ok(layout) = Layout::array::<T>(len) else {
598        return Err(AllocError);
599    };
600
601    match bump.allocate(layout) {
602        Ok(ptr) => Ok(ptr.cast()),
603        Err(err) => Err(err),
604    }
605}
606
607#[inline]
608#[expect(clippy::unnecessary_wraps)]
609unsafe fn shrink_slice<T>(bump: impl BumpAllocator, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
610    unsafe {
611        Some(
612            bump.shrink(
613                ptr.cast(),
614                Layout::array::<T>(old_len).unwrap_unchecked(),
615                Layout::array::<T>(new_len).unwrap_unchecked(),
616            )
617            .unwrap_unchecked()
618            .cast(),
619        )
620    }
621}
622
623fn is_upwards_allocating(bump: &impl BumpAllocator) -> bool {
624    let chunk = bump.checkpoint().chunk;
625    let header = chunk.addr();
626    let end = unsafe { chunk.as_ref() }.end.addr();
627    end > header
628}
629
630#[inline(always)]
631#[cfg(feature = "panic-on-alloc")]
632fn prepare_slice_allocation<T>(bump: impl BumpAllocator, min_cap: usize) -> NonNull<[T]> {
633    let Ok(layout) = Layout::array::<T>(min_cap) else {
634        capacity_overflow()
635    };
636
637    match bump.prepare_allocation(layout) {
638        Ok(range) => {
639            // NB: We can't use `offset_from_unsigned`, because the size is not a multiple of `T`'s.
640            let cap = unsafe { non_null::byte_offset_from_unsigned(range.end, range.start) } / T::SIZE;
641
642            let ptr = if is_upwards_allocating(&bump) {
643                range.start.cast::<T>()
644            } else {
645                unsafe { range.end.cast::<T>().sub(cap) }
646            };
647
648            NonNull::slice_from_raw_parts(ptr.cast(), cap)
649        }
650        Err(AllocError) => handle_alloc_error(layout),
651    }
652}
653
654#[inline(always)]
655fn try_prepare_slice_allocation<T>(bump: impl BumpAllocator, len: usize) -> Result<NonNull<[T]>, AllocError> {
656    let Ok(layout) = Layout::array::<T>(len) else {
657        return Err(AllocError);
658    };
659
660    match bump.prepare_allocation(layout) {
661        Ok(range) => {
662            // NB: We can't use `offset_from_unsigned`, because the size is not a multiple of `T`'s.
663            let cap = unsafe { non_null::byte_offset_from_unsigned(range.end, range.start) } / T::SIZE;
664
665            let ptr = if is_upwards_allocating(&bump) {
666                range.start.cast::<T>()
667            } else {
668                unsafe { range.end.cast::<T>().sub(cap) }
669            };
670
671            Ok(NonNull::slice_from_raw_parts(ptr.cast(), cap))
672        }
673        Err(err) => Err(err),
674    }
675}
676
677#[inline(always)]
678unsafe fn allocate_prepared_slice<T>(bump: impl BumpAllocator, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
679    unsafe {
680        let range = non_null::cast_range(ptr..ptr.add(cap));
681        let layout = Layout::from_size_align_unchecked(core::mem::size_of::<T>() * len, T::ALIGN);
682        let data = bump.allocate_prepared(layout, range).cast();
683        NonNull::slice_from_raw_parts(data, len)
684    }
685}
686
687#[inline(always)]
688unsafe fn allocate_prepared_slice_rev<T>(bump: impl BumpAllocator, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
689    unsafe {
690        let range = non_null::cast_range(ptr.sub(cap)..ptr);
691        let layout = Layout::from_size_align_unchecked(core::mem::size_of::<T>() * len, T::ALIGN);
692        let data = bump.allocate_prepared_rev(layout, range).cast();
693        NonNull::slice_from_raw_parts(data, len)
694    }
695}
696
697unsafe impl<B: BumpAllocatorExt + ?Sized> BumpAllocatorExt for &B {
698    type Stats<'b>
699        = B::Stats<'b>
700    where
701        Self: 'b;
702
703    #[inline(always)]
704    fn stats(&self) -> Self::Stats<'_> {
705        B::stats(self)
706    }
707
708    #[inline(always)]
709    #[cfg(feature = "panic-on-alloc")]
710    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
711        B::allocate_layout(self, layout)
712    }
713
714    #[inline(always)]
715    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
716        B::try_allocate_layout(self, layout)
717    }
718
719    #[inline(always)]
720    #[cfg(feature = "panic-on-alloc")]
721    fn allocate_sized<T>(&self) -> NonNull<T> {
722        B::allocate_sized(self)
723    }
724
725    #[inline(always)]
726    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
727        B::try_allocate_sized(self)
728    }
729
730    #[inline(always)]
731    #[cfg(feature = "panic-on-alloc")]
732    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
733        B::allocate_slice(self, len)
734    }
735
736    #[inline(always)]
737    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
738        B::try_allocate_slice(self, len)
739    }
740
741    #[inline(always)]
742    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
743        unsafe { B::shrink_slice(self, ptr, old_len, new_len) }
744    }
745
746    #[inline(always)]
747    #[cfg(feature = "panic-on-alloc")]
748    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
749        B::prepare_slice_allocation(self, len)
750    }
751
752    #[inline(always)]
753    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
754        B::try_prepare_slice_allocation(self, len)
755    }
756
757    #[inline(always)]
758    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
759        unsafe { B::allocate_prepared_slice(self, ptr, len, cap) }
760    }
761
762    #[inline(always)]
763    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
764        unsafe { B::allocate_prepared_slice_rev(self, ptr, len, cap) }
765    }
766}
767
768unsafe impl<B: BumpAllocatorExt + ?Sized> BumpAllocatorExt for &mut B {
769    type Stats<'b>
770        = B::Stats<'b>
771    where
772        Self: 'b;
773
774    #[inline(always)]
775    fn stats(&self) -> Self::Stats<'_> {
776        B::stats(self)
777    }
778
779    #[inline(always)]
780    #[cfg(feature = "panic-on-alloc")]
781    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
782        B::allocate_layout(self, layout)
783    }
784
785    #[inline(always)]
786    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
787        B::try_allocate_layout(self, layout)
788    }
789
790    #[inline(always)]
791    #[cfg(feature = "panic-on-alloc")]
792    fn allocate_sized<T>(&self) -> NonNull<T> {
793        B::allocate_sized(self)
794    }
795
796    #[inline(always)]
797    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
798        B::try_allocate_sized(self)
799    }
800
801    #[inline(always)]
802    #[cfg(feature = "panic-on-alloc")]
803    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
804        B::allocate_slice(self, len)
805    }
806
807    #[inline(always)]
808    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
809        B::try_allocate_slice(self, len)
810    }
811
812    #[inline(always)]
813    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
814        unsafe { B::shrink_slice(self, ptr, old_len, new_len) }
815    }
816
817    #[inline(always)]
818    #[cfg(feature = "panic-on-alloc")]
819    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
820        B::prepare_slice_allocation(self, len)
821    }
822
823    #[inline(always)]
824    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
825        B::try_prepare_slice_allocation(self, len)
826    }
827
828    #[inline(always)]
829    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
830        unsafe { B::allocate_prepared_slice(self, ptr, len, cap) }
831    }
832
833    #[inline(always)]
834    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
835        unsafe { B::allocate_prepared_slice_rev(self, ptr, len, cap) }
836    }
837}
838
839unsafe impl<B: BumpAllocatorExt> BumpAllocatorExt for WithoutDealloc<B> {
840    type Stats<'b>
841        = B::Stats<'b>
842    where
843        Self: 'b;
844
845    #[inline(always)]
846    fn stats(&self) -> Self::Stats<'_> {
847        B::stats(&self.0)
848    }
849
850    #[inline(always)]
851    #[cfg(feature = "panic-on-alloc")]
852    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
853        B::allocate_layout(&self.0, layout)
854    }
855
856    #[inline(always)]
857    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
858        B::try_allocate_layout(&self.0, layout)
859    }
860
861    #[inline(always)]
862    #[cfg(feature = "panic-on-alloc")]
863    fn allocate_sized<T>(&self) -> NonNull<T> {
864        B::allocate_sized(&self.0)
865    }
866
867    #[inline(always)]
868    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
869        B::try_allocate_sized(&self.0)
870    }
871
872    #[inline(always)]
873    #[cfg(feature = "panic-on-alloc")]
874    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
875        B::allocate_slice(&self.0, len)
876    }
877
878    #[inline(always)]
879    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
880        B::try_allocate_slice(&self.0, len)
881    }
882
883    #[inline(always)]
884    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
885        unsafe { B::shrink_slice(&self.0, ptr, old_len, new_len) }
886    }
887
888    #[inline(always)]
889    #[cfg(feature = "panic-on-alloc")]
890    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
891        B::prepare_slice_allocation(&self.0, len)
892    }
893
894    #[inline(always)]
895    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
896        B::try_prepare_slice_allocation(&self.0, len)
897    }
898
899    #[inline(always)]
900    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
901        unsafe { B::allocate_prepared_slice(&self.0, ptr, len, cap) }
902    }
903
904    #[inline(always)]
905    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
906        unsafe { B::allocate_prepared_slice_rev(&self.0, ptr, len, cap) }
907    }
908}
909
910unsafe impl<B: BumpAllocatorExt> BumpAllocatorExt for WithoutShrink<B> {
911    type Stats<'b>
912        = B::Stats<'b>
913    where
914        Self: 'b;
915
916    #[inline(always)]
917    fn stats(&self) -> Self::Stats<'_> {
918        B::stats(&self.0)
919    }
920
921    #[inline(always)]
922    #[cfg(feature = "panic-on-alloc")]
923    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
924        B::allocate_layout(&self.0, layout)
925    }
926
927    #[inline(always)]
928    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
929        B::try_allocate_layout(&self.0, layout)
930    }
931
932    #[inline(always)]
933    #[cfg(feature = "panic-on-alloc")]
934    fn allocate_sized<T>(&self) -> NonNull<T> {
935        B::allocate_sized(&self.0)
936    }
937
938    #[inline(always)]
939    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
940        B::try_allocate_sized(&self.0)
941    }
942
943    #[inline(always)]
944    #[cfg(feature = "panic-on-alloc")]
945    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
946        B::allocate_slice(&self.0, len)
947    }
948
949    #[inline(always)]
950    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
951        B::try_allocate_slice(&self.0, len)
952    }
953
954    #[inline(always)]
955    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
956        _ = (ptr, old_len, new_len);
957        None
958    }
959
960    #[inline(always)]
961    #[cfg(feature = "panic-on-alloc")]
962    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
963        B::prepare_slice_allocation(&self.0, len)
964    }
965
966    #[inline(always)]
967    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
968        B::try_prepare_slice_allocation(&self.0, len)
969    }
970
971    #[inline(always)]
972    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
973        unsafe { B::allocate_prepared_slice(&self.0, ptr, len, cap) }
974    }
975
976    #[inline(always)]
977    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
978        unsafe { B::allocate_prepared_slice_rev(&self.0, ptr, len, cap) }
979    }
980}
981
982unsafe impl<A, const MIN_ALIGN: usize, const UP: bool, const GUARANTEED_ALLOCATED: bool, const DEALLOCATES: bool>
983    BumpAllocatorExt for BumpScope<'_, A, MIN_ALIGN, UP, GUARANTEED_ALLOCATED, DEALLOCATES>
984where
985    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
986    A: BaseAllocator<GUARANTEED_ALLOCATED>,
987{
988    type Stats<'b>
989        = Stats<'b, A, UP, GUARANTEED_ALLOCATED>
990    where
991        Self: 'b;
992
993    #[inline(always)]
994    fn stats(&self) -> Self::Stats<'_> {
995        BumpScope::stats(self)
996    }
997
998    #[inline(always)]
999    #[cfg(feature = "panic-on-alloc")]
1000    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
1001        self.alloc_layout(layout)
1002    }
1003
1004    #[inline(always)]
1005    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
1006        self.try_alloc_layout(layout)
1007    }
1008
1009    #[inline(always)]
1010    #[cfg(feature = "panic-on-alloc")]
1011    fn allocate_sized<T>(&self) -> NonNull<T> {
1012        panic_on_error(self.do_alloc_sized())
1013    }
1014
1015    #[inline(always)]
1016    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
1017        self.do_alloc_sized()
1018    }
1019
1020    #[inline(always)]
1021    #[cfg(feature = "panic-on-alloc")]
1022    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
1023        panic_on_error(self.do_alloc_slice(len))
1024    }
1025
1026    #[inline(always)]
1027    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
1028        self.do_alloc_slice(len)
1029    }
1030
1031    #[inline]
1032    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
1033        if !DEALLOCATES {
1034            return None;
1035        }
1036
1037        let old_ptr = ptr.cast::<u8>();
1038        let old_size = old_len * T::SIZE; // we already allocated that amount so this can't overflow
1039        let new_size = new_len * T::SIZE; // its less than the capacity so this can't overflow
1040
1041        // Adapted from `Allocator::shrink`.
1042        unsafe {
1043            let is_last_and_allocated = self.chunk.get().is_allocated()
1044                && if UP {
1045                    old_ptr.as_ptr().add(old_size) == self.chunk.get().pos().as_ptr()
1046                } else {
1047                    old_ptr == self.chunk.get().pos()
1048                };
1049
1050            // if that's not the last allocation, there is nothing we can do
1051            if !is_last_and_allocated {
1052                return None;
1053            }
1054
1055            if UP {
1056                let end = old_ptr.addr().get() + new_size;
1057
1058                // Up-aligning a pointer inside a chunk by `MIN_ALIGN` never overflows.
1059                let new_pos = up_align_usize_unchecked(end, MIN_ALIGN);
1060
1061                self.chunk.get().guaranteed_allocated_unchecked().set_pos_addr(new_pos);
1062                Some(old_ptr.cast())
1063            } else {
1064                let old_addr = old_ptr.addr();
1065                let old_addr_old_end = NonZeroUsize::new_unchecked(old_addr.get() + old_size);
1066
1067                let new_addr = bump_down(old_addr_old_end, new_size, T::ALIGN.max(MIN_ALIGN));
1068                let new_addr = NonZeroUsize::new_unchecked(new_addr);
1069                let old_addr_new_end = NonZeroUsize::new_unchecked(old_addr.get() + new_size);
1070
1071                let new_ptr = old_ptr.with_addr(new_addr);
1072                let overlaps = old_addr_new_end > new_addr;
1073
1074                if overlaps {
1075                    old_ptr.copy_to(new_ptr, new_size);
1076                } else {
1077                    old_ptr.copy_to_nonoverlapping(new_ptr, new_size);
1078                }
1079
1080                self.chunk.get().guaranteed_allocated_unchecked().set_pos(new_ptr);
1081                Some(new_ptr.cast())
1082            }
1083        }
1084    }
1085
1086    #[inline(always)]
1087    #[cfg(feature = "panic-on-alloc")]
1088    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
1089        panic_on_error(BumpScope::generic_prepare_slice_allocation::<_, T>(self, len))
1090    }
1091
1092    #[inline(always)]
1093    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
1094        BumpScope::generic_prepare_slice_allocation::<_, T>(self, len)
1095    }
1096
1097    #[inline(always)]
1098    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
1099        unsafe { BumpScope::use_prepared_slice_allocation(self, ptr, len, cap) }
1100    }
1101
1102    #[inline(always)]
1103    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
1104        unsafe { BumpScope::use_prepared_slice_allocation_rev(self, ptr, len, cap) }
1105    }
1106}
1107
1108unsafe impl<A, const MIN_ALIGN: usize, const UP: bool, const GUARANTEED_ALLOCATED: bool, const DEALLOCATES: bool>
1109    BumpAllocatorExt for Bump<A, MIN_ALIGN, UP, GUARANTEED_ALLOCATED, DEALLOCATES>
1110where
1111    MinimumAlignment<MIN_ALIGN>: SupportedMinimumAlignment,
1112    A: BaseAllocator<GUARANTEED_ALLOCATED>,
1113{
1114    type Stats<'b>
1115        = Stats<'b, A, UP, GUARANTEED_ALLOCATED>
1116    where
1117        Self: 'b;
1118
1119    #[inline(always)]
1120    fn stats(&self) -> Self::Stats<'_> {
1121        self.as_scope().stats()
1122    }
1123
1124    #[inline(always)]
1125    #[cfg(feature = "panic-on-alloc")]
1126    fn allocate_layout(&self, layout: Layout) -> NonNull<u8> {
1127        self.as_scope().allocate_layout(layout)
1128    }
1129
1130    #[inline(always)]
1131    fn try_allocate_layout(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
1132        self.as_scope().try_allocate_layout(layout)
1133    }
1134
1135    #[inline(always)]
1136    #[cfg(feature = "panic-on-alloc")]
1137    fn allocate_sized<T>(&self) -> NonNull<T> {
1138        self.as_scope().allocate_sized()
1139    }
1140
1141    #[inline(always)]
1142    fn try_allocate_sized<T>(&self) -> Result<NonNull<T>, AllocError> {
1143        self.as_scope().try_allocate_sized()
1144    }
1145
1146    #[inline(always)]
1147    #[cfg(feature = "panic-on-alloc")]
1148    fn allocate_slice<T>(&self, len: usize) -> NonNull<T> {
1149        self.as_scope().allocate_slice(len)
1150    }
1151
1152    #[inline(always)]
1153    fn try_allocate_slice<T>(&self, len: usize) -> Result<NonNull<T>, AllocError> {
1154        self.as_scope().try_allocate_slice(len)
1155    }
1156
1157    #[inline(always)]
1158    unsafe fn shrink_slice<T>(&self, ptr: NonNull<T>, old_len: usize, new_len: usize) -> Option<NonNull<T>> {
1159        unsafe { self.as_scope().shrink_slice(ptr, old_len, new_len) }
1160    }
1161
1162    #[inline(always)]
1163    #[cfg(feature = "panic-on-alloc")]
1164    fn prepare_slice_allocation<T>(&self, len: usize) -> NonNull<[T]> {
1165        self.as_scope().prepare_slice_allocation(len)
1166    }
1167
1168    #[inline(always)]
1169    fn try_prepare_slice_allocation<T>(&self, len: usize) -> Result<NonNull<[T]>, AllocError> {
1170        self.as_scope().try_prepare_slice_allocation(len)
1171    }
1172
1173    #[inline(always)]
1174    unsafe fn allocate_prepared_slice<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
1175        unsafe { self.as_scope().use_prepared_slice_allocation(ptr, len, cap) }
1176    }
1177
1178    #[inline(always)]
1179    unsafe fn allocate_prepared_slice_rev<T>(&self, ptr: NonNull<T>, len: usize, cap: usize) -> NonNull<[T]> {
1180        unsafe { self.as_scope().use_prepared_slice_allocation_rev(ptr, len, cap) }
1181    }
1182}
1183
1184#[cold]
1185#[inline(never)]
1186#[cfg(feature = "panic-on-alloc")]
1187pub(crate) const fn invalid_slice_layout() -> ! {
1188    panic!("invalid slice layout");
1189}
bump_scope/traits/bump_allocator_ext.rs

bump_scope/traits/
bump_allocator_ext.rs
