fix has_more_rbsp_data, PicParameterSetExtra::read

Cargo.toml

@@ -11,7 +11,7 @@ repository = "https://github.com/dholroyd/h264-reader"
 edition = "2018"
 
 [dependencies]
-bitstream-io = "1.0"
+bitstream-io = "1.1"
 memchr = "2.1.1"
 rfc6381-codec = "0.1"
 log = "0.4"

src/nal/pps.rs

@@ -172,7 +172,7 @@ pub struct PicParameterSetExtra {
 }
 impl PicParameterSetExtra {
     fn read(r: &mut RbspBitReader<'_>, sps: &sps::SeqParameterSet) -> Result<Option<PicParameterSetExtra>,PpsError> {
-        Ok(if r.has_more_rbsp_data() {
+        Ok(if r.has_more_rbsp_data("pic_parameter_set_extra")? {
             let transform_8x8_mode_flag = r.read_bool()?;
             Some(PicParameterSetExtra {
                 transform_8x8_mode_flag,
@@ -231,7 +231,7 @@ impl PicParameterSet {
             .map_err(PpsError::BadPicParamSetId)?;
         let seq_parameter_set_id = ParamSetId::from_u32(r.read_ue_named("seq_parameter_set_id")?)
             .map_err(PpsError::BadSeqParamSetId)?;
-        let _seq_parameter_set = ctx.sps_by_id(seq_parameter_set_id)
+        let seq_parameter_set = ctx.sps_by_id(seq_parameter_set_id)
             .ok_or_else(|| PpsError::UnknownSeqParamSetId(seq_parameter_set_id))?;
         Ok(PicParameterSet {
             pic_parameter_set_id,
@@ -249,7 +249,7 @@ impl PicParameterSet {
             deblocking_filter_control_present_flag: r.read_bool()?,
             constrained_intra_pred_flag: r.read_bool()?,
             redundant_pic_cnt_present_flag: r.read_bool()?,
-            extension: None, // TODO: buggy?: PicParameterSetExtra::read(&mut r, seq_parameter_set)?,
+            extension: PicParameterSetExtra::read(&mut r, seq_parameter_set)?,
         })
     }
 

src/nal/sei/mod.rs

@@ -366,10 +366,6 @@ impl<R: SeiIncrementalPayloadReader> NalHandler for SeiHeaderReader<R> {
         match self.state {
             SeiHeaderState::Begin => (),
             SeiHeaderState::End => panic!("SeiHeaderReader already ended and end() called again"),
-            SeiHeaderState::PayloadSize { payload_type: HeaderType::ReservedSeiMessage(0x80), payload_size: 0 } => {
-                // TODO: this is a bit of a hack to ignore rbsp_trailing_bits (which will always
-                //       be 0b10000000 in an SEI payload since SEI messages are byte-aligned).
-            },
             SeiHeaderState::PayloadType { .. } => {
                 error!("End of SEI data encountered while reading SEI payloadType");
                 self.reader.reset(ctx);

src/rbsp.rs

@@ -248,9 +248,24 @@ impl<'buf> RbspBitReader<'buf> {
         self.reader.read(u32::from(bit_count)).map_err( |e| RbspBitReaderError::ReaderError(e) )
     }
 
-    pub fn has_more_rbsp_data(&mut self) -> bool {
+    /// Returns true if positioned before the RBSP trailing bits.
+    ///
+    /// This matches the definition of `more_rbsp_data()` in Rec. ITU-T H.264
+    /// (03/2010) section 7.2.
+    pub fn has_more_rbsp_data(&mut self, name: &'static str) -> Result<bool, RbspBitReaderError> {
         // BitReader returns its reader iff at an aligned position.
-        self.reader.reader().map(|r| (r.position() as usize) < r.get_ref().len()).unwrap_or(true)
+        //self.reader.reader().map(|r| (r.position() as usize) < r.get_ref().len()).unwrap_or(true)
+        let mut throwaway = self.reader.clone();
+        let r = (move || {
+            throwaway.skip(1)?;
+            throwaway.read_unary1()?;
+            Ok::<_, std::io::Error>(())
+        })();
+        match r {
+            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => Ok(false),
+            Err(e) => Err(RbspBitReaderError::ReaderErrorFor(name, e)),
+            Ok(_) => Ok(true),
+        }
     }
 
     fn golomb_to_signed(val: u32) -> i32 {
@@ -365,15 +380,20 @@ mod tests {
 
     #[test]
     fn bitreader_has_more_data() {
-        let mut reader = RbspBitReader::new(&[0x12, 0x34]);
-        assert!(reader.has_more_rbsp_data());
+        // Should work when the end bit is byte-aligned.
+        let mut reader = RbspBitReader::new(&[0x12, 0x80]);
+        assert!(reader.has_more_rbsp_data("call 1").unwrap());
+        assert_eq!(reader.read_u8(8).unwrap(), 0x12);
+        assert!(!reader.has_more_rbsp_data("call 2").unwrap());
+
+        // and when it's not.
+        let mut reader = RbspBitReader::new(&[0x18]);
+        assert!(reader.has_more_rbsp_data("call 3").unwrap());
         assert_eq!(reader.read_u8(4).unwrap(), 0x1);
-        assert!(reader.has_more_rbsp_data()); // unaligned, backing reader not at EOF
-        assert_eq!(reader.read_u8(4).unwrap(), 0x2);
-        assert!(reader.has_more_rbsp_data()); // aligned, backing reader not at EOF
-        assert_eq!(reader.read_u8(4).unwrap(), 0x3);
-        assert!(reader.has_more_rbsp_data()); // unaligned, backing reader at EOF
-        assert_eq!(reader.read_u8(4).unwrap(), 0x4);
-        assert!(!reader.has_more_rbsp_data()); // aligned, backing reader at EOF
+        assert!(!reader.has_more_rbsp_data("call 4").unwrap());
+
+        // should also work when there are cabac-zero-words.
+        let mut reader = RbspBitReader::new(&[0x80, 0x00, 0x00]);
+        assert!(!reader.has_more_rbsp_data("at end with cabac-zero-words").unwrap());
     }
 }