Skip to main content

polars_utils/
array.rs

1use std::mem::ManuallyDrop;
2
3#[repr(C)]
4struct ArrayPair<T, const NUM_LEFT: usize, const NUM_RIGHT: usize>([T; NUM_LEFT], [T; NUM_RIGHT]);
5
6pub fn try_map<T, U, const N: usize>(
7    array: [T; N],
8    f: impl FnMut(T) -> Option<U>,
9) -> Option<[U; N]> {
10    let mut array = array.map(f);
11
12    if array.iter().any(Option::is_none) {
13        return None;
14    }
15
16    Some(std::array::from_fn(|n| array[n].take().unwrap()))
17}
18
19/// Concatenate 2 arrays.
20pub fn array_concat<T, const NUM_LEFT: usize, const NUM_RIGHT: usize, const NUM_TOTAL: usize>(
21    left: [T; NUM_LEFT],
22    right: [T; NUM_RIGHT],
23) -> [T; NUM_TOTAL] {
24    const {
25        assert!(NUM_LEFT + NUM_RIGHT == NUM_TOTAL);
26    }
27
28    unsafe { std::mem::transmute_copy(&ManuallyDrop::new(ArrayPair(left, right))) }
29}
30
31/// Split an array to 2 arrays.
32pub fn array_split<T, const NUM_LEFT: usize, const NUM_RIGHT: usize, const NUM_TOTAL: usize>(
33    array: [T; NUM_TOTAL],
34) -> ([T; NUM_LEFT], [T; NUM_RIGHT]) {
35    const {
36        assert!(NUM_LEFT + NUM_RIGHT == NUM_TOTAL);
37    }
38
39    let ArrayPair::<T, NUM_LEFT, NUM_RIGHT>(l, r) =
40        unsafe { std::mem::transmute_copy(&ManuallyDrop::new(array)) };
41
42    (l, r)
43}