Skip to main content

bump_scope/bump_box/
slice_initializer.rs

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