zoomvtools 2.0.0

Video motion vector analysis utilities in pure Rust
Documentation
#[cfg(test)]
mod tests;

use std::num::NonZeroUsize;

use anyhow::{Result, anyhow};

#[derive(Debug, Clone)]
/// Owned storage for a single frame plane and its stride metadata.
pub struct FramePlane<T> {
    data: Box<[T]>,
    width: NonZeroUsize,
    height: NonZeroUsize,
    stride: NonZeroUsize,
}

impl<T> FramePlane<T> {
    /// Creates an owned plane from a backing buffer and logical dimensions.
    ///
    /// Returns an error when `data.len()` is smaller than `stride * height`.
    #[inline]
    pub fn new(
        data: Box<[T]>,
        width: NonZeroUsize,
        height: NonZeroUsize,
        stride: NonZeroUsize,
    ) -> Result<Self> {
        let required_len = stride.get() * height.get();
        if data.len() < required_len {
            return Err(anyhow!(
                "plane data is shorter than stride * height: {} < {}",
                data.len(),
                required_len
            ));
        }

        Ok(Self {
            data,
            width,
            height,
            stride,
        })
    }

    /// Returns the full backing storage, including any stride padding.
    #[must_use]
    #[inline]
    pub fn data(&self) -> &[T] {
        &self.data
    }

    /// Returns the logical plane width in pixels.
    #[must_use]
    #[inline]
    pub const fn width(&self) -> NonZeroUsize {
        self.width
    }

    /// Returns the logical plane height in pixels.
    #[must_use]
    #[inline]
    pub const fn height(&self) -> NonZeroUsize {
        self.height
    }

    /// Returns the number of samples between adjacent rows.
    #[must_use]
    #[inline]
    pub const fn stride(&self) -> NonZeroUsize {
        self.stride
    }
}

#[derive(Debug, Clone)]
/// Owned frame data stored as up to three Y/U/V planes.
pub struct Frame<T> {
    planes: [Option<FramePlane<T>>; 3],
}

impl<T> Frame<T> {
    /// Creates an owned frame from optional Y, U, and V planes.
    #[must_use]
    #[inline]
    pub const fn new(planes: [Option<FramePlane<T>>; 3]) -> Self {
        Self { planes }
    }

    /// Returns the requested plane if it is present.
    #[must_use]
    #[inline]
    pub fn plane(&self, plane: usize) -> Option<&FramePlane<T>> {
        self.planes.get(plane).and_then(Option::as_ref)
    }

    /// Returns the raw Y/U/V plane array.
    #[must_use]
    #[inline]
    pub const fn planes(&self) -> &[Option<FramePlane<T>>; 3] {
        &self.planes
    }

    /// Returns per-plane stride metadata for the owned frame.
    ///
    /// Returns an error if the luma plane is missing.
    #[inline]
    pub fn pitch(&self) -> Result<PlaneSizeTuple> {
        Ok((
            self.plane(0)
                .ok_or_else(|| anyhow!("owned frame is missing luma plane"))?
                .stride(),
            self.plane(1).map(FramePlane::stride),
            self.plane(2).map(FramePlane::stride),
        ))
    }

    /// Borrows this frame as immutable plane references.
    #[must_use]
    #[inline]
    pub fn as_planes(&self) -> FramePlanes<'_, T> {
        FramePlanes::new(
            self.planes
                .each_ref()
                .map(|plane| plane.as_ref().map(|plane| PlaneRef::new(plane.data()))),
        )
    }

    /// Borrows this frame as a [`FrameView`] with matching pitch metadata.
    #[inline]
    pub fn as_view(&self) -> Result<FrameView<'_, T>> {
        Ok(FrameView::new(self.as_planes(), self.pitch()?))
    }
}

#[derive(Debug)]
/// Borrowed frame planes paired with per-plane pitch information.
pub struct FrameView<'a, T> {
    planes: FramePlanes<'a, T>,
    pitch: PlaneSizeTuple,
}

impl<'a, T> FrameView<'a, T> {
    /// Creates a borrowed frame view from plane slices and per-plane pitch.
    /// CONTRACT: `planes` and `pitch` must describe the same logical plane layout.
    /// Chroma pitch entries must be present if and only if the matching plane exists.
    #[must_use]
    #[inline]
    pub fn new(planes: FramePlanes<'a, T>, pitch: PlaneSizeTuple) -> Self {
        debug_assert!(
            planes.planes[0].is_some(),
            "FrameView requires a luma plane"
        );
        debug_assert_eq!(
            planes.planes[1].is_some(),
            pitch.1.is_some(),
            "FrameView planes and pitch must agree for chroma U"
        );
        debug_assert_eq!(
            planes.planes[2].is_some(),
            pitch.2.is_some(),
            "FrameView planes and pitch must agree for chroma V"
        );

        Self { planes, pitch }
    }

    /// Returns the stored per-plane pitch tuple.
    #[must_use]
    #[inline]
    pub const fn pitch(&self) -> PlaneSizeTuple {
        self.pitch
    }

    /// Returns the number of leading planes present in this view.
    #[must_use]
    #[inline]
    pub fn plane_count(&self) -> usize {
        self.planes
            .planes
            .iter()
            .rposition(Option::is_some)
            .map_or(0, |last_plane| last_plane + 1)
    }

    /// Returns the borrowed plane set backing this view.
    #[must_use]
    #[inline]
    pub const fn planes(&self) -> &FramePlanes<'a, T> {
        &self.planes
    }

    /// Returns the requested plane slice.
    #[inline]
    pub fn plane(&self, plane: usize) -> Result<&'a [T]> {
        self.planes.plane(plane)
    }

    /// Returns the stride for the requested plane.
    #[inline]
    pub fn pitch_for_plane(&self, plane: usize) -> Result<NonZeroUsize> {
        let pitch = match plane {
            0 => Some(self.pitch.0),
            1 => self.pitch.1,
            2 => self.pitch.2,
            _ => None,
        };

        pitch.ok_or_else(|| anyhow!("requested plane {plane} is not available"))
    }
}

#[derive(Debug, Clone, Copy)]
/// Borrowed reference to one frame plane.
pub struct PlaneRef<'a, T> {
    data: &'a [T],
}

impl<'a, T> PlaneRef<'a, T> {
    /// Wraps a borrowed plane slice.
    #[must_use]
    #[inline]
    pub const fn new(data: &'a [T]) -> Self {
        Self { data }
    }

    /// Returns the borrowed plane slice.
    #[must_use]
    #[inline]
    pub const fn data(&self) -> &'a [T] {
        self.data
    }
}

#[derive(Debug)]
/// Borrowed Y/U/V plane slices without stride metadata.
pub struct FramePlanes<'a, T> {
    planes: [Option<PlaneRef<'a, T>>; 3],
}

impl<'a, T> FramePlanes<'a, T> {
    /// Creates a borrowed plane set from optional Y, U, and V planes.
    #[must_use]
    #[inline]
    pub const fn new(planes: [Option<PlaneRef<'a, T>>; 3]) -> Self {
        Self { planes }
    }

    /// Returns the requested plane slice.
    #[inline]
    pub fn plane(&self, plane: usize) -> Result<&'a [T]> {
        self.planes
            .get(plane)
            .and_then(|plane| plane.as_ref().map(PlaneRef::data))
            .ok_or_else(|| anyhow!("requested plane {plane} is not available"))
    }
}

#[derive(Debug)]
/// Mutable borrowed Y/U/V plane slices stored behind stable raw pointers.
pub struct FramePlanesMut<'a, T> {
    planes: [Option<(*mut T, usize)>; 3],
    _marker: std::marker::PhantomData<&'a mut [T]>,
}

impl<'a, T> FramePlanesMut<'a, T> {
    /// Creates a mutable plane set from optional Y, U, and V slices.
    #[must_use]
    #[inline]
    pub fn new(planes: [Option<&'a mut [T]>; 3]) -> Self {
        let planes = planes.map(|plane| plane.map(|plane| (plane.as_mut_ptr(), plane.len())));
        Self {
            planes,
            _marker: std::marker::PhantomData,
        }
    }

    /// Returns an immutable view of the requested plane.
    #[inline]
    pub fn plane(&self, plane: usize) -> Result<&[T]> {
        let (ptr, len) = self
            .planes
            .get(plane)
            .and_then(|plane| *plane)
            .ok_or_else(|| anyhow!("requested plane {plane} is not available"))?;

        // SAFETY: `ptr,len` originate from a live mutable slice borrowed for `'a`.
        Ok(unsafe { std::slice::from_raw_parts(ptr.cast_const(), len) })
    }

    /// Returns a mutable view of the requested plane.
    #[inline]
    pub fn plane_mut(&mut self, plane: usize) -> Result<&mut [T]> {
        let (ptr, len) = self
            .planes
            .get(plane)
            .and_then(|plane| *plane)
            .ok_or_else(|| anyhow!("requested plane {plane} is not available"))?;

        // SAFETY: `ptr,len` originate from a unique mutable slice borrowed for `'a`.
        Ok(unsafe { std::slice::from_raw_parts_mut(ptr, len) })
    }

    /// Returns immutable and mutable views over the same plane backing storage.
    ///
    /// # Safety
    ///
    /// The caller must ensure the returned immutable and mutable slices are only used on
    /// non-overlapping logical regions of the same underlying plane.
    #[inline]
    pub unsafe fn plane_split(&mut self, plane: usize) -> Result<(&[T], &mut [T])> {
        let (ptr, len) = self
            .planes
            .get(plane)
            .and_then(|plane| *plane)
            .ok_or_else(|| anyhow!("requested plane {plane} is not available"))?;

        // SAFETY: Caller guarantees the immutable and mutable views are used on non-overlapping
        // logical regions, matching the previous VapourSynth-specific helper contract.
        unsafe {
            Ok((
                std::slice::from_raw_parts(ptr.cast_const(), len),
                std::slice::from_raw_parts_mut(ptr, len),
            ))
        }
    }
}

/// Per-plane stride tuple in Y/U/V order.
///
/// Chroma entries are `None` when the matching plane is absent.
pub type PlaneSizeTuple = (NonZeroUsize, Option<NonZeroUsize>, Option<NonZeroUsize>);