zoomvtools 1.1.1

//! Group of Frames code for managing a hierarchical frame structure

#[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, FrameView},
    mv_frame::MVFrame,
    mv_plane::{MVPlane, plane_height_luma, plane_super_offset, plane_width_luma},
    params::{MVPlaneSet, ReduceFilter, Subpel, SubpelMethod},
    util::{Pixel, PlaneSizeTuple},
};

pub(crate) struct BorrowedSuperFrame<'a, T> {
    view: &'a FrameView<'a, T>,
    gof: MVGroupOfFrames,
}

impl<'a, T> BorrowedSuperFrame<'a, T> {
    #[inline]
    pub(crate) fn new(
        view: &'a FrameView<'a, T>,
        level_count: usize,
        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,
    ) -> Result<Self> {
        let plane_count = if yuv_mode.contains(MVPlaneSet::UVPLANES) {
            view.plane_count()
        } else {
            1
        };
        let mut gof = MVGroupOfFrames::new(
            level_count,
            width,
            height,
            pel,
            hpad,
            vpad,
            yuv_mode,
            x_ratio_uv,
            y_ratio_uv,
            bits_per_sample,
            view.pitch(),
            plane_count,
        )?;
        gof.update(view.pitch());

        Ok(Self { view, gof })
    }

    #[inline]
    pub(crate) fn plane(&self, plane: usize) -> Result<&'a [T]> {
        self.view.plane(plane)
    }

    #[inline]
    pub(crate) const fn view(&self) -> &'a FrameView<'a, T> {
        self.view
    }

    #[inline]
    pub(crate) fn mv_plane(&self, plane: usize) -> &MVPlane {
        semisafe_get(&semisafe_get(&self.gof.frames, 0).planes, plane)
    }
}

#[derive(Debug, Clone)]
pub struct MVGroupOfFrames {
    level_count: usize,
    width: [NonZeroUsize; 3],
    height: [NonZeroUsize; 3],
    pel: Subpel,
    hpad: [usize; 3],
    vpad: [usize; 3],
    x_ratio_uv: NonZeroU8,
    y_ratio_uv: NonZeroU8,
    pub frames: Box<[MVFrame]>,
}

impl MVGroupOfFrames {
    #[inline]
    pub fn new(
        level_count: usize,
        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,
        pitch: PlaneSizeTuple,
        plane_count: usize,
    ) -> 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 mut this = Self {
            level_count,
            width: [width, chroma_width, chroma_width],
            height: [height, chroma_height, chroma_height],
            pel,
            hpad: [hpad, chroma_hpad, chroma_hpad],
            vpad: [vpad, chroma_vpad, chroma_vpad],
            x_ratio_uv,
            y_ratio_uv,
            frames: Default::default(),
        };

        let mut frames = Vec::with_capacity(level_count);

        for i in 0..level_count {
            let width_i = plane_width_luma(this.width[0], i, this.x_ratio_uv, this.hpad[0]);
            let height_i = plane_height_luma(this.height[0], i, this.y_ratio_uv, this.vpad[0]);
            let mut plane_offsets = SmallVec::with_capacity(3);
            for plane_idx in 0..plane_count {
                if (yuv_mode.bits() & (1 << plane_idx)) == 0 {
                    continue;
                }

                let offset = plane_super_offset(
                    plane_idx > 0,
                    *semisafe_get(&this.height, plane_idx),
                    i,
                    this.pel,
                    *semisafe_get(&this.vpad, 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!(),
                    },
                    this.y_ratio_uv,
                );
                plane_offsets.push(offset);
            }

            frames.push(MVFrame::new(
                width_i,
                height_i,
                if i == 0 { pel } else { Subpel::Full },
                this.hpad[0],
                this.vpad[0],
                yuv_mode,
                this.x_ratio_uv,
                this.y_ratio_uv,
                bits_per_sample,
                &plane_offsets,
                pitch,
            )?);
        }

        this.frames = frames.into_boxed_slice();

        Ok(this)
    }

    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip_all, name = "mv_gof::reduce", fields(self.level_count))
    )]
    #[inline]
    pub fn reduce<T: Pixel>(
        &mut self,
        mode: MVPlaneSet,
        filter: ReduceFilter,
        frame: &mut FramePlanesMut<'_, T>,
    ) {
        for i in 0..(self.level_count - 1) {
            let next_idx = i + 1;
            let source_frame = semisafe_get(&self.frames, i).clone();
            source_frame.reduce_to::<T>(
                semisafe_get_mut(&mut self.frames, next_idx),
                mode,
                filter,
                frame,
            );
            semisafe_get_mut(&mut self.frames, next_idx).pad::<T>(MVPlaneSet::YUVPLANES, frame);
        }
    }

    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip_all, name = "mv_gof::pad")
    )]
    #[inline]
    pub fn pad<T: Pixel>(&mut self, mode: MVPlaneSet, frame: &mut FramePlanesMut<'_, T>) {
        semisafe_get_mut(&mut self.frames, 0).pad::<T>(mode, frame);
    }

    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip_all, name = "mv_gof::refine")
    )]
    #[inline]
    pub fn refine<T: Pixel>(
        &mut self,
        mode: MVPlaneSet,
        subpel: SubpelMethod,
        frame: &mut FramePlanesMut<'_, T>,
    ) {
        semisafe_get_mut(&mut self.frames, 0).refine::<T>(mode, subpel, frame);
    }

    #[inline]
    pub fn update(&mut self, src_pitch: PlaneSizeTuple) {
        let mut plane_offsets = [0; 3];
        for i in 0..self.level_count {
            for plane in 0..3 {
                if let Some(pitch) = match plane {
                    0 => Some(src_pitch.0),
                    1 => src_pitch.1,
                    2 => src_pitch.2,
                    _ => unreachable!(),
                } {
                    *semisafe_get_mut(&mut plane_offsets, plane) = plane_super_offset(
                        plane > 0,
                        *semisafe_get(&self.height, plane),
                        i,
                        self.pel,
                        *semisafe_get(&self.vpad, plane),
                        pitch,
                        self.y_ratio_uv,
                    );
                }
            }

            semisafe_get_mut(&mut self.frames, i).update(plane_offsets, src_pitch);
        }
    }
}