X-Git-Url: https://git.nihav.org/?a=blobdiff_plain;f=src%2Fio%2Fcodebook.rs;h=78ef5ae0d142164106696f478916905da2514477;hb=b74ff9fac35d41737d71d97227fad233aa4a4b49;hp=89641f24e7a62bbddac2c5bccef50d26246bc4e8;hpb=06fd8c8865256ab5348996a23afeeb512bc45835;p=nihav.git

diff --git a/src/io/codebook.rs b/src/io/codebook.rs
index 89641f2..78ef5ae 100644
--- a/src/io/codebook.rs
+++ b/src/io/codebook.rs
@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 use std::cmp::{max, min};
-use io::bitreader::BitReader;
+use super::bitreader::BitReader;
 
 #[derive(Debug)]
 pub enum CodebookError {
@@ -158,7 +158,7 @@ fn build_esc_lut(table: &mut Vec<u32>,
 
     for code in &bucket.codes {
         let bits = code.bits;
-        if code.bits < MAX_LUT_BITS {
+        if code.bits <= MAX_LUT_BITS {
             fill_lut(table, mode, bucket.offset, code.code, bits,
                      maxlen, code.idx as u32, false);
         } else {
@@ -181,7 +181,7 @@ fn build_esc_lut(table: &mut Vec<u32>,
     for (_, sec_bucket) in &escape_list {
         build_esc_lut(table, mode, sec_bucket)?;
     }
-    
+
     Ok(())
 }
 
@@ -195,7 +195,12 @@ impl<S: Copy> Codebook<S> {
         let mut symidx: usize = 0;
         for i in 0..cb.len() {
             let bits = cb.bits(i);
-            if bits > 0 { nnz = nnz + 1; }
+            if bits > 0 {
+                nnz = nnz + 1;
+                if cb.code(i) >= (1 << bits) {
+                    return Err(CodebookError::InvalidCodebook);
+                }
+            }
             maxbits = max(bits, maxbits);
             if bits > MAX_LUT_BITS {
                 let code = cb.code(i);
@@ -309,10 +314,29 @@ impl CodebookDescReader<u32> for ShortCodebookDescReader {
     fn len(&mut self) -> usize { self.data.len() }
 }
 
+pub struct TableCodebookDescReader<CodeType:'static, IndexType:'static> {
+    bits:       &'static [u8],
+    codes:      &'static [CodeType],
+    idx_map:    fn(usize) -> IndexType,
+}
+
+impl<'a, CodeType, IndexType> TableCodebookDescReader<CodeType, IndexType> {
+    pub fn new(codes: &'static [CodeType], bits: &'static [u8], idx_map: fn(usize) -> IndexType) -> Self {
+        Self { bits, codes, idx_map }
+    }
+}
+impl<CodeType: Copy+Into<u32>, IndexType> CodebookDescReader<IndexType> for TableCodebookDescReader<CodeType, IndexType>
+{
+    fn bits(&mut self, idx: usize) -> u8  { self.bits[idx] }
+    fn code(&mut self, idx: usize) -> u32 { self.codes[idx].into() }
+    fn sym (&mut self, idx: usize) -> IndexType { (self.idx_map)(idx) }
+    fn len(&mut self) -> usize { self.bits.len() }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
-    use io::bitreader::*;
+    use crate::io::bitreader::*;
 
     #[test]
     fn test_cb() {