use std::fmt::Debug;
macro_rules! simd_chunks {
($expr:expr $(,)?) => {
$crate::utils::simd_chunks_helper::<T, _, N, _>([$expr]).map(|[result]| result)
};
($($expr:expr),* $(,)?) => {
$crate::utils::simd_chunks_helper::<T, _, N, _>([$($expr),*])
};
}
pub(crate) use simd_chunks;
#[doc(hidden)]
pub fn simd_chunks_helper<T, const I: usize, const O: usize, const N: usize>(
arrays: [[T; I]; N],
) -> impl Iterator<Item = [[T; O]; N]>
where
T: Debug + Copy + 'static,
{
(0..I).step_by(2).map(move |offset| {
std::array::from_fn(|i| {
let mut iter = arrays[i].into_iter().cycle().skip(offset);
std::array::from_fn(|_| iter.next().unwrap())
})
})
}
#[test]
fn test_simd_chunks() {
type T = u32;
const N: usize = 8;
assert_eq!(
simd_chunks!(
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
)
.collect::<Vec<_>>(),
vec![
[[1, 2, 3, 4, 5, 6, 7, 8], [11, 12, 13, 14, 15, 16, 17, 18]],
[[3, 4, 5, 6, 7, 8, 9, 10], [13, 14, 15, 16, 17, 18, 19, 20]],
[[5, 6, 7, 8, 9, 10, 1, 2], [15, 16, 17, 18, 19, 20, 11, 12]],
[[7, 8, 9, 10, 1, 2, 3, 4], [17, 18, 19, 20, 11, 12, 13, 14]],
[[9, 10, 1, 2, 3, 4, 5, 6], [19, 20, 11, 12, 13, 14, 15, 16]],
]
);
assert_eq!(
simd_chunks!([1, 2, 3], [4, 5, 6]).collect::<Vec<_>>(),
vec![
[[1, 2, 3, 1, 2, 3, 1, 2], [4, 5, 6, 4, 5, 6, 4, 5]],
[[3, 1, 2, 3, 1, 2, 3, 1], [6, 4, 5, 6, 4, 5, 6, 4]],
]
);
}