1use std::io;
14
15use bitstream_io::{BitRead, BitReader, BitWrite, BitWriter, LittleEndian};
16
17pub 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(); bw.write::<16, _>(32).unwrap(); 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
37pub 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
58pub 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
95pub 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 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}