zoomvtools 1.1.1

#[cfg(test)]
mod tests;

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

use anyhow::{Result, anyhow, bail};
use semisafe::slice::get as semisafe_get;
use semisafe::slice::get_mut as semisafe_get_mut;

use crate::{
    filters::analyse::SuperClipInfo,
    frame::{FramePlanesMut, FrameView, PlaneSizeTuple},
    mv_gof::MVGroupOfFrames,
    mv_plane::MVPlane,
    params::Subpel,
    util::{Pixel, vs_bitblt},
    video::{ColorFamily, Resolution, SampleType, VideoFormat, VideoInfo},
};

/// Extracts the finest subpixel plane from a super clip.
#[derive(Debug, Clone)]
pub struct Finest {
    info: VideoInfo,
    super_info: SuperClipInfo,
    width: NonZeroUsize,
    x_ratio_uv: NonZeroU8,
    y_ratio_uv: NonZeroU8,
}

impl Finest {
    /// Builds a finest-level extractor for the supplied super-clip metadata.
    ///
    /// # Errors
    /// Returns an error if the clip format or super metadata are incompatible.
    #[inline]
    pub fn new(info: VideoInfo, super_info: SuperClipInfo) -> Result<Self> {
        let format = info.format;
        if format.bits_per_sample.get() > 16
            || format.sample_type != SampleType::Integer
            || ![ColorFamily::Yuv, ColorFamily::Gray].contains(&format.color_family)
            || format.sub_sampling_w > 1
            || format.sub_sampling_h > 1
        {
            bail!(
                "Finest: input clip must be GRAY, 420, 422, 440, or 444, up to 16 bits, with constant dimensions."
            );
        }

        let width = NonZeroUsize::new(
            info.resolution
                .width
                .checked_sub(super_info.hpad * 2)
                .ok_or_else(|| anyhow!("Finest: parameters from super clip appear to be wrong."))?,
        )
        .ok_or_else(|| anyhow!("Finest: parameters from super clip appear to be wrong."))?;

        Ok(Self {
            info,
            super_info,
            width,
            x_ratio_uv: NonZeroU8::new(1 << format.sub_sampling_w)
                .expect("subsampling ratio should never be zero"),
            y_ratio_uv: NonZeroU8::new(1 << format.sub_sampling_h)
                .expect("subsampling ratio should never be zero"),
        })
    }

    /// Returns the frame size produced by [`Self::render_frame`].
    #[must_use]
    #[inline]
    pub fn output_resolution(&self) -> Resolution {
        let pel = usize::from(u8::from(self.super_info.pel));
        Resolution {
            width: (self.width.get() + self.super_info.hpad * 2) * pel,
            height: (self.super_info.height.get() + self.super_info.vpad * 2) * pel,
        }
    }

    /// Renders the finest level into `output`.
    ///
    /// # Errors
    /// Returns an error if the input frame or output buffers do not match the stored geometry.
    #[inline]
    pub fn render_frame<T: Pixel>(
        &self,
        src: &FrameView<'_, T>,
        output: &mut FramePlanesMut<'_, T>,
        output_pitch: PlaneSizeTuple,
    ) -> Result<()> {
        let plane_count = self.info.format.plane_count();
        if self.super_info.pel == Subpel::Full {
            for plane in 0..plane_count {
                let src_plane = src.plane(plane)?;
                let dst_plane = output.plane_mut(plane)?;
                vs_bitblt(
                    dst_plane,
                    pitch_for_plane(output_pitch, plane)?,
                    src_plane,
                    src.pitch_for_plane(plane)?,
                    plane_width(self.info.format, self.output_resolution(), plane)?,
                    plane_height(self.info.format, self.output_resolution(), plane)?,
                );
            }

            return Ok(());
        }

        let src_strides = (
            src.pitch_for_plane(0)?,
            (plane_count > 1)
                .then(|| src.pitch_for_plane(1))
                .transpose()?,
            (plane_count > 2)
                .then(|| src.pitch_for_plane(2))
                .transpose()?,
        );
        let gof = MVGroupOfFrames::new(
            self.super_info.levels,
            self.width,
            self.super_info.height,
            self.super_info.pel,
            self.super_info.hpad,
            self.super_info.vpad,
            self.super_info.mode_yuv,
            self.x_ratio_uv,
            self.y_ratio_uv,
            self.info.format.bits_per_sample,
            src_strides,
            plane_count,
        )?;
        let finest_planes = &semisafe_get(&gof.frames, 0).planes;

        for plane in 0..plane_count {
            if (self.super_info.mode_yuv.bits() & (1 << plane)) == 0 {
                continue;
            }

            let mv_plane = semisafe_get(finest_planes, plane);
            let src_plane = src.plane(plane)?;
            let dst_plane = output.plane_mut(plane)?;

            match self.super_info.pel {
                Subpel::Half => merge_plane_to_big::<T, 2>(
                    dst_plane,
                    pitch_for_plane(output_pitch, plane)?,
                    src_plane,
                    mv_plane,
                ),
                Subpel::Quarter => merge_plane_to_big::<T, 4>(
                    dst_plane,
                    pitch_for_plane(output_pitch, plane)?,
                    src_plane,
                    mv_plane,
                ),
                Subpel::Full => unreachable!(),
            }
        }

        Ok(())
    }
}

#[inline]
fn plane_width(format: VideoFormat, resolution: Resolution, plane: usize) -> Result<NonZeroUsize> {
    let width = match plane {
        0 => resolution.width,
        1 | 2 => resolution.width >> usize::from(format.sub_sampling_w),
        _ => bail!("requested plane {plane} is not available"),
    };

    NonZeroUsize::new(width).ok_or_else(|| anyhow!("requested plane {plane} has zero width"))
}

#[inline]
fn plane_height(format: VideoFormat, resolution: Resolution, plane: usize) -> Result<NonZeroUsize> {
    let height = match plane {
        0 => resolution.height,
        1 | 2 => resolution.height >> usize::from(format.sub_sampling_h),
        _ => bail!("requested plane {plane} is not available"),
    };

    NonZeroUsize::new(height).ok_or_else(|| anyhow!("requested plane {plane} has zero height"))
}

#[inline]
fn pitch_for_plane(pitch: PlaneSizeTuple, plane: usize) -> Result<NonZeroUsize> {
    match plane {
        0 => Ok(pitch.0),
        1 => pitch
            .1
            .ok_or_else(|| anyhow!("requested plane {plane} is not available")),
        2 => pitch
            .2
            .ok_or_else(|| anyhow!("requested plane {plane} is not available")),
        _ => bail!("requested plane {plane} is not available"),
    }
}

#[inline]
fn merge_plane_to_big<T: Pixel, const PEL: usize>(
    dst: &mut [T],
    dst_pitch: NonZeroUsize,
    src: &[T],
    plane: &MVPlane,
) {
    let dst_pitch = dst_pitch.get();

    for y in 0..plane.padded_height.get() {
        for sub_y in 0..PEL {
            let dst_row = (y * PEL + sub_y) * dst_pitch;
            for x in 0..plane.padded_width.get() {
                for sub_x in 0..PEL {
                    let src_offset = plane.get_absolute_pix_offset_pel::<PEL>(
                        (x * PEL + sub_x) as i32,
                        (y * PEL + sub_y) as i32,
                    );
                    *semisafe_get_mut(dst, dst_row + x * PEL + sub_x) =
                        *semisafe_get(src, src_offset);
                }
            }
        }
    }
}