polars_utils/
vec.rs

1use bytemuck::Pod;
2
3use crate::with_drop::WithDrop;
4
5/// Re-uses the memory for a vec while clearing it. Allows casting the type of
6/// the vec at the same time. The stdlib specializes collect() to re-use the
7/// memory.
8pub fn reuse_vec<T, U>(v: Vec<T>) -> Vec<U> {
9    const {
10        assert!(core::mem::size_of::<T>() == core::mem::size_of::<U>());
11        assert!(core::mem::align_of::<T>() == core::mem::align_of::<U>());
12    }
13    v.into_iter().filter_map(|_| None).collect()
14}
15
16pub trait PushUnchecked<T> {
17    /// Will push an item and not check if there is enough capacity
18    ///
19    /// # Safety
20    /// Caller must ensure the array has enough capacity to hold `T`.
21    unsafe fn push_unchecked(&mut self, value: T);
22}
23
24impl<T> PushUnchecked<T> for Vec<T> {
25    #[inline]
26    unsafe fn push_unchecked(&mut self, value: T) {
27        unsafe {
28            let len = self.len();
29            debug_assert!(self.capacity() > len);
30            let end = self.as_mut_ptr().add(len);
31            std::ptr::write(end, value);
32            self.set_len(len + 1);
33        }
34    }
35}
36
37pub fn with_cast_mut_vec<T: Pod, U: Pod, R, F: FnOnce(&mut Vec<U>) -> R>(
38    v: &mut Vec<T>,
39    f: F,
40) -> R {
41    let mut vu = WithDrop::new(bytemuck::cast_vec::<T, U>(core::mem::take(v)), |vu| {
42        *v = bytemuck::cast_vec::<U, T>(vu)
43    });
44    f(&mut vu)
45}
polars_utils/vec.rs

polars_utils/
vec.rs
