ivf/
lib.rs

1// Copyright (c) 2001-2016, Alliance for Open Media. All rights reserved
2// Copyright (c) 2017-2022, The rav1e contributors. All rights reserved
3//
4// This source code is subject to the terms of the BSD 2 Clause License and
5// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6// was not distributed with this source code in the LICENSE file, you can
7// obtain it at www.aomedia.org/license/software. If the Alliance for Open
8// Media Patent License 1.0 was not distributed with this source code in the
9// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10
11//! Simple ivf muxer
12
13use std::io;
14
15use bitstream_io::{BitRead, BitReader, BitWrite, BitWriter, LittleEndian};
16
17/// # Panics
18///
19/// - If header cannot be written to output file.
20pub fn write_ivf_header(
21  output_file: &mut dyn io::Write, width: usize, height: usize,
22  framerate_num: usize, framerate_den: usize,
23) {
24  let mut bw = BitWriter::endian(output_file, LittleEndian);
25  bw.write_bytes(b"DKIF").unwrap();
26  bw.write::<16, _>(0).unwrap(); // version
27  bw.write::<16, _>(32).unwrap(); // version
28  bw.write_bytes(b"AV01").unwrap();
29  bw.write::<16, _>(width as u16).unwrap();
30  bw.write::<16, _>(height as u16).unwrap();
31  bw.write::<32, _>(framerate_num as u32).unwrap();
32  bw.write::<32, _>(framerate_den as u32).unwrap();
33  bw.write::<32, _>(0).unwrap();
34  bw.write::<32, _>(0).unwrap();
35}
36
37/// # Panics
38///
39/// - If frame cannot be written to output file.
40pub fn write_ivf_frame(
41  output_file: &mut dyn io::Write, pts: u64, data: &[u8],
42) {
43  let mut bw = BitWriter::endian(output_file, LittleEndian);
44  bw.write::<32, _>(data.len() as u32).unwrap();
45  bw.write::<64, _>(pts).unwrap();
46  bw.write_bytes(data).unwrap();
47}
48
49#[derive(Debug, PartialEq, Eq)]
50pub struct Header {
51  pub tag: [u8; 4],
52  pub w: u16,
53  pub h: u16,
54  pub timebase_num: u32,
55  pub timebase_den: u32,
56}
57
58/// # Errors
59///
60/// - Returns `io::Error` if packet cannot be read
61/// - Returns `io::ErrorKind::InvalidData` if header signature is invalid
62pub fn read_header(r: &mut dyn io::Read) -> io::Result<Header> {
63  let mut br = BitReader::endian(r, LittleEndian);
64
65  let mut signature = [0u8; 4];
66  let mut tag = [0u8; 4];
67
68  br.read_bytes(&mut signature)?;
69
70  if &signature != b"DKIF" {
71    return Err(io::ErrorKind::InvalidData.into());
72  }
73
74  let _v0: u16 = br.read::<16, _>()?;
75  let _v1: u16 = br.read::<16, _>()?;
76  br.read_bytes(&mut tag)?;
77
78  let w: u16 = br.read::<16, _>()?;
79  let h: u16 = br.read::<16, _>()?;
80
81  let timebase_den: u32 = br.read::<32, _>()?;
82  let timebase_num: u32 = br.read::<32, _>()?;
83
84  let _: u32 = br.read::<32, _>()?;
85  let _: u32 = br.read::<32, _>()?;
86
87  Ok(Header { tag, w, h, timebase_num, timebase_den })
88}
89
90pub struct Packet {
91  pub data: Box<[u8]>,
92  pub pts: u64,
93}
94
95/// # Errors
96///
97/// - Returns `io::Error` if packet cannot be read
98pub fn read_packet(r: &mut dyn io::Read) -> io::Result<Packet> {
99  let mut br = BitReader::endian(r, LittleEndian);
100
101  let len: u32 = br.read::<32, _>()?;
102  let pts: u64 = br.read::<64, _>()?;
103  let mut buf = vec![0u8; len as usize];
104
105  br.read_bytes(&mut buf)?;
106
107  Ok(Packet { data: buf.into_boxed_slice(), pts })
108}
109
110#[cfg(test)]
111mod tests {
112  use std::io::{BufReader, ErrorKind::InvalidData};
113
114  use crate::{read_header, read_packet};
115
116  #[test]
117  fn read_invalid_headers() {
118    // Invalid magic.
119    let mut br = BufReader::new(&b"FIKD"[..]);
120    let result = read_header(&mut br).map_err(|e| e.kind());
121    let expected = Err(InvalidData);
122    assert_eq!(result, expected);
123  }
124
125  #[test]
126  fn read_valid_headers() {
127    let bytes: [u8; 32] = [
128      0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x41, 0x56, 0x30, 0x31,
129      0x80, 0x07, 0x38, 0x04, 0x18, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
130      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
131    ];
132
133    let mut br = BufReader::new(&bytes[..]);
134    let header = read_header(&mut br).unwrap();
135    assert_eq!(header.tag, [0x41, 0x56, 0x30, 0x31]);
136    assert_eq!(header.w, 1920);
137    assert_eq!(header.h, 1080);
138    assert_eq!(header.timebase_num, 1);
139    assert_eq!(header.timebase_den, 24);
140  }
141
142  #[test]
143  fn read_valid_packet() {
144    let bytes: [u8; 13] = [
145      0x01, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
146      0x01,
147    ];
148    let mut br = BufReader::new(&bytes[..]);
149    let packet = read_packet(&mut br).unwrap();
150    assert_eq!(packet.pts, 3u64);
151  }
152}