/rust/registry/src/index.crates.io-1949cf8c6b5b557f/y4m-0.8.0/src/lib.rs

Source
//! # YUV4MPEG2 (.y4m) Encoder/Decoder
#![deny(missing_docs)]

use std::fmt;
use std::io;
use std::io::Read;
use std::io::Write;
use std::num;
use std::str;

const MAX_PARAMS_SIZE: usize = 1024;
const FILE_MAGICK: &[u8] = b"YUV4MPEG2 ";
const FRAME_MAGICK: &[u8] = b"FRAME";
const TERMINATOR: u8 = 0x0A;
const FIELD_SEP: u8 = b' ';
const RATIO_SEP: u8 = b':';

/// Both encoding and decoding errors.
#[derive(Debug)]
pub enum Error {
    /// End of the file. Technically not an error, but it's easier to process
    /// that way.
    EOF,
    /// Bad input parameters provided.
    BadInput,
    /// Unknown colorspace (possibly just unimplemented).
    UnknownColorspace,
    /// Error while parsing the file/frame header.
    // TODO(Kagami): Better granularity of parse errors.
    ParseError(ParseError),
    /// Error while reading/writing the file.
    IoError(io::Error),
    /// Out of memory (limits exceeded).
    OutOfMemory,
}

impl std::error::Error for crate::Error {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match *self {
            Error::EOF => None,
            Error::BadInput => None,
            Error::UnknownColorspace => None,
            Error::ParseError(ref err) => Some(err),
            Error::IoError(ref err) => Some(err),
            Error::OutOfMemory => None,
        }
    }
}

impl fmt::Display for crate::Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::EOF => write!(f, "End of file"),
            Error::BadInput => write!(f, "Bad input parameters provided"),
            Error::UnknownColorspace => write!(f, "Bad input parameters provided"),
            Error::ParseError(ref err) => err.fmt(f),
            Error::IoError(ref err) => err.fmt(f),
            Error::OutOfMemory => write!(f, "Out of memory (limits exceeded)"),
        }
    }
}

/// Granular ParseError Definiations
pub enum ParseError {
    /// Error reading y4m header
    InvalidY4M,
    /// Error parsing int
    Int,
    /// Error parsing UTF8
    Utf8,
    /// General Parsing Error
    General,
}

impl std::error::Error for crate::ParseError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match *self {
            ParseError::InvalidY4M => None,
            ParseError::Int => None,
            ParseError::Utf8 => None,
            ParseError::General => None,
        }
    }
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            ParseError::InvalidY4M => write!(f, "Error parsing y4m header"),
            ParseError::Int => write!(f, "Error parsing Int"),
            ParseError::Utf8 => write!(f, "Error parsing UTF8"),
            ParseError::General => write!(f, "General parsing error"),
        }
    }
}

impl fmt::Debug for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            ParseError::InvalidY4M => write!(f, "Error parsing y4m header"),
            ParseError::Int => write!(f, "Error parsing Int"),
            ParseError::Utf8 => write!(f, "Error parsing UTF8"),
            ParseError::General => write!(f, "General parsing error"),
        }
    }
}

macro_rules! parse_error {
    ($p:expr) => {
        return Err(Error::ParseError($p))
    };
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        match err.kind() {
            io::ErrorKind::UnexpectedEof => Error::EOF,
            _ => Error::IoError(err),
        }
    }
}

impl From<num::ParseIntError> for Error {
    fn from(_: num::ParseIntError) -> Error {
        Error::ParseError(ParseError::Int)
    }
}

impl From<str::Utf8Error> for Error {
    fn from(_: str::Utf8Error) -> Error {
        Error::ParseError(ParseError::Utf8)
    }
}

trait EnhancedRead {
    fn read_until(&mut self, ch: u8, buf: &mut [u8]) -> Result<usize, Error>;
}

impl<R: Read> EnhancedRead for R {
    // Current implementation does one `read` call per byte. This might be a
    // bit slow for long headers but it simplifies things: we don't need to
    // check whether start of the next frame is already read and so on.
    fn read_until(&mut self, ch: u8, buf: &mut [u8]) -> Result<usize, Error> {
        let mut collected = 0;
        while collected < buf.len() {
            let chunk_size = self.read(&mut buf[collected..=collected])?;
            if chunk_size == 0 {
                return Err(Error::EOF);
            }
            if buf[collected] == ch {
                return Ok(collected);
            }
            collected += chunk_size;
        }
        parse_error!(ParseError::General)
    }
}

fn parse_bytes(buf: &[u8]) -> Result<usize, Error> {
    // A bit kludgy but seems like there is no other way.
    Ok(str::from_utf8(buf)?.parse()?)
}

/// A newtype wrapper around Vec<u8> to ensure validity as a vendor extension.
#[derive(Debug, Clone)]
pub struct VendorExtensionString(Vec<u8>);

impl VendorExtensionString {
    /// Create a new vendor extension string.
    ///
    /// For example, setting to `b"COLORRANGE=FULL"` sets the interpretation of
    /// the YUV values to cover the full range (rather a limited "studio swing"
    /// range).
    ///
    /// The argument `x_option` must not contain a space (b' ') character,
    /// otherwise [Error::BadInput] is returned.
    pub fn new(value: Vec<u8>) -> Result<VendorExtensionString, Error> {
        if value.contains(&b' ') {
            return Err(Error::BadInput);
        }
        Ok(VendorExtensionString(value))
    }
    /// Get the vendor extension string.
    pub fn value(&self) -> &[u8] {
        self.0.as_slice()
    }
}

/// Simple ratio structure since stdlib lacks one.
#[derive(Debug, Clone, Copy)]
pub struct Ratio {
    /// Numerator.
    pub num: usize,
    /// Denominator.
    pub den: usize,
}

impl Ratio {
    /// Create a new ratio.
    pub fn new(num: usize, den: usize) -> Ratio {
        Ratio { num, den }
    }

    /// Parse a ratio from a byte slice.
    pub fn parse(value: &[u8]) -> Result<Ratio, Error> {
        let parts: Vec<_> = value.splitn(2, |&b| b == RATIO_SEP).collect();
        if parts.len() != 2 {
            parse_error!(ParseError::General)
        }
        let num = parse_bytes(parts[0])?;
        let den = parse_bytes(parts[1])?;
        Ok(Ratio::new(num, den))
    }
}

impl fmt::Display for Ratio {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.num, self.den)
    }
}

/// Colorspace (color model/pixel format). Only subset of them is supported.
///
/// From libavformat/yuv4mpegenc.c:
///
/// > yuv4mpeg can only handle yuv444p, yuv422p, yuv420p, yuv411p and gray8
/// pixel formats. And using 'strict -1' also yuv444p9, yuv422p9, yuv420p9,
/// yuv444p10, yuv422p10, yuv420p10, yuv444p12, yuv422p12, yuv420p12,
/// yuv444p14, yuv422p14, yuv420p14, yuv444p16, yuv422p16, yuv420p16, gray9,
/// gray10, gray12 and gray16 pixel formats.
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub enum Colorspace {
    /// Grayscale only, 8-bit.
    Cmono,
    /// Grayscale only, 12-bit.
    Cmono12,
    /// 4:2:0 with coincident chroma planes, 8-bit.
    C420,
    /// 4:2:0 with coincident chroma planes, 10-bit.
    C420p10,
    /// 4:2:0 with coincident chroma planes, 12-bit.
    C420p12,
    /// 4:2:0 with biaxially-displaced chroma planes, 8-bit.
    C420jpeg,
    /// 4:2:0 with vertically-displaced chroma planes, 8-bit.
    C420paldv,
    /// Found in some files. Same as `C420`.
    C420mpeg2,
    /// 4:2:2, 8-bit.
    C422,
    /// 4:2:2, 10-bit.
    C422p10,
    /// 4:2:2, 12-bit.
    C422p12,
    /// 4:4:4, 8-bit.
    C444,
    /// 4:4:4, 10-bit.
    C444p10,
    /// 4:4:4, 12-bit.
    C444p12,
}

impl Colorspace {
    /// Return the bit depth per sample
    #[inline]
    pub fn get_bit_depth(self) -> usize {
        match self {
            Colorspace::Cmono
            | Colorspace::C420
            | Colorspace::C422
            | Colorspace::C444
            | Colorspace::C420jpeg
            | Colorspace::C420paldv
            | Colorspace::C420mpeg2 => 8,
            Colorspace::C420p10 | Colorspace::C422p10 | Colorspace::C444p10 => 10,
            Colorspace::Cmono12
            | Colorspace::C420p12
            | Colorspace::C422p12
            | Colorspace::C444p12 => 12,
        }
    }

    /// Return the number of bytes in a sample
    #[inline]
    pub fn get_bytes_per_sample(self) -> usize {
        if self.get_bit_depth() <= 8 {
            1
        } else {
            2
        }
    }
}

fn get_plane_sizes(width: usize, height: usize, colorspace: Colorspace) -> (usize, usize, usize) {
    let y_plane_size = width * height * colorspace.get_bytes_per_sample();

    let c420_chroma_size =
        ((width + 1) / 2) * ((height + 1) / 2) * colorspace.get_bytes_per_sample();
    let c422_chroma_size = ((width + 1) / 2) * height * colorspace.get_bytes_per_sample();

    let c420_sizes = (y_plane_size, c420_chroma_size, c420_chroma_size);
    let c422_sizes = (y_plane_size, c422_chroma_size, c422_chroma_size);
    let c444_sizes = (y_plane_size, y_plane_size, y_plane_size);

    match colorspace {
        Colorspace::Cmono | Colorspace::Cmono12 => (y_plane_size, 0, 0),
        Colorspace::C420
        | Colorspace::C420p10
        | Colorspace::C420p12
        | Colorspace::C420jpeg
        | Colorspace::C420paldv
        | Colorspace::C420mpeg2 => c420_sizes,
        Colorspace::C422 | Colorspace::C422p10 | Colorspace::C422p12 => c422_sizes,
        Colorspace::C444 | Colorspace::C444p10 | Colorspace::C444p12 => c444_sizes,
    }
}

/// Limits on the resources `Decoder` is allowed to use.
#[derive(Clone, Copy, Debug)]
pub struct Limits {
    /// Maximum allowed size of frame buffer, default is 1 GiB.
    pub bytes: usize,
}

impl Default for Limits {
    fn default() -> Limits {
        Limits {
            bytes: 1024 * 1024 * 1024,
        }
    }
}

/// YUV4MPEG2 decoder.
pub struct Decoder<R: Read> {
    reader: R,
    params_buf: Vec<u8>,
    frame_buf: Vec<u8>,
    raw_params: Vec<u8>,
    width: usize,
    height: usize,
    framerate: Ratio,
    pixel_aspect: Ratio,
    colorspace: Colorspace,
    y_len: usize,
    u_len: usize,
}

impl<R: Read> Decoder<R> {
    /// Create a new decoder instance.
    pub fn new(reader: R) -> Result<Decoder<R>, Error> {
        Decoder::new_with_limits(reader, Limits::default())
    }

    /// Create a new decoder instance with custom limits.
    pub fn new_with_limits(mut reader: R, limits: Limits) -> Result<Decoder<R>, Error> {
        let mut params_buf = vec![0; MAX_PARAMS_SIZE];
        let end_params_pos = reader.read_until(TERMINATOR, &mut params_buf)?;
        if end_params_pos < FILE_MAGICK.len() || !params_buf.starts_with(FILE_MAGICK) {
            parse_error!(ParseError::InvalidY4M)
        }
        let raw_params = params_buf[FILE_MAGICK.len()..end_params_pos].to_owned();
        let mut width = 0;
        let mut height = 0;
        // Framerate is actually required per spec, but let's be a bit more
        // permissive as per ffmpeg behavior.
        let mut framerate = Ratio::new(25, 1);
        let mut pixel_aspect = Ratio::new(1, 1);
        let mut colorspace = None;
        // We shouldn't convert it to string because encoding is unspecified.
        for param in raw_params.split(|&b| b == FIELD_SEP) {
            if param.is_empty() {
                continue;
            }
            let (name, value) = (param[0], &param[1..]);
            // TODO(Kagami): interlacing, comment.
            match name {
                b'W' => width = parse_bytes(value)?,
                b'H' => height = parse_bytes(value)?,
                b'F' => framerate = Ratio::parse(value)?,
                b'A' => pixel_aspect = Ratio::parse(value)?,
                b'C' => {
                    colorspace = match value {
                        b"mono" => Some(Colorspace::Cmono),
                        b"mono12" => Some(Colorspace::Cmono12),
                        b"420" => Some(Colorspace::C420),
                        b"420p10" => Some(Colorspace::C420p10),
                        b"420p12" => Some(Colorspace::C420p12),
                        b"422" => Some(Colorspace::C422),
                        b"422p10" => Some(Colorspace::C422p10),
                        b"422p12" => Some(Colorspace::C422p12),
                        b"444" => Some(Colorspace::C444),
                        b"444p10" => Some(Colorspace::C444p10),
                        b"444p12" => Some(Colorspace::C444p12),
                        b"420jpeg" => Some(Colorspace::C420jpeg),
                        b"420paldv" => Some(Colorspace::C420paldv),
                        b"420mpeg2" => Some(Colorspace::C420mpeg2),
                        _ => return Err(Error::UnknownColorspace),
                    }
                }
                _ => {}
            }
        }
        let colorspace = colorspace.unwrap_or(Colorspace::C420);
        if width == 0 || height == 0 {
            parse_error!(ParseError::General)
        }
        let (y_len, u_len, v_len) = get_plane_sizes(width, height, colorspace);
        let frame_size = y_len + u_len + v_len;
        if frame_size > limits.bytes {
            return Err(Error::OutOfMemory);
        }
        let frame_buf = vec![0; frame_size];
        Ok(Decoder {
            reader,
            params_buf,
            frame_buf,
            raw_params,
            width,
            height,
            framerate,
            pixel_aspect,
            colorspace,
            y_len,
            u_len,
        })
    }

    /// Iterate over frames. End of input is indicated by `Error::EOF`.
    pub fn read_frame(&mut self) -> Result<Frame, Error> {
        let end_params_pos = self.reader.read_until(TERMINATOR, &mut self.params_buf)?;
        if end_params_pos < FRAME_MAGICK.len() || !self.params_buf.starts_with(FRAME_MAGICK) {
            parse_error!(ParseError::InvalidY4M)
        }
        // We don't parse frame params currently but user has access to them.
        let start_params_pos = FRAME_MAGICK.len();
        let raw_params = if end_params_pos - start_params_pos > 0 {
            // Check for extra space.
            if self.params_buf[start_params_pos] != FIELD_SEP {
                parse_error!(ParseError::InvalidY4M)
            }
            Some(self.params_buf[start_params_pos + 1..end_params_pos].to_owned())
        } else {
            None
        };
        self.reader.read_exact(&mut self.frame_buf)?;
        Ok(Frame::new(
            [
                &self.frame_buf[0..self.y_len],
                &self.frame_buf[self.y_len..self.y_len + self.u_len],
                &self.frame_buf[self.y_len + self.u_len..],
            ],
            raw_params,
        ))
    }

    /// Return file width.
    #[inline]
    pub fn get_width(&self) -> usize {
        self.width
    }
    /// Return file height.
    #[inline]
    pub fn get_height(&self) -> usize {
        self.height
    }
    /// Return file framerate.
    #[inline]
    pub fn get_framerate(&self) -> Ratio {
        self.framerate
    }
    /// Return file pixel aspect.
    #[inline]
    pub fn get_pixel_aspect(&self) -> Ratio {
        self.pixel_aspect
    }
    /// Return file colorspace.
    ///
    /// **NOTE:** normally all .y4m should have colorspace param, but there are
    /// files encoded without that tag and it's unclear what should we do in
    /// that case. Currently C420 is implied by default as per ffmpeg behavior.
    #[inline]
    pub fn get_colorspace(&self) -> Colorspace {
        self.colorspace
    }
    /// Return file raw parameters.
    #[inline]
    pub fn get_raw_params(&self) -> &[u8] {
        &self.raw_params
    }
    /// Return the bit depth per sample
    #[inline]
    pub fn get_bit_depth(&self) -> usize {
        self.colorspace.get_bit_depth()
    }
    /// Return the number of bytes in a sample
    #[inline]
    pub fn get_bytes_per_sample(&self) -> usize {
        self.colorspace.get_bytes_per_sample()
    }
}

/// A single frame.
#[derive(Debug)]
pub struct Frame<'f> {
    planes: [&'f [u8]; 3],
    raw_params: Option<Vec<u8>>,
}

impl<'f> Frame<'f> {
    /// Create a new frame with optional parameters.
    /// No heap allocations are made.
    pub fn new(planes: [&'f [u8]; 3], raw_params: Option<Vec<u8>>) -> Frame<'f> {
        Frame { planes, raw_params }
    }

    /// Create a new frame from data in 16-bit format.
    pub fn from_u16(planes: [&'f [u16]; 3], raw_params: Option<Vec<u8>>) -> Frame<'f> {
        Frame::new(
            [
                unsafe {
                    std::slice::from_raw_parts::<u8>(
                        planes[0].as_ptr() as *const u8,
                        planes[0].len() * 2,
                    )
                },
                unsafe {
                    std::slice::from_raw_parts::<u8>(
                        planes[1].as_ptr() as *const u8,
                        planes[1].len() * 2,
                    )
                },
                unsafe {
                    std::slice::from_raw_parts::<u8>(
                        planes[2].as_ptr() as *const u8,
                        planes[2].len() * 2,
                    )
                },
            ],
            raw_params,
        )
    }

    /// Return Y (first) plane.
    #[inline]
    pub fn get_y_plane(&self) -> &[u8] {
        self.planes[0]
    }
    /// Return U (second) plane. Empty in case of grayscale.
    #[inline]
    pub fn get_u_plane(&self) -> &[u8] {
        self.planes[1]
    }
    /// Return V (third) plane. Empty in case of grayscale.
    #[inline]
    pub fn get_v_plane(&self) -> &[u8] {
        self.planes[2]
    }
    /// Return frame raw parameters if any.
    #[inline]
    pub fn get_raw_params(&self) -> Option<&[u8]> {
        self.raw_params.as_ref().map(|v| &v[..])
    }
}

/// Encoder builder. Allows to set y4m file parameters using builder pattern.
// TODO(Kagami): Accept all known tags and raw params.
#[derive(Debug)]
pub struct EncoderBuilder {
    width: usize,
    height: usize,
    framerate: Ratio,
    pixel_aspect: Ratio,
    colorspace: Colorspace,
    vendor_extensions: Vec<Vec<u8>>,
}

impl EncoderBuilder {
    /// Create a new encoder builder.
    pub fn new(width: usize, height: usize, framerate: Ratio) -> EncoderBuilder {
        EncoderBuilder {
            width,
            height,
            framerate,
            pixel_aspect: Ratio::new(1, 1),
            colorspace: Colorspace::C420,
            vendor_extensions: vec![],
        }
    }

    /// Specify file colorspace.
    pub fn with_colorspace(mut self, colorspace: Colorspace) -> Self {
        self.colorspace = colorspace;
        self
    }

    /// Specify file pixel aspect.
    pub fn with_pixel_aspect(mut self, pixel_aspect: Ratio) -> Self {
        self.pixel_aspect = pixel_aspect;
        self
    }

    /// Add vendor extension.
    pub fn append_vendor_extension(mut self, x_option: VendorExtensionString) -> Self {
        self.vendor_extensions.push(x_option.0);
        self
    }

    /// Write header to the stream and create encoder instance.
    pub fn write_header<W: Write>(self, mut writer: W) -> Result<Encoder<W>, Error> {
        // XXX(Kagami): Beware that FILE_MAGICK already contains space.
        writer.write_all(FILE_MAGICK)?;
        write!(
            writer,
            "W{} H{} F{}",
            self.width, self.height, self.framerate
        )?;
        if self.pixel_aspect.num != 1 || self.pixel_aspect.den != 1 {
            write!(writer, " A{}", self.pixel_aspect)?;
        }
        for x_option in self.vendor_extensions.iter() {
            write!(writer, " X")?;
            writer.write_all(x_option)?;
        }
        write!(writer, " {:?}", self.colorspace)?;
        writer.write_all(&[TERMINATOR])?;
        let (y_len, u_len, v_len) = get_plane_sizes(self.width, self.height, self.colorspace);
        Ok(Encoder {
            writer,
            y_len,
            u_len,
            v_len,
        })
    }
}

/// YUV4MPEG2 encoder.
pub struct Encoder<W: Write> {
    writer: W,
    y_len: usize,
    u_len: usize,
    v_len: usize,
}

impl<W: Write> Encoder<W> {
    /// Write next frame to the stream.
    pub fn write_frame(&mut self, frame: &Frame) -> Result<(), Error> {
        if frame.get_y_plane().len() != self.y_len
            || frame.get_u_plane().len() != self.u_len
            || frame.get_v_plane().len() != self.v_len
        {
            return Err(Error::BadInput);
        }
        self.writer.write_all(FRAME_MAGICK)?;
        if let Some(params) = frame.get_raw_params() {
            self.writer.write_all(&[FIELD_SEP])?;
            self.writer.write_all(params)?;
        }
        self.writer.write_all(&[TERMINATOR])?;
        self.writer.write_all(frame.get_y_plane())?;
        self.writer.write_all(frame.get_u_plane())?;
        self.writer.write_all(frame.get_v_plane())?;
        Ok(())
    }
}

/// Create a new decoder instance. Alias for `Decoder::new`.
pub fn decode<R: Read>(reader: R) -> Result<Decoder<R>, Error> {
    Decoder::new(reader)
}

/// Create a new encoder builder. Alias for `EncoderBuilder::new`.
pub fn encode(width: usize, height: usize, framerate: Ratio) -> EncoderBuilder {
    EncoderBuilder::new(width, height, framerate)
}

Coverage Report

Created: 2026-02-26 07:34

Line	Count	Source
1		//! # YUV4MPEG2 (.y4m) Encoder/Decoder
2		#![deny(missing_docs)]
3
4		use std::fmt;
5		use std::io;
6		use std::io::Read;
7		use std::io::Write;
8		use std::num;
9		use std::str;
10
11		const MAX_PARAMS_SIZE: usize = 1024;
12		const FILE_MAGICK: &[u8] = b"YUV4MPEG2 ";
13		const FRAME_MAGICK: &[u8] = b"FRAME";
14		const TERMINATOR: u8 = 0x0A;
15		const FIELD_SEP: u8 = b' ';
16		const RATIO_SEP: u8 = b':';
17
18		/// Both encoding and decoding errors.
19		#[derive(Debug)]
20		pub enum Error {
21		/// End of the file. Technically not an error, but it's easier to process
22		/// that way.
23		EOF,
24		/// Bad input parameters provided.
25		BadInput,
26		/// Unknown colorspace (possibly just unimplemented).
27		UnknownColorspace,
28		/// Error while parsing the file/frame header.
29		// TODO(Kagami): Better granularity of parse errors.
30		ParseError(ParseError),
31		/// Error while reading/writing the file.
32		IoError(io::Error),
33		/// Out of memory (limits exceeded).
34		OutOfMemory,
35		}
36
37		impl std::error::Error for crate::Error {
38	0	fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
39	0	match *self {
40	0	Error::EOF => None,
41	0	Error::BadInput => None,
42	0	Error::UnknownColorspace => None,
43	0	Error::ParseError(ref err) => Some(err),
44	0	Error::IoError(ref err) => Some(err),
45	0	Error::OutOfMemory => None,
46		}
47	0	}
48		}
49
50		impl fmt::Display for crate::Error {
51	0	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
52	0	match *self {
53	0	Error::EOF => write!(f, "End of file"),
54	0	Error::BadInput => write!(f, "Bad input parameters provided"),
55	0	Error::UnknownColorspace => write!(f, "Bad input parameters provided"),
56	0	Error::ParseError(ref err) => err.fmt(f),
57	0	Error::IoError(ref err) => err.fmt(f),
58	0	Error::OutOfMemory => write!(f, "Out of memory (limits exceeded)"),
59		}
60	0	}
61		}
62
63		/// Granular ParseError Definiations
64		pub enum ParseError {
65		/// Error reading y4m header
66		InvalidY4M,
67		/// Error parsing int
68		Int,
69		/// Error parsing UTF8
70		Utf8,
71		/// General Parsing Error
72		General,
73		}
74
75		impl std::error::Error for crate::ParseError {
76	0	fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
77	0	match *self {
78	0	ParseError::InvalidY4M => None,
79	0	ParseError::Int => None,
80	0	ParseError::Utf8 => None,
81	0	ParseError::General => None,
82		}
83	0	}
84		}
85
86		impl fmt::Display for ParseError {
87	0	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
88	0	match self {
89	0	ParseError::InvalidY4M => write!(f, "Error parsing y4m header"),
90	0	ParseError::Int => write!(f, "Error parsing Int"),
91	0	ParseError::Utf8 => write!(f, "Error parsing UTF8"),
92	0	ParseError::General => write!(f, "General parsing error"),
93		}
94	0	}
95		}
96
97		impl fmt::Debug for ParseError {
98	0	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
99	0	match self {
100	0	ParseError::InvalidY4M => write!(f, "Error parsing y4m header"),
101	0	ParseError::Int => write!(f, "Error parsing Int"),
102	0	ParseError::Utf8 => write!(f, "Error parsing UTF8"),
103	0	ParseError::General => write!(f, "General parsing error"),
104		}
105	0	}
106		}
107
108		macro_rules! parse_error {
109		($p:expr) => {
110		return Err(Error::ParseError($p))
111		};
112		}
113
114		impl From<io::Error> for Error {
115	0	fn from(err: io::Error) -> Error {
116	0	match err.kind() {
117	0	io::ErrorKind::UnexpectedEof => Error::EOF,
118	0	_ => Error::IoError(err),
119		}
120	0	}
121		}
122
123		impl From<num::ParseIntError> for Error {
124	0	fn from(_: num::ParseIntError) -> Error {
125	0	Error::ParseError(ParseError::Int)
126	0	}
127		}
128
129		impl From<str::Utf8Error> for Error {
130	0	fn from(_: str::Utf8Error) -> Error {
131	0	Error::ParseError(ParseError::Utf8)
132	0	}
133		}
134
135		trait EnhancedRead {
136		fn read_until(&mut self, ch: u8, buf: &mut [u8]) -> Result<usize, Error>;
137		}
138
139		impl<R: Read> EnhancedRead for R {
140		// Current implementation does one `read` call per byte. This might be a
141		// bit slow for long headers but it simplifies things: we don't need to
142		// check whether start of the next frame is already read and so on.
143	0	fn read_until(&mut self, ch: u8, buf: &mut [u8]) -> Result<usize, Error> {
144	0	let mut collected = 0;
145	0	while collected < buf.len() {
146	0	let chunk_size = self.read(&mut buf[collected..=collected])?;
147	0	if chunk_size == 0 {
148	0	return Err(Error::EOF);
149	0	}
150	0	if buf[collected] == ch {
151	0	return Ok(collected);
152	0	}
153	0	collected += chunk_size;
154		}
155	0	parse_error!(ParseError::General)
156	0	}
157		}
158
159	0	fn parse_bytes(buf: &[u8]) -> Result<usize, Error> {
160		// A bit kludgy but seems like there is no other way.
161	0	Ok(str::from_utf8(buf)?.parse()?)
162	0	}
163
164		/// A newtype wrapper around Vec<u8> to ensure validity as a vendor extension.
165		#[derive(Debug, Clone)]
166		pub struct VendorExtensionString(Vec<u8>);
167
168		impl VendorExtensionString {
169		/// Create a new vendor extension string.
170		///
171		/// For example, setting to `b"COLORRANGE=FULL"` sets the interpretation of
172		/// the YUV values to cover the full range (rather a limited "studio swing"
173		/// range).
174		///
175		/// The argument `x_option` must not contain a space (b' ') character,
176		/// otherwise [Error::BadInput] is returned.
177	0	pub fn new(value: Vec<u8>) -> Result<VendorExtensionString, Error> {
178	0	if value.contains(&b' ') {
179	0	return Err(Error::BadInput);
180	0	}
181	0	Ok(VendorExtensionString(value))
182	0	}
183		/// Get the vendor extension string.
184	0	pub fn value(&self) -> &[u8] {
185	0	self.0.as_slice()
186	0	}
187		}
188
189		/// Simple ratio structure since stdlib lacks one.
190		#[derive(Debug, Clone, Copy)]
191		pub struct Ratio {
192		/// Numerator.
193		pub num: usize,
194		/// Denominator.
195		pub den: usize,
196		}
197
198		impl Ratio {
199		/// Create a new ratio.
200	0	pub fn new(num: usize, den: usize) -> Ratio {
201	0	Ratio { num, den }
202	0	}
203
204		/// Parse a ratio from a byte slice.
205	0	pub fn parse(value: &[u8]) -> Result<Ratio, Error> {
206	0	let parts: Vec<_> = value.splitn(2, \|&b\| b == RATIO_SEP).collect();
207	0	if parts.len() != 2 {
208	0	parse_error!(ParseError::General)
209	0	}
210	0	let num = parse_bytes(parts[0])?;
211	0	let den = parse_bytes(parts[1])?;
212	0	Ok(Ratio::new(num, den))
213	0	}
214		}
215
216		impl fmt::Display for Ratio {
217	0	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
218	0	write!(f, "{}:{}", self.num, self.den)
219	0	}
220		}
221
222		/// Colorspace (color model/pixel format). Only subset of them is supported.
223		///
224		/// From libavformat/yuv4mpegenc.c:
225		///
226		/// > yuv4mpeg can only handle yuv444p, yuv422p, yuv420p, yuv411p and gray8
227		/// pixel formats. And using 'strict -1' also yuv444p9, yuv422p9, yuv420p9,
228		/// yuv444p10, yuv422p10, yuv420p10, yuv444p12, yuv422p12, yuv420p12,
229		/// yuv444p14, yuv422p14, yuv420p14, yuv444p16, yuv422p16, yuv420p16, gray9,
230		/// gray10, gray12 and gray16 pixel formats.
231		#[derive(Debug, Clone, Copy)]
232		#[non_exhaustive]
233		pub enum Colorspace {
234		/// Grayscale only, 8-bit.
235		Cmono,
236		/// Grayscale only, 12-bit.
237		Cmono12,
238		/// 4:2:0 with coincident chroma planes, 8-bit.
239		C420,
240		/// 4:2:0 with coincident chroma planes, 10-bit.
241		C420p10,
242		/// 4:2:0 with coincident chroma planes, 12-bit.
243		C420p12,
244		/// 4:2:0 with biaxially-displaced chroma planes, 8-bit.
245		C420jpeg,
246		/// 4:2:0 with vertically-displaced chroma planes, 8-bit.
247		C420paldv,
248		/// Found in some files. Same as `C420`.
249		C420mpeg2,
250		/// 4:2:2, 8-bit.
251		C422,
252		/// 4:2:2, 10-bit.
253		C422p10,
254		/// 4:2:2, 12-bit.
255		C422p12,
256		/// 4:4:4, 8-bit.
257		C444,
258		/// 4:4:4, 10-bit.
259		C444p10,
260		/// 4:4:4, 12-bit.
261		C444p12,
262		}
263
264		impl Colorspace {
265		/// Return the bit depth per sample
266		#[inline]
267	0	pub fn get_bit_depth(self) -> usize {
268	0	match self {
269		Colorspace::Cmono
270		\| Colorspace::C420
271		\| Colorspace::C422
272		\| Colorspace::C444
273		\| Colorspace::C420jpeg
274		\| Colorspace::C420paldv
275	0	\| Colorspace::C420mpeg2 => 8,
276	0	Colorspace::C420p10 \| Colorspace::C422p10 \| Colorspace::C444p10 => 10,
277		Colorspace::Cmono12
278		\| Colorspace::C420p12
279		\| Colorspace::C422p12
280	0	\| Colorspace::C444p12 => 12,
281		}
282	0	}
283
284		/// Return the number of bytes in a sample
285		#[inline]
286	0	pub fn get_bytes_per_sample(self) -> usize {
287	0	if self.get_bit_depth() <= 8 {
288	0	1
289		} else {
290	0	2
291		}
292	0	}
293		}
294
295	0	fn get_plane_sizes(width: usize, height: usize, colorspace: Colorspace) -> (usize, usize, usize) {
296	0	let y_plane_size = width * height * colorspace.get_bytes_per_sample();
297
298	0	let c420_chroma_size =
299	0	((width + 1) / 2) * ((height + 1) / 2) * colorspace.get_bytes_per_sample();
300	0	let c422_chroma_size = ((width + 1) / 2) * height * colorspace.get_bytes_per_sample();
301
302	0	let c420_sizes = (y_plane_size, c420_chroma_size, c420_chroma_size);
303	0	let c422_sizes = (y_plane_size, c422_chroma_size, c422_chroma_size);
304	0	let c444_sizes = (y_plane_size, y_plane_size, y_plane_size);
305
306	0	match colorspace {
307	0	Colorspace::Cmono \| Colorspace::Cmono12 => (y_plane_size, 0, 0),
308		Colorspace::C420
309		\| Colorspace::C420p10
310		\| Colorspace::C420p12
311		\| Colorspace::C420jpeg
312		\| Colorspace::C420paldv
313	0	\| Colorspace::C420mpeg2 => c420_sizes,
314	0	Colorspace::C422 \| Colorspace::C422p10 \| Colorspace::C422p12 => c422_sizes,
315	0	Colorspace::C444 \| Colorspace::C444p10 \| Colorspace::C444p12 => c444_sizes,
316		}
317	0	}
318
319		/// Limits on the resources `Decoder` is allowed to use.
320		#[derive(Clone, Copy, Debug)]
321		pub struct Limits {
322		/// Maximum allowed size of frame buffer, default is 1 GiB.
323		pub bytes: usize,
324		}
325
326		impl Default for Limits {
327	0	fn default() -> Limits {
328	0	Limits {
329	0	bytes: 1024 * 1024 * 1024,
330	0	}
331	0	}
332		}
333
334		/// YUV4MPEG2 decoder.
335		pub struct Decoder<R: Read> {
336		reader: R,
337		params_buf: Vec<u8>,
338		frame_buf: Vec<u8>,
339		raw_params: Vec<u8>,
340		width: usize,
341		height: usize,
342		framerate: Ratio,
343		pixel_aspect: Ratio,
344		colorspace: Colorspace,
345		y_len: usize,
346		u_len: usize,
347		}
348
349		impl<R: Read> Decoder<R> {
350		/// Create a new decoder instance.
351	0	pub fn new(reader: R) -> Result<Decoder<R>, Error> {
352	0	Decoder::new_with_limits(reader, Limits::default())
353	0	}
354
355		/// Create a new decoder instance with custom limits.
356	0	pub fn new_with_limits(mut reader: R, limits: Limits) -> Result<Decoder<R>, Error> {
357	0	let mut params_buf = vec![0; MAX_PARAMS_SIZE];
358	0	let end_params_pos = reader.read_until(TERMINATOR, &mut params_buf)?;
359	0	if end_params_pos < FILE_MAGICK.len() \|\| !params_buf.starts_with(FILE_MAGICK) {
360	0	parse_error!(ParseError::InvalidY4M)
361	0	}
362	0	let raw_params = params_buf[FILE_MAGICK.len()..end_params_pos].to_owned();
363	0	let mut width = 0;
364	0	let mut height = 0;
365		// Framerate is actually required per spec, but let's be a bit more
366		// permissive as per ffmpeg behavior.
367	0	let mut framerate = Ratio::new(25, 1);
368	0	let mut pixel_aspect = Ratio::new(1, 1);
369	0	let mut colorspace = None;
370		// We shouldn't convert it to string because encoding is unspecified.
371	0	for param in raw_params.split(\|&b\| b == FIELD_SEP) {
372	0	if param.is_empty() {
373	0	continue;
374	0	}
375	0	let (name, value) = (param[0], &param[1..]);
376		// TODO(Kagami): interlacing, comment.
377	0	match name {
378	0	b'W' => width = parse_bytes(value)?,
379	0	b'H' => height = parse_bytes(value)?,
380	0	b'F' => framerate = Ratio::parse(value)?,
381	0	b'A' => pixel_aspect = Ratio::parse(value)?,
382		b'C' => {
383	0	colorspace = match value {
384	0	b"mono" => Some(Colorspace::Cmono),
385	0	b"mono12" => Some(Colorspace::Cmono12),
386	0	b"420" => Some(Colorspace::C420),
387	0	b"420p10" => Some(Colorspace::C420p10),
388	0	b"420p12" => Some(Colorspace::C420p12),
389	0	b"422" => Some(Colorspace::C422),
390	0	b"422p10" => Some(Colorspace::C422p10),
391	0	b"422p12" => Some(Colorspace::C422p12),
392	0	b"444" => Some(Colorspace::C444),
393	0	b"444p10" => Some(Colorspace::C444p10),
394	0	b"444p12" => Some(Colorspace::C444p12),
395	0	b"420jpeg" => Some(Colorspace::C420jpeg),
396	0	b"420paldv" => Some(Colorspace::C420paldv),
397	0	b"420mpeg2" => Some(Colorspace::C420mpeg2),
398	0	_ => return Err(Error::UnknownColorspace),
399		}
400		}
401	0	_ => {}
402		}
403		}
404	0	let colorspace = colorspace.unwrap_or(Colorspace::C420);
405	0	if width == 0 \|\| height == 0 {
406	0	parse_error!(ParseError::General)
407	0	}
408	0	let (y_len, u_len, v_len) = get_plane_sizes(width, height, colorspace);
409	0	let frame_size = y_len + u_len + v_len;
410	0	if frame_size > limits.bytes {
411	0	return Err(Error::OutOfMemory);
412	0	}
413	0	let frame_buf = vec![0; frame_size];
414	0	Ok(Decoder {
415	0	reader,
416	0	params_buf,
417	0	frame_buf,
418	0	raw_params,
419	0	width,
420	0	height,
421	0	framerate,
422	0	pixel_aspect,
423	0	colorspace,
424	0	y_len,
425	0	u_len,
426	0	})
427	0	}
428
429		/// Iterate over frames. End of input is indicated by `Error::EOF`.
430	0	pub fn read_frame(&mut self) -> Result<Frame, Error> {
431	0	let end_params_pos = self.reader.read_until(TERMINATOR, &mut self.params_buf)?;
432	0	if end_params_pos < FRAME_MAGICK.len() \|\| !self.params_buf.starts_with(FRAME_MAGICK) {
433	0	parse_error!(ParseError::InvalidY4M)
434	0	}
435		// We don't parse frame params currently but user has access to them.
436	0	let start_params_pos = FRAME_MAGICK.len();
437	0	let raw_params = if end_params_pos - start_params_pos > 0 {
438		// Check for extra space.
439	0	if self.params_buf[start_params_pos] != FIELD_SEP {
440	0	parse_error!(ParseError::InvalidY4M)
441	0	}
442	0	Some(self.params_buf[start_params_pos + 1..end_params_pos].to_owned())
443		} else {
444	0	None
445		};
446	0	self.reader.read_exact(&mut self.frame_buf)?;
447	0	Ok(Frame::new(
448	0	[
449	0	&self.frame_buf[0..self.y_len],
450	0	&self.frame_buf[self.y_len..self.y_len + self.u_len],
451	0	&self.frame_buf[self.y_len + self.u_len..],
452	0	],
453	0	raw_params,
454	0	))
455	0	}
456
457		/// Return file width.
458		#[inline]
459	0	pub fn get_width(&self) -> usize {
460	0	self.width
461	0	}
462		/// Return file height.
463		#[inline]
464	0	pub fn get_height(&self) -> usize {
465	0	self.height
466	0	}
467		/// Return file framerate.
468		#[inline]
469	0	pub fn get_framerate(&self) -> Ratio {
470	0	self.framerate
471	0	}
472		/// Return file pixel aspect.
473		#[inline]
474	0	pub fn get_pixel_aspect(&self) -> Ratio {
475	0	self.pixel_aspect
476	0	}
477		/// Return file colorspace.
478		///
479		/// NOTE: normally all .y4m should have colorspace param, but there are
480		/// files encoded without that tag and it's unclear what should we do in
481		/// that case. Currently C420 is implied by default as per ffmpeg behavior.
482		#[inline]
483	0	pub fn get_colorspace(&self) -> Colorspace {
484	0	self.colorspace
485	0	}
486		/// Return file raw parameters.
487		#[inline]
488	0	pub fn get_raw_params(&self) -> &[u8] {
489	0	&self.raw_params
490	0	}
491		/// Return the bit depth per sample
492		#[inline]
493	0	pub fn get_bit_depth(&self) -> usize {
494	0	self.colorspace.get_bit_depth()
495	0	}
496		/// Return the number of bytes in a sample
497		#[inline]
498	0	pub fn get_bytes_per_sample(&self) -> usize {
499	0	self.colorspace.get_bytes_per_sample()
500	0	}
501		}
502
503		/// A single frame.
504		#[derive(Debug)]
505		pub struct Frame<'f> {
506		planes: [&'f [u8]; 3],
507		raw_params: Option<Vec<u8>>,
508		}
509
510		impl<'f> Frame<'f> {
511		/// Create a new frame with optional parameters.
512		/// No heap allocations are made.
513	0	pub fn new(planes: [&'f [u8]; 3], raw_params: Option<Vec<u8>>) -> Frame<'f> {
514	0	Frame { planes, raw_params }
515	0	}
516
517		/// Create a new frame from data in 16-bit format.
518	0	pub fn from_u16(planes: [&'f [u16]; 3], raw_params: Option<Vec<u8>>) -> Frame<'f> {
519	0	Frame::new(
520	0	[
521	0	unsafe {
522	0	std::slice::from_raw_parts::<u8>(
523	0	planes[0].as_ptr() as *const u8,
524	0	planes[0].len() * 2,
525	0	)
526	0	},
527	0	unsafe {
528	0	std::slice::from_raw_parts::<u8>(
529	0	planes[1].as_ptr() as *const u8,
530	0	planes[1].len() * 2,
531	0	)
532	0	},
533	0	unsafe {
534	0	std::slice::from_raw_parts::<u8>(
535	0	planes[2].as_ptr() as *const u8,
536	0	planes[2].len() * 2,
537	0	)
538	0	},
539	0	],
540	0	raw_params,
541		)
542	0	}
543
544		/// Return Y (first) plane.
545		#[inline]
546	0	pub fn get_y_plane(&self) -> &[u8] {
547	0	self.planes[0]
548	0	}
549		/// Return U (second) plane. Empty in case of grayscale.
550		#[inline]
551	0	pub fn get_u_plane(&self) -> &[u8] {
552	0	self.planes[1]
553	0	}
554		/// Return V (third) plane. Empty in case of grayscale.
555		#[inline]
556	0	pub fn get_v_plane(&self) -> &[u8] {
557	0	self.planes[2]
558	0	}
559		/// Return frame raw parameters if any.
560		#[inline]
561	0	pub fn get_raw_params(&self) -> Option<&[u8]> {
562	0	self.raw_params.as_ref().map(\|v\| &v[..])
563	0	}
564		}
565
566		/// Encoder builder. Allows to set y4m file parameters using builder pattern.
567		// TODO(Kagami): Accept all known tags and raw params.
568		#[derive(Debug)]
569		pub struct EncoderBuilder {
570		width: usize,
571		height: usize,
572		framerate: Ratio,
573		pixel_aspect: Ratio,
574		colorspace: Colorspace,
575		vendor_extensions: Vec<Vec<u8>>,
576		}
577
578		impl EncoderBuilder {
579		/// Create a new encoder builder.
580	0	pub fn new(width: usize, height: usize, framerate: Ratio) -> EncoderBuilder {
581	0	EncoderBuilder {
582	0	width,
583	0	height,
584	0	framerate,
585	0	pixel_aspect: Ratio::new(1, 1),
586	0	colorspace: Colorspace::C420,
587	0	vendor_extensions: vec![],
588	0	}
589	0	}
590
591		/// Specify file colorspace.
592	0	pub fn with_colorspace(mut self, colorspace: Colorspace) -> Self {
593	0	self.colorspace = colorspace;
594	0	self
595	0	}
596
597		/// Specify file pixel aspect.
598	0	pub fn with_pixel_aspect(mut self, pixel_aspect: Ratio) -> Self {
599	0	self.pixel_aspect = pixel_aspect;
600	0	self
601	0	}
602
603		/// Add vendor extension.
604	0	pub fn append_vendor_extension(mut self, x_option: VendorExtensionString) -> Self {
605	0	self.vendor_extensions.push(x_option.0);
606	0	self
607	0	}
608
609		/// Write header to the stream and create encoder instance.
610	0	pub fn write_header<W: Write>(self, mut writer: W) -> Result<Encoder<W>, Error> {
611		// XXX(Kagami): Beware that FILE_MAGICK already contains space.
612	0	writer.write_all(FILE_MAGICK)?;
613	0	write!(
614	0	writer,
615	0	"W{} H{} F{}",
616		self.width, self.height, self.framerate
617	0	)?;
618	0	if self.pixel_aspect.num != 1 \|\| self.pixel_aspect.den != 1 {
619	0	write!(writer, " A{}", self.pixel_aspect)?;
620	0	}
621	0	for x_option in self.vendor_extensions.iter() {
622	0	write!(writer, " X")?;
623	0	writer.write_all(x_option)?;
624		}
625	0	write!(writer, " {:?}", self.colorspace)?;
626	0	writer.write_all(&[TERMINATOR])?;
627	0	let (y_len, u_len, v_len) = get_plane_sizes(self.width, self.height, self.colorspace);
628	0	Ok(Encoder {
629	0	writer,
630	0	y_len,
631	0	u_len,
632	0	v_len,
633	0	})
634	0	}
635		}
636
637		/// YUV4MPEG2 encoder.
638		pub struct Encoder<W: Write> {
639		writer: W,
640		y_len: usize,
641		u_len: usize,
642		v_len: usize,
643		}
644
645		impl<W: Write> Encoder<W> {
646		/// Write next frame to the stream.
647	0	pub fn write_frame(&mut self, frame: &Frame) -> Result<(), Error> {
648	0	if frame.get_y_plane().len() != self.y_len
649	0	\|\| frame.get_u_plane().len() != self.u_len
650	0	\|\| frame.get_v_plane().len() != self.v_len
651		{
652	0	return Err(Error::BadInput);
653	0	}
654	0	self.writer.write_all(FRAME_MAGICK)?;
655	0	if let Some(params) = frame.get_raw_params() {
656	0	self.writer.write_all(&[FIELD_SEP])?;
657	0	self.writer.write_all(params)?;
658	0	}
659	0	self.writer.write_all(&[TERMINATOR])?;
660	0	self.writer.write_all(frame.get_y_plane())?;
661	0	self.writer.write_all(frame.get_u_plane())?;
662	0	self.writer.write_all(frame.get_v_plane())?;
663	0	Ok(())
664	0	}
665		}
666
667		/// Create a new decoder instance. Alias for `Decoder::new`.
668	0	pub fn decode<R: Read>(reader: R) -> Result<Decoder<R>, Error> {
669	0	Decoder::new(reader)
670	0	}
671
672		/// Create a new encoder builder. Alias for `EncoderBuilder::new`.
673	0	pub fn encode(width: usize, height: usize, framerate: Ratio) -> EncoderBuilder {
674	0	EncoderBuilder::new(width, height, framerate)
675	0	}