polars_utils/
vec.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2use std::mem::MaybeUninit;
3
4use bytemuck::Pod;
5use num_traits::Zero;
6
7use crate::with_drop::WithDrop;
8
9pub trait IntoRawParts<T> {
10    fn into_raw_parts(self) -> (*mut T, usize, usize);
11
12    // doesn't take ownership
13    fn raw_parts(&self) -> (*mut T, usize, usize);
14}
15
16impl<T> IntoRawParts<T> for Vec<T> {
17    fn into_raw_parts(self) -> (*mut T, usize, usize) {
18        let mut me = std::mem::ManuallyDrop::new(self);
19        (me.as_mut_ptr(), me.len(), me.capacity())
20    }
21
22    fn raw_parts(&self) -> (*mut T, usize, usize) {
23        (self.as_ptr() as *mut T, self.len(), self.capacity())
24    }
25}
26
27/// Fill current allocation if > 0
28/// otherwise realloc
29pub trait ResizeFaster<T: Copy> {
30    fn fill_or_alloc(&mut self, new_len: usize, value: T);
31}
32
33impl<T: Copy + Zero + PartialEq> ResizeFaster<T> for Vec<T> {
34    fn fill_or_alloc(&mut self, new_len: usize, value: T) {
35        if self.capacity() == 0 {
36            // it is faster to allocate zeroed
37            // so if the capacity is 0, we alloc (value might be 0)
38            *self = vec![value; new_len]
39        } else {
40            // first clear then reserve so that the reserve doesn't have
41            // to memcpy in case it needs to realloc.
42            self.clear();
43            self.reserve(new_len);
44
45            // // init the uninit values
46            let spare = &mut self.spare_capacity_mut()[..new_len];
47            let init_value = MaybeUninit::new(value);
48            spare.fill(init_value);
49            unsafe { self.set_len(new_len) }
50        }
51    }
52}
53pub trait PushUnchecked<T> {
54    /// Will push an item and not check if there is enough capacity
55    ///
56    /// # Safety
57    /// Caller must ensure the array has enough capacity to hold `T`.
58    unsafe fn push_unchecked(&mut self, value: T);
59}
60
61impl<T> PushUnchecked<T> for Vec<T> {
62    #[inline]
63    unsafe fn push_unchecked(&mut self, value: T) {
64        debug_assert!(self.capacity() > self.len());
65        let end = self.as_mut_ptr().add(self.len());
66        std::ptr::write(end, value);
67        self.set_len(self.len() + 1);
68    }
69}
70
71pub trait CapacityByFactor {
72    fn with_capacity_by_factor(original_len: usize, factor: f64) -> Self;
73}
74
75impl<T> CapacityByFactor for Vec<T> {
76    fn with_capacity_by_factor(original_len: usize, factor: f64) -> Self {
77        let cap = (original_len as f64 * factor) as usize;
78        Vec::with_capacity(cap)
79    }
80}
81
82// Trait to convert a Vec.
83// The reason for this is to reduce code-generation. Conversion functions that are named
84// functions should only generate the conversion loop once.
85pub trait ConvertVec<Out> {
86    type ItemIn;
87
88    fn convert_owned<F: FnMut(Self::ItemIn) -> Out>(self, f: F) -> Vec<Out>;
89
90    fn convert<F: FnMut(&Self::ItemIn) -> Out>(&self, f: F) -> Vec<Out>;
91}
92
93impl<T, Out> ConvertVec<Out> for Vec<T> {
94    type ItemIn = T;
95
96    fn convert_owned<F: FnMut(Self::ItemIn) -> Out>(self, f: F) -> Vec<Out> {
97        self.into_iter().map(f).collect()
98    }
99
100    fn convert<F: FnMut(&Self::ItemIn) -> Out>(&self, f: F) -> Vec<Out> {
101        self.iter().map(f).collect()
102    }
103}
104
105pub fn with_cast_mut_vec<T: Pod, U: Pod, R, F: FnOnce(&mut Vec<U>) -> R>(
106    v: &mut Vec<T>,
107    f: F,
108) -> R {
109    let mut vu = WithDrop::new(bytemuck::cast_vec::<T, U>(core::mem::take(v)), |vu| {
110        *v = bytemuck::cast_vec::<U, T>(vu)
111    });
112    f(&mut vu)
113}
polars_utils/vec.rs

polars_utils/
vec.rs
