zoomvtools 1.1.0

Video motion vector analysis utilities in pure Rust
Documentation
#[cfg(test)]
use std::mem::size_of;
use std::num::NonZeroUsize;

use crate::util::Pixel;
use semisafe::option::unwrap as semisafe_opt_unwrap;
use semisafe::slice::get as semisafe_get;
use semisafe::slice::get_mut as semisafe_get_mut;

pub unsafe fn degrain<T: Pixel, const RADIUS: usize, const WIDTH: usize, const HEIGHT: usize>(
    dest: *mut u8,
    dest_stride_bytes: NonZeroUsize,
    src: *const u8,
    src_stride_bytes: NonZeroUsize,
    refs: &[*const u8],
    refs_strides_bytes: &[NonZeroUsize],
    w_src: i32,
    w_refs: &[i32],
) {
    let dest: *mut T = dest.cast();
    let src: *const T = src.cast();
    let refs: Box<[*const T]> = refs.iter().map(|ref_| ref_.cast()).collect();

    // SAFETY: size of T is non-zero and byte stride must be divisible by it
    let dest_stride_pixels =
        unsafe { NonZeroUsize::new_unchecked(dest_stride_bytes.get() / size_of::<T>()) };
    // SAFETY: size of T is non-zero and byte stride must be divisible by it
    let src_stride_pixels =
        unsafe { NonZeroUsize::new_unchecked(src_stride_bytes.get() / size_of::<T>()) };
    let ref_strides_pixels = refs_strides_bytes
        .iter()
        .map(|stride| {
            // SAFETY: size of T is non-zero and byte stride must be divisible by it
            unsafe { NonZeroUsize::new_unchecked(stride.get() / size_of::<T>()) }
        })
        .collect::<Box<[_]>>();

    let mut src_offset = 0;
    let mut dest_offset = 0;
    let mut ref_offsets = vec![0; refs.len()];

    for _y in 0..HEIGHT {
        for x in 0..WIDTH {
            let src = src.add(src_offset + x);
            let dest = dest.add(dest_offset + x);

            let src_i32: i32 = (*src).as_();
            let mut sum = 128 + src_i32 * w_src;

            for r in 0..(RADIUS * 2) {
                let ref_ = semisafe_get(&refs, r);
                let ref_ = ref_.add(semisafe_get(&ref_offsets, r) + x);
                let ref_i32: i32 = (*ref_).as_();
                sum += ref_i32 * semisafe_get(w_refs, r);
            }

            *dest = semisafe_opt_unwrap(T::from(sum >> 8));
        }

        dest_offset += dest_stride_pixels.get();
        src_offset += src_stride_pixels.get();
        for r in 0..(RADIUS * 2) {
            *semisafe_get_mut(&mut ref_offsets, r) += semisafe_get(&ref_strides_pixels, r).get();
        }
    }
}

#[cfg(test)]
pub(super) fn degrain_test<T: Pixel>(
    dest: &mut Vec<T>,
    width: NonZeroUsize,
    height: NonZeroUsize,
    src: &[T],
    src_stride_pixels: NonZeroUsize,
    refs: &[&[T]],
    refs_strides_pixels: &[NonZeroUsize],
    w_src: i32,
    w_refs: &[i32],
) {
    let radius = refs.len() / 2;

    // SAFETY: size_of::<T>() cannot be 0
    let stride_bytes =
        // SAFETY: size_of::<T>() cannot be 0
        unsafe { NonZeroUsize::new_unchecked(src_stride_pixels.get().saturating_mul(size_of::<T>())) };

    let refs_ptrs = refs
        .iter()
        .map(|ref_| ref_.as_ptr().cast())
        .collect::<Box<[_]>>();
    let refs_strides_bytes = refs_strides_pixels
        .iter()
        .map(|stride| {
            // SAFETY: size_of::<T>() cannot be 0
            // SAFETY: size_of::<T>() cannot be 0
            unsafe { NonZeroUsize::new_unchecked(stride.get().saturating_mul(size_of::<T>())) }
        })
        .collect::<Box<[_]>>();

    let func = match radius {
        1 => super::select_degrain_rust::<T, 1>(width, height),
        2 => super::select_degrain_rust::<T, 2>(width, height),
        3 => super::select_degrain_rust::<T, 3>(width, height),
        4 => super::select_degrain_rust::<T, 4>(width, height),
        5 => super::select_degrain_rust::<T, 5>(width, height),
        6 => super::select_degrain_rust::<T, 6>(width, height),
        _ => unreachable!("unsupported degrain radius"),
    };

    // SAFETY: pointers and strides are derived from owned slices, and function is selected for
    // the requested radius/width/height.
    unsafe {
        func(
            dest.as_mut_ptr().cast(),
            stride_bytes,
            src.as_ptr().cast(),
            stride_bytes,
            &refs_ptrs,
            &refs_strides_bytes,
            w_src,
            w_refs,
        );
    }
}