use core::cmp::min;
use zenpixels::{InPlacePixels, Orientation, PixelBuffer, PixelSlice, PixelSliceMut};
use crate::error::ConvertError;
use archmage::prelude::*;
use magetypes::simd::generic::f32x4 as GenericF32x4;
const TILE: u32 = 32;
#[must_use]
pub fn apply_orientation(src: PixelSlice<'_>, orientation: Orientation) -> PixelBuffer {
let (ow, oh) = orientation.output_dimensions(src.width(), src.rows());
let desc = src.descriptor();
let mut out = PixelBuffer::new(ow, oh, desc);
apply_orientation_into(src, orientation, out.as_slice_mut())
.expect("apply_orientation: freshly allocated buffer matches output geometry");
out
}
pub fn apply_orientation_into(
src: PixelSlice<'_>,
orientation: Orientation,
mut dst: PixelSliceMut<'_>,
) -> Result<(), ConvertError> {
let w = src.width();
let h = src.rows();
let bpp = src.descriptor().bytes_per_pixel();
let (ow, oh) = orientation.output_dimensions(w, h);
let dst_bpp = dst.descriptor().bytes_per_pixel();
if dst.width() != ow || dst.rows() != oh || dst_bpp != bpp {
return Err(ConvertError::BufferSize {
expected: ow as usize * oh as usize * bpp,
actual: dst.width() as usize * dst.rows() as usize * dst_bpp,
});
}
if w == 0 || h == 0 || bpp == 0 {
return Ok(());
}
{
match orientation {
Orientation::Identity => {
for y in 0..h {
dst.row_mut(y).copy_from_slice(src.row(y));
}
}
Orientation::FlipV => {
for y in 0..h {
dst.row_mut(y).copy_from_slice(src.row(h - 1 - y));
}
}
Orientation::FlipH => {
for y in 0..h {
reverse_row(src.row(y), dst.row_mut(y), w as usize, bpp);
}
}
Orientation::Rotate180 => {
for y in 0..h {
reverse_row(src.row(h - 1 - y), dst.row_mut(y), w as usize, bpp);
}
}
Orientation::Transpose
| Orientation::Rotate90
| Orientation::Rotate270
| Orientation::Transverse
| _ => {
do_transpose(&src, &mut dst, orientation, w, h, bpp);
}
}
}
Ok(())
}
const MAX_INPLACE_BPP: usize = 16;
pub fn apply_orientation_in_place(
dst: &mut PixelBuffer,
orientation: Orientation,
) -> Result<(), ConvertError> {
let bpp = dst.descriptor().bytes_per_pixel();
if bpp == 0 || bpp > MAX_INPLACE_BPP {
return Err(ConvertError::BufferSize {
expected: MAX_INPLACE_BPP,
actual: bpp,
});
}
if !matches!(
orientation,
Orientation::Identity
| Orientation::FlipH
| Orientation::FlipV
| Orientation::Rotate180
| Orientation::Transpose
| Orientation::Rotate90
| Orientation::Rotate270
| Orientation::Transverse
) {
return Err(ConvertError::BufferSize {
expected: MAX_INPLACE_BPP,
actual: 0,
});
}
dst.transform_in_place(|px| orient_in_place_impl(px, orientation));
Ok(())
}
fn orient_in_place_impl(px: InPlacePixels<'_>, orientation: Orientation) -> PixelSliceMut<'_> {
let InPlacePixels {
bytes,
width: w,
rows: h,
stride: in_stride,
descriptor: desc,
color,
..
} = px;
let bpp = desc.bytes_per_pixel();
let (ow, oh) = orientation.output_dimensions(w, h);
let tight = w as usize * bpp;
let out_stride = ow as usize * bpp;
let out_len = out_stride * oh as usize;
fn rewrap<'b>(
bytes: &'b mut [u8],
ow: u32,
oh: u32,
out_stride: usize,
desc: zenpixels::PixelDescriptor,
color: Option<alloc::sync::Arc<zenpixels::ColorContext>>,
) -> PixelSliceMut<'b> {
let out = PixelSliceMut::new(bytes, ow, oh, out_stride, desc)
.expect("oriented in-place geometry is always valid");
match color {
Some(c) => out.with_color_context(c),
None => out,
}
}
if w == 0 || h == 0 {
return rewrap(&mut bytes[..out_len], ow, oh, out_stride, desc, color);
}
if in_stride != tight {
for y in 1..h as usize {
bytes.copy_within(y * in_stride..y * in_stride + tight, y * tight);
}
}
let content = &mut bytes[..tight * h as usize];
match orientation {
Orientation::Identity => {}
Orientation::FlipH => inplace_flip_h(content, w, h, bpp),
Orientation::FlipV => inplace_flip_v(content, w, h, bpp),
Orientation::Rotate180 => inplace_reverse_elements(content, bpp),
Orientation::Transpose => inplace_transpose(content, w, h, bpp),
Orientation::Rotate90 => {
inplace_transpose(content, w, h, bpp);
inplace_flip_h(content, ow, oh, bpp); }
Orientation::Rotate270 => {
inplace_transpose(content, w, h, bpp);
inplace_flip_v(content, ow, oh, bpp); }
Orientation::Transverse => {
inplace_transpose(content, w, h, bpp);
inplace_reverse_elements(content, bpp); }
_ => {}
}
rewrap(&mut bytes[..out_len], ow, oh, out_stride, desc, color)
}
fn inplace_flip_h(a: &mut [u8], w: u32, h: u32, bpp: usize) {
let w = w as usize;
let row_len = w * bpp;
for y in 0..h as usize {
let row = &mut a[y * row_len..y * row_len + row_len];
let (mut lo, mut hi) = (0usize, w - 1);
while lo < hi {
let (al, ah) = (lo * bpp, hi * bpp);
for k in 0..bpp {
row.swap(al + k, ah + k);
}
lo += 1;
hi -= 1;
}
}
}
fn inplace_flip_v(a: &mut [u8], w: u32, h: u32, bpp: usize) {
let row_len = w as usize * bpp;
let h = h as usize;
let (mut top, mut bot) = (0usize, h - 1);
while top < bot {
let split = bot * row_len;
let (head, tail) = a.split_at_mut(split);
head[top * row_len..top * row_len + row_len].swap_with_slice(&mut tail[..row_len]);
top += 1;
bot -= 1;
}
}
fn inplace_reverse_elements(a: &mut [u8], bpp: usize) {
let n = a.len() / bpp;
if n < 2 {
return;
}
let (mut lo, mut hi) = (0usize, n - 1);
while lo < hi {
let (al, ah) = (lo * bpp, hi * bpp);
for k in 0..bpp {
a.swap(al + k, ah + k);
}
lo += 1;
hi -= 1;
}
}
fn inplace_transpose(a: &mut [u8], w: u32, h: u32, bpp: usize) {
if w == h {
let n = w as usize;
for i in 0..n {
for j in (i + 1)..n {
let (p, q) = ((i * n + j) * bpp, (j * n + i) * bpp);
for k in 0..bpp {
a.swap(p + k, q + k);
}
}
}
return;
}
let (w, h) = (w as usize, h as usize);
let n = w * h;
if n <= 1 {
return;
}
let mn1 = n - 1;
let mut moved = alloc::vec![false; n];
moved[0] = true;
moved[mn1] = true;
let mut tmp = [0u8; MAX_INPLACE_BPP];
let mut start = 1;
while start < mn1 {
if moved[start] {
start += 1;
continue;
}
tmp[..bpp].copy_from_slice(&a[start * bpp..start * bpp + bpp]);
let mut cur = start;
loop {
moved[cur] = true;
let prev = (cur * w) % mn1; if prev == start {
break;
}
a.copy_within(prev * bpp..prev * bpp + bpp, cur * bpp);
cur = prev;
}
a[cur * bpp..cur * bpp + bpp].copy_from_slice(&tmp[..bpp]);
start += 1;
}
}
#[cfg(feature = "__bench_orient")]
#[doc(hidden)]
#[must_use]
pub fn __bench_apply_orientation_scalar(
src: PixelSlice<'_>,
orientation: Orientation,
) -> PixelBuffer {
let w = src.width();
let h = src.rows();
let desc = src.descriptor();
let bpp = desc.bytes_per_pixel();
let (ow, oh) = orientation.output_dimensions(w, h);
let mut out = PixelBuffer::new(ow, oh, desc);
if w == 0 || h == 0 || bpp == 0 {
return out;
}
{
let mut dst = out.as_slice_mut();
transpose_blocked(&src, &mut dst, orientation, w, h, bpp);
}
out
}
#[inline]
fn reverse_row(s: &[u8], d: &mut [u8], width: usize, bpp: usize) {
for x in 0..width {
let si = (width - 1 - x) * bpp;
let di = x * bpp;
d[di..di + bpp].copy_from_slice(&s[si..si + bpp]);
}
}
#[inline]
#[allow(clippy::too_many_arguments)] fn scatter_pixel(
s: &[u8],
dst: &mut PixelSliceMut<'_>,
orientation: Orientation,
sx: u32,
sy: u32,
w: u32,
h: u32,
bpp: usize,
) {
let (dx, dy) = orientation.forward_map(sx, sy, w, h);
let si = sx as usize * bpp;
let di = dx as usize * bpp;
dst.row_mut(dy)[di..di + bpp].copy_from_slice(&s[si..si + bpp]);
}
fn do_transpose(
src: &PixelSlice<'_>,
dst: &mut PixelSliceMut<'_>,
orientation: Orientation,
w: u32,
h: u32,
bpp: usize,
) {
if bpp == 4
&& matches!(
orientation,
Orientation::Transpose
| Orientation::Rotate90
| Orientation::Rotate270
| Orientation::Transverse
)
{
incant!(
transpose4_simd(src, dst, orientation, w, h),
[v3, neon, wasm128, scalar]
);
return;
}
transpose_blocked(src, dst, orientation, w, h, bpp);
}
fn transpose_blocked(
src: &PixelSlice<'_>,
dst: &mut PixelSliceMut<'_>,
orientation: Orientation,
w: u32,
h: u32,
bpp: usize,
) {
let mut tile_y = 0;
while tile_y < h {
let y_end = min(tile_y + TILE, h);
let mut tile_x = 0;
while tile_x < w {
let x_end = min(tile_x + TILE, w);
for sy in tile_y..y_end {
let s = src.row(sy);
for sx in tile_x..x_end {
scatter_pixel(s, dst, orientation, sx, sy, w, h, bpp);
}
}
tile_x += TILE;
}
tile_y += TILE;
}
}
#[allow(clippy::too_many_arguments)] fn transpose_edges(
src: &PixelSlice<'_>,
dst: &mut PixelSliceMut<'_>,
orientation: Orientation,
w: u32,
h: u32,
bpp: usize,
full_w: u32,
full_h: u32,
) {
for sy in 0..full_h {
let s = src.row(sy);
for sx in full_w..w {
scatter_pixel(s, dst, orientation, sx, sy, w, h, bpp);
}
}
for sy in full_h..h {
let s = src.row(sy);
for sx in 0..w {
scatter_pixel(s, dst, orientation, sx, sy, w, h, bpp);
}
}
}
#[inline]
fn tile_dest(
orientation: Orientation,
bx: u32,
by: u32,
r: u32,
w: u32,
h: u32,
) -> (u32, u32, bool) {
match orientation {
Orientation::Transpose => (bx + r, by, false),
Orientation::Rotate90 => (bx + r, h - 4 - by, true),
Orientation::Rotate270 => (w - 1 - bx - r, by, false),
Orientation::Transverse => (w - 1 - bx - r, h - 4 - by, true),
_ => unreachable!("tile_dest only handles the four transposing orientations"),
}
}
#[magetypes(v3, neon, wasm128, scalar)]
fn transpose4_simd(
token: Token,
src: &PixelSlice<'_>,
dst: &mut PixelSliceMut<'_>,
orientation: Orientation,
w: u32,
h: u32,
) {
#[allow(non_camel_case_types)]
type f32x4 = GenericF32x4<Token>;
let full_w = w & !3; let full_h = h & !3;
let mut by = 0;
while by < full_h {
let mut bx = 0;
while bx < full_w {
let xb = bx as usize * 4;
let f0: [f32; 4] =
bytemuck::cast::<[u8; 16], _>(src.row(by)[xb..xb + 16].try_into().unwrap());
let f1: [f32; 4] =
bytemuck::cast::<[u8; 16], _>(src.row(by + 1)[xb..xb + 16].try_into().unwrap());
let f2: [f32; 4] =
bytemuck::cast::<[u8; 16], _>(src.row(by + 2)[xb..xb + 16].try_into().unwrap());
let f3: [f32; 4] =
bytemuck::cast::<[u8; 16], _>(src.row(by + 3)[xb..xb + 16].try_into().unwrap());
let mut rows = [
f32x4::load(token, &f0),
f32x4::load(token, &f1),
f32x4::load(token, &f2),
f32x4::load(token, &f3),
];
f32x4::transpose_4x4(&mut rows);
for r in 0..4u32 {
let mut lanes = [0f32; 4];
rows[r as usize].store(&mut lanes);
let (drow, dcol, rev) = tile_dest(orientation, bx, by, r, w, h);
if rev {
lanes.reverse();
}
let bytes: [u8; 16] = bytemuck::cast(lanes);
let db = dcol as usize * 4;
dst.row_mut(drow)[db..db + 16].copy_from_slice(&bytes);
}
bx += 4;
}
by += 4;
}
transpose_edges(src, dst, orientation, w, h, 4, full_w, full_h);
}
#[cfg(test)]
mod tests {
use super::*;
use zenpixels::PixelDescriptor;
fn slice<'a>(data: &'a [u8], w: u32, h: u32, desc: PixelDescriptor) -> PixelSlice<'a> {
PixelSlice::new(data, w, h, w as usize * desc.bytes_per_pixel(), desc).unwrap()
}
const SRC_3X2: [u8; 6] = [0, 1, 2, 3, 4, 5];
fn expected_3x2(o: Orientation) -> (u32, u32, Vec<u8>) {
match o {
Orientation::Identity => (3, 2, vec![0, 1, 2, 3, 4, 5]),
Orientation::FlipH => (3, 2, vec![2, 1, 0, 5, 4, 3]),
Orientation::FlipV => (3, 2, vec![3, 4, 5, 0, 1, 2]),
Orientation::Rotate180 => (3, 2, vec![5, 4, 3, 2, 1, 0]),
Orientation::Transpose => (2, 3, vec![0, 3, 1, 4, 2, 5]),
Orientation::Rotate90 => (2, 3, vec![3, 0, 4, 1, 5, 2]),
Orientation::Rotate270 => (2, 3, vec![2, 5, 1, 4, 0, 3]),
Orientation::Transverse => (2, 3, vec![5, 2, 4, 1, 3, 0]),
_ => unreachable!("non-exhaustive Orientation in test oracle"),
}
}
#[test]
fn all_orientations_match_hand_derived_oracle_gray8() {
let desc = PixelDescriptor::GRAY8;
for &o in &Orientation::ALL {
let out = apply_orientation(slice(&SRC_3X2, 3, 2, desc), o);
let (ew, eh, ebytes) = expected_3x2(o);
assert_eq!((out.width(), out.height()), (ew, eh), "{o:?} dims");
let s = out.as_slice();
for y in 0..eh {
let got = s.row(y);
let exp = &ebytes[y as usize * ew as usize..][..ew as usize];
assert_eq!(got, exp, "{o:?} row {y}");
}
}
}
#[test]
fn all_orientations_match_oracle_rgba8() {
let desc = PixelDescriptor::RGBA8;
let mut src = Vec::new();
for v in 0u8..6 {
src.extend_from_slice(&[v, v + 64, v + 128, 255]);
}
for &o in &Orientation::ALL {
let out = apply_orientation(slice(&src, 3, 2, desc), o);
let (ew, eh, gray) = expected_3x2(o);
assert_eq!((out.width(), out.height()), (ew, eh), "{o:?} dims");
let s = out.as_slice();
for y in 0..eh {
let got = s.row(y);
for x in 0..ew {
let v = gray[(y * ew + x) as usize];
let exp = [v, v + 64, v + 128, 255];
assert_eq!(&got[x as usize * 4..][..4], &exp, "{o:?} px ({x},{y})");
}
}
}
}
fn fill(n: usize) -> Vec<u8> {
let mut v = Vec::with_capacity(n);
let mut s = 0x9e3779b9u32;
for _ in 0..n {
s = s.wrapping_mul(1664525).wrapping_add(1013904223);
v.push((s >> 24) as u8);
}
v
}
#[test]
fn roundtrip_orientation_then_inverse_is_identity() {
for &desc in &[
PixelDescriptor::GRAY8, PixelDescriptor::GRAYA8, PixelDescriptor::RGB8, PixelDescriptor::RGBA8, PixelDescriptor::RGBAF32, ] {
let bpp = desc.bytes_per_pixel();
for &(w, h) in &[(1u32, 1u32), (17, 13), (33, 31), (64, 48)] {
let data = fill(w as usize * h as usize * bpp);
for &o in &Orientation::ALL {
let once = apply_orientation(slice(&data, w, h, desc), o);
let back = apply_orientation(once.as_slice(), o.inverse());
assert_eq!(
(back.width(), back.height()),
(w, h),
"{o:?} {desc:?} {w}x{h}"
);
for y in 0..h {
let exp = &data[y as usize * w as usize * bpp..][..w as usize * bpp];
assert_eq!(
back.as_slice().row(y),
exp,
"{o:?} {desc:?} {w}x{h} row {y}"
);
}
}
}
}
}
#[test]
fn compose_matches_sequential_application() {
let desc = PixelDescriptor::RGBA8;
let (w, h) = (11u32, 7u32);
let data = fill(w as usize * h as usize * 4);
for &a in &Orientation::ALL {
for &b in &Orientation::ALL {
let seq =
apply_orientation(apply_orientation(slice(&data, w, h, desc), a).as_slice(), b);
let fused = apply_orientation(slice(&data, w, h, desc), a.then(b));
assert_eq!(
(seq.width(), seq.height()),
(fused.width(), fused.height()),
"{a:?}.then({b:?}) dims"
);
for y in 0..seq.height() {
assert_eq!(
seq.as_slice().row(y),
fused.as_slice().row(y),
"{a:?}.then({b:?}) row {y}"
);
}
}
}
}
#[test]
fn handles_strided_source() {
let desc = PixelDescriptor::RGBA8;
let (w, h) = (5u32, 4u32);
let tight_stride = w as usize * 4;
let padded_stride = tight_stride + 12;
let tight = fill(tight_stride * h as usize);
let mut padded = vec![0xABu8; padded_stride * h as usize];
for y in 0..h as usize {
padded[y * padded_stride..y * padded_stride + tight_stride]
.copy_from_slice(&tight[y * tight_stride..][..tight_stride]);
}
for &o in &Orientation::ALL {
let tight_slice = PixelSlice::new(&tight, w, h, tight_stride, desc).unwrap();
let padded_slice = PixelSlice::new(&padded, w, h, padded_stride, desc).unwrap();
let a = apply_orientation(tight_slice, o);
let b = apply_orientation(padded_slice, o);
for y in 0..a.height() {
assert_eq!(a.as_slice().row(y), b.as_slice().row(y), "{o:?} row {y}");
}
}
}
#[test]
fn simd_transpose_matches_scalar_reference_rgba8() {
let desc = PixelDescriptor::RGBA8;
let dims = [
(8u32, 8u32),
(16, 16),
(64, 48),
(17, 13),
(9, 7),
(12, 4),
(4, 12),
(3, 3),
(1, 1),
(5, 5),
];
for &(w, h) in &dims {
let data = fill(w as usize * h as usize * 4);
for &o in &[
Orientation::Transpose,
Orientation::Rotate90,
Orientation::Rotate270,
Orientation::Transverse,
] {
let got = apply_orientation(slice(&data, w, h, desc), o);
let (ow, oh) = o.output_dimensions(w, h);
let mut reference = PixelBuffer::new(ow, oh, desc);
{
let src = slice(&data, w, h, desc);
let mut d = reference.as_slice_mut();
transpose_blocked(&src, &mut d, o, w, h, 4);
}
for y in 0..oh {
assert_eq!(
got.as_slice().row(y),
reference.as_slice().row(y),
"{o:?} {w}x{h} row {y}"
);
}
}
}
}
#[test]
fn into_writes_caller_buffer_and_is_reusable() {
let desc = PixelDescriptor::RGBA8;
let (w, h) = (17u32, 13u32);
let data = fill(w as usize * h as usize * 4);
let (ow, oh) = Orientation::Rotate90.output_dimensions(w, h);
let mut target = PixelBuffer::new(ow, oh, desc);
for &o in &[
Orientation::Rotate90,
Orientation::Rotate270,
Orientation::Transverse,
Orientation::Transpose,
] {
apply_orientation_into(slice(&data, w, h, desc), o, target.as_slice_mut())
.expect("into should accept a correctly-sized buffer");
let want = apply_orientation(slice(&data, w, h, desc), o);
for y in 0..oh {
assert_eq!(
target.as_slice().row(y),
want.as_slice().row(y),
"{o:?} row {y}"
);
}
}
}
#[test]
fn into_rejects_wrong_sized_dst() {
let desc = PixelDescriptor::RGBA8;
let (w, h) = (8u32, 6u32);
let data = fill(w as usize * h as usize * 4);
let mut wrong = PixelBuffer::new(w, h, desc); let result = apply_orientation_into(
slice(&data, w, h, desc),
Orientation::Rotate90,
wrong.as_slice_mut(),
);
assert!(
matches!(result, Err(ConvertError::BufferSize { .. })),
"expected BufferSize, got {result:?}"
);
}
#[test]
fn in_place_matches_out_of_place() {
let descs = [
PixelDescriptor::GRAY8,
PixelDescriptor::GRAYA8,
PixelDescriptor::RGB8,
PixelDescriptor::RGBA8,
PixelDescriptor::RGBAF32,
];
let dims = [
(1u32, 1u32),
(2, 2),
(4, 4),
(8, 8),
(32, 32),
(3, 5),
(5, 3),
(17, 13),
(13, 17),
(16, 9),
(9, 16),
(7, 1),
(1, 7),
];
for &desc in &descs {
let bpp = desc.bytes_per_pixel();
for &(w, h) in &dims {
let data = fill(w as usize * h as usize * bpp);
for &o in &Orientation::ALL {
let want = apply_orientation(slice(&data, w, h, desc), o);
let mut buf = PixelBuffer::new(w, h, desc);
{
let mut s = buf.as_slice_mut();
for y in 0..h {
s.row_mut(y).copy_from_slice(
&data[y as usize * w as usize * bpp..][..w as usize * bpp],
);
}
}
apply_orientation_in_place(&mut buf, o)
.expect("in_place should accept bpp ≤ 16");
assert_eq!(
(buf.width(), buf.height()),
(want.width(), want.height()),
"{o:?} {desc:?} {w}x{h} dims"
);
for y in 0..buf.height() {
assert_eq!(
buf.as_slice().row(y),
want.as_slice().row(y),
"{o:?} {desc:?} {w}x{h} row {y}"
);
}
}
}
}
}
}