zoomvtools 1.1.0

Video motion vector analysis utilities in pure Rust
Documentation
use std::num::NonZeroUsize;

use anyhow::{Result, anyhow};

#[derive(Debug, Clone, Copy)]
pub struct PlaneRef<'a, T> {
    data: &'a [T],
}

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

    #[must_use]
    #[inline]
    pub const fn data(&self) -> &'a [T] {
        self.data
    }
}

#[derive(Debug)]
pub struct FramePlanes<'a, T> {
    planes: [Option<PlaneRef<'a, T>>; 3],
}

impl<'a, T> FramePlanes<'a, T> {
    #[must_use]
    #[inline]
    pub const fn new(planes: [Option<PlaneRef<'a, T>>; 3]) -> Self {
        Self { planes }
    }

    #[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)]
pub struct FramePlanesMut<'a, T> {
    planes: [Option<(*mut T, usize)>; 3],
    _marker: std::marker::PhantomData<&'a mut [T]>,
}

impl<'a, T> FramePlanesMut<'a, T> {
    #[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,
        }
    }

    #[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) })
    }

    #[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) })
    }

    /// # 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),
            ))
        }
    }
}

pub type PlaneSizeTuple = (NonZeroUsize, Option<NonZeroUsize>, Option<NonZeroUsize>);

#[cfg(test)]
mod tests {
    #![allow(clippy::indexing_slicing, reason = "allow in test files")]
    #![allow(clippy::unwrap_used, reason = "allow in test files")]

    use super::{FramePlanes, FramePlanesMut, PlaneRef};

    #[test]
    fn frame_planes_returns_requested_plane() {
        let y = [1u8, 2, 3, 4];
        let planes = FramePlanes::new([Some(PlaneRef::new(&y)), None, None]);

        assert_eq!(planes.plane(0).unwrap(), &y);
    }

    #[test]
    fn frame_planes_reports_missing_plane() {
        let planes = FramePlanes::<u8>::new([None, None, None]);

        let err = planes.plane(0).unwrap_err();
        assert!(
            err.to_string()
                .contains("requested plane 0 is not available")
        );
    }

    #[test]
    fn frame_planes_mut_can_write_plane() {
        let mut y = [0u8; 4];
        let mut planes = FramePlanesMut::new([Some(&mut y), None, None]);

        planes.plane_mut(0).unwrap().copy_from_slice(&[5, 6, 7, 8]);

        assert_eq!(planes.plane(0).unwrap(), &[5, 6, 7, 8]);
    }

    #[test]
    fn frame_planes_mut_split_exposes_shared_backing() {
        let mut y = [10u8, 20, 30, 40];
        let mut planes = FramePlanesMut::new([Some(&mut y), None, None]);

        // SAFETY: Test reads and writes disjoint logical usage through the returned views.
        let (src, dst) = unsafe { planes.plane_split(0).unwrap() };
        assert_eq!(src, &[10, 20, 30, 40]);
        dst[2] = 99;

        assert_eq!(planes.plane(0).unwrap(), &[10, 20, 99, 40]);
    }
}