zoomvtools 1.1.0

#[cfg(test)]
mod tests;

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

use anyhow::Result;
use semisafe::slice::get as semisafe_get;
use semisafe::slice::get_mut as semisafe_get_mut;
use smallvec::SmallVec;

use crate::{
    frame::FramePlanesMut,
    mv_plane::MVPlane,
    params::{MVPlaneSet, ReduceFilter, Subpel, SubpelMethod},
    util::{Pixel, PlaneSizeTuple},
};

#[derive(Debug, Clone)]
pub struct MVFrame {
    pub planes: SmallVec<[MVPlane; 3]>,
    pub yuv_mode: MVPlaneSet,
}

impl MVFrame {
    #[inline]
    pub fn new(
        width: NonZeroUsize,
        height: NonZeroUsize,
        pel: Subpel,
        hpad: usize,
        vpad: usize,
        yuv_mode: MVPlaneSet,
        x_ratio_uv: NonZeroU8,
        y_ratio_uv: NonZeroU8,
        bits_per_sample: NonZeroU8,
        plane_offsets: &SmallVec<[usize; 3]>,
        pitch: PlaneSizeTuple,
    ) -> Result<Self> {
        // SAFETY: Width must be at least the value of its ratio
        let chroma_width =
            unsafe { NonZeroUsize::new_unchecked(width.get() / x_ratio_uv.get() as usize) };
        // SAFETY: Height must be at least the value of its ratio
        let chroma_height =
            unsafe { NonZeroUsize::new_unchecked(height.get() / y_ratio_uv.get() as usize) };
        let chroma_hpad = hpad / x_ratio_uv.get() as usize;
        let chroma_vpad = vpad / y_ratio_uv.get() as usize;

        let width = [width, chroma_width, chroma_width];
        let height = [height, chroma_height, chroma_height];
        let hpad = [hpad, chroma_hpad, chroma_hpad];
        let vpad = [vpad, chroma_vpad, chroma_vpad];

        let mut planes = SmallVec::new();
        for plane_idx in 0..3 {
            if (yuv_mode.bits() & (1 << plane_idx)) == 0 {
                continue;
            }

            let plane = MVPlane::new(
                *semisafe_get(&width, plane_idx),
                *semisafe_get(&height, plane_idx),
                pel,
                *semisafe_get(&hpad, plane_idx),
                *semisafe_get(&vpad, plane_idx),
                bits_per_sample,
                *semisafe_get(plane_offsets, plane_idx),
                match plane_idx {
                    0 => pitch.0,
                    1 => pitch.1.expect("should be set here"),
                    2 => pitch.2.expect("should be set here"),
                    _ => unreachable!(),
                },
            )?;
            planes.push(plane);
        }

        Ok(Self { planes, yuv_mode })
    }

    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip_all, name = "mv_frame::reduce_to")
    )]
    pub(crate) fn reduce_to<T: Pixel>(
        &self,
        reduced_frame: &mut MVFrame,
        mode: MVPlaneSet,
        filter: ReduceFilter,
        frame: &mut FramePlanesMut<'_, T>,
    ) {
        for i in 0..3 {
            if let Some(plane) = self.planes.get(i)
                && (mode.bits() & (1 << i)) > 0
            {
                let reduced_pitch = semisafe_get(&reduced_frame.planes, i).stride;
                let (width, height) = (
                    semisafe_get(&reduced_frame.planes, i).width,
                    semisafe_get(&reduced_frame.planes, i).height,
                );
                // Use the new helper function to avoid cloning the source data
                // SAFETY: The windows inside each plane are set up so that they do not overlap.
                unsafe {
                    let (src, dest) = frame
                        .plane_split(i)
                        .expect("Super: plane should exist but does not");
                    plane.reduce_to::<T>(
                        semisafe_get_mut(&mut reduced_frame.planes, i),
                        filter,
                        dest,
                        src,
                        reduced_pitch,
                        semisafe_get(&self.planes, i).stride,
                        width,
                        height,
                    );
                }
            }
        }
    }

    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip_all, name = "mv_frame::pad")
    )]
    pub(crate) fn pad<T: Pixel>(&mut self, mode: MVPlaneSet, frame: &mut FramePlanesMut<'_, T>) {
        for i in 0..3 {
            if let Some(plane) = self.planes.get_mut(i)
                && (mode.bits() & (1 << i)) > 0
            {
                plane.pad(
                    frame
                        .plane_mut(i)
                        .expect("Super: source plane should exist but does not"),
                );
            }
        }
    }

    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip_all, name = "mv_frame::refine")
    )]
    pub(crate) fn refine<T: Pixel>(
        &mut self,
        mode: MVPlaneSet,
        subpel: SubpelMethod,
        frame: &mut FramePlanesMut<'_, T>,
    ) {
        for i in 0..3 {
            if let Some(plane) = self.planes.get_mut(i)
                && (mode.bits() & (1 << i)) > 0
            {
                plane.refine::<T>(
                    subpel,
                    frame
                        .plane_mut(i)
                        .expect("Super: source plane should exist but does not"),
                );
            }
        }
    }

    pub(crate) fn update(&mut self, plane_offsets: [usize; 3], src_pitch: PlaneSizeTuple) {
        for i in 0..3 {
            if let Some(pitch) = match i {
                0 => Some(src_pitch.0),
                1 => src_pitch.1,
                2 => src_pitch.2,
                _ => unreachable!(),
            } {
                semisafe_get_mut(&mut self.planes, i)
                    .update(*semisafe_get(&plane_offsets, i), pitch);
            }
        }
    }
}