zoomvtools 1.1.1

#[cfg(test)]
mod tests;

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

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

use crate::{
    filters::analyse::SuperClipInfo,
    frame::{FramePlanesMut, FrameView, PlaneSizeTuple},
    mv_gof::MVGroupOfFrames,
    mv_plane::{plane_height_luma, plane_super_offset, plane_width_luma},
    params::{MVPlaneSet, ReduceFilter, Subpel, SubpelMethod},
    util::Pixel,
    video::{ColorFamily, Resolution, SampleType, VideoInfo},
};

/// Options used to construct [`Super`].
#[derive(Debug, Clone, Copy)]
pub struct SuperOptions {
    /// Horizontal padding, in source pixels.
    pub hpad: usize,
    /// Vertical padding, in source pixels.
    pub vpad: usize,
    /// Subpixel precision stored in the output clip.
    pub pel: Subpel,
    /// Requested number of pyramid levels; `0` selects the maximum valid count.
    pub levels: usize,
    /// Includes chroma planes in the super clip when `true`.
    pub chroma: bool,
    /// Subpixel refinement kernel.
    pub sharp: SubpelMethod,
    /// Reduction filter used for coarser pyramid levels.
    pub rfilter: ReduceFilter,
}

/// Optional refined subpixel clip passed to [`Super::render_frame`].
pub struct SuperPelClip<'a, T> {
    /// Borrowed frame containing externally prepared subpixel data.
    pub frame: &'a FrameView<'a, T>,
    /// Indicates whether `frame` already includes super-style padding.
    pub is_padded: bool,
}

/// Super-clip builder used as input to other motion filters.
#[derive(Debug, Clone)]
pub struct Super {
    info: VideoInfo,
    super_info: SuperClipInfo,
    sharp: SubpelMethod,
    rfilter: ReduceFilter,
    output_resolution: Resolution,
    x_ratio_uv: NonZeroU8,
    y_ratio_uv: NonZeroU8,
}

impl Super {
    /// Builds a super-clip generator from the source clip metadata.
    ///
    /// # Errors
    /// Returns an error if the clip format or requested options are unsupported.
    #[inline]
    pub fn new(info: VideoInfo, options: SuperOptions) -> Result<Self> {
        let format = info.format;
        if format.bits_per_sample.get() > 16 {
            bail!("Super: input clip must be 8-16 bits");
        }
        if format.sample_type != SampleType::Integer {
            bail!("Super: input clip must be integer format");
        }
        if ![ColorFamily::Yuv, ColorFamily::Gray].contains(&format.color_family)
            || format.sub_sampling_w > 1
            || format.sub_sampling_h > 1
        {
            bail!("Super: input clip must be GRAY, 420, 422, 440, or 444");
        }

        let width = NonZeroUsize::new(info.resolution.width).ok_or_else(|| {
            anyhow::anyhow!("Super: variable resolution input clips are not supported")
        })?;
        let height = NonZeroUsize::new(info.resolution.height).ok_or_else(|| {
            anyhow::anyhow!("Super: variable resolution input clips are not supported")
        })?;
        let chroma = if format.color_family == ColorFamily::Gray {
            false
        } else {
            options.chroma
        };
        let mode_yuv = if chroma {
            MVPlaneSet::YUVPLANES
        } else {
            MVPlaneSet::YPLANE
        };
        let x_ratio_uv = NonZeroU8::new(1 << format.sub_sampling_w)
            .expect("subsampling ratio should never be zero");
        let y_ratio_uv = NonZeroU8::new(1 << format.sub_sampling_h)
            .expect("subsampling ratio should never be zero");

        let mut levels_max = 0;
        while plane_height_luma(height, levels_max, y_ratio_uv, options.vpad).get()
            >= y_ratio_uv.get() as usize * 2
            && plane_width_luma(width, levels_max, x_ratio_uv, options.hpad).get()
                >= x_ratio_uv.get() as usize * 2
        {
            levels_max += 1;
        }

        let levels = if options.levels == 0 || options.levels > levels_max {
            levels_max
        } else {
            options.levels
        };

        let mut super_width = width.saturating_add(2 * options.hpad);
        let mut super_height = NonZeroUsize::new(
            plane_super_offset(
                false,
                height,
                levels,
                options.pel,
                options.vpad,
                super_width,
                y_ratio_uv,
            ) / super_width,
        )
        .expect("super height should be non-zero when levels are valid");
        if y_ratio_uv.get() == 2 && super_height.get() & 1 > 0 {
            super_height = super_height.saturating_add(1);
        }
        if x_ratio_uv.get() == 2 && super_width.get() & 1 > 0 {
            super_width = super_width.saturating_add(1);
        }

        Ok(Self {
            info,
            super_info: SuperClipInfo {
                height,
                hpad: options.hpad,
                vpad: options.vpad,
                pel: options.pel,
                mode_yuv,
                levels,
            },
            sharp: options.sharp,
            rfilter: options.rfilter,
            output_resolution: Resolution {
                width: super_width.get(),
                height: super_height.get(),
            },
            x_ratio_uv,
            y_ratio_uv,
        })
    }

    /// Returns the metadata needed to reuse this super clip with other filters.
    #[must_use]
    #[inline]
    pub const fn super_info(&self) -> SuperClipInfo {
        self.super_info
    }

    /// Returns the frame size produced by [`Self::render_frame`].
    #[must_use]
    #[inline]
    pub const fn output_resolution(&self) -> Resolution {
        self.output_resolution
    }

    /// Renders one super frame into `output`.
    ///
    /// When `pel_clip` is provided, its finest level replaces the internally refined pel data.
    ///
    /// # Errors
    /// Returns an error if the source frame, optional pel clip, or output buffers do not match
    /// the filter configuration.
    #[inline]
    pub fn render_frame<T: Pixel>(
        &self,
        src: &FrameView<'_, T>,
        pel_clip: Option<SuperPelClip<'_, T>>,
        output: &mut FramePlanesMut<'_, T>,
        output_pitch: PlaneSizeTuple,
    ) -> Result<()> {
        let plane_count = self.info.format.plane_count();
        for plane in 0..plane_count {
            output.plane_mut(plane)?.fill(T::from_u32_or_max_value(0));
        }

        let mut src_gof = MVGroupOfFrames::new(
            self.super_info.levels,
            NonZeroUsize::new(self.info.resolution.width).expect("validated in constructor"),
            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,
            output_pitch,
            plane_count,
        )?;

        for plane in 0..plane_count {
            if let Some(plane_ref) = semisafe_get_mut(&mut src_gof.frames, 0)
                .planes
                .get_mut(plane)
            {
                plane_ref.fill_plane(
                    src.plane(plane)?,
                    src.pitch_for_plane(plane)?,
                    output.plane_mut(plane)?,
                );
            }
        }

        src_gof.reduce::<T>(self.super_info.mode_yuv, self.rfilter, output);
        src_gof.pad::<T>(self.super_info.mode_yuv, output);

        if let Some(pel_clip) = pel_clip {
            let src_frames = semisafe_get_mut(&mut src_gof.frames, 0);
            for plane in 0..plane_count {
                let src_plane = semisafe_get_mut(&mut src_frames.planes, plane);
                if !(self.super_info.mode_yuv & plane_mode(plane)).is_empty() {
                    src_plane.refine_ext(
                        pel_clip.frame.plane(plane)?,
                        pel_clip.frame.pitch_for_plane(plane)?,
                        pel_clip.is_padded,
                        output.plane_mut(plane)?,
                    );
                }
            }
        } else {
            src_gof.refine::<T>(self.super_info.mode_yuv, self.sharp, output);
        }

        Ok(())
    }
}

const fn plane_mode(plane: usize) -> MVPlaneSet {
    match plane {
        0 => MVPlaneSet::YPLANE,
        1 => MVPlaneSet::UPLANE,
        2 => MVPlaneSet::VPLANE,
        _ => MVPlaneSet::empty(),
    }
}