zoomvtools 1.1.1

#![allow(clippy::unwrap_used, reason = "allow in test files")]
#![allow(clippy::undocumented_unsafe_blocks, reason = "allow in test files")]
#![allow(clippy::indexing_slicing, reason = "allow in test files")]
#![allow(unused_unsafe)]

use std::num::{NonZeroU8, NonZeroUsize};

use num_traits::clamp;
use pastey::paste;
use safefma::Fma;

use super::DctHelper;
use crate::util::round_ties_to_even;

fn make_helper(size_x: usize, size_y: usize, bits_per_sample: u8) -> DctHelper {
    DctHelper::new(
        NonZeroUsize::new(size_x).unwrap(),
        NonZeroUsize::new(size_y).unwrap(),
        NonZeroU8::new(bits_per_sample).unwrap(),
    )
    .unwrap()
}

fn expected_ac<T>(helper: &DctHelper, value: f32) -> T
where
    T: crate::util::Pixel,
{
    let sqrt_2_div_2 = (2f32).sqrt() / 2.0;
    let pixel_max = (1 << usize::from(helper.bits_per_sample.get())) - 1;
    let pixel_half = 1 << (usize::from(helper.bits_per_sample.get()) - 1);
    let integ = round_ties_to_even(value * sqrt_2_div_2) as i32;

    T::from(clamp(
        (integ >> helper.dct_shift) + pixel_half,
        0,
        pixel_max,
    ))
    .unwrap()
}

fn expected_dc<T>(helper: &DctHelper, value: f32) -> T
where
    T: crate::util::Pixel,
{
    let pixel_max = (1 << usize::from(helper.bits_per_sample.get())) - 1;
    let pixel_half = 1 << (usize::from(helper.bits_per_sample.get()) - 1);
    let integ = round_ties_to_even(value * 0.5) as i32;

    T::from(clamp(
        (integ >> helper.dct_shift0) + pixel_half,
        0,
        pixel_max,
    ))
    .unwrap()
}

#[test]
fn pixels_to_float_src_u8_contiguous_pitch_copies_block() {
    let mut helper = make_helper(4, 4, 8);
    let src = (0u8..16).collect::<Vec<_>>();

    helper.pixels_to_float_src(&src, NonZeroUsize::new(4).unwrap());

    let expected = (0..16).map(|x| x as f32).collect::<Vec<_>>();
    assert_eq!(helper.src.as_ref(), expected.as_slice());
}

#[test]
fn pixels_to_float_src_u8_non_trivial_pitch_copies_only_active_region() {
    let mut helper = make_helper(4, 4, 8);
    let src_pitch = NonZeroUsize::new(6).unwrap();
    let src = vec![
        1u8, 2, 3, 4, 99, 99, 5, 6, 7, 8, 99, 99, 9, 10, 11, 12, 99, 99, 13, 14, 15, 16, 99, 99,
    ];

    helper.pixels_to_float_src(&src, src_pitch);

    let expected = (1..=16).map(|x| x as f32).collect::<Vec<_>>();
    assert_eq!(helper.src.as_ref(), expected.as_slice());
}

#[test]
fn pixels_to_float_src_u8_boundary_values_preserved() {
    let mut helper = make_helper(4, 4, 8);
    let src = vec![0u8, 1, 2, 255, 3, 4, 5, 6, 7, 8, 9, 10, 255, 11, 12, 0];

    helper.pixels_to_float_src(&src, NonZeroUsize::new(4).unwrap());

    assert_eq!(helper.src[0], 0.0);
    assert_eq!(helper.src[3], 255.0);
    assert_eq!(helper.src[12], 255.0);
    assert_eq!(helper.src[15], 0.0);
}

#[test]
fn pixels_to_float_src_u16_full_range_preserved() {
    let mut helper = make_helper(8, 8, 16);
    let mut src = (0..64)
        .map(|i| u16::try_from(i * 1024).unwrap())
        .collect::<Vec<_>>();
    src[0] = 0;
    src[1] = u16::MAX;
    src[63] = 32768;

    helper.pixels_to_float_src(&src, NonZeroUsize::new(8).unwrap());

    assert_eq!(helper.src[0], 0.0);
    assert_eq!(helper.src[1], f32::from(u16::MAX));
    assert_eq!(helper.src[63], 32768.0);
}

#[test]
fn pixels_to_float_src_rectangular_block_uses_row_stride() {
    let mut helper = make_helper(8, 4, 8);
    let src_pitch = NonZeroUsize::new(10).unwrap();
    let mut src = vec![255u8; src_pitch.get() * 4];

    for y in 0..4 {
        for x in 0..8 {
            src[y * src_pitch.get() + x] = u8::try_from(y * 8 + x).unwrap();
        }
    }

    helper.pixels_to_float_src(&src, src_pitch);

    let expected = (0..32).map(|x| x as f32).collect::<Vec<_>>();
    assert_eq!(helper.src.as_ref(), expected.as_slice());
}

#[test]
fn float_src_to_pixels_u8_zero_input_returns_half_range() {
    let helper = make_helper(4, 4, 8);
    let mut dst = vec![0u8; 16];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(4).unwrap());

    assert_eq!(dst, vec![128u8; 16]);
}

#[test]
fn float_src_to_pixels_u8_positive_ac_clamps_to_max() {
    let mut helper = make_helper(4, 4, 8);
    helper.src_dct[1] = 8192.0;
    let mut dst = vec![0u8; 16];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(4).unwrap());

    assert_eq!(dst[1], 255);
}

#[test]
fn float_src_to_pixels_u8_negative_ac_clamps_to_zero() {
    let mut helper = make_helper(4, 4, 8);
    helper.src_dct[1] = -8192.0;
    let mut dst = vec![0u8; 16];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(4).unwrap());

    assert_eq!(dst[1], 0);
}

#[test]
fn float_src_to_pixels_u8_dc_uses_dct_shift0_override() {
    let mut helper = make_helper(8, 8, 8);
    helper.src_dct[0] = 256.0;
    helper.src_dct[1] = 256.0;
    let mut dst = vec![0u8; 64];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(8).unwrap());

    assert_eq!(dst[0], expected_dc::<u8>(&helper, 256.0));
    assert_eq!(dst[1], expected_ac::<u8>(&helper, 256.0));
    assert_ne!(dst[0], dst[1]);
}

#[test]
fn float_src_to_pixels_u8_rounds_ties_to_even() {
    let mut helper = make_helper(4, 4, 8);
    let tie = 0.5 / ((2f32).sqrt() / 2.0);
    helper.src_dct[1] = 2.5 * tie;
    helper.src_dct[2] = 3.5 * tie;
    helper.src_dct[3] = -0.5 * tie;
    let mut dst = vec![0u8; 16];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(4).unwrap());

    assert_eq!(dst[1], expected_ac::<u8>(&helper, helper.src_dct[1]));
    assert_eq!(dst[2], expected_ac::<u8>(&helper, helper.src_dct[2]));
    assert_eq!(dst[3], expected_ac::<u8>(&helper, helper.src_dct[3]));
}

#[test]
fn float_src_to_pixels_u16_10bit_clamps_range() {
    let mut helper = make_helper(4, 4, 10);
    helper.src_dct[1] = 16384.0;
    helper.src_dct[2] = -16384.0;
    let mut dst = vec![0u16; 16];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(4).unwrap());

    assert_eq!(dst[0], 512);
    assert_eq!(dst[1], 1023);
    assert_eq!(dst[2], 0);
}

#[test]
fn float_src_to_pixels_u16_16bit_clamps_range() {
    let mut helper = make_helper(4, 4, 16);
    helper.src_dct[1] = 1_048_576.0;
    helper.src_dct[2] = -1_048_576.0;
    let mut dst = vec![0u16; 16];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(4).unwrap());

    assert_eq!(dst[0], 32768);
    assert_eq!(dst[1], 65535);
    assert_eq!(dst[2], 0);
}

#[test]
fn float_src_to_pixels_non_trivial_dst_pitch_preserves_padding() {
    let mut helper = make_helper(4, 4, 8);
    helper.src_dct[1] = 64.0;
    let dst_pitch = NonZeroUsize::new(6).unwrap();
    let mut dst = vec![77u8; dst_pitch.get() * 4];

    helper.float_src_to_pixels(&mut dst, dst_pitch);

    for y in 0..4 {
        assert_eq!(dst[y * dst_pitch.get() + 4], 77);
        assert_eq!(dst[y * dst_pitch.get() + 5], 77);
    }
}

#[test]
fn float_src_to_pixels_rectangular_block_writes_expected_values() {
    let mut helper = make_helper(8, 4, 8);
    helper.src_dct[8] = 32.0;
    helper.src_dct[31] = -32.0;
    let mut dst = vec![0u8; 32];

    helper.float_src_to_pixels(&mut dst, NonZeroUsize::new(8).unwrap());

    assert_eq!(dst[8], expected_ac::<u8>(&helper, 32.0));
    assert_eq!(dst[31], expected_ac::<u8>(&helper, -32.0));
}

const DCT_FLOAT_SRC_TO_PIXELS_CASES: &[(usize, usize, u8)] = &[
    (4, 4, 8),
    (4, 4, 10),
    (4, 4, 16),
    (8, 4, 8),
    (8, 8, 8),
    (8, 8, 10),
    (8, 8, 16),
    (16, 8, 8),
    (16, 16, 10),
    (32, 16, 16),
    (32, 32, 8),
];

fn prepare_helper_with_coeffs(size_x: usize, size_y: usize, bits: u8, coeffs: &[f32]) -> DctHelper {
    let mut helper = super::DctHelper::new(
        NonZeroUsize::new(size_x).unwrap(),
        NonZeroUsize::new(size_y).unwrap(),
        NonZeroU8::new(bits).unwrap(),
    )
    .unwrap();
    assert_eq!(coeffs.len(), size_x * size_y);
    helper.src_dct.copy_from_slice(coeffs);
    helper
}

macro_rules! float_src_to_pixels_tests {
    ($module:ident) => {
        paste! {
            #[test]
            fn [<float_src_to_pixels_zero_coeffs_u8_ $module>]() {
                for &(w, h, bits) in DCT_FLOAT_SRC_TO_PIXELS_CASES.iter().filter(|&&(_, _, b)| b == 8) {
                    let coeffs = vec![0.0f32; w * h];
                    let helper = prepare_helper_with_coeffs(w, h, bits, &coeffs);
                    let mut dst = vec![0u8; w * h];
                    let dst_pitch = NonZeroUsize::new(w).unwrap();

                    verify_asm!($module, float_src_to_pixels(
                        &mut dst,
                        dst_pitch,
                        &helper.src_dct,
                        helper.size_x,
                        helper.size_y,
                        helper.bits_per_sample,
                        helper.dct_shift,
                        helper.dct_shift0,
                    ));

                    assert!(dst.iter().all(|&p| p == 128), "case {w}x{h}@{bits}bpp");
                }
            }

            #[test]
            fn [<float_src_to_pixels_zero_coeffs_u16_ $module>]() {
                for &(w, h, bits) in DCT_FLOAT_SRC_TO_PIXELS_CASES.iter().filter(|&&(_, _, b)| b != 8) {
                    let coeffs = vec![0.0f32; w * h];
                    let helper = prepare_helper_with_coeffs(w, h, bits, &coeffs);
                    let mut dst = vec![0u16; w * h];
                    let dst_pitch = NonZeroUsize::new(w).unwrap();
                    let expected_half = 1u16 << (bits - 1);

                    verify_asm!($module, float_src_to_pixels(
                        &mut dst,
                        dst_pitch,
                        &helper.src_dct,
                        helper.size_x,
                        helper.size_y,
                        helper.bits_per_sample,
                        helper.dct_shift,
                        helper.dct_shift0,
                    ));

                    assert!(
                        dst.iter().all(|&p| p == expected_half),
                        "case {w}x{h}@{bits}bpp"
                    );
                }
            }

            #[test]
            fn [<float_src_to_pixels_mixed_ac_u8_ $module>]() {
                for &(w, h, bits) in DCT_FLOAT_SRC_TO_PIXELS_CASES.iter().filter(|&&(_, _, b)| b == 8) {
                    let mut coeffs = vec![0.0f32; w * h];
                    for (i, c) in coeffs.iter_mut().enumerate() {
                        let v = (i as i32 - (w * h) as i32 / 2) as f32 * 12.5;
                        *c = v;
                    }
                    let helper = prepare_helper_with_coeffs(w, h, bits, &coeffs);
                    let mut dst = vec![0u8; w * h];
                    let dst_pitch = NonZeroUsize::new(w).unwrap();

                    verify_asm!($module, float_src_to_pixels(
                        &mut dst,
                        dst_pitch,
                        &helper.src_dct,
                        helper.size_x,
                        helper.size_y,
                        helper.bits_per_sample,
                        helper.dct_shift,
                        helper.dct_shift0,
                    ));
                }
            }

            #[test]
            fn [<float_src_to_pixels_mixed_ac_u16_ $module>]() {
                for &(w, h, bits) in DCT_FLOAT_SRC_TO_PIXELS_CASES.iter().filter(|&&(_, _, b)| b != 8) {
                    let mut coeffs = vec![0.0f32; w * h];
                    for (i, c) in coeffs.iter_mut().enumerate() {
                        let v = (i as i32 - (w * h) as i32 / 2) as f32 * 97.25;
                        *c = v;
                    }
                    let helper = prepare_helper_with_coeffs(w, h, bits, &coeffs);
                    let mut dst = vec![0u16; w * h];
                    let dst_pitch = NonZeroUsize::new(w).unwrap();

                    verify_asm!($module, float_src_to_pixels(
                        &mut dst,
                        dst_pitch,
                        &helper.src_dct,
                        helper.size_x,
                        helper.size_y,
                        helper.bits_per_sample,
                        helper.dct_shift,
                        helper.dct_shift0,
                    ));
                }
            }

            #[test]
            fn [<float_src_to_pixels_clamp_saturation_u8_ $module>]() {
                let w = 8usize;
                let h = 8usize;
                let bits = 8u8;
                let mut coeffs = vec![0.0f32; w * h];
                for (i, c) in coeffs.iter_mut().enumerate().skip(1) {
                    *c = if i % 2 == 0 { 8192.0 } else { -8192.0 };
                }
                let helper = prepare_helper_with_coeffs(w, h, bits, &coeffs);
                let mut dst = vec![0u8; w * h];
                let dst_pitch = NonZeroUsize::new(w).unwrap();

                verify_asm!($module, float_src_to_pixels(
                    &mut dst,
                    dst_pitch,
                    &helper.src_dct,
                    helper.size_x,
                    helper.size_y,
                    helper.bits_per_sample,
                    helper.dct_shift,
                    helper.dct_shift0,
                ));
            }

            #[test]
            fn [<float_src_to_pixels_non_trivial_pitch_u8_ $module>]() {
                let w = 8usize;
                let h = 8usize;
                let bits = 8u8;
                let pitch = w + 4;
                let mut coeffs = vec![0.0f32; w * h];
                for (i, c) in coeffs.iter_mut().enumerate() {
                    *c = (i as f32).fma(3.5, -32.0);
                }
                let helper = prepare_helper_with_coeffs(w, h, bits, &coeffs);
                let mut dst = vec![0xaau8; pitch * h];
                let dst_pitch = NonZeroUsize::new(pitch).unwrap();

                verify_asm!($module, float_src_to_pixels(
                    &mut dst,
                    dst_pitch,
                    &helper.src_dct,
                    helper.size_x,
                    helper.size_y,
                    helper.bits_per_sample,
                    helper.dct_shift,
                    helper.dct_shift0,
                ));

                for y in 0..h {
                    for x in w..pitch {
                        assert_eq!(dst[y * pitch + x], 0xaa, "padding modified at ({x},{y})");
                    }
                }
            }
        }
    };
}

float_src_to_pixels_tests!(rust);

#[cfg(all(target_arch = "x86_64", feature = "avx2"))]
float_src_to_pixels_tests!(avx2);