From: Kostya Shishkov Date: Sun, 30 Apr 2017 12:37:44 +0000 (+0200) Subject: add codebook reader X-Git-Url: https://git.nihav.org/?a=commitdiff_plain;h=4667915ab74fdf5bb14acbb486543090232bc74b;p=nihav.git add codebook reader --- diff --git a/src/io/codebook.rs b/src/io/codebook.rs new file mode 100644 index 0000000..f31e813 --- /dev/null +++ b/src/io/codebook.rs @@ -0,0 +1,183 @@ +use io::bitreader::BitReader; + +#[derive(Debug)] +pub enum CodebookError { + InvalidCodebook, + MemoryError, + InvalidCode, +} + +type CodebookResult = Result; + +pub struct FullCodebookDesc { + code: u32, + bits: u8, + sym: S, +} + +pub struct ShortCodebookDesc { + code: u32, + bits: u8, +} + +pub trait CodebookDescReader { + fn bits(&mut self, idx: usize) -> u8; + fn code(&mut self, idx: usize) -> u32; + fn sym (&mut self, idx: usize) -> S; + fn len (&mut self) -> usize; +} + +#[allow(dead_code)] +pub struct Codebook { + table: Vec, + syms: Vec, + lut_bits: u8, +} + +pub trait CodebookReader { + fn read_cb(&mut self, cb: &Codebook) -> CodebookResult; +} + +impl Codebook { +//todo allow add escapes + pub fn new(cb: &mut CodebookDescReader) -> CodebookResult { + let mut maxbits = 0; + let mut nnz = 0; + for i in 0..cb.len() { + let bits = cb.bits(i); + if bits > 0 { nnz = nnz + 1; } + if bits > maxbits { + maxbits = bits; + } + } + if maxbits == 0 { return Err(CodebookError::InvalidCodebook); } + + let mut table: Vec = Vec::new(); + let mut syms: Vec = Vec::new(); + let tab_len = 1 << maxbits; + table.reserve(tab_len); + if table.capacity() < tab_len { return Err(CodebookError::MemoryError); } + table.resize(tab_len, 0xFF); + syms.reserve(nnz); + if syms.capacity() < nnz { return Err(CodebookError::MemoryError); } + + let mut symidx: u32 = 0; + for i in 0..cb.len() { + let bits = cb.bits(i); + if bits == 0 { continue; } + let code = cb.code(i) << (maxbits - bits); + let fill_len = 1 << (maxbits - bits); + for j in 0..fill_len { + let idx = (code + j) as usize; + table[idx] = (symidx << 8) | (bits as u32); + } + symidx = symidx + 1; + } + + for i in 0..cb.len() { + if cb.bits(i) > 0 { + syms.push(cb.sym(i)); + } + } + + Ok(Codebook { table: table, syms: syms, lut_bits: maxbits }) + } +} + +impl<'a, S: Copy> CodebookReader for BitReader<'a> { + #[allow(unused_variables)] + fn read_cb(&mut self, cb: &Codebook) -> CodebookResult { + let lut_idx = self.peek(cb.lut_bits) as usize; + let bits = cb.table[lut_idx] & 0xFF; + let idx = (cb.table[lut_idx] >> 8) as usize; + if bits == 0xFF || (bits as isize) > self.left() { + return Err(CodebookError::InvalidCode); + } + if let Err(_) = self.skip(bits) {} + let sym = cb.syms[idx]; + return Ok(sym) + } +} + +pub struct FullCodebookDescReader { + data: Vec>, +} + +impl FullCodebookDescReader { + pub fn new(data: Vec>) -> Self { + FullCodebookDescReader { data: data } + } +} + +impl CodebookDescReader for FullCodebookDescReader { + fn bits(&mut self, idx: usize) -> u8 { self.data[idx].bits } + fn code(&mut self, idx: usize) -> u32 { self.data[idx].code } + fn sym (&mut self, idx: usize) -> S { self.data[idx].sym } + fn len(&mut self) -> usize { self.data.len() } +} + +pub struct ShortCodebookDescReader { + data: Vec, +} + +impl ShortCodebookDescReader { + pub fn new(data: Vec>) -> Self { + ShortCodebookDescReader { data: data } + } +} + +impl CodebookDescReader for ShortCodebookDescReader { + fn bits(&mut self, idx: usize) -> u8 { self.data[idx].bits } + fn code(&mut self, idx: usize) -> u32 { self.data[idx].code } + fn sym (&mut self, idx: usize) -> u32 { idx as u32 } + fn len(&mut self) -> usize { self.data.len() } +} + +#[cfg(test)] +mod test { + use super::*; + use io::bitreader::*; + + #[test] + fn test_cb() { + const BITS: [u8; 2] = [0b01011011, 0b10111100]; + let cb_desc: Vec> = vec!( + FullCodebookDesc { code: 0b0, bits: 1, sym: 16 }, + FullCodebookDesc { code: 0b10, bits: 2, sym: -3 }, + FullCodebookDesc { code: 0b110, bits: 3, sym: 42 }, + FullCodebookDesc { code: 0b1110, bits: 4, sym: -42 } + ); + let buf = &BITS; + let mut br = BitReader::new(buf, buf.len(), BitReaderMode::BE); + let mut cfr = FullCodebookDescReader::new(cb_desc); + let cb = Codebook::new(&mut cfr).unwrap(); + assert_eq!(br.read_cb(&cb).unwrap(), 16); + assert_eq!(br.read_cb(&cb).unwrap(), -3); + assert_eq!(br.read_cb(&cb).unwrap(), 42); + assert_eq!(br.read_cb(&cb).unwrap(), -42); + let ret = br.read_cb(&cb); + if let Err(e) = ret { + assert_eq!(e as i32, CodebookError::InvalidCode as i32); + } else { + assert_eq!(0, 1); + } + + let scb_desc: Vec = vec!( + ShortCodebookDesc { code: 0b0, bits: 1 }, + ShortCodebookDesc { code: 0, bits: 0 }, + ShortCodebookDesc { code: 0b10, bits: 2 }, + ShortCodebookDesc { code: 0, bits: 0 }, + ShortCodebookDesc { code: 0, bits: 0 }, + ShortCodebookDesc { code: 0b110, bits: 3 }, + ShortCodebookDesc { code: 0, bits: 0 }, + ShortCodebookDesc { code: 0b1110, bits: 4 } + ); + let mut br2 = BitReader::new(buf, buf.len(), BitReaderMode::BE); + let mut cfr = ShortCodebookDescReader::new(scb_desc); + let cb = Codebook::new(&mut cfr).unwrap(); + assert_eq!(br2.read_cb(&cb).unwrap(), 0); + assert_eq!(br2.read_cb(&cb).unwrap(), 2); + assert_eq!(br2.read_cb(&cb).unwrap(), 5); + assert_eq!(br2.read_cb(&cb).unwrap(), 7); + } +} diff --git a/src/io/mod.rs b/src/io/mod.rs index cb908c9..17f9c98 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -1,2 +1,5 @@ pub mod bitreader; +pub mod codebook; + pub mod byteio; +