X-Git-Url: https://git.nihav.org/?a=blobdiff_plain;f=nihav-core%2Fsrc%2Fio%2Fcodebook.rs;h=d835ce31501af39b963a0e6985a09299ca0da690;hb=0b257d9fde8ee0cc24e15a63544b100a0b6da52d;hp=9882610ab0d102d2be487b56ca5abc4494fd4a4a;hpb=fd47a9b679a723f08a76a7f83e3be18f87cef63c;p=nihav.git

diff --git a/nihav-core/src/io/codebook.rs b/nihav-core/src/io/codebook.rs
index 9882610..d835ce3 100644
--- a/nihav-core/src/io/codebook.rs
+++ b/nihav-core/src/io/codebook.rs
@@ -16,7 +16,7 @@
 //!             ShortCodebookDesc { code: 0b01,   bits: 2 },
 //!             ShortCodebookDesc { code: 0b1,    bits: 1 });
 //! let mut cr = ShortCodebookDescReader::new(cb_desc);
-//! let cb = Codebook::new(&mut cr, CodebookMode::LSB).unwrap();
+//! let cb = Codebook::new(&mut cr, CodebookMode::MSB).unwrap();
 //! ```
 //!
 //! Create a codebook using more flexible [`TableCodebookDescReader`] approach.
@@ -127,9 +127,9 @@ pub trait CodebookDescReader<S> {
 /// The codebook structure for code reading.
 #[allow(dead_code)]
 pub struct Codebook<S> {
-    table: Vec<u32>,
-    syms:  Vec<S>,
-    lut_bits: u8,
+    pub table: Vec<u32>,
+    pub syms:  Vec<S>,
+    pub lut_bits: u8,
 }
 
 /// Trait allowing bitreader to use codebook for decoding bit sequences.
@@ -140,11 +140,11 @@ pub trait CodebookReader<S> {
     fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S>;
 }
 
-const TABLE_FILL_VALUE: u32 = 0x7F;
+pub const TABLE_FILL_VALUE: u32 = 0x7F;
 const MAX_LUT_BITS: u8 = 10;
 
 fn fill_lut_msb(table: &mut Vec<u32>, off: usize,
-                code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) {
+                code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) -> CodebookResult<()> {
     if !esc {
         let fill_len  = lut_bits - bits;
         let fill_size = 1 << fill_len;
@@ -152,16 +152,19 @@ fn fill_lut_msb(table: &mut Vec<u32>, off: usize,
         let lut_value = (symidx << 8) | u32::from(bits);
         for j in 0..fill_size {
             let idx = (fill_code + j) as usize;
+            if table[idx + off] != TABLE_FILL_VALUE { return Err(CodebookError::InvalidCodebook); }
             table[idx + off] = lut_value;
         }
     } else {
         let idx = (code as usize) + off;
+        if table[idx] != TABLE_FILL_VALUE { return Err(CodebookError::InvalidCodebook); }
         table[idx] = (symidx << 8) | 0x80 | u32::from(bits);
     }
+    Ok(())
 }
 
 fn fill_lut_lsb(table: &mut Vec<u32>, off: usize,
-                code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) {
+                code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) -> CodebookResult<()> {
     if !esc {
         let fill_len  = lut_bits - bits;
         let fill_size = 1 << fill_len;
@@ -169,21 +172,24 @@ fn fill_lut_lsb(table: &mut Vec<u32>, off: usize,
         let step = lut_bits - fill_len;
         for j in 0..fill_size {
             let idx = (fill_code + (j << step)) as usize;
+            if table[idx + off] != TABLE_FILL_VALUE { return Err(CodebookError::InvalidCodebook); }
             table[idx + off] = (symidx << 8) | u32::from(bits);
         }
     } else {
         let idx = (code as usize) + off;
+        if table[idx] != TABLE_FILL_VALUE { return Err(CodebookError::InvalidCodebook); }
         table[idx] = (symidx << 8) | 0x80 | u32::from(bits);
     }
+    Ok(())
 }
 
 fn fill_lut(table: &mut Vec<u32>, mode: CodebookMode,
-            off: usize, code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) -> bool {
+            off: usize, code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) -> CodebookResult<bool> {
     match mode {
-        CodebookMode::MSB => fill_lut_msb(table, off, code, bits, lut_bits, symidx, esc),
-        CodebookMode::LSB => fill_lut_lsb(table, off, code, bits, lut_bits, symidx, esc),
+        CodebookMode::MSB => fill_lut_msb(table, off, code, bits, lut_bits, symidx, esc)?,
+        CodebookMode::LSB => fill_lut_lsb(table, off, code, bits, lut_bits, symidx, esc)?,
     };
-    bits > lut_bits
+    Ok(bits > lut_bits)
 }
 
 fn resize_table(table: &mut Vec<u32>, bits: u8) -> CodebookResult<u32> {
@@ -254,7 +260,7 @@ fn build_esc_lut(table: &mut Vec<u32>,
         let bits = code.bits;
         if code.bits <= MAX_LUT_BITS {
             fill_lut(table, mode, bucket.offset, code.code, bits,
-                     maxlen, code.idx as u32, false);
+                     maxlen, code.idx as u32, false)?;
         } else {
             let ckey = extract_lut_part(code.code, bits, MAX_LUT_BITS, mode);
             let cval = extract_esc_part(code.code, bits, MAX_LUT_BITS, mode);
@@ -268,7 +274,7 @@ fn build_esc_lut(table: &mut Vec<u32>,
         let maxlen = min(sec_bucket.maxlen, MAX_LUT_BITS);
         let new_off = resize_table(table, maxlen)?;
         fill_lut(table, mode, cur_offset, key, maxlen,
-                 MAX_LUT_BITS, new_off, true);
+                 MAX_LUT_BITS, new_off, true)?;
         sec_bucket.offset = new_off as usize;
     }
 
@@ -322,7 +328,7 @@ impl<S: Copy> Codebook<S> {
             let code = cb.code(i);
             if bits == 0 { continue; }
             if bits <= MAX_LUT_BITS {
-                fill_lut(&mut table, mode, 0, code, bits, maxbits, symidx, false);
+                fill_lut(&mut table, mode, 0, code, bits, maxbits, symidx, false)?;
             } else {
                 let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode) as usize;
                 if table[ckey] == TABLE_FILL_VALUE {
@@ -330,7 +336,7 @@ impl<S: Copy> Codebook<S> {
                     if let Some(bucket) = escape_list.get_mut(&key) {
                         let maxlen = min(bucket.maxlen, MAX_LUT_BITS);
                         let new_off = resize_table(&mut table, maxlen)?;
-                        fill_lut(&mut table, mode, 0, key, maxlen, MAX_LUT_BITS, new_off, true);
+                        fill_lut(&mut table, mode, 0, key, maxlen, MAX_LUT_BITS, new_off, true)?;
                         bucket.offset = new_off as usize;
                     }
                 }
@@ -364,10 +370,10 @@ impl<'a, S: Copy> CodebookReader<S> for BitReader<'a> {
             let bits = cb.table[lut_idx] & 0x7F;
             esc  = (cb.table[lut_idx] & 0x80) != 0;
             idx  = (cb.table[lut_idx] >> 8) as usize;
-            if (bits as isize) > self.left() {
+            let skip_bits = if esc { u32::from(lut_bits) } else { bits };
+            if (skip_bits as isize) > self.left() {
                 return Err(CodebookError::InvalidCode);
             }
-            let skip_bits = if esc { u32::from(lut_bits) } else { bits };
             self.skip(skip_bits as u32).unwrap();
             lut_bits = bits as u8;
         }