use crate::error::{Error, Result};
pub fn unshuffle_into(data: &[u8], element_size: usize, out: &mut [u8]) -> Result<()> {
if out.len() != data.len() {
return Err(crate::error::Error::InvalidFormat(
"shuffle output length mismatch".into(),
));
}
if element_size <= 1 || data.is_empty() {
out.copy_from_slice(data);
return Ok(());
}
let n_elements = data.len() / element_size;
let grouped = n_elements * element_size;
match element_size {
8 => {
let p0 = &data[0..n_elements];
let p1 = &data[n_elements..2 * n_elements];
let p2 = &data[2 * n_elements..3 * n_elements];
let p3 = &data[3 * n_elements..4 * n_elements];
let p4 = &data[4 * n_elements..5 * n_elements];
let p5 = &data[5 * n_elements..6 * n_elements];
let p6 = &data[6 * n_elements..7 * n_elements];
let p7 = &data[7 * n_elements..8 * n_elements];
for (i, elem) in out[..grouped].chunks_exact_mut(8).enumerate() {
elem[0] = p0[i];
elem[1] = p1[i];
elem[2] = p2[i];
elem[3] = p3[i];
elem[4] = p4[i];
elem[5] = p5[i];
elem[6] = p6[i];
elem[7] = p7[i];
}
}
4 => {
let p0 = &data[0..n_elements];
let p1 = &data[n_elements..2 * n_elements];
let p2 = &data[2 * n_elements..3 * n_elements];
let p3 = &data[3 * n_elements..4 * n_elements];
for (i, elem) in out[..grouped].chunks_exact_mut(4).enumerate() {
elem[0] = p0[i];
elem[1] = p1[i];
elem[2] = p2[i];
elem[3] = p3[i];
}
}
2 => {
let p0 = &data[0..n_elements];
let p1 = &data[n_elements..2 * n_elements];
for (i, elem) in out[..grouped].chunks_exact_mut(2).enumerate() {
elem[0] = p0[i];
elem[1] = p1[i];
}
}
_ => {
for i in 0..n_elements {
for j in 0..element_size {
let dst = shuffle_index(i, element_size, j)?;
let src = shuffle_index(j, n_elements, i)?;
out[dst] = data[src];
}
}
}
}
out[grouped..].copy_from_slice(&data[grouped..]);
Ok(())
}
pub fn shuffle_into(data: &[u8], element_size: usize, out: &mut [u8]) -> Result<()> {
if out.len() != data.len() {
return Err(crate::error::Error::InvalidFormat(
"shuffle output length mismatch".into(),
));
}
if element_size <= 1 || data.is_empty() {
out.copy_from_slice(data);
return Ok(());
}
let n_elements = data.len() / element_size;
for i in 0..n_elements {
for j in 0..element_size {
let dst = shuffle_index(j, n_elements, i)?;
let src = shuffle_index(i, element_size, j)?;
out[dst] = data[src];
}
}
let grouped = n_elements * element_size;
out[grouped..].copy_from_slice(&data[grouped..]);
Ok(())
}
pub fn is_noop(data_len: usize, element_size: usize) -> bool {
element_size <= 1 || data_len / element_size <= 1
}
pub fn set_local_shuffle(element_size: usize) -> Result<usize> {
if element_size == 0 {
return Err(Error::InvalidFormat(
"shuffle filter element size is zero".into(),
));
}
Ok(element_size)
}
pub fn filter_shuffle_into(
data: &[u8],
element_size: usize,
reverse: bool,
out: &mut [u8],
) -> Result<()> {
let element_size = set_local_shuffle(element_size)?;
if reverse {
unshuffle_into(data, element_size, out)
} else {
shuffle_into(data, element_size, out)
}
}
fn shuffle_index(base: usize, stride: usize, offset: usize) -> Result<usize> {
base.checked_mul(stride)
.and_then(|value| value.checked_add(offset))
.ok_or_else(|| Error::InvalidFormat("shuffle index overflow".into()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_shuffle_roundtrip() {
let data = vec![1, 2, 3, 4, 5, 6, 7, 8]; let mut shuffled = vec![0; data.len()];
shuffle_into(&data, 4, &mut shuffled).unwrap();
let mut unshuffled = vec![0; shuffled.len()];
unshuffle_into(&shuffled, 4, &mut unshuffled).unwrap();
assert_eq!(unshuffled, data);
}
#[test]
fn test_shuffle_roundtrip_preserves_trailing_bytes() {
let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
let mut shuffled = vec![0; data.len()];
shuffle_into(&data, 4, &mut shuffled).unwrap();
let mut unshuffled = vec![0; shuffled.len()];
unshuffle_into(&shuffled, 4, &mut unshuffled).unwrap();
assert_eq!(unshuffled, data);
}
#[test]
fn shuffle_index_rejects_overflow() {
let err = shuffle_index(usize::MAX, 2, 0).unwrap_err();
assert!(
err.to_string().contains("shuffle index overflow"),
"unexpected error: {err}"
);
}
#[test]
fn shuffle_noop_detects_single_or_smaller_element_payloads() {
assert!(is_noop(0, 4));
assert!(is_noop(1, 1));
assert!(is_noop(3, 4));
assert!(is_noop(4, 4));
assert!(!is_noop(8, 4));
}
}