refactor(av1): split OBU parsing in decoder (S3776)

crates/av1/src/decoder.rs

@@ -405,49 +405,35 @@ impl Decoder {
         r
     }
 
+    /// Whether an OBU belongs to the selected operating point (spec 7.5: OBUs of other layers are dropped).
+    fn in_operating_point(&self, obu: &ObuUnit) -> bool {
+        if obu.obu_type != OBU_SEQUENCE_HEADER_T
+            && obu.obu_type != OBU_TEMPORAL_DELIMITER_T
+            && obu.ext
+            && let Some(seq) = &self.seq
+            && seq.op_idc != 0
+        {
+            let idc = seq.op_idc;
+            let in_t = (idc >> obu.temporal_id) & 1;
+            let in_s = (idc >> (obu.spatial_id + 8)) & 1;
+            if in_t == 0 || in_s == 0 {
+                return false;
+            }
+        }
+        true
+    }
+
     fn parse_obus(&mut self, data: &[u8], out: &mut Vec<(Output, usize, u32)>) -> Result<()> {
         let mut tu = 0usize;
         let mut pos = 0;
         while pos < data.len() {
-            let h = data[pos];
+            let obu = parse_obu_unit(data, pos)?;
-            if h & 0x80 != 0 {
+            pos = obu.end;
-                return Err(Error::Invalid("obu_forbidden_bit"));
+            if !self.in_operating_point(&obu) {
+                continue;
             }
-            let obu_type = ((h >> 3) & 0xf) as u32;
+            let (payload, spatial_id) = (obu.payload, obu.spatial_id);
-            let ext = h & 4 != 0;
+            match obu.obu_type {
-            let has_size = h & 2 != 0;
-            let mut p = pos + 1;
-            let (mut temporal_id, mut spatial_id) = (0, 0);
-            if ext {
-                let e = *data.get(p).ok_or(Error::Truncated)?;
-                temporal_id = (e >> 5) as u32;
-                spatial_id = ((e >> 3) & 3) as u32;
-                p += 1;
-            }
-            let size = if has_size {
-                let (v, n) = leb128(data.get(p..).ok_or(Error::Truncated)?)?;
-                p += n;
-                v as usize
-            } else {
-                data.len().checked_sub(p).ok_or(Error::Truncated)?
-            };
-            let end = p.checked_add(size).filter(|&e| e <= data.len()).ok_or(Error::Truncated)?;
-            let payload = &data[p..end];
-            pos = end;
-            if obu_type != OBU_SEQUENCE_HEADER_T
-                && obu_type != OBU_TEMPORAL_DELIMITER_T
-                && ext
-                && let Some(seq) = &self.seq
-                && seq.op_idc != 0
-            {
-                let idc = seq.op_idc;
-                let in_t = (idc >> temporal_id) & 1;
-                let in_s = (idc >> (spatial_id + 8)) & 1;
-                if in_t == 0 || in_s == 0 {
-                    continue;
-                }
-            }
-            match obu_type {
                 OBU_SEQUENCE_HEADER_T => {
                     let s = SequenceHeader::parse(payload)?;
                     self.seq = Some(s);
@@ -457,27 +443,11 @@ impl Decoder {
                     tu += 1;
                 }
                 OBU_FRAME_HEADER_T | OBU_REDUNDANT_FRAME_HEADER_T | OBU_FRAME_T => {
-                    if self.seen_frame_header {
+                    if let Some(o) = self.frame_obu(&obu)? {
-                        // frame_header_copy(): identical to the active header
-                        if obu_type != OBU_FRAME_T {
-                            continue;
-                        }
-                    }
-                    let mut r = BitReader::new(payload);
-                    let shown = self.frame_header(&mut r, temporal_id, spatial_id)?;
-                    if let Some(o) = shown {
                         if out.iter().filter(|(o, _, _)| matches!(o, Output::Existing { .. })).count() >= 32 {
                             return Err(Error::Invalid("too many frame outputs in chunk"));
                         }
                         out.push((o, tu, spatial_id));
-                        continue;
-                    }
-                    if obu_type == OBU_FRAME_T {
-                        r.byte_align();
-                        let rest = payload.get(r.byte_pos()..).ok_or(Error::Truncated)?;
-                        if let Some(o) = self.tile_group(rest)? {
-                            out.push((o, tu, spatial_id));
-                        }
                     }
                 }
                 OBU_TILE_GROUP_T => {
@@ -492,6 +462,22 @@ impl Decoder {
         Ok(())
     }
 
+    /// A frame header or frame OBU: returns the picture it completes, if any.
+    fn frame_obu(&mut self, obu: &ObuUnit) -> Result<Option<Output>> {
+        // frame_header_copy(): identical to the active header
+        if self.seen_frame_header && obu.obu_type != OBU_FRAME_T {
+            return Ok(None);
+        }
+        let mut r = BitReader::new(obu.payload);
+        let shown = self.frame_header(&mut r, obu.temporal_id, obu.spatial_id)?;
+        if shown.is_some() || obu.obu_type != OBU_FRAME_T {
+            return Ok(shown);
+        }
+        r.byte_align();
+        let rest = obu.payload.get(r.byte_pos()..).ok_or(Error::Truncated)?;
+        self.tile_group(rest)
+    }
+
     /// frame_header_obu( ): returns the output for show_existing_frame.
     fn frame_header(&mut self, r: &mut BitReader, temporal_id: u32, spatial_id: u32) -> Result<Option<Output>> {
         let seq = self.seq.clone().ok_or(Error::Invalid("frame before sequence header"))?;
@@ -1017,6 +1003,43 @@ mod tests {
     }
 }
 
+/// One OBU: its header fields, payload and the position just past it.
+struct ObuUnit<'a> {
+    obu_type: u32,
+    ext: bool,
+    temporal_id: u32,
+    spatial_id: u32,
+    payload: &'a [u8],
+    end: usize,
+}
+
+fn parse_obu_unit(data: &[u8], pos: usize) -> Result<ObuUnit<'_>> {
+    let h = data[pos];
+    if h & 0x80 != 0 {
+        return Err(Error::Invalid("obu_forbidden_bit"));
+    }
+    let obu_type = ((h >> 3) & 0xf) as u32;
+    let ext = h & 4 != 0;
+    let has_size = h & 2 != 0;
+    let mut p = pos + 1;
+    let (mut temporal_id, mut spatial_id) = (0, 0);
+    if ext {
+        let e = *data.get(p).ok_or(Error::Truncated)?;
+        temporal_id = (e >> 5) as u32;
+        spatial_id = ((e >> 3) & 3) as u32;
+        p += 1;
+    }
+    let size = if has_size {
+        let (v, n) = leb128(data.get(p..).ok_or(Error::Truncated)?)?;
+        p += n;
+        v as usize
+    } else {
+        data.len().checked_sub(p).ok_or(Error::Truncated)?
+    };
+    let end = p.checked_add(size).filter(|&e| e <= data.len()).ok_or(Error::Truncated)?;
+    Ok(ObuUnit { obu_type, ext, temporal_id, spatial_id, payload: &data[p..end], end })
+}
+
 #[cfg(test)]
 mod guard_tests {
     use super::*;