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];
-            if h & 0x80 != 0 {
-                return Err(Error::Invalid("obu_forbidden_bit"));
+            let obu = parse_obu_unit(data, pos)?;
+            pos = obu.end;
+            if !self.in_operating_point(&obu) {
+                continue;
             }
-            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)?;
-            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 {
+            let (payload, spatial_id) = (obu.payload, obu.spatial_id);
+            match obu.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 {
-                        // 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 let Some(o) = self.frame_obu(&obu)? {
                         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::*;

crates/h264/src/slicedec.rs

@@ -665,38 +665,33 @@ impl<'a> SliceDecoder<'a> {
             return Ok(());
         }
         for p in 0..4 {
-            self.derive_sub_mb_motion(p, n)?;
-        }
-        Ok(())
-    }
-
-    /// Motion of one 8x8 sub-macroblock `p` of a P8x8/B8x8 MB.
-    fn derive_sub_mb_motion(&mut self, p: usize, n: [usize; 2]) -> Result<()> {
-        let sub = self.cur.sub[p];
-        let (x8, y8) = ((p & 1) * 2, (p >> 1) * 2);
-        if sub.direct {
-            return self.derive_direct(1 << p);
-        }
-        let mut refs = [-1i8; 2];
-        for (l, r) in refs.iter_mut().enumerate() {
-            if sub.pred & (1 << l) != 0 {
-                let v = self.cur.ref_idx[l][p];
-                ensure!((v as usize) < n[l], "ref_idx {v} out of range (list {l}, size {})", n[l]);
-                *r = v;
+            let sub = self.cur.sub[p];
+            let (x8, y8) = ((p & 1) * 2, (p >> 1) * 2);
+            if sub.direct {
+                self.derive_direct(1 << p)?;
+                continue;
             }
-        }
-        for j in 0..sub.shape.num_parts() {
-            let (sx, sy, w, h) = sub.shape.rect(j);
-            let (x, y) = (x8 + sx, y8 + sy);
-            let mut mvs = [[0i16; 2]; 2];
-            for l in 0..2 {
-                if refs[l] >= 0 {
-                    let mvp = self.mvp(l, x, y, w, refs[l], 0, 0);
-                    let d = self.cur.mvd[l][p * 4 + j];
-                    mvs[l] = [mvp[0].wrapping_add(d[0]), mvp[1].wrapping_add(d[1])];
+            let mut refs = [-1i8; 2];
+            for (l, r) in refs.iter_mut().enumerate() {
+                if sub.pred & (1 << l) != 0 {
+                    let v = self.cur.ref_idx[l][p];
+                    ensure!((v as usize) < n[l], "ref_idx {v} out of range (list {l}, size {})", n[l]);
+                    *r = v;
                 }
             }
-            self.set_motion(x, y, w, h, refs, mvs);
+            for j in 0..sub.shape.num_parts() {
+                let (sx, sy, w, h) = sub.shape.rect(j);
+                let (x, y) = (x8 + sx, y8 + sy);
+                let mut mvs = [[0i16; 2]; 2];
+                for l in 0..2 {
+                    if refs[l] >= 0 {
+                        let mvp = self.mvp(l, x, y, w, refs[l], 0, 0);
+                        let d = self.cur.mvd[l][p * 4 + j];
+                        mvs[l] = [mvp[0].wrapping_add(d[0]), mvp[1].wrapping_add(d[1])];
+                    }
+                }
+                self.set_motion(x, y, w, h, refs, mvs);
+            }
         }
         Ok(())
     }
@@ -1385,11 +1380,6 @@ impl<'a> SliceDecoder<'a> {
             cavlc::residual_block(r, &mut lv, 0, 15, 16, nc)?;
             self.put_luma_dc(&lv);
         }
-        self.residual_cavlc_luma(r, cbp, kind, t8)?;
-        self.residual_cavlc_chroma(r, cbp)
-    }
-
-    fn residual_cavlc_luma(&mut self, r: &mut BitReader, cbp: u8, kind: MbKind, t8: bool) -> Result<()> {
         for b8 in 0..4 {
             let (bx8, by8) = ((b8 & 1) * 2, (b8 >> 1) * 2);
             if cbp & (1 << b8) == 0 {
@@ -1423,10 +1413,6 @@ impl<'a> SliceDecoder<'a> {
                 }
             }
         }
-        Ok(())
-    }
-
-    fn residual_cavlc_chroma(&mut self, r: &mut BitReader, cbp: u8) -> Result<()> {
         let cbp_c = cbp >> 4;
         if cbp_c & 3 != 0 {
             for c in 0..2 {