fn transpose_2d<T: Copy>(rows: usize, cols: usize, row_major: &[T]) -> Vec<T> {
let n = rows
.checked_mul(cols)
.expect("transpose: rows * cols overflows usize");
assert_eq!(
row_major.len(),
n,
"transpose: source length {} does not match {rows}x{cols} = {n}",
row_major.len()
);
let mut out: Vec<T> = Vec::with_capacity(n);
if n == 0 {
return out;
}
const BLK: usize = 32;
let dst = out.as_mut_ptr();
for cb in (0..cols).step_by(BLK) {
let c_end = (cb + BLK).min(cols);
for rb in (0..rows).step_by(BLK) {
let r_end = (rb + BLK).min(rows);
for r in rb..r_end {
let src_row_base = r * cols;
for c in cb..c_end {
let value = row_major[src_row_base + c];
unsafe {
dst.add(c * rows + r).write(value);
}
}
}
}
}
unsafe {
out.set_len(n);
}
out
}
#[inline]
pub(crate) fn transpose_scalars<T: Copy>(rows: usize, cols: usize, row_major: &[T]) -> Vec<T> {
transpose_2d(rows, cols, row_major)
}
#[inline]
pub(crate) fn transpose_pairs<T: Copy>(
rows: usize,
cols: usize,
row_major: &[(T, T)],
) -> Vec<(T, T)> {
transpose_2d(rows, cols, row_major)
}
#[cfg(test)]
mod tests {
use super::*;
fn transpose_reference<T: Copy>(rows: usize, cols: usize, row_major: &[T]) -> Vec<T> {
let mut out = Vec::with_capacity(rows * cols);
for c in 0..cols {
for r in 0..rows {
out.push(row_major[r * cols + c]);
}
}
out
}
const SHAPES: &[(usize, usize)] = &[
(0, 0),
(0, 5),
(5, 0),
(1, 1),
(1, 7),
(7, 1),
(3, 4),
(32, 32),
(33, 32),
(32, 33),
(33, 47),
(65, 3),
];
#[test]
fn scalars_match_the_reference_transpose() {
for &(rows, cols) in SHAPES {
let src: Vec<u32> = (0..(rows * cols) as u32).collect();
assert_eq!(
transpose_scalars(rows, cols, &src),
transpose_reference(rows, cols, &src),
"shape {rows}x{cols}"
);
}
}
#[test]
fn pairs_match_the_reference_transpose() {
for &(rows, cols) in SHAPES {
let src: Vec<(f64, f64)> = (0..rows * cols).map(|i| (i as f64, -(i as f64))).collect();
assert_eq!(
transpose_pairs(rows, cols, &src),
transpose_reference(rows, cols, &src),
"shape {rows}x{cols}"
);
}
}
#[test]
fn transposing_twice_restores_the_original() {
let (rows, cols) = (33usize, 47usize);
let src: Vec<u64> = (0..(rows * cols) as u64).map(|i| i * 7 + 1).collect();
let once = transpose_scalars(rows, cols, &src);
assert_eq!(transpose_scalars(cols, rows, &once), src);
}
#[test]
fn an_empty_matrix_yields_an_empty_vec() {
assert!(transpose_scalars::<u8>(0, 0, &[]).is_empty());
assert!(transpose_pairs::<f64>(4, 0, &[]).is_empty());
}
const WRAPS_TO_FOUR: usize = usize::MAX / 4 + 2;
#[test]
fn the_wrapping_shape_really_does_wrap_to_the_source_length() {
assert_eq!(WRAPS_TO_FOUR.wrapping_mul(4), 4);
assert!(WRAPS_TO_FOUR.checked_mul(4).is_none());
}
#[test]
#[should_panic(expected = "overflows usize")]
fn a_shape_whose_product_wraps_is_refused() {
let src = vec![0u8; 4];
let _ = transpose_scalars(WRAPS_TO_FOUR, 4, &src);
}
#[test]
#[should_panic(expected = "overflows usize")]
fn a_matrix_whose_product_wraps_is_refused() {
let _ = crate::mat::Matrix::from_row_major(WRAPS_TO_FOUR, 4, vec![0u8; 4]);
}
}