nihav-core/src/io/codebook.rs

   1 use std::collections::HashMap;
   2 use std::cmp::{max, min};
   3 use super::bitreader::BitReader;
   4
   5 #[derive(Debug)]
   6 pub enum CodebookError {
   7     InvalidCodebook,
   8     MemoryError,
   9     InvalidCode,
  10 }
  11
  12 #[derive(Debug, Copy, Clone)]
  13 pub enum CodebookMode {
  14     MSB,
  15     LSB,
  16 }
  17
  18 type CodebookResult<T> = Result<T, CodebookError>;
  19
  20 pub struct FullCodebookDesc<S> {
  21     code: u32,
  22     bits: u8,
  23     sym:  S,
  24 }
  25
  26 pub struct ShortCodebookDesc {
  27     code: u32,
  28     bits: u8,
  29 }
  30
  31 pub trait CodebookDescReader<S> {
  32     fn bits(&mut self, idx: usize) -> u8;
  33     fn code(&mut self, idx: usize) -> u32;
  34     fn sym (&mut self, idx: usize) -> S;
  35     fn len (&mut self)             -> usize;
  36 }
  37
  38 #[allow(dead_code)]
  39 pub struct Codebook<S> {
  40     table: Vec<u32>,
  41     syms:  Vec<S>,
  42     lut_bits: u8,
  43 }
  44
  45 pub trait CodebookReader<S> {
  46     fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S>;
  47 }
  48
  49 const TABLE_FILL_VALUE: u32 = 0x7F;
  50 const MAX_LUT_BITS: u8 = 10;
  51
  52 fn fill_lut_msb(table: &mut Vec<u32>, off: usize,
  53                 code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) {
  54     if !esc {
  55         let fill_len  = lut_bits - bits;
  56         let fill_size = 1 << fill_len;
  57         let fill_code = code << (lut_bits - bits);
  58         let lut_value = (symidx << 8) | (bits as u32);
  59         for j in 0..fill_size {
  60             let idx = (fill_code + j) as usize;
  61             table[idx + off] = lut_value;
  62         }
  63     } else {
  64         let idx = (code as usize) + off;
  65         table[idx] = (symidx << 8) | 0x80 | (bits as u32);
  66     }
  67 }
  68
  69 fn fill_lut_lsb(table: &mut Vec<u32>, off: usize,
  70                 code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) {
  71     if !esc {
  72         let fill_len  = lut_bits - bits;
  73         let fill_size = 1 << fill_len;
  74         let fill_code = code;
  75         let step = lut_bits - fill_len;
  76         for j in 0..fill_size {
  77             let idx = (fill_code + (j << step)) as usize;
  78             table[idx + off] = (symidx << 8) | (bits as u32);
  79         }
  80     } else {
  81         let idx = (code as usize) + off;
  82         table[idx] = (symidx << 8) | 0x80 | (bits as u32);
  83     }
  84 }
  85
  86 fn fill_lut(table: &mut Vec<u32>, mode: CodebookMode,
  87             off: usize, code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) -> bool {
  88     match mode {
  89         CodebookMode::MSB => fill_lut_msb(table, off, code, bits, lut_bits, symidx, esc),
  90         CodebookMode::LSB => fill_lut_lsb(table, off, code, bits, lut_bits, symidx, esc),
  91     };
  92     bits > lut_bits
  93 }
  94
  95 fn resize_table(table: &mut Vec<u32>, bits: u8) -> CodebookResult<u32> {
  96     let add_size = (1 << bits) as usize;
  97     table.reserve(add_size);
  98     let cur_off = table.len() as u32;
  99     let new_size = table.len() + add_size;
 100     if table.capacity() < new_size { return Err(CodebookError::MemoryError); }
 101     table.resize(new_size, TABLE_FILL_VALUE);
 102     Ok(cur_off)
 103 }
 104
 105
 106 fn extract_lut_part(code: u32, bits: u8, lut_bits: u8, mode: CodebookMode) -> u32 {
 107     match mode {
 108         CodebookMode::MSB => code >> (bits - lut_bits),
 109         CodebookMode::LSB => code & ((1 << lut_bits) - 1),
 110     }
 111 }
 112
 113 fn extract_esc_part(code: u32, bits: u8, lut_bits: u8, mode: CodebookMode) -> u32 {
 114     match mode {
 115         CodebookMode::MSB => code & ((1 << (bits - lut_bits)) - 1),
 116         CodebookMode::LSB => code >> lut_bits,
 117     }
 118 }
 119
 120 #[derive(Clone,Copy)]
 121 struct Code {
 122     code: u32,
 123     bits: u8,
 124     idx:  usize,
 125 }
 126
 127 struct CodeBucket {
 128     maxlen: u8,
 129     offset: usize,
 130     codes:  Vec<Code>,
 131 }
 132
 133 impl CodeBucket {
 134     fn new() -> Self {
 135         CodeBucket { maxlen: 0, offset: 0, codes: Vec::new() }
 136     }
 137     fn add_code(&mut self, c: Code) {
 138         if c.bits > self.maxlen { self.maxlen = c.bits; }
 139         self.codes.push(c);
 140     }
 141 }
 142
 143 type EscapeCodes = HashMap<u32, CodeBucket>;
 144
 145 fn add_esc_code(cc: &mut EscapeCodes, key: u32, code: u32, bits: u8, idx: usize) {
 146     if !cc.contains_key(&key) { cc.insert(key, CodeBucket::new()); }
 147     let b = cc.get_mut(&key);
 148     if let Some(bucket) = b {
 149         bucket.add_code(Code {code: code, bits: bits, idx: idx });
 150     } else { panic!("no bucket when expected!"); }
 151 }
 152
 153 fn build_esc_lut(table: &mut Vec<u32>,
 154                  mode: CodebookMode,
 155                  bucket: &CodeBucket) -> CodebookResult<()> {
 156     let mut escape_list: EscapeCodes = HashMap::new();
 157     let maxlen = if bucket.maxlen > MAX_LUT_BITS { MAX_LUT_BITS } else { bucket.maxlen };
 158
 159     for code in &bucket.codes {
 160         let bits = code.bits;
 161         if code.bits <= MAX_LUT_BITS {
 162             fill_lut(table, mode, bucket.offset, code.code, bits,
 163                      maxlen, code.idx as u32, false);
 164         } else {
 165             let ckey = extract_lut_part(code.code, bits, MAX_LUT_BITS, mode);
 166             let cval = extract_esc_part(code.code, bits, MAX_LUT_BITS, mode);
 167             add_esc_code(&mut escape_list, ckey, cval, bits - MAX_LUT_BITS, code.idx);
 168         }
 169     }
 170
 171     let cur_offset = bucket.offset;
 172     for (ckey, sec_bucket) in &mut escape_list {
 173         let key = *ckey as u32;
 174         let maxlen = min(sec_bucket.maxlen, MAX_LUT_BITS);
 175         let new_off = resize_table(table, maxlen)?;
 176         fill_lut(table, mode, cur_offset, key, maxlen,
 177                  MAX_LUT_BITS, new_off, true);
 178         sec_bucket.offset = new_off as usize;
 179     }
 180
 181     for (_, sec_bucket) in &escape_list {
 182         build_esc_lut(table, mode, sec_bucket)?;
 183     }
 184
 185     Ok(())
 186 }
 187
 188 impl<S: Copy> Codebook<S> {
 189
 190     pub fn new(cb: &mut CodebookDescReader<S>, mode: CodebookMode) -> CodebookResult<Self> {
 191         let mut maxbits = 0;
 192         let mut nnz = 0;
 193         let mut escape_list: EscapeCodes = HashMap::new();
 194
 195         let mut symidx: usize = 0;
 196         for i in 0..cb.len() {
 197             let bits = cb.bits(i);
 198             if bits > 0 {
 199                 nnz = nnz + 1;
 200                 if cb.code(i) >= (1 << bits) {
 201                     return Err(CodebookError::InvalidCodebook);
 202                 }
 203             }
 204             maxbits = max(bits, maxbits);
 205             if bits > MAX_LUT_BITS {
 206                 let code = cb.code(i);
 207                 let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode);
 208                 let cval = extract_esc_part(code, bits, MAX_LUT_BITS, mode);
 209                 add_esc_code(&mut escape_list, ckey, cval, bits - MAX_LUT_BITS, symidx);
 210             }
 211             if bits > 0 { symidx = symidx + 1; }
 212         }
 213         if maxbits == 0 { return Err(CodebookError::InvalidCodebook); }
 214
 215         if maxbits > MAX_LUT_BITS { maxbits = MAX_LUT_BITS; }
 216
 217         let tab_len = 1 << maxbits;
 218         let mut table: Vec<u32> = Vec::with_capacity(tab_len);
 219         let mut syms:  Vec<S>   = Vec::with_capacity(nnz);
 220         if table.capacity() < tab_len { return Err(CodebookError::MemoryError); }
 221         if syms.capacity()  < nnz     { return Err(CodebookError::MemoryError); }
 222         table.resize(tab_len, TABLE_FILL_VALUE);
 223
 224         let mut symidx: u32 = 0;
 225         for i in 0..cb.len() {
 226             let bits = cb.bits(i);
 227             let code = cb.code(i);
 228             if bits == 0 { continue; }
 229             if bits <= MAX_LUT_BITS {
 230                 fill_lut(&mut table, mode, 0, code, bits, maxbits, symidx, false);
 231             } else {
 232                 let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode) as usize;
 233                 if table[ckey] == TABLE_FILL_VALUE {
 234                     let key = ckey as u32;
 235                     if let Some(bucket) = escape_list.get_mut(&key) {
 236                         let maxlen = min(bucket.maxlen, MAX_LUT_BITS);
 237                         let new_off = resize_table(&mut table, maxlen)?;
 238                         fill_lut(&mut table, mode, 0, key, maxlen, MAX_LUT_BITS, new_off, true);
 239                         bucket.offset = new_off as usize;
 240                     }
 241                 }
 242             }
 243             symidx = symidx + 1;
 244         }
 245
 246         for (_, bucket) in &escape_list {
 247             build_esc_lut(&mut table, mode, &bucket)?;
 248         }
 249
 250         for i in 0..cb.len() {
 251             if cb.bits(i) > 0 {
 252                 syms.push(cb.sym(i));
 253             }
 254         }
 255
 256         Ok(Codebook { table: table, syms: syms, lut_bits: maxbits })
 257     }
 258 }
 259
 260 impl<'a, S: Copy> CodebookReader<S> for BitReader<'a> {
 261     #[allow(unused_variables)]
 262     fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S> {
 263         let mut esc = true;
 264         let mut idx = 0;
 265         let mut lut_bits = cb.lut_bits;
 266         while esc {
 267             let lut_idx = (self.peek(lut_bits) as usize) + (idx as usize);
 268             if cb.table[lut_idx] == TABLE_FILL_VALUE { return Err(CodebookError::InvalidCode); }
 269             let bits = cb.table[lut_idx] & 0x7F;
 270             esc  = (cb.table[lut_idx] & 0x80) != 0;
 271             idx  = (cb.table[lut_idx] >> 8) as usize;
 272             if (bits as isize) > self.left() {
 273                 return Err(CodebookError::InvalidCode);
 274             }
 275             let skip_bits = if esc { lut_bits as u32 } else { bits };
 276             if let Err(_) = self.skip(skip_bits as u32) {}
 277             lut_bits = bits as u8;
 278         }
 279         Ok(cb.syms[idx])
 280     }
 281 }
 282
 283 pub struct FullCodebookDescReader<S> {
 284     data: Vec<FullCodebookDesc<S>>,
 285 }
 286
 287 impl<S> FullCodebookDescReader<S> {
 288     pub fn new(data: Vec<FullCodebookDesc<S>>) -> Self {
 289         FullCodebookDescReader { data: data }
 290     }
 291 }
 292
 293 impl<S: Copy> CodebookDescReader<S> for FullCodebookDescReader<S> {
 294     fn bits(&mut self, idx: usize) -> u8  { self.data[idx].bits }
 295     fn code(&mut self, idx: usize) -> u32 { self.data[idx].code }
 296     fn sym (&mut self, idx: usize) -> S   { self.data[idx].sym  }
 297     fn len(&mut self) -> usize { self.data.len() }
 298 }
 299
 300 pub struct ShortCodebookDescReader {
 301     data: Vec<ShortCodebookDesc>,
 302 }
 303
 304 impl ShortCodebookDescReader {
 305     pub fn new(data: Vec<ShortCodebookDesc<>>) -> Self {
 306         ShortCodebookDescReader { data: data }
 307     }
 308 }
 309
 310 impl CodebookDescReader<u32> for ShortCodebookDescReader {
 311     fn bits(&mut self, idx: usize) -> u8  { self.data[idx].bits }
 312     fn code(&mut self, idx: usize) -> u32 { self.data[idx].code }
 313     fn sym (&mut self, idx: usize) -> u32 { idx as u32 }
 314     fn len(&mut self) -> usize { self.data.len() }
 315 }
 316
 317 pub struct TableCodebookDescReader<CodeType:'static, IndexType:'static> {
 318     bits:       &'static [u8],
 319     codes:      &'static [CodeType],
 320     idx_map:    fn(usize) -> IndexType,
 321 }
 322
 323 impl<'a, CodeType, IndexType> TableCodebookDescReader<CodeType, IndexType> {
 324     pub fn new(codes: &'static [CodeType], bits: &'static [u8], idx_map: fn(usize) -> IndexType) -> Self {
 325         Self { bits, codes, idx_map }
 326     }
 327 }
 328 impl<CodeType: Copy+Into<u32>, IndexType> CodebookDescReader<IndexType> for TableCodebookDescReader<CodeType, IndexType>
 329 {
 330     fn bits(&mut self, idx: usize) -> u8  { self.bits[idx] }
 331     fn code(&mut self, idx: usize) -> u32 { self.codes[idx].into() }
 332     fn sym (&mut self, idx: usize) -> IndexType { (self.idx_map)(idx) }
 333     fn len(&mut self) -> usize { self.bits.len() }
 334 }
 335
 336 #[cfg(test)]
 337 mod test {
 338     use super::*;
 339     use crate::io::bitreader::*;
 340
 341     #[test]
 342     fn test_cb() {
 343         const BITS: [u8; 2] = [0b01011011, 0b10111100];
 344         let cb_desc: Vec<FullCodebookDesc<i8>> = vec!(
 345             FullCodebookDesc { code: 0b0,    bits: 1, sym:  16 },
 346             FullCodebookDesc { code: 0b10,   bits: 2, sym:  -3 },
 347             FullCodebookDesc { code: 0b110,  bits: 3, sym:  42 },
 348             FullCodebookDesc { code: 0b1110, bits: 4, sym: -42 }
 349         );
 350         let buf = &BITS;
 351         let mut br = BitReader::new(buf, buf.len(), BitReaderMode::BE);
 352         let mut cfr = FullCodebookDescReader::new(cb_desc);
 353         let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap();
 354         assert_eq!(br.read_cb(&cb).unwrap(),  16);
 355         assert_eq!(br.read_cb(&cb).unwrap(),  -3);
 356         assert_eq!(br.read_cb(&cb).unwrap(),  42);
 357         assert_eq!(br.read_cb(&cb).unwrap(), -42);
 358         let ret = br.read_cb(&cb);
 359         if let Err(e) = ret {
 360             assert_eq!(e as i32, CodebookError::InvalidCode as i32);
 361         } else {
 362             assert_eq!(0, 1);
 363         }
 364
 365         let scb_desc: Vec<ShortCodebookDesc> = vec!(
 366             ShortCodebookDesc { code: 0b0,    bits: 1 },
 367             ShortCodebookDesc { code: 0,      bits: 0 },
 368             ShortCodebookDesc { code: 0b10,   bits: 2 },
 369             ShortCodebookDesc { code: 0,      bits: 0 },
 370             ShortCodebookDesc { code: 0,      bits: 0 },
 371             ShortCodebookDesc { code: 0b110,  bits: 3 },
 372             ShortCodebookDesc { code: 0,      bits: 0 },
 373             ShortCodebookDesc { code: 0b11100, bits: 5 },
 374             ShortCodebookDesc { code: 0b11101, bits: 5 },
 375             ShortCodebookDesc { code: 0b1111010, bits: 7 },
 376             ShortCodebookDesc { code: 0b1111011, bits: 7 },
 377             ShortCodebookDesc { code: 0b1111110, bits: 7 },
 378             ShortCodebookDesc { code: 0b11111111, bits: 8 }
 379         );
 380         let mut br2 = BitReader::new(buf, buf.len(), BitReaderMode::BE);
 381         let mut cfr = ShortCodebookDescReader::new(scb_desc);
 382         let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap();
 383         assert_eq!(br2.read_cb(&cb).unwrap(), 0);
 384         assert_eq!(br2.read_cb(&cb).unwrap(), 2);
 385         assert_eq!(br2.read_cb(&cb).unwrap(), 5);
 386         assert_eq!(br2.read_cb(&cb).unwrap(), 8);
 387
 388         assert_eq!(reverse_bits(0b0000_0101_1011_1011_1101_1111_0111_1111, 32),
 389                                 0b1111_1110_1111_1011_1101_1101_1010_0000);
 390
 391         const BITS_LE: [u8; 3] = [0b11101111, 0b01110010, 0b01];
 392         let buf = &BITS_LE;
 393         let scble_desc: Vec<ShortCodebookDesc> = vec!(
 394             ShortCodebookDesc { code: 0b00,   bits: 2 },
 395             ShortCodebookDesc { code: 0,      bits: 0 },
 396             ShortCodebookDesc { code: 0b01,   bits: 2 },
 397             ShortCodebookDesc { code: 0,      bits: 0 },
 398             ShortCodebookDesc { code: 0,      bits: 0 },
 399             ShortCodebookDesc { code: 0b011,  bits: 3 },
 400             ShortCodebookDesc { code: 0,      bits: 0 },
 401             ShortCodebookDesc { code: 0b10111, bits: 5 },
 402             ShortCodebookDesc { code: 0b00111, bits: 5 },
 403             ShortCodebookDesc { code: 0b0101111, bits: 7 },
 404             ShortCodebookDesc { code: 0b0111111, bits: 7 },
 405             ShortCodebookDesc { code: 0b1011101111, bits: 10 }
 406         );
 407         let mut brl = BitReader::new(buf, buf.len(), BitReaderMode::LE);
 408         let mut cfr = ShortCodebookDescReader::new(scble_desc);
 409         let cb = Codebook::new(&mut cfr, CodebookMode::LSB).unwrap();
 410         assert_eq!(brl.read_cb(&cb).unwrap(), 11);
 411         assert_eq!(brl.read_cb(&cb).unwrap(), 0);
 412         assert_eq!(brl.read_cb(&cb).unwrap(), 7);
 413         assert_eq!(brl.read_cb(&cb).unwrap(), 0);
 414     }
 415 }