vec_strings/
small_array_box.rs

1use std::mem::{ManuallyDrop, MaybeUninit};
2use std::ptr::NonNull;
3use std::slice::{from_raw_parts, from_raw_parts_mut};
4
5use std::iter::IntoIterator;
6use std::iter::{ExactSizeIterator, Iterator};
7
8use std::fmt::{self, Debug};
9use std::ops::{Deref, DerefMut};
10
11use std::cmp::{Eq, PartialEq};
12
13pub(crate) union SmallArrayBoxInner<T, const INLINE_LEN: usize> {
14    ptr: NonNull<T>,
15    pub(crate) inline_storage: ManuallyDrop<[MaybeUninit<T>; INLINE_LEN]>,
16}
17
18/// * `INLINE_LEN` - Number of elements that can be stored inline.
19pub struct SmallArrayBox<T, const INLINE_LEN: usize> {
20    pub(crate) storage: SmallArrayBoxInner<T, INLINE_LEN>,
21    pub(crate) len: usize,
22}
23
24unsafe impl<T: Send, const INLINE_LEN: usize> Send for SmallArrayBox<T, INLINE_LEN> {}
25unsafe impl<T: Sync, const INLINE_LEN: usize> Sync for SmallArrayBox<T, INLINE_LEN> {}
26
27impl<T, const INLINE_LEN: usize> Default for SmallArrayBox<T, INLINE_LEN> {
28    fn default() -> Self {
29        Self::new_empty()
30    }
31}
32
33impl<T, const INLINE_LEN: usize> From<Box<[T]>> for SmallArrayBox<T, INLINE_LEN> {
34    fn from(boxed: Box<[T]>) -> Self {
35        Self::from_box(boxed)
36    }
37}
38
39impl<T, const INLINE_LEN: usize> From<Vec<T>> for SmallArrayBox<T, INLINE_LEN> {
40    fn from(vec: Vec<T>) -> Self {
41        if vec.len() <= INLINE_LEN {
42            Self::new(vec)
43        } else {
44            vec.into_boxed_slice().into()
45        }
46    }
47}
48
49impl<T: Clone, const INLINE_LEN: usize> From<&[T]> for SmallArrayBox<T, INLINE_LEN> {
50    fn from(slice: &[T]) -> Self {
51        Self::new(slice.iter().cloned())
52    }
53}
54
55impl<T: Clone, const INLINE_LEN: usize> Clone for SmallArrayBox<T, INLINE_LEN> {
56    fn clone(&self) -> Self {
57        Self::new(self.iter().cloned())
58    }
59}
60
61impl<T, const INLINE_LEN: usize> SmallArrayBox<T, INLINE_LEN> {
62    pub(crate) fn uninit_inline_storage() -> Self {
63        Self {
64            storage: SmallArrayBoxInner {
65                // Safety:
66                //
67                // It is safe because the array contains `MaybeUninit<T>`.
68                inline_storage: ManuallyDrop::new(unsafe { MaybeUninit::uninit().assume_init() }),
69            },
70            len: 0,
71        }
72    }
73
74    pub const fn new_empty() -> Self {
75        Self {
76            storage: SmallArrayBoxInner {
77                ptr: NonNull::dangling(),
78            },
79            len: 0,
80        }
81    }
82
83    pub fn new<I>(iter: impl IntoIterator<IntoIter = I>) -> Self
84    where
85        I: Iterator<Item = T> + ExactSizeIterator,
86    {
87        let iter = iter.into_iter();
88
89        let len = iter.len();
90
91        if len <= INLINE_LEN {
92            let mut this = Self::uninit_inline_storage();
93
94            let inline_storage = unsafe { this.storage.inline_storage.deref_mut() };
95
96            iter.zip(inline_storage).for_each(|(src, dst)| {
97                *dst = MaybeUninit::new(src);
98            });
99
100            this.len = len;
101
102            this
103        } else {
104            let vec: Vec<T> = iter.collect();
105            let array_ptr = Box::into_raw(vec.into_boxed_slice());
106            let slice = unsafe { &mut *array_ptr };
107            let ptr = unsafe { NonNull::new_unchecked(slice.as_mut_ptr()) };
108
109            Self {
110                storage: SmallArrayBoxInner { ptr },
111                len,
112            }
113        }
114    }
115
116    pub fn from_box(boxed: Box<[T]>) -> Self {
117        let len = boxed.len();
118
119        if len <= INLINE_LEN {
120            let vec: Vec<T> = boxed.into();
121            Self::new(vec)
122        } else {
123            let array_ptr = Box::into_raw(boxed);
124            let slice = unsafe { &mut *array_ptr };
125            let ptr = unsafe { NonNull::new_unchecked(slice.as_mut_ptr()) };
126
127            debug_assert_eq!(slice.len(), len);
128
129            Self {
130                storage: SmallArrayBoxInner { ptr },
131                len,
132            }
133        }
134    }
135
136    pub fn into_boxed_slice(self) -> Box<[T]> {
137        let len = self.len;
138
139        let mut this = ManuallyDrop::new(self);
140
141        if len <= INLINE_LEN {
142            let inline_storage = unsafe { &mut this.storage.inline_storage };
143            let mut vec = Vec::with_capacity(len);
144
145            for elem in inline_storage[..len].iter_mut() {
146                let ptr = elem.as_mut_ptr();
147                vec.push(unsafe { ptr.read() });
148            }
149
150            debug_assert_eq!(vec.len(), len);
151
152            vec.into_boxed_slice()
153        } else {
154            let ptr = unsafe { this.storage.ptr }.as_ptr();
155            let slice = unsafe { from_raw_parts_mut(ptr, len) };
156            unsafe { Box::from_raw(slice) }
157        }
158    }
159}
160
161impl<T, const INLINE_LEN: usize> Deref for SmallArrayBox<T, INLINE_LEN> {
162    type Target = [T];
163
164    fn deref(&self) -> &Self::Target {
165        let len = self.len;
166
167        if len <= INLINE_LEN {
168            let inline_storage = unsafe { self.storage.inline_storage.deref() };
169            unsafe { &*(&inline_storage[..len] as *const _ as *const [T]) }
170        } else {
171            let ptr = unsafe { self.storage.ptr }.as_ptr();
172            unsafe { from_raw_parts(ptr, len) }
173        }
174    }
175}
176
177impl<T, const INLINE_LEN: usize> DerefMut for SmallArrayBox<T, INLINE_LEN> {
178    fn deref_mut(&mut self) -> &mut Self::Target {
179        let len = self.len;
180
181        if len <= INLINE_LEN {
182            let inline_storage = unsafe { self.storage.inline_storage.deref_mut() };
183            unsafe { &mut *(&mut inline_storage[..len] as *mut _ as *mut [T]) }
184        } else {
185            let ptr = unsafe { self.storage.ptr }.as_ptr();
186            unsafe { from_raw_parts_mut(ptr, len) }
187        }
188    }
189}
190
191impl<T, const INLINE_LEN: usize> Drop for SmallArrayBox<T, INLINE_LEN> {
192    fn drop(&mut self) {
193        let len = self.len;
194
195        if len <= INLINE_LEN {
196            let inline_storage = unsafe { self.storage.inline_storage.deref_mut() };
197
198            inline_storage[..len].iter_mut().for_each(|elem| {
199                let ptr = elem.as_mut_ptr();
200                unsafe {
201                    ptr.drop_in_place();
202                }
203            });
204        } else {
205            let ptr = unsafe { self.storage.ptr }.as_ptr();
206            let slice = unsafe { from_raw_parts_mut(ptr, len) };
207            drop(unsafe { Box::from_raw(slice) });
208        }
209    }
210}
211
212impl<T: Debug, const INLINE_LEN: usize> Debug for SmallArrayBox<T, INLINE_LEN> {
213    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
214        write!(f, "{:#?}", self.deref())
215    }
216}
217
218impl<T: PartialEq, const INLINE_LEN: usize> PartialEq for SmallArrayBox<T, INLINE_LEN> {
219    fn eq(&self, other: &Self) -> bool {
220        self.deref().eq(other.deref())
221    }
222}
223
224impl<T: Eq, const INLINE_LEN: usize> Eq for SmallArrayBox<T, INLINE_LEN> {}
225
226#[cfg(test)]
227mod tests {
228    type SmallArrayBox = super::SmallArrayBox<u8, 8>;
229
230    use std::ops::{Deref, DerefMut};
231    use std::ptr;
232
233    fn assert_ptr_eq(x: *const [u8], y: *const [u8]) {
234        assert!(ptr::eq(x, y));
235    }
236
237    #[test]
238    fn test_new_empty() {
239        let mut empty_array = SmallArrayBox::new_empty();
240
241        let empty: &[u8] = &[];
242
243        assert_eq!(empty_array.deref(), empty);
244        assert_eq!(empty_array.deref_mut(), empty);
245        assert_ptr_eq(empty_array.deref(), empty_array.deref_mut());
246
247        let boxed = empty_array.into_boxed_slice();
248
249        assert_eq!(&*boxed, empty);
250    }
251
252    #[test]
253    fn test_new() {
254        let vec: Vec<u8> = (0..100).collect();
255
256        for len in 0..vec.len() {
257            let slice = &vec[..len];
258
259            let mut array = SmallArrayBox::new(slice.iter().copied());
260
261            assert_eq!(array.deref(), slice);
262            assert_eq!(array.deref_mut(), slice);
263            assert_ptr_eq(array.deref(), array.deref_mut());
264
265            let boxed = array.into_boxed_slice();
266
267            assert_eq!(&*boxed, slice);
268        }
269    }
270
271    #[test]
272    fn test_from_box() {
273        let vec: Vec<u8> = (0..100).collect();
274
275        for len in 0..vec.len() {
276            let slice = &vec[..len];
277
278            let vec: Vec<u8> = slice.to_vec();
279
280            let mut array = SmallArrayBox::from_box(vec.into_boxed_slice());
281
282            assert_eq!(array.deref(), slice);
283            assert_eq!(array.deref_mut(), slice);
284            assert_ptr_eq(array.deref(), array.deref_mut());
285
286            let boxed = array.into_boxed_slice();
287
288            assert_eq!(&*boxed, slice);
289        }
290    }
291}