falcon_ffmpeg/testing/
av1.rs

1//! AV1 test clips in the containers falcon meets them in. The build has no
2//! AV1 encoder, so the pictures come from an IVF file encoded elsewhere
3//! (`testdata/av1.ivf`); these functions only rewrap its temporal units
4//! as WebM, Matroska or MP4, writing every box and element by hand.
5
6/// An AV1 stream: one temporal unit (the OBUs of one picture) per frame.
7#[derive(Clone, Debug)]
8pub struct Av1Clip {
9    pub width: u32,
10    pub height: u32,
11    pub fps: u32,
12    pub frames: Vec<Vec<u8>>,
13}
14
15/// OBU types used below.
16const OBU_SEQUENCE_HEADER: u8 = 1;
17
18/// Reads a LEB128 number from the front of `b`: the value and its length.
19fn leb128(b: &[u8]) -> Option<(usize, usize)> {
20    let mut v = 0usize;
21    for (i, byte) in b.iter().enumerate().take(8) {
22        v |= usize::from(byte & 0x7f) << (7 * i);
23        if byte & 0x80 == 0 {
24            return Some((v, i + 1));
25        }
26    }
27    None
28}
29
30/// The OBUs of a temporal unit in the low-overhead format (each with a
31/// size field), as (type, whole OBU, payload).
32fn obus(tu: &[u8]) -> Vec<(u8, &[u8], &[u8])> {
33    let mut out = Vec::new();
34    let mut p = 0;
35    while p < tu.len() {
36        let header = tu[p];
37        let kind = (header >> 3) & 0x0f;
38        let ext = usize::from(header >> 2 & 1);
39        let Some((size, n)) = leb128(&tu[p + 1 + ext..]) else { break };
40        let start = p + 1 + ext + n;
41        let end = (start + size).min(tu.len());
42        out.push((kind, &tu[p..end], &tu[start..end]));
43        p = end;
44    }
45    out
46}
47
48impl Av1Clip {
49    /// Parses an IVF file (32-byte header, then a 12-byte header before
50    /// each frame) holding AV1 at a whole number of frames per second.
51    pub fn from_ivf(b: &[u8]) -> Option<Self> {
52        let u16_at = |p: usize| u16::from_le_bytes([b[p], b[p + 1]]);
53        let u32_at = |p: usize| u32::from_le_bytes(b[p..p + 4].try_into().unwrap());
54        if b.len() < 32 || &b[0..4] != b"DKIF" || &b[8..12] != b"AV01" {
55            return None;
56        }
57        let (width, height) = (u32::from(u16_at(12)), u32::from(u16_at(14)));
58        let (rate, scale) = (u32_at(16), u32_at(20).max(1));
59        let mut frames = Vec::new();
60        let mut p = usize::from(u16_at(6));
61        while p + 12 <= b.len() {
62            let size = u32_at(p) as usize;
63            frames.push(b.get(p + 12..p + 12 + size)?.to_vec());
64            p += 12 + size;
65        }
66        Some(Av1Clip { width, height, fps: rate / scale, frames })
67    }
68
69    /// Whether frame `i` starts a new sequence: it carries a sequence
70    /// header, which encoders (rav1e among them) repeat at each keyframe.
71    pub fn is_key(&self, i: usize) -> bool {
72        obus(&self.frames[i]).iter().any(|(kind, _, _)| *kind == OBU_SEQUENCE_HEADER)
73    }
74
75    /// The `av1C` record (AV1 codec configuration) that MP4 and Matroska
76    /// carry: profile, level and tier read from the sequence header, the
77    /// colour fields for 8-bit 4:2:0, and the sequence header OBU itself.
78    pub fn av1c(&self) -> Vec<u8> {
79        let (obu, payload) = obus(&self.frames[0])
80            .into_iter()
81            .find(|(kind, _, _)| *kind == OBU_SEQUENCE_HEADER)
82            .map(|(_, obu, payload)| (obu, payload))
83            .expect("the first frame carries the sequence header");
84        let bit = |n: usize| payload[n / 8] >> (7 - n % 8) & 1;
85        let bits = |from: usize, count: usize| (from..from + count).fold(0u8, |v, n| v << 1 | bit(n));
86        let profile = bits(0, 3);
87        // Without timing info: one bit for the display delay flag, five for
88        // the operating point count, twelve for the first one's idc.
89        let reduced = bit(4) == 1;
90        let (level, tier) = if reduced {
91            (bits(5, 5), 0)
92        } else {
93            assert_eq!(bit(5), 0, "timing info is not parsed");
94            let level = bits(25, 5);
95            (level, if level > 7 { bit(30) } else { 0 })
96        };
97        let mut v = vec![0x81, profile << 5 | level, tier << 7 | 0b0000_1100, 0];
98        v.extend_from_slice(obu);
99        v
100    }
101
102    fn duration_ms(&self) -> f64 {
103        self.frames.len() as f64 * 1000.0 / f64::from(self.fps)
104    }
105
106    fn frame_ms(&self, i: usize) -> u64 {
107        i as u64 * 1000 / u64::from(self.fps)
108    }
109
110    /// The clip as a WebM (`webm`) or Matroska file: one cluster per
111    /// keyframe, with cues pointing at them.
112    pub fn matroska(&self, webm: bool) -> Vec<u8> {
113        let doc: &[u8] = if webm { b"webm" } else { b"matroska" };
114        let header = el(
115            0x1A45DFA3,
116            &[uint(0x4286, 1), uint(0x42F7, 1), uint(0x42F2, 4), uint(0x42F3, 8), el(0x4282, doc), uint(0x4287, 4), uint(0x4285, 2)]
117                .concat(),
118        );
119        let info = el(
120            0x1549A966,
121            &[uint(0x2AD7B1, 1_000_000), el(0x4489, &self.duration_ms().to_be_bytes()), el(0x4D80, b"falcon tests"), el(0x5741, b"falcon tests")]
122                .concat(),
123        );
124        let video = el(0xE0, &[uint(0xB0, u64::from(self.width)), uint(0xBA, u64::from(self.height))].concat());
125        let track = el(
126            0xAE,
127            &[
128                uint(0xD7, 1),
129                uint(0x73C5, 1),
130                uint(0x83, 1),
131                el(0x86, b"V_AV1"),
132                el(0x63A2, &self.av1c()),
133                uint(0x23E383, 1_000_000_000 / u64::from(self.fps)),
134                video,
135            ]
136            .concat(),
137        );
138        let tracks = el(0x1654AE6B, &track);
139
140        let mut clusters = Vec::new();
141        let mut starts = Vec::new();
142        let mut i = 0;
143        while i < self.frames.len() {
144            let start_ms = self.frame_ms(i);
145            starts.push((start_ms, clusters.len()));
146            let mut body = uint(0xE7, start_ms);
147            loop {
148                let mut block = vec![0x81];
149                block.extend_from_slice(&((self.frame_ms(i) - start_ms) as i16).to_be_bytes());
150                block.push(if self.is_key(i) { 0x80 } else { 0 });
151                block.extend_from_slice(&self.frames[i]);
152                body.extend(el(0xA3, &block));
153                i += 1;
154                if i == self.frames.len() || self.is_key(i) {
155                    break;
156                }
157            }
158            clusters.extend(el(0x1F43B675, &body));
159        }
160        // Cluster positions count from the start of the segment's data.
161        let before = info.len() + tracks.len();
162        let cues: Vec<u8> = starts
163            .iter()
164            .map(|&(ms, at)| el(0xBB, &[uint(0xB3, ms), el(0xB7, &[uint(0xF7, 1), uint(0xF1, (before + at) as u64)].concat())].concat()))
165            .collect::<Vec<_>>()
166            .concat();
167        let segment = el(0x18538067, &[info, tracks, clusters, el(0x1C53BB6B, &cues)].concat());
168        [header, segment].concat()
169    }
170
171    /// The clip as an MP4 file: `moov` first, then one chunk of samples
172    /// in `mdat`, with keyframes listed in `stss`.
173    pub fn mp4(&self) -> Vec<u8> {
174        const TIMESCALE: u32 = 12_800;
175        let delta = TIMESCALE / self.fps;
176        let n = self.frames.len() as u32;
177        let ms = self.duration_ms() as u32;
178        let ftyp = mp4_box(b"ftyp", &[&b"isom"[..], &be32(0x200), b"isom", b"iso6", b"av01", b"mp41"].concat());
179        let matrix: Vec<u8> = [0x0001_0000u32, 0, 0, 0, 0x0001_0000, 0, 0, 0, 0x4000_0000].iter().flat_map(|v| v.to_be_bytes()).collect();
180
181        let moov = |chunk_offset: u32| {
182            let mvhd = full_box(
183                b"mvhd",
184                0,
185                &[&be32(0)[..], &be32(0), &be32(1000), &be32(ms), &be32(0x0001_0000), &[1, 0], &[0; 10], &matrix, &[0; 24], &be32(2)].concat(),
186            );
187            let tkhd = full_box(
188                b"tkhd",
189                3,
190                &[
191                    &be32(0)[..],
192                    &be32(0),
193                    &be32(1),
194                    &be32(0),
195                    &be32(ms),
196                    &[0; 8],
197                    &[0; 8],
198                    &matrix,
199                    &be32(self.width << 16),
200                    &be32(self.height << 16),
201                ]
202                .concat(),
203            );
204            let mdhd = full_box(b"mdhd", 0, &[&be32(0)[..], &be32(0), &be32(TIMESCALE), &be32(n * delta), &[0x55, 0xC4, 0, 0]].concat());
205            let hdlr = full_box(b"hdlr", 0, &[&be32(0)[..], b"vide", &[0; 12], b"VideoHandler\0"].concat());
206            let vmhd = full_box(b"vmhd", 1, &[0; 8]);
207            let dinf = mp4_box(b"dinf", &full_box(b"dref", 0, &[&be32(1)[..], &full_box(b"url ", 1, &[])].concat()));
208            let entry = mp4_box(
209                b"av01",
210                &[
211                    &[0u8; 6][..],
212                    &[0, 1],
213                    &[0; 16],
214                    &(self.width as u16).to_be_bytes(),
215                    &(self.height as u16).to_be_bytes(),
216                    &be32(0x0048_0000),
217                    &be32(0x0048_0000),
218                    &be32(0),
219                    &[0, 1],
220                    &[0; 32],
221                    &[0, 0x18, 0xff, 0xff],
222                    &mp4_box(b"av1C", &self.av1c()),
223                ]
224                .concat(),
225            );
226            let stsd = full_box(b"stsd", 0, &[&be32(1)[..], &entry].concat());
227            let stts = full_box(b"stts", 0, &[be32(1), be32(n), be32(delta)].concat());
228            let keys: Vec<u32> = (0..self.frames.len()).filter(|&i| self.is_key(i)).map(|i| i as u32 + 1).collect();
229            let stss = full_box(b"stss", 0, &[be32(keys.len() as u32), keys.iter().flat_map(|k| k.to_be_bytes()).collect()].concat());
230            let stsc = full_box(b"stsc", 0, &[be32(1), be32(1), be32(n), be32(1)].concat());
231            let sizes: Vec<u8> = self.frames.iter().flat_map(|f| (f.len() as u32).to_be_bytes()).collect();
232            let stsz = full_box(b"stsz", 0, &[be32(0), be32(n), sizes].concat());
233            let stco = full_box(b"stco", 0, &[be32(1), be32(chunk_offset)].concat());
234            let stbl = mp4_box(b"stbl", &[stsd, stts, stss, stsc, stsz, stco].concat());
235            let minf = mp4_box(b"minf", &[vmhd, dinf, stbl].concat());
236            let mdia = mp4_box(b"mdia", &[mdhd, hdlr, minf].concat());
237            let trak = mp4_box(b"trak", &[tkhd, mdia].concat());
238            mp4_box(b"moov", &[mvhd, trak].concat())
239        };
240        // The offset doesn't change moov's size, so measure with a dummy.
241        let offset = (ftyp.len() + moov(0).len() + 8) as u32;
242        let mdat = mp4_box(b"mdat", &self.frames.concat());
243        [ftyp, moov(offset), mdat].concat()
244    }
245}
246
247/// A Matroska element: ID, size (as an eight-byte vint) and data.
248pub(super) fn el(id: u32, data: &[u8]) -> Vec<u8> {
249    let id = id.to_be_bytes();
250    let skip = id.iter().take_while(|b| **b == 0).count();
251    let mut v = id[skip..].to_vec();
252    v.push(0x01);
253    v.extend_from_slice(&(data.len() as u64).to_be_bytes()[1..]);
254    v.extend_from_slice(data);
255    v
256}
257
258/// A Matroska unsigned integer element, in eight bytes.
259pub(super) fn uint(id: u32, value: u64) -> Vec<u8> {
260    el(id, &value.to_be_bytes())
261}
262
263fn be32(v: u32) -> Vec<u8> {
264    v.to_be_bytes().to_vec()
265}
266
267fn mp4_box(kind: &[u8; 4], data: &[u8]) -> Vec<u8> {
268    let mut v = be32(8 + data.len() as u32);
269    v.extend_from_slice(kind);
270    v.extend_from_slice(data);
271    v
272}
273
274/// A box with a version (0 here) and 24 bits of flags.
275fn full_box(kind: &[u8; 4], flags: u32, data: &[u8]) -> Vec<u8> {
276    mp4_box(kind, &[&be32(flags & 0x00ff_ffff)[..], data].concat())
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282
283    fn clip() -> Av1Clip {
284        Av1Clip::from_ivf(include_bytes!("../../testdata/av1.ivf")).unwrap()
285    }
286
287    #[test]
288    fn reads_the_ivf_fixture() {
289        let c = clip();
290        assert_eq!((c.width, c.height, c.fps, c.frames.len()), (64, 48, 25, 25));
291        let keys: Vec<usize> = (0..c.frames.len()).filter(|&i| c.is_key(i)).collect();
292        assert_eq!(keys, [0, 13]);
293    }
294
295    #[test]
296    fn av1c_carries_the_sequence_header() {
297        let v = clip().av1c();
298        assert_eq!(v[0], 0x81, "marker and version 1");
299        assert_eq!(v[1] >> 5, 0, "main profile");
300        assert_eq!(v[4] >> 3 & 0x0f, OBU_SEQUENCE_HEADER);
301    }
302
303    #[test]
304    fn containers_have_their_signatures() {
305        let c = clip();
306        let mkv = c.matroska(true);
307        assert_eq!(&mkv[0..4], [0x1A, 0x45, 0xDF, 0xA3]);
308        assert!(mkv.windows(5).any(|w| w == b"V_AV1"));
309        let mp4 = c.mp4();
310        assert_eq!(&mp4[4..8], b"ftyp");
311        let total: usize = c.frames.iter().map(Vec::len).sum();
312        assert_eq!(&mp4[mp4.len() - total - 4..mp4.len() - total], b"mdat");
313    }
314}
falcon_ffmpeg/testing/av1.rs

falcon_ffmpeg/testing/
av1.rs
