polars_utils/
idx_vec.rs

1use std::fmt::{Debug, Formatter};
2use std::mem::ManuallyDrop;
3use std::ops::{Deref, DerefMut};
4
5use crate::index::{IdxSize, NonZeroIdxSize};
6
7pub type IdxVec = UnitVec<IdxSize>;
8
9union PointerOrValue<T> {
10    ptr: *mut T,
11    value: ManuallyDrop<T>,
12}
13
14/// A type logically equivalent to `Vec<T>`, but which does not do a
15/// memory allocation until at least two elements have been pushed, storing the
16/// first element inside the UnitVec directly.
17///
18/// Uses IdxSize internally to store lengths, will panic if trying to reserve
19/// for more elements.
20pub struct UnitVec<T> {
21    len: IdxSize,
22    capacity: NonZeroIdxSize,
23    data: PointerOrValue<T>,
24}
25
26unsafe impl<T: Send + Sync> Send for UnitVec<T> {}
27unsafe impl<T: Send + Sync> Sync for UnitVec<T> {}
28
29impl<T> UnitVec<T> {
30    #[inline(always)]
31    fn data_ptr_mut(&mut self) -> *mut T {
32        if self.is_inline() {
33            unsafe { &mut *self.data.value }
34        } else {
35            unsafe { self.data.ptr }
36        }
37    }
38
39    #[inline(always)]
40    fn data_ptr(&self) -> *const T {
41        if self.is_inline() {
42            unsafe { &*self.data.value }
43        } else {
44            unsafe { self.data.ptr }
45        }
46    }
47
48    #[inline]
49    pub fn new() -> Self {
50        Self {
51            len: 0,
52            capacity: NonZeroIdxSize::new(1).unwrap(),
53            data: PointerOrValue {
54                ptr: std::ptr::null_mut(),
55            },
56        }
57    }
58
59    #[inline(always)]
60    pub fn is_inline(&self) -> bool {
61        self.capacity.get() == 1
62    }
63
64    #[inline(always)]
65    pub fn len(&self) -> usize {
66        self.len as usize
67    }
68
69    #[inline(always)]
70    pub fn is_empty(&self) -> bool {
71        self.len == 0
72    }
73
74    #[inline(always)]
75    pub fn capacity(&self) -> usize {
76        self.capacity.get() as usize
77    }
78
79    #[inline(always)]
80    pub fn clear(&mut self) {
81        if std::mem::needs_drop::<T>() {
82            while self.len > 0 {
83                self.pop();
84            }
85        } else {
86            self.len = 0;
87        }
88    }
89
90    #[inline(always)]
91    pub fn push(&mut self, idx: T) {
92        if self.len == self.capacity.get() {
93            self.reserve(1);
94        }
95
96        unsafe { self.push_unchecked(idx) }
97    }
98
99    #[inline(always)]
100    /// # Safety
101    /// Caller must ensure that `UnitVec` has enough capacity.
102    pub unsafe fn push_unchecked(&mut self, idx: T) {
103        unsafe {
104            self.data_ptr_mut().add(self.len as usize).write(idx);
105            self.len += 1;
106        }
107    }
108
109    #[inline]
110    pub fn pop(&mut self) -> Option<T> {
111        if self.len == 0 {
112            None
113        } else {
114            unsafe {
115                self.len -= 1;
116                Some(self.data_ptr().add(self.len as usize).read())
117            }
118        }
119    }
120
121    #[cold]
122    #[inline(never)]
123    pub fn reserve(&mut self, additional: usize) {
124        let new_len = self
125            .len
126            .checked_add(additional.try_into().unwrap())
127            .unwrap();
128        if new_len > self.capacity.get() {
129            let double = self.capacity.get() * 2;
130            self.realloc(double.max(new_len).max(8));
131        }
132    }
133
134    /// # Panics
135    /// Panics if `new_cap <= 1` or `new_cap < self.len`
136    fn realloc(&mut self, mut new_cap: IdxSize) {
137        assert!(new_cap > 1 && new_cap >= self.len);
138        unsafe {
139            let mut me = std::mem::ManuallyDrop::new(Vec::with_capacity(new_cap as usize));
140            new_cap = me.capacity().try_into().unwrap();
141            let buffer = me.as_mut_ptr();
142            std::ptr::copy(self.data_ptr(), buffer, self.len as usize);
143            self.dealloc();
144            self.data = PointerOrValue { ptr: buffer };
145            self.capacity = NonZeroIdxSize::new(new_cap).unwrap();
146        }
147    }
148
149    unsafe fn dealloc(&mut self) {
150        unsafe {
151            if !self.is_inline() {
152                drop(Vec::from_raw_parts(
153                    self.data.ptr.cast::<ManuallyDrop<T>>(),
154                    self.len as usize,
155                    self.capacity(),
156                ));
157            }
158        }
159    }
160
161    pub fn with_capacity(capacity: usize) -> Self {
162        if capacity <= 1 {
163            Self::new()
164        } else {
165            let mut me = std::mem::ManuallyDrop::new(Vec::with_capacity(capacity));
166            let cap = me.capacity().try_into().unwrap();
167            let ptr = me.as_mut_ptr();
168            Self {
169                len: 0,
170                capacity: NonZeroIdxSize::new(cap).unwrap(),
171                data: PointerOrValue { ptr },
172            }
173        }
174    }
175
176    #[inline]
177    pub fn iter(&self) -> std::slice::Iter<'_, T> {
178        self.as_slice().iter()
179    }
180
181    #[inline]
182    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
183        self.as_mut_slice().iter_mut()
184    }
185
186    #[inline]
187    pub fn as_slice(&self) -> &[T] {
188        self.as_ref()
189    }
190
191    #[inline]
192    pub fn as_mut_slice(&mut self) -> &mut [T] {
193        self.as_mut()
194    }
195}
196
197impl<T: Copy> UnitVec<T> {
198    pub fn retain(&mut self, mut f: impl FnMut(T) -> bool) {
199        let mut i = 0;
200        for j in 0..self.len() {
201            if f(self[j]) {
202                self[i] = self[j];
203                i += 1;
204            }
205        }
206
207        if i == 0 {
208            *self = Self::new();
209        } else if i == 1 {
210            *self = Self::from_slice(&[self[0]]);
211        } else {
212            self.len = i as IdxSize;
213        }
214    }
215}
216
217impl<T: Clone> UnitVec<T> {
218    pub fn from_slice(sl: &[T]) -> Self {
219        if sl.len() <= 1 {
220            let mut new = UnitVec::new();
221            if let Some(v) = sl.first() {
222                new.push(v.clone())
223            }
224            new
225        } else {
226            sl.to_vec().into()
227        }
228    }
229}
230
231impl<T> Extend<T> for UnitVec<T> {
232    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
233        let iter = iter.into_iter();
234        self.reserve(iter.size_hint().0);
235        for v in iter {
236            self.push(v)
237        }
238    }
239}
240
241impl<T> Drop for UnitVec<T> {
242    fn drop(&mut self) {
243        self.clear();
244        unsafe { self.dealloc() }
245    }
246}
247
248impl<T: Clone> Clone for UnitVec<T> {
249    fn clone(&self) -> Self {
250        Self::from_iter(self.iter().cloned())
251    }
252}
253
254impl<T: Debug> Debug for UnitVec<T> {
255    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
256        write!(f, "UnitVec: {:?}", self.as_slice())
257    }
258}
259
260impl<T> Default for UnitVec<T> {
261    fn default() -> Self {
262        Self::new()
263    }
264}
265
266impl<T> Deref for UnitVec<T> {
267    type Target = [T];
268
269    #[inline(always)]
270    fn deref(&self) -> &Self::Target {
271        self.as_slice()
272    }
273}
274
275impl<T> DerefMut for UnitVec<T> {
276    #[inline(always)]
277    fn deref_mut(&mut self) -> &mut Self::Target {
278        self.as_mut_slice()
279    }
280}
281
282impl<T> AsRef<[T]> for UnitVec<T> {
283    #[inline(always)]
284    fn as_ref(&self) -> &[T] {
285        unsafe { std::slice::from_raw_parts(self.data_ptr(), self.len as usize) }
286    }
287}
288
289impl<T> AsMut<[T]> for UnitVec<T> {
290    #[inline(always)]
291    fn as_mut(&mut self) -> &mut [T] {
292        unsafe { std::slice::from_raw_parts_mut(self.data_ptr_mut(), self.len as usize) }
293    }
294}
295
296impl<T: PartialEq> PartialEq for UnitVec<T> {
297    fn eq(&self, other: &Self) -> bool {
298        self.as_slice() == other.as_slice()
299    }
300}
301
302impl<T: Eq> Eq for UnitVec<T> {}
303
304impl<T> FromIterator<T> for UnitVec<T> {
305    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
306        let mut iter = iter.into_iter();
307
308        let Some(first) = iter.next() else {
309            return Self::new();
310        };
311
312        let Some(second) = iter.next() else {
313            let mut out = Self::new();
314            out.push(first);
315            return out;
316        };
317
318        let mut vec = Vec::with_capacity(iter.size_hint().0 + 2);
319        vec.push(first);
320        vec.push(second);
321        vec.extend(iter);
322        Self::from(vec)
323    }
324}
325
326impl<T> IntoIterator for UnitVec<T> {
327    type Item = T;
328
329    type IntoIter = IntoIter<T>;
330
331    fn into_iter(mut self) -> Self::IntoIter {
332        if self.is_inline() {
333            IntoIter::Inline(self.pop().into_iter())
334        } else {
335            IntoIter::External(Vec::from(self).into_iter())
336        }
337    }
338}
339
340pub enum IntoIter<T> {
341    Inline(std::option::IntoIter<T>),
342    External(std::vec::IntoIter<T>),
343}
344
345impl<T> Iterator for IntoIter<T> {
346    type Item = T;
347
348    fn next(&mut self) -> Option<Self::Item> {
349        match self {
350            IntoIter::Inline(it) => it.next(),
351            IntoIter::External(it) => it.next(),
352        }
353    }
354
355    fn size_hint(&self) -> (usize, Option<usize>) {
356        match self {
357            IntoIter::Inline(it) => it.size_hint(),
358            IntoIter::External(it) => it.size_hint(),
359        }
360    }
361}
362
363impl<T, const N: usize> From<[T; N]> for UnitVec<T> {
364    fn from(value: [T; N]) -> Self {
365        UnitVec::from_iter(value)
366    }
367}
368
369impl<T> ExactSizeIterator for IntoIter<T> {}
370
371impl<T> From<Vec<T>> for UnitVec<T> {
372    fn from(mut value: Vec<T>) -> Self {
373        if value.capacity() <= 1 {
374            let mut new = UnitVec::new();
375            if let Some(v) = value.pop() {
376                new.push(v)
377            }
378            new
379        } else {
380            let mut me = std::mem::ManuallyDrop::new(value);
381            UnitVec {
382                data: PointerOrValue {
383                    ptr: me.as_mut_ptr(),
384                },
385                capacity: NonZeroIdxSize::new(me.capacity().try_into().unwrap()).unwrap(),
386                len: me.len().try_into().unwrap(),
387            }
388        }
389    }
390}
391
392impl<T> From<UnitVec<T>> for Vec<T> {
393    fn from(mut value: UnitVec<T>) -> Self {
394        if value.is_inline() {
395            let mut out = Vec::with_capacity(value.len());
396            if let Some(item) = value.pop() {
397                out.push(item);
398            }
399            out
400        } else {
401            // SAFETY: when not inline, the data points to a buffer allocated by a Vec.
402            let out = unsafe {
403                Vec::from_raw_parts(value.data.ptr, value.len as usize, value.capacity())
404            };
405            // Prevent deallocating the buffer
406            std::mem::forget(value);
407            out
408        }
409    }
410}
411
412#[macro_export]
413macro_rules! unitvec {
414    () => {{
415        $crate::idx_vec::UnitVec::new()
416    }};
417    ($elem:expr; $n:expr) => {{
418        let mut new = $crate::idx_vec::UnitVec::new();
419        for _ in 0..$n {
420            new.push($elem)
421        }
422        new
423    }};
424    ($elem:expr) => {{
425        let mut new = $crate::idx_vec::UnitVec::new();
426        let v = $elem;
427        // SAFETY: first element always fits.
428        unsafe { new.push_unchecked(v) };
429        new
430    }};
431    ($($x:expr),+ $(,)?) => {{
432        vec![$($x),+].into()
433    }};
434}
435
436mod tests {
437
438    #[test]
439    #[should_panic]
440    fn test_unitvec_realloc_zero() {
441        super::UnitVec::<usize>::new().realloc(0);
442    }
443
444    #[test]
445    #[should_panic]
446    fn test_unitvec_realloc_one() {
447        super::UnitVec::<usize>::new().realloc(1);
448    }
449
450    #[test]
451    #[should_panic]
452    fn test_untivec_realloc_lt_len() {
453        super::UnitVec::<usize>::from([1, 2]).realloc(1)
454    }
455
456    #[test]
457    fn test_unitvec_clone() {
458        {
459            let v = unitvec![1usize];
460            assert_eq!(v, v.clone());
461        }
462
463        for n in [
464            26903816120209729usize,
465            42566276440897687,
466            44435161834424652,
467            49390731489933083,
468            51201454727649242,
469            83861672190814841,
470            92169290527847622,
471            92476373900398436,
472            95488551309275459,
473            97499984126814549,
474        ] {
475            let v = unitvec![n];
476            assert_eq!(v, v.clone());
477        }
478    }
479
480    #[test]
481    fn test_unitvec_repeat_n() {
482        assert_eq!(unitvec![5; 3].as_slice(), &[5, 5, 5])
483    }
484}
polars_utils/idx_vec.rs

polars_utils/
idx_vec.rs
