Skip to main content

otf_pixels_codec_avif/av1/
still.rs

1//! The still-picture front-end driver: OBUs to parsed headers.
2//!
3//! This ties the layers below together the way an AVIF still is actually laid
4//! out. The sequence header arrives in the `av1C` box's configuration OBUs; the
5//! coded frame arrives as the primary item's data, as an `OBU_FRAME` (a frame
6//! header immediately followed by its tile group) or as a separate
7//! `OBU_FRAME_HEADER` plus one or more `OBU_TILE_GROUP`s. Either way this
8//! returns the two parsed headers and the byte ranges of the tile groups that
9//! follow them, which is where P2 stops and reconstruction picks up.
10
11use super::bits::BitReader;
12use super::frame::FrameHeader;
13use super::obu::{Obu, ObuType};
14use super::seq::SequenceHeader;
15use core::ops::Range;
16use otf_pixels_core::{PixelsError, Result};
17
18/// The parsed headers of a still picture, plus a locator for its tile data.
19#[derive(Debug, Clone)]
20pub struct StillPicture {
21    /// The sequence header governing the frame.
22    pub sequence: SequenceHeader,
23    /// The frame (uncompressed) header.
24    pub frame: FrameHeader,
25    /// Each tile group's bytes, in order, as ranges of the `frame_data` passed
26    /// to [`StillPicture::parse`].
27    ///
28    /// For an `OBU_FRAME` the first range lies inside that OBU's payload, past
29    /// the byte-aligned frame header; each `OBU_TILE_GROUP` adds its payload.
30    pub tile_groups: Vec<Range<usize>>,
31}
32
33impl StillPicture {
34    /// The tile groups' bytes, ready for [`super::decode_still`].
35    ///
36    /// # Errors
37    ///
38    /// Returns [`PixelsError::Malformed`] if a range lies outside
39    /// `frame_data`, which means it is not the buffer these were parsed from.
40    pub fn tile_group_data<'a>(&self, frame_data: &'a [u8]) -> Result<Vec<&'a [u8]>> {
41        self.tile_groups
42            .iter()
43            .map(|range| {
44                frame_data.get(range.clone()).ok_or_else(|| {
45                    PixelsError::malformed("avif", "tile group lies outside the coded frame")
46                })
47            })
48            .collect()
49    }
50}
51
52/// Parse the sequence header out of a run of configuration OBUs.
53///
54/// The `av1C` config OBUs must contain exactly the sequence header (and may
55/// contain metadata or padding around it). The first sequence header found
56/// wins.
57pub fn sequence_header_from_config(config_obus: &[u8]) -> Result<SequenceHeader> {
58    for obu in Obu::parse_stream(config_obus)? {
59        if obu.header.kind == ObuType::SequenceHeader {
60            let mut reader = BitReader::new(obu.payload);
61            return SequenceHeader::parse(&mut reader);
62        }
63    }
64    Err(PixelsError::malformed(
65        "avif",
66        "the av1C configuration OBUs contain no sequence header",
67    ))
68}
69
70impl StillPicture {
71    /// Parse the headers of a coded still picture.
72    ///
73    /// `config_obus` is the `av1C` configuration OBU run; `frame_data` is the
74    /// coded frame. A sequence header repeated in `frame_data` overrides the
75    /// configuration one, as the spec allows.
76    pub fn parse(config_obus: &[u8], frame_data: &[u8]) -> Result<Self> {
77        let mut sequence = sequence_header_from_config(config_obus).ok();
78
79        let obus = Obu::parse_stream(frame_data)?;
80        let mut frame = None;
81        let mut tile_groups = Vec::new();
82
83        for obu in &obus {
84            // Each OBU records where its payload begins in frame_data, so the
85            // tile-data locator stays absolute.
86            let payload_offset = obu.payload_start;
87
88            match obu.header.kind {
89                ObuType::SequenceHeader => {
90                    let mut reader = BitReader::new(obu.payload);
91                    sequence = Some(SequenceHeader::parse(&mut reader)?);
92                }
93                ObuType::Frame => {
94                    let seq = sequence.as_ref().ok_or_else(missing_sequence)?;
95                    let mut reader = BitReader::new(obu.payload);
96                    let header = FrameHeader::parse(
97                        &mut reader,
98                        seq,
99                        obu.header.temporal_id,
100                        obu.header.spatial_id,
101                    )?;
102                    reader.byte_alignment()?;
103                    let consumed = reader.byte_position().min(obu.payload.len());
104                    frame = Some(header);
105                    tile_groups.clear();
106                    tile_groups.push(payload_offset + consumed..payload_offset + obu.payload.len());
107                }
108                // A redundant copy repeats a header already in hand, so it
109                // matters only if the original was lost; it must not discard
110                // the tile groups already collected.
111                ObuType::RedundantFrameHeader if frame.is_some() => {}
112                ObuType::FrameHeader | ObuType::RedundantFrameHeader => {
113                    let seq = sequence.as_ref().ok_or_else(missing_sequence)?;
114                    let mut reader = BitReader::new(obu.payload);
115                    let header = FrameHeader::parse(
116                        &mut reader,
117                        seq,
118                        obu.header.temporal_id,
119                        obu.header.spatial_id,
120                    )?;
121                    // Tile data arrives in the OBU_TILE_GROUPs that follow.
122                    frame = Some(header);
123                    tile_groups.clear();
124                }
125                ObuType::TileGroup if frame.is_some() => {
126                    tile_groups.push(payload_offset..payload_offset + obu.payload.len());
127                }
128                _ => {}
129            }
130        }
131
132        let sequence = sequence.ok_or_else(missing_sequence)?;
133        let frame = frame.ok_or_else(|| {
134            PixelsError::malformed("avif", "the coded frame contains no frame header")
135        })?;
136        if tile_groups.is_empty() {
137            return Err(PixelsError::malformed(
138                "avif",
139                "the coded frame has a header but no tile group",
140            ));
141        }
142
143        Ok(Self {
144            sequence,
145            frame,
146            tile_groups,
147        })
148    }
149}
150
151fn missing_sequence() -> PixelsError {
152    PixelsError::malformed(
153        "avif",
154        "a frame header was reached before any sequence header",
155    )
156}
157
158#[cfg(test)]
159#[allow(
160    clippy::unwrap_used,
161    clippy::indexing_slicing,
162    clippy::panic,
163    reason = "tests operate on known-good values and assert shapes directly"
164)]
165mod tests {
166    use super::*;
167
168    /// Bit-level builder shared in shape with the seq/frame tests.
169    struct Bldr {
170        bits: Vec<u8>,
171    }
172    impl Bldr {
173        fn new() -> Self {
174            Self { bits: Vec::new() }
175        }
176        fn put(&mut self, value: u32, n: u32) -> &mut Self {
177            for i in (0..n).rev() {
178                self.bits.push(((value >> i) & 1) as u8);
179            }
180            self
181        }
182        fn flag(&mut self, b: bool) -> &mut Self {
183            self.put(u32::from(b), 1)
184        }
185        fn pack(&self) -> Vec<u8> {
186            let mut out = vec![0_u8; self.bits.len().div_ceil(8)];
187            for (i, &bit) in self.bits.iter().enumerate() {
188                if bit != 0 {
189                    out[i / 8] |= 1 << (7 - (i % 8));
190                }
191            }
192            out
193        }
194    }
195
196    /// Wrap payload bytes in a low-overhead OBU of the given type.
197    fn obu(kind: u8, payload: &[u8]) -> Vec<u8> {
198        let mut out = vec![(kind << 3) | 0b0000_0010];
199        assert!(payload.len() < 128);
200        out.push(payload.len() as u8);
201        out.extend_from_slice(payload);
202        out
203    }
204
205    /// The reduced-still-picture sequence header bytes for 8-bit 4:2:0.
206    fn seq_bytes(width: u32, height: u32) -> Vec<u8> {
207        let mut b = Bldr::new();
208        b.put(0, 3).flag(true).flag(true).put(1, 5);
209        b.put(15, 4)
210            .put(15, 4)
211            .put(width - 1, 16)
212            .put(height - 1, 16);
213        b.flag(false).flag(false).flag(false);
214        b.flag(false).flag(false).flag(false);
215        b.flag(false).flag(false).flag(false);
216        b.flag(false).put(0, 2).flag(false);
217        b.flag(false);
218        b.pack()
219    }
220
221    /// The uncompressed header bytes for the reduced-still key frame used above,
222    /// followed by whatever trailing tile bytes the caller wants.
223    fn frame_bytes(tile: &[u8]) -> Vec<u8> {
224        let mut b = Bldr::new();
225        b.flag(false); // disable_cdf_update
226        b.flag(false); // allow_screen_content_tools
227        b.flag(false); // render_and_frame_size_different
228        b.flag(true); // uniform_tile_spacing
229        b.put(100, 8); // base_q_idx
230        b.flag(false).flag(false).flag(false).flag(false); // deltas + qmatrix
231        b.flag(false); // segmentation_enabled
232        b.flag(false); // delta_q_present
233        b.put(0, 6).put(0, 6).put(0, 3).flag(false); // loop filter
234        b.flag(false); // tx_mode_select -> Largest
235        b.flag(false); // reduced_tx_set
236        let mut bytes = b.pack();
237        // The header is byte-aligned by construction here (its bit count is a
238        // multiple of 8 after packing rounds up); append the tile payload.
239        bytes.extend_from_slice(tile);
240        bytes
241    }
242
243    #[test]
244    fn parses_config_and_frame_into_headers_and_tile_locator() {
245        let config = obu(1, &seq_bytes(64, 64)); // OBU_SEQUENCE_HEADER
246        let tile = [0xDE, 0xAD, 0xBE, 0xEF];
247        let frame_payload = frame_bytes(&tile);
248        let frame_stream = obu(6, &frame_payload); // OBU_FRAME
249
250        let still = StillPicture::parse(&config, &frame_stream).unwrap();
251        assert_eq!(still.sequence.max_frame_width, 64);
252        assert_eq!(still.frame.frame_width, 64);
253        assert_eq!(still.frame.quantization.base_q_idx, 100);
254        assert_eq!(still.frame.tile_info.count(), 1);
255        // The tile locator points at the trailing bytes of the OBU_FRAME.
256        let located = still.tile_group_data(&frame_stream).unwrap();
257        assert_eq!(located, [&tile[..]]);
258    }
259
260    #[test]
261    fn a_separate_frame_header_and_tile_group_are_joined() {
262        let config = obu(1, &seq_bytes(32, 32));
263        let header_only = frame_bytes(&[]);
264        let mut stream = obu(3, &header_only); // OBU_FRAME_HEADER
265        let tile = [1, 2, 3];
266        stream.extend(obu(4, &tile)); // OBU_TILE_GROUP
267        stream.extend(obu(4, &[4, 5])); // a second OBU_TILE_GROUP
268
269        let still = StillPicture::parse(&config, &stream).unwrap();
270        assert_eq!(still.frame.frame_width, 32);
271        let located = still.tile_group_data(&stream).unwrap();
272        assert_eq!(located, [&tile[..], &[4, 5][..]]);
273    }
274
275    #[test]
276    fn a_redundant_frame_header_keeps_the_tile_groups_already_seen() {
277        let config = obu(1, &seq_bytes(32, 32));
278        let header = frame_bytes(&[]);
279        let mut stream = obu(3, &header); // OBU_FRAME_HEADER
280        stream.extend(obu(4, &[1, 2])); // OBU_TILE_GROUP
281        stream.extend(obu(7, &header)); // OBU_REDUNDANT_FRAME_HEADER
282        stream.extend(obu(4, &[3])); // OBU_TILE_GROUP
283
284        let still = StillPicture::parse(&config, &stream).unwrap();
285        let located = still.tile_group_data(&stream).unwrap();
286        assert_eq!(located, [&[1, 2][..], &[3][..]]);
287    }
288
289    #[test]
290    fn a_frame_header_without_tile_data_is_malformed() {
291        let config = obu(1, &seq_bytes(32, 32));
292        let stream = obu(3, &frame_bytes(&[])); // OBU_FRAME_HEADER alone
293        let error = StillPicture::parse(&config, &stream).unwrap_err();
294        assert!(error.to_string().contains("no tile group"), "{error}");
295    }
296
297    #[test]
298    fn an_in_band_sequence_header_overrides_the_config_one() {
299        // Config says 64x64; the frame stream carries a 40x40 sequence header
300        // ahead of the frame, which must win.
301        let config = obu(1, &seq_bytes(64, 64));
302        let mut stream = obu(1, &seq_bytes(40, 40));
303        stream.extend(obu(6, &frame_bytes(&[0x00])));
304        let still = StillPicture::parse(&config, &stream).unwrap();
305        assert_eq!(still.sequence.max_frame_width, 40);
306        assert_eq!(still.frame.frame_width, 40);
307    }
308
309    #[test]
310    fn a_frame_with_no_sequence_header_anywhere_is_rejected() {
311        let stream = obu(6, &frame_bytes(&[0x00]));
312        let err = StillPicture::parse(&[], &stream).unwrap_err();
313        assert!(err.to_string().contains("sequence header"), "{err}");
314    }
315
316    #[test]
317    fn config_without_a_sequence_header_is_reported() {
318        let config = obu(2, &[]); // a temporal delimiter, no seq header
319        let err = sequence_header_from_config(&config).unwrap_err();
320        assert!(err.to_string().contains("no sequence header"), "{err}");
321    }
322}