bump_scope/bump_box/
slice_initializer.rs

1use core::{
2    marker::PhantomData,
3    mem::{self, ManuallyDrop, MaybeUninit},
4    ptr::NonNull,
5};
6
7use crate::{
8    BumpBox, SizedTypeProperties,
9    polyfill::{non_null, pointer},
10};
11
12/// Allows for initializing a `BumpBox<[MaybeUninit<T>]>` by pushing values.
13/// On drop, this will drop the values that have been pushed so far.
14pub(crate) struct BumpBoxSliceInitializer<'a, T> {
15    pos: NonNull<T>,
16
17    start: NonNull<T>,
18    end: NonNull<T>, // if T is a ZST this is ptr + len
19
20    /// First field marks the lifetime.
21    /// Second field marks ownership over T. (<https://doc.rust-lang.org/nomicon/phantom-data.html#generic-parameters-and-drop-checking>)
22    marker: PhantomData<(&'a (), T)>,
23}
24
25impl<T> Drop for BumpBoxSliceInitializer<'_, T> {
26    fn drop(&mut self) {
27        unsafe {
28            let to_drop_len = self.init_len();
29            let to_drop = NonNull::slice_from_raw_parts(self.start, to_drop_len);
30            to_drop.as_ptr().drop_in_place();
31        }
32    }
33}
34
35impl<'a, T> BumpBoxSliceInitializer<'a, T> {
36    #[inline(always)]
37    pub(crate) fn new(slice: BumpBox<'a, [MaybeUninit<T>]>) -> Self {
38        if T::IS_ZST {
39            let start = NonNull::dangling();
40            let end = unsafe { non_null::wrapping_byte_add(start, slice.len()) };
41
42            return Self {
43                pos: start,
44                start,
45                end,
46                marker: PhantomData,
47            };
48        }
49
50        let len = slice.len();
51        let slice = slice.into_raw();
52
53        unsafe {
54            let start = slice.cast::<T>();
55            let end = start.add(len);
56
57            Self {
58                pos: start,
59                start,
60                end,
61                marker: PhantomData,
62            }
63        }
64    }
65
66    #[inline(always)]
67    fn init_len(&self) -> usize {
68        if T::IS_ZST {
69            self.pos.addr().get().wrapping_sub(self.start.addr().get())
70        } else {
71            unsafe { non_null::offset_from_unsigned(self.pos, self.start) }
72        }
73    }
74
75    #[inline(always)]
76    fn len(&self) -> usize {
77        if T::IS_ZST {
78            self.end.addr().get().wrapping_sub(self.start.addr().get())
79        } else {
80            unsafe { non_null::offset_from_unsigned(self.end, self.start) }
81        }
82    }
83
84    #[inline(always)]
85    pub(crate) fn is_full(&self) -> bool {
86        self.pos == self.end
87    }
88
89    #[inline(always)]
90    pub(crate) fn push(&mut self, value: T) {
91        self.push_with(|| value);
92    }
93
94    #[inline(always)]
95    pub(crate) fn push_with(&mut self, f: impl FnOnce() -> T) {
96        assert!(!self.is_full());
97        unsafe { self.push_with_unchecked(f) }
98    }
99
100    #[inline(always)]
101    pub(crate) unsafe fn push_unchecked(&mut self, value: T) {
102        unsafe { self.push_with_unchecked(|| value) };
103    }
104
105    #[inline(always)]
106    pub(crate) unsafe fn push_with_unchecked(&mut self, f: impl FnOnce() -> T) {
107        debug_assert!(!self.is_full());
108
109        unsafe {
110            if T::IS_ZST {
111                mem::forget(f());
112                self.pos = non_null::wrapping_byte_add(self.pos, 1);
113            } else {
114                pointer::write_with(self.pos.as_ptr(), f);
115                self.pos = self.pos.add(1);
116            }
117        }
118    }
119
120    #[inline(always)]
121    pub(crate) fn into_init(self) -> BumpBox<'a, [T]> {
122        assert!(self.is_full());
123        unsafe { self.into_init_unchecked() }
124    }
125
126    #[inline(always)]
127    pub(crate) unsafe fn into_init_unchecked(self) -> BumpBox<'a, [T]> {
128        unsafe {
129            let this = ManuallyDrop::new(self);
130            debug_assert!(this.is_full());
131            let len = this.len();
132            let slice = NonNull::slice_from_raw_parts(this.start, len);
133            BumpBox::from_raw(slice)
134        }
135    }
136}