zoomvtools 1.1.1

#[cfg(test)]
mod tests;

use std::num::NonZeroUsize;

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

use crate::{
    analysis::MVAnalysisData,
    fake::group_of_planes::FakeGroupOfPlanes,
    frame::{FramePlanesMut, FrameView, PlaneSizeTuple},
    params::Subpel,
    resize::SimpleResize,
    util::Pixel,
    video::{ColorFamily, SampleType, VideoInfo},
};

/// Output sampling mode used by [`Flow`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FlowMode {
    /// Fetches interpolated pixels directly from the reference clip.
    Fetch,
    /// Shifts source pixels by the motion field.
    Shift,
}

/// Options used to construct [`Flow`].
#[derive(Debug, Clone, Copy)]
pub struct FlowOptions {
    /// Position within the motion field in the `0.0..=100.0` range.
    pub time: f64,
    /// Pixel sampling mode used while warping.
    pub mode: FlowMode,
}

/// One-way motion-warping filter.
pub struct Flow {
    vectors_data: MVAnalysisData,
    mode: FlowMode,
    time256: i32,
    blk_x_padded: NonZeroUsize,
    blk_y_padded: NonZeroUsize,
    width_uv: Option<NonZeroUsize>,
    height_uv: Option<NonZeroUsize>,
    h_padding_uv: usize,
    v_padding_uv: usize,
    vector_pitch_y: NonZeroUsize,
    vector_pitch_uv: Option<NonZeroUsize>,
    pixel_max: u32,
    upsizer: SimpleResize,
    upsizer_uv: Option<SimpleResize>,
}

#[derive(Debug)]
struct ResizedVectorPair {
    vx_full: Vec<i16>,
    vy_full: Vec<i16>,
}

impl Flow {
    /// Builds a flow-warping filter from validated clip and vector metadata.
    ///
    /// # Errors
    /// Returns an error if the clip format, vector metadata, or options are invalid.
    #[inline]
    pub fn new(
        info: VideoInfo,
        vectors_data: MVAnalysisData,
        options: FlowOptions,
    ) -> Result<Self> {
        ensure!(
            (0.0..=100.0).contains(&options.time),
            "Flow: time must be between 0.0 and 100.0 (inclusive)."
        );

        let format = info.format;
        if format.bits_per_sample.get() > 16 {
            bail!("Flow: input clip must be 8-16 bits");
        }
        if format.sample_type != SampleType::Integer {
            bail!("Flow: input clip must be integer super_format");
        }
        if ![ColorFamily::Yuv, ColorFamily::Gray].contains(&format.color_family)
            || format.sub_sampling_w > 1
            || format.sub_sampling_h > 1
        {
            bail!("Flow: input clip must be GRAY, 420, 422, 440, or 444");
        }

        ensure!(
            info.resolution.width == vectors_data.width.get()
                && info.resolution.height == vectors_data.height.get(),
            "Flow: wrong source or super clip frame size."
        );
        let expected_x_ratio = 1usize << usize::from(format.sub_sampling_w);
        let expected_y_ratio = 1usize << usize::from(format.sub_sampling_h);
        ensure!(
            vectors_data.x_ratio_uv.get() as usize == expected_x_ratio
                && vectors_data.y_ratio_uv.get() as usize == expected_y_ratio,
            "Flow: input clip subsampling does not match vector metadata."
        );

        let mut blk_x_padded = vectors_data.blk_x.get();
        while blk_x_padded * (vectors_data.blk_size_x.get() - vectors_data.overlap_x)
            + vectors_data.overlap_x
            < vectors_data.width.get()
        {
            blk_x_padded += 1;
        }

        let mut blk_y_padded = vectors_data.blk_y.get();
        while blk_y_padded * (vectors_data.blk_size_y.get() - vectors_data.overlap_y)
            + vectors_data.overlap_y
            < vectors_data.height.get()
        {
            blk_y_padded += 1;
        }

        let blk_x_padded = NonZeroUsize::new(blk_x_padded)
            .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
        let blk_y_padded = NonZeroUsize::new(blk_y_padded)
            .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;

        let width_padded = NonZeroUsize::new(
            blk_x_padded.get() * (vectors_data.blk_size_x.get() - vectors_data.overlap_x)
                + vectors_data.overlap_x,
        )
        .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
        let height_padded = NonZeroUsize::new(
            blk_y_padded.get() * (vectors_data.blk_size_y.get() - vectors_data.overlap_y)
                + vectors_data.overlap_y,
        )
        .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;

        let vector_pitch_y = NonZeroUsize::new((width_padded.get() + 15) & !15)
            .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
        let upsizer = SimpleResize::new(
            width_padded,
            height_padded,
            blk_x_padded,
            blk_y_padded,
            vectors_data.width,
            vectors_data.height,
            vectors_data.pel,
        );

        let include_chroma = format.color_family == ColorFamily::Yuv;
        let x_ratio_uv = NonZeroUsize::from(vectors_data.x_ratio_uv);
        let y_ratio_uv = NonZeroUsize::from(vectors_data.y_ratio_uv);
        let (width_uv, height_uv, vector_pitch_uv, upsizer_uv) = if include_chroma {
            let width_uv = NonZeroUsize::new(vectors_data.width.get() / x_ratio_uv.get())
                .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
            let height_uv = NonZeroUsize::new(vectors_data.height.get() / y_ratio_uv.get())
                .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
            let width_padded_uv = NonZeroUsize::new(width_padded.get() / x_ratio_uv.get())
                .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
            let height_padded_uv = NonZeroUsize::new(height_padded.get() / y_ratio_uv.get())
                .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
            let vector_pitch_uv = NonZeroUsize::new((width_padded_uv.get() + 15) & !15)
                .ok_or_else(|| anyhow!("Flow: wrong source or super clip frame size."))?;
            let upsizer_uv = SimpleResize::new(
                width_padded_uv,
                height_padded_uv,
                blk_x_padded,
                blk_y_padded,
                width_uv,
                height_uv,
                vectors_data.pel,
            );

            (
                Some(width_uv),
                Some(height_uv),
                Some(vector_pitch_uv),
                Some(upsizer_uv),
            )
        } else {
            (None, None, None, None)
        };

        Ok(Self {
            vectors_data,
            mode: options.mode,
            time256: (options.time * 256.0 / 100.0) as i32,
            blk_x_padded,
            blk_y_padded,
            width_uv,
            height_uv,
            h_padding_uv: vectors_data.h_padding / x_ratio_uv.get(),
            v_padding_uv: vectors_data.v_padding / y_ratio_uv.get(),
            vector_pitch_y,
            vector_pitch_uv,
            pixel_max: (1u32 << format.bits_per_sample.get()) - 1,
            upsizer,
            upsizer_uv,
        })
    }

    /// Returns whether the vector field passes the supplied scene-change thresholds.
    #[must_use]
    #[inline]
    pub fn vectors_are_usable(
        &self,
        vectors: &FakeGroupOfPlanes,
        thscd1: u64,
        thscd2: u64,
    ) -> bool {
        vectors.is_usable(thscd1, thscd2)
    }

    /// Renders one warped frame into `output`.
    ///
    /// `field_shift` is applied to vertical vectors in subpixel units.
    ///
    /// # Errors
    /// Returns an error if the frame, vectors, or output buffers do not match the filter
    /// configuration.
    #[inline]
    pub fn render_frame<T: Pixel>(
        &self,
        reference: &FrameView<'_, T>,
        output: &mut FramePlanesMut<'_, T>,
        output_pitch: PlaneSizeTuple,
        vectors: &FakeGroupOfPlanes,
        field_shift: i16,
    ) -> Result<()> {
        let (mut vx_small_y, mut vy_small_y) = make_vector_small_masks(
            vectors,
            self.vectors_data.blk_x.get(),
            self.vectors_data.blk_y.get(),
            self.blk_x_padded.get(),
            self.blk_y_padded.get(),
        );
        check_and_pad_small_vectors(
            &mut vx_small_y,
            &mut vy_small_y,
            self.blk_x_padded.get(),
            self.blk_y_padded.get(),
            self.vectors_data.blk_x.get(),
            self.vectors_data.blk_y.get(),
        );

        if field_shift != 0 {
            for value in &mut vy_small_y {
                *value += field_shift;
            }
        }

        let vectors_y = resize_vector_pair(
            &self.upsizer,
            self.blk_x_padded,
            self.vector_pitch_y,
            &vx_small_y,
            &vy_small_y,
        );

        {
            let ref_stride = reference.pitch_for_plane(0)?;
            let dst_stride = pitch_for_plane(output_pitch, 0)?;
            let ref_origin_offset = reference_origin_offset(
                ref_stride,
                self.vectors_data.h_padding,
                self.vectors_data.v_padding,
                self.vectors_data.pel,
            );
            let ref_plane = reference.plane(0)?;
            let dst_plane = output.plane_mut(0)?;

            if self.mode == FlowMode::Shift {
                dst_plane.fill(T::from_u32_or_max_value(self.pixel_max));
            }

            flow_plane(
                dst_plane,
                dst_stride.get(),
                ref_plane,
                ref_stride.get(),
                ref_origin_offset,
                &vectors_y.vx_full,
                &vectors_y.vy_full,
                self.vector_pitch_y.get(),
                self.vectors_data.width.get(),
                self.vectors_data.height.get(),
                self.time256,
                self.vectors_data.pel,
                self.mode,
            );
        }

        if let Some(upsizer_uv) = &self.upsizer_uv {
            let vx_small_uv =
                vector_small_mask_y_to_uv(&vx_small_y, self.vectors_data.x_ratio_uv.get() as usize);
            let vy_small_uv =
                vector_small_mask_y_to_uv(&vy_small_y, self.vectors_data.y_ratio_uv.get() as usize);
            let vectors_uv = resize_vector_pair(
                upsizer_uv,
                self.blk_x_padded,
                self.vector_pitch_uv.expect("chroma pitch must exist"),
                &vx_small_uv,
                &vy_small_uv,
            );

            for plane in [1_usize, 2] {
                let ref_stride = reference.pitch_for_plane(plane)?;
                let dst_stride = pitch_for_plane(output_pitch, plane)?;
                let ref_origin_offset = reference_origin_offset(
                    ref_stride,
                    self.h_padding_uv,
                    self.v_padding_uv,
                    self.vectors_data.pel,
                );
                let ref_plane = reference.plane(plane)?;
                let dst_plane = output.plane_mut(plane)?;

                if self.mode == FlowMode::Shift {
                    dst_plane.fill(T::from_u32_or_max_value(self.pixel_max));
                }

                flow_plane(
                    dst_plane,
                    dst_stride.get(),
                    ref_plane,
                    ref_stride.get(),
                    ref_origin_offset,
                    &vectors_uv.vx_full,
                    &vectors_uv.vy_full,
                    self.vector_pitch_uv.expect("chroma pitch must exist").get(),
                    self.width_uv.expect("chroma width must exist").get(),
                    self.height_uv.expect("chroma height must exist").get(),
                    self.time256,
                    self.vectors_data.pel,
                    self.mode,
                );
            }
        }

        Ok(())
    }
}

fn flow_plane<T: Pixel>(
    dst: &mut [T],
    dst_stride: usize,
    reference: &[T],
    ref_stride: usize,
    ref_origin_offset: isize,
    vx_full: &[i16],
    vy_full: &[i16],
    vector_pitch: usize,
    width: usize,
    height: usize,
    time256: i32,
    pel: Subpel,
    mode: FlowMode,
) {
    match mode {
        FlowMode::Fetch => flow_fetch(
            dst,
            dst_stride,
            reference,
            ref_stride,
            ref_origin_offset,
            vx_full,
            vy_full,
            vector_pitch,
            width,
            height,
            time256,
            pel,
        ),
        FlowMode::Shift => flow_shift(
            dst,
            dst_stride,
            reference,
            ref_stride,
            ref_origin_offset,
            vx_full,
            vy_full,
            vector_pitch,
            width,
            height,
            time256,
            pel,
        ),
    }
}

fn flow_fetch<T: Pixel>(
    dst: &mut [T],
    dst_stride: usize,
    reference: &[T],
    ref_stride: usize,
    ref_origin_offset: isize,
    vx_full: &[i16],
    vy_full: &[i16],
    vector_pitch: usize,
    width: usize,
    height: usize,
    time256: i32,
    pel: Subpel,
) {
    let pel_log = pel.log();

    for y in 0..height {
        for x in 0..width {
            let vx = ((*semisafe_get(vx_full, y * vector_pitch + x) as i32) * time256 + 128) >> 8;
            let vy = ((*semisafe_get(vy_full, y * vector_pitch + x) as i32) * time256 + 128) >> 8;
            let row_base = ((y * ref_stride) << pel_log) as isize;
            let src_index = ref_origin_offset
                + row_base
                + vy as isize * ref_stride as isize
                + vx as isize
                + (x << pel_log) as isize;
            let src_index =
                usize::try_from(src_index).expect("flow fetch index must stay within padded plane");
            *semisafe_get_mut(dst, y * dst_stride + x) = *semisafe_get(reference, src_index);
        }
    }
}

fn flow_shift<T: Pixel>(
    dst: &mut [T],
    dst_stride: usize,
    reference: &[T],
    ref_stride: usize,
    ref_origin_offset: isize,
    vx_full: &[i16],
    vy_full: &[i16],
    vector_pitch: usize,
    width: usize,
    height: usize,
    time256: i32,
    pel: Subpel,
) {
    let pel_log = pel.log();
    let rounding = 128_i32 << pel_log;
    let shift = 8 + pel_log;

    for y in 0..height {
        for x in 0..width {
            let vx = (-(*semisafe_get(vx_full, y * vector_pitch + x) as i32) * time256 + rounding)
                >> shift;
            let vy = (-(*semisafe_get(vy_full, y * vector_pitch + x) as i32) * time256 + rounding)
                >> shift;
            let href = y as i32 + vy;
            let wref = x as i32 + vx;
            if (0..height as i32).contains(&href) && (0..width as i32).contains(&wref) {
                let dst_index = href as usize * dst_stride + wref as usize;
                let src_index = ref_origin_offset
                    + ((y * ref_stride) << pel_log) as isize
                    + (x << pel_log) as isize;
                let src_index = usize::try_from(src_index)
                    .expect("flow shift source index must stay within padded plane");
                *semisafe_get_mut(dst, dst_index) = *semisafe_get(reference, src_index);
            }
        }
    }
}

fn make_vector_small_masks(
    fake_gop: &FakeGroupOfPlanes,
    blk_x: usize,
    blk_y: usize,
    blk_x_padded: usize,
    blk_y_padded: usize,
) -> (Vec<i16>, Vec<i16>) {
    let mut vx = vec![0; blk_x_padded * blk_y_padded];
    let mut vy = vec![0; blk_x_padded * blk_y_padded];

    for by in 0..blk_y {
        for bx in 0..blk_x {
            let block = fake_gop.get_block(0, bx + by * blk_x);
            *semisafe_get_mut(&mut vx, bx + by * blk_x_padded) = block.vector.x as i16;
            *semisafe_get_mut(&mut vy, bx + by * blk_x_padded) = block.vector.y as i16;
        }
    }

    (vx, vy)
}

fn check_and_pad_small_vectors(
    vx_small: &mut [i16],
    vy_small: &mut [i16],
    blk_x_padded: usize,
    blk_y_padded: usize,
    blk_x: usize,
    blk_y: usize,
) {
    if blk_x_padded > blk_x {
        for row in 0..blk_y {
            let vx_right = (*semisafe_get(vx_small, row * blk_x_padded + blk_x - 1)).min(0);
            let vy_right = *semisafe_get(vy_small, row * blk_x_padded + blk_x - 1);
            for col in blk_x..blk_x_padded {
                *semisafe_get_mut(vx_small, row * blk_x_padded + col) = vx_right;
                *semisafe_get_mut(vy_small, row * blk_x_padded + col) = vy_right;
            }
        }
    }

    if blk_y_padded > blk_y {
        for col in 0..blk_x_padded {
            let vx_bottom = *semisafe_get(vx_small, blk_x_padded * (blk_y - 1) + col);
            let vy_bottom = (*semisafe_get(vy_small, blk_x_padded * (blk_y - 1) + col)).min(0);
            for row in blk_y..blk_y_padded {
                *semisafe_get_mut(vx_small, blk_x_padded * row + col) = vx_bottom;
                *semisafe_get_mut(vy_small, blk_x_padded * row + col) = vy_bottom;
            }
        }
    }
}

fn vector_small_mask_y_to_uv(v_small_y: &[i16], ratio_uv: usize) -> Vec<i16> {
    v_small_y
        .iter()
        .map(|value| if ratio_uv == 2 { *value >> 1 } else { *value })
        .collect()
}

fn resize_vector_pair(
    upsizer: &SimpleResize,
    blk_stride: NonZeroUsize,
    vec_pitch: NonZeroUsize,
    vx_small: &[i16],
    vy_small: &[i16],
) -> ResizedVectorPair {
    let vx_full = upsizer.resize_i16_to_vec(vec_pitch, vx_small, blk_stride, true);
    let vy_full = upsizer.resize_i16_to_vec(vec_pitch, vy_small, blk_stride, false);
    ResizedVectorPair { vx_full, vy_full }
}

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 1 is not available")),
        2 => pitch
            .2
            .ok_or_else(|| anyhow!("requested plane 2 is not available")),
        _ => bail!("requested plane {plane} is not available"),
    }
}

#[inline]
const fn reference_origin_offset(
    stride: NonZeroUsize,
    h_padding: usize,
    v_padding: usize,
    pel: Subpel,
) -> isize {
    (stride.get() * v_padding * pel as usize + h_padding * pel as usize) as isize
}