Handover of final codec fixes

src/codec/av1/parser.rs

@@ -59,7 +59,8 @@ pub const FEATURE_SIGNED: [bool; SEG_LVL_MAX] = [true, true, true, true, true, f
 // Same as Segmentation_Feature_Max in the specification. See 5.9.14
 pub const FEATURE_MAX: [i32; SEG_LVL_MAX] = [255, 63, 63, 63, 63, 7, 0, 0];
 
-pub enum ParsedObu<'a> {
+/// Tells what should be done with the OBU after [`Parser::read_obu`] is called.
+pub enum ObuAction<'a> {
     /// We should process the OBU normally.
     Process(Obu<'a>),
     /// We should drop this OBU and advance to the next one. The u32 is how much
@@ -117,21 +118,63 @@ impl ObuHeader {
     }
 }
 
+/// Contains the OBU header and a reference to its data. The OBU itself hasn't been parsed yet.
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct Obu<'a> {
     /// The OBU header.
     pub header: ObuHeader,
-    /// The data backing the OBU.
-    pub data: Cow<'a, [u8]>,
-    /// Where the OBU data starts, after the size has been read.
-    pub start_offset: usize,
-    /// The OBU size as per the specification after `start_offset`.
-    pub size: usize,
+    /// Amount of bytes from the input consumed to parse this OBU.
+    pub bytes_used: usize,
+    /// The slice backing the OBU.
+    data: Cow<'a, [u8]>,
 }
 
 impl<'a> AsRef<[u8]> for Obu<'a> {
     fn as_ref(&self) -> &[u8] {
-        &self.data[self.start_offset..self.start_offset + self.size]
+        self.data.as_ref()
+    }
+}
+
+/// A fully parsed OBU, with additional data when relevant.
+pub enum ParsedObu<'a> {
+    Reserved,
+    SequenceHeader(Rc<SequenceHeaderObu>),
+    TemporalDelimiter,
+    FrameHeader(FrameHeaderObu),
+    TileGroup(TileGroupObu<'a>),
+    Metadata,
+    Frame(FrameObu<'a>),
+    RedundantFrameHeader,
+    TileList,
+    Reserved2,
+    Reserved3,
+    Reserved4,
+    Reserved5,
+    Reserved6,
+    Reserved7,
+    Padding,
+}
+
+impl<'a> ParsedObu<'a> {
+    pub fn obu_type(&self) -> ObuType {
+        match self {
+            ParsedObu::Reserved => ObuType::Reserved,
+            ParsedObu::SequenceHeader(_) => ObuType::SequenceHeader,
+            ParsedObu::TemporalDelimiter => ObuType::TemporalDelimiter,
+            ParsedObu::FrameHeader(_) => ObuType::FrameHeader,
+            ParsedObu::TileGroup(_) => ObuType::TileGroup,
+            ParsedObu::Metadata => ObuType::Metadata,
+            ParsedObu::Frame(_) => ObuType::Frame,
+            ParsedObu::RedundantFrameHeader => ObuType::RedundantFrameHeader,
+            ParsedObu::TileList => ObuType::TileList,
+            ParsedObu::Reserved2 => ObuType::Reserved2,
+            ParsedObu::Reserved3 => ObuType::Reserved3,
+            ParsedObu::Reserved4 => ObuType::Reserved4,
+            ParsedObu::Reserved5 => ObuType::Reserved5,
+            ParsedObu::Reserved6 => ObuType::Reserved6,
+            ParsedObu::Reserved7 => ObuType::Reserved7,
+            ParsedObu::Padding => ObuType::Padding,
+        }
     }
 }
 
@@ -1453,7 +1496,7 @@ impl Parser {
         // We can't have that in parse_obu as per the spec, because the reader
         // is not initialized on our design at that point, so move the check to
         // inside this function.
-        if obu.size == 0
+        if obu.data.len() == 0
             || matches!(
                 obu.header.obu_type,
                 ObuType::TileList | ObuType::TileGroup | ObuType::Frame
@@ -1493,7 +1536,7 @@ impl Parser {
     /// format.
     ///
     /// `None` may eventually be returned if the OBU is to be dropped.
-    pub fn parse_obu<'a>(&mut self, data: &'a [u8]) -> anyhow::Result<ParsedObu<'a>> {
+    pub fn read_obu<'a>(&mut self, data: &'a [u8]) -> anyhow::Result<ObuAction<'a>> {
         if data.is_empty() {
             return Err(anyhow!("Empty data"));
         }
@@ -1519,7 +1562,7 @@ impl Parser {
             let obu_length = reader.current_annexb_obu_length(annexb_state)?;
             match obu_length {
                 Some(length) => length,
-                None => return Ok(ParsedObu::Drop(reader.consumed(0))),
+                None => return Ok(ObuAction::Drop(reader.consumed(0))),
             }
         } else {
             0
@@ -1574,15 +1617,14 @@ impl Parser {
             let in_spatial_layer = ((self.operating_point_idc >> (header.spatial_id + 8)) & 1) != 0;
             if !in_temporal_layer || !in_spatial_layer {
                 log::debug!("Dropping obu as per drop_obu() in the specification",);
-                return Ok(ParsedObu::Drop(reader.consumed(0)));
+                return Ok(ObuAction::Drop(reader.consumed(0)));
             }
         }
 
-        Ok(ParsedObu::Process(Obu {
+        Ok(ObuAction::Process(Obu {
             header,
-            data: Cow::from(&data[..start_offset + obu_size]),
-            start_offset,
-            size: obu_size,
+            data: Cow::from(&data[start_offset..start_offset + obu_size]),
+            bytes_used: start_offset + obu_size,
         }))
     }
 
@@ -1721,29 +1763,12 @@ impl Parser {
         Ok(())
     }
 
-    pub fn parse_temporal_delimiter_obu(&mut self, obu: &Obu) -> anyhow::Result<()> {
-        if !matches!(obu.header.obu_type, ObuType::TemporalDelimiter) {
-            return Err(anyhow!(
-                "Expected a TemporalDelimiterOBU, got {:?}",
-                obu.header.obu_type
-            ));
-        }
-
+    fn parse_temporal_delimiter_obu(&mut self) -> anyhow::Result<()> {
         self.seen_frame_header = false;
         Ok(())
     }
 
-    pub fn parse_sequence_header_obu(
-        &mut self,
-        obu: &Obu,
-    ) -> anyhow::Result<Rc<SequenceHeaderObu>> {
-        if !matches!(obu.header.obu_type, ObuType::SequenceHeader) {
-            return Err(anyhow!(
-                "Expected a SequenceHeaderOBU, got {:?}",
-                obu.header.obu_type
-            ));
-        }
-
+    fn parse_sequence_header_obu(&mut self, obu: &Obu) -> anyhow::Result<Rc<SequenceHeaderObu>> {
         let mut s = SequenceHeaderObu {
             obu_header: obu.header.clone(),
             ..Default::default()
@@ -1919,7 +1944,7 @@ impl Parser {
     /// header, so we must save them now, as they will not be parsed from the
     /// bitstream. We also save some internal parser state which will be useful
     /// later.
-    fn load_reference_frame(&mut self, fh: &mut FrameHeaderObu) -> anyhow::Result<()> {
+    fn load_reference_frame(&self, fh: &mut FrameHeaderObu) -> anyhow::Result<()> {
         let rf = &self.ref_info[fh.frame_to_show_map_idx as usize];
 
         // Section 6.8.1: It is a requirement of bitstream conformance that a
@@ -3552,7 +3577,7 @@ impl Parser {
         Ok(fh)
     }
 
-    pub fn parse_tile_group_obu<'a>(&mut self, obu: Obu<'a>) -> anyhow::Result<TileGroupObu<'a>> {
+    fn parse_tile_group_obu<'a>(&mut self, obu: Obu<'a>) -> anyhow::Result<TileGroupObu<'a>> {
         let mut tg = TileGroupObu {
             obu,
             ..Default::default()
@@ -3637,20 +3662,15 @@ impl Parser {
         Ok(tg)
     }
 
-    pub fn parse_frame_obu<'a>(&mut self, obu: Obu<'a>) -> anyhow::Result<FrameObu<'a>> {
-        if !matches!(obu.header.obu_type, ObuType::Frame) {
-            return Err(anyhow!(
-                "Expected a FrameOBU, got {:?}",
-                obu.header.obu_type
-            ));
-        }
-
+    fn parse_frame_obu<'a>(&mut self, obu: Obu<'a>) -> anyhow::Result<FrameObu<'a>> {
         let frame_header_obu = self.parse_frame_header_obu(&obu)?;
         let obu = Obu {
             header: obu.header,
-            data: obu.data,
-            start_offset: obu.start_offset + frame_header_obu.header_bytes,
-            size: obu.size - frame_header_obu.header_bytes,
+            data: match obu.data {
+                Cow::Borrowed(d) => Cow::Borrowed(&d[frame_header_obu.header_bytes..]),
+                Cow::Owned(d) => Cow::Owned(d[frame_header_obu.header_bytes..].to_owned()),
+            },
+            bytes_used: obu.bytes_used,
         };
         let tile_group_obu = self.parse_tile_group_obu(obu)?;
 
@@ -3758,6 +3778,34 @@ impl Parser {
             Some(helpers::floor_log2((self.operating_point_idc >> 8) as u32))
         }
     }
+
+    /// Fully parse an OBU.
+    pub fn parse_obu<'a>(&mut self, obu: Obu<'a>) -> anyhow::Result<ParsedObu<'a>> {
+        match obu.header.obu_type {
+            ObuType::Reserved => Ok(ParsedObu::Reserved),
+            ObuType::SequenceHeader => self
+                .parse_sequence_header_obu(&obu)
+                .map(ParsedObu::SequenceHeader),
+            ObuType::TemporalDelimiter => self
+                .parse_temporal_delimiter_obu()
+                .map(|_| ParsedObu::TemporalDelimiter),
+            ObuType::FrameHeader => self
+                .parse_frame_header_obu(&obu)
+                .map(ParsedObu::FrameHeader),
+            ObuType::TileGroup => self.parse_tile_group_obu(obu).map(ParsedObu::TileGroup),
+            ObuType::Metadata => Ok(ParsedObu::Metadata),
+            ObuType::Frame => self.parse_frame_obu(obu).map(ParsedObu::Frame),
+            ObuType::RedundantFrameHeader => Ok(ParsedObu::RedundantFrameHeader),
+            ObuType::TileList => Ok(ParsedObu::TileList),
+            ObuType::Reserved2 => Ok(ParsedObu::Reserved2),
+            ObuType::Reserved3 => Ok(ParsedObu::Reserved3),
+            ObuType::Reserved4 => Ok(ParsedObu::Reserved4),
+            ObuType::Reserved5 => Ok(ParsedObu::Reserved5),
+            ObuType::Reserved6 => Ok(ParsedObu::Reserved6),
+            ObuType::Reserved7 => Ok(ParsedObu::Reserved7),
+            ObuType::Padding => Ok(ParsedObu::Padding),
+        }
+    }
 }
 
 impl Default for Parser {
@@ -3821,7 +3869,7 @@ impl Clone for Parser {
 
 #[cfg(test)]
 mod tests {
-    use crate::codec::av1::parser::{ParsedObu, Parser, StreamFormat};
+    use crate::codec::av1::parser::{ObuAction, Parser, StreamFormat};
     use crate::utils::IvfIterator;
 
     use super::ObuType;
@@ -3849,16 +3897,16 @@ mod tests {
         for packet in ivf_iter {
             let mut consumed = 0;
 
-            while let Ok(obu) = parser.parse_obu(&packet[consumed..]) {
+            while let Ok(obu) = parser.read_obu(&packet[consumed..]) {
                 let obu = match obu {
-                    ParsedObu::Process(obu) => obu,
+                    ObuAction::Process(obu) => obu,
                     // This OBU should be dropped.
-                    ParsedObu::Drop(length) => {
+                    ObuAction::Drop(length) => {
                         consumed += usize::try_from(length).unwrap();
                         continue;
                     }
                 };
-                consumed += obu.data.len();
+                consumed += obu.bytes_used;
                 num_obus += 1;
             }
         }
@@ -3877,14 +3925,14 @@ mod tests {
         let mut ivf_iter = IvfIterator::new(STREAM_TEST_25_FPS);
         let packet = ivf_iter.next().unwrap();
 
-        parser.parse_obu(packet).unwrap();
+        parser.read_obu(packet).unwrap();
         assert!(matches!(parser.stream_format, StreamFormat::LowOverhead));
 
         let mut parser = Parser::default();
         let mut ivf_iter = IvfIterator::new(STREAM_ANNEXB);
         let packet = ivf_iter.next().unwrap();
 
-        parser.parse_obu(packet).unwrap();
+        parser.read_obu(packet).unwrap();
         assert!(matches!(parser.stream_format, StreamFormat::AnnexB { .. }));
     }
 
@@ -3898,17 +3946,17 @@ mod tests {
         for packet in ivf_iter {
             let mut consumed = 0;
 
-            while let Ok(obu) = parser.parse_obu(&packet[consumed..]) {
+            while let Ok(obu) = parser.read_obu(&packet[consumed..]) {
                 let obu = match obu {
-                    ParsedObu::Process(obu) => obu,
+                    ObuAction::Process(obu) => obu,
                     // This OBU should be dropped.
-                    ParsedObu::Drop(length) => {
+                    ObuAction::Drop(length) => {
                         consumed += usize::try_from(length).unwrap();
                         continue;
                     }
                 };
                 assert!(matches!(parser.stream_format, StreamFormat::LowOverhead));
-                consumed += obu.data.len();
+                consumed += obu.bytes_used;
             }
         }
 
@@ -3919,17 +3967,17 @@ mod tests {
         for packet in ivf_iter {
             let mut consumed = 0;
 
-            while let Ok(obu) = parser.parse_obu(&packet[consumed..]) {
+            while let Ok(obu) = parser.read_obu(&packet[consumed..]) {
                 let obu = match obu {
-                    ParsedObu::Process(obu) => obu,
+                    ObuAction::Process(obu) => obu,
                     // This OBU should be dropped.
-                    ParsedObu::Drop(length) => {
+                    ObuAction::Drop(length) => {
                         consumed += usize::try_from(length).unwrap();
                         continue;
                     }
                 };
                 assert!(matches!(parser.stream_format, StreamFormat::AnnexB { .. }));
-                consumed += obu.data.len();
+                consumed += obu.bytes_used;
                 num_obus += 1;
             }
         }
@@ -3957,17 +4005,17 @@ mod tests {
         for packet in ivf_iter {
             let mut consumed = 0;
 
-            while let Ok(obu) = parser.parse_obu(&packet[consumed..]) {
+            while let Ok(obu) = parser.read_obu(&packet[consumed..]) {
                 let obu = match obu {
-                    ParsedObu::Process(obu) => obu,
+                    ObuAction::Process(obu) => obu,
                     // This OBU should be dropped.
-                    ParsedObu::Drop(length) => {
+                    ObuAction::Drop(length) => {
                         consumed += usize::try_from(length).unwrap();
                         continue;
                     }
                 };
 
-                let data_len = obu.data.len();
+                let data_len = obu.bytes_used;
 
                 match obu.header.obu_type {
                     ObuType::SequenceHeader => {