core/io: document codebook module
[nihav.git] / nihav-core / src / io / codebook.rs
1 //! Codebook support for bitstream reader.
2 //!
3 //! Codebook is a set of unique bit strings and values assigned to them.
4 //! Since there are many ways to define codebook, this implementation employs [`CodebookDescReader`] trait to provide codebook generator with the codes.
5 //! Please also pay attention to the codebook creation mode: if bitstream reader reads bits starting from most significant bit first then you should use [`MSB`] mode and [`LSB`] mode otherwise.
6 //!
7 //! # Examples
8 //!
9 //! Create a codebook from arrays with codeword descriptions:
10 //! ```
11 //! use nihav_core::io::codebook::{ShortCodebookDesc, ShortCodebookDescReader, Codebook, CodebookMode};
12 //!
13 //! let cb_desc: Vec<ShortCodebookDesc> = vec!(
14 //! ShortCodebookDesc { code: 0b00, bits: 2 },
15 //! ShortCodebookDesc { code: 0, bits: 0 },
16 //! ShortCodebookDesc { code: 0b01, bits: 2 },
17 //! ShortCodebookDesc { code: 0b1, bits: 1 });
18 //! let mut cr = ShortCodebookDescReader::new(cb_desc);
19 //! let cb = Codebook::new(&mut cr, CodebookMode::LSB).unwrap();
20 //! ```
21 //!
22 //! Create a codebook using more flexible [`TableCodebookDescReader`] approach.
23 //! This will create a codebook for the following set: `1` -> -2, `01` -> -1, `001` -> 0, `0001` -> 1, `00001` -> 2.
24 //! ```
25 //! use nihav_core::io::codebook::{TableCodebookDescReader, Codebook, CodebookMode};
26 //!
27 //! fn map_cb_index(index: usize) -> i16 { (index as i16) - 2 }
28 //! const CB_BITS: [u8; 5] = [ 1, 2, 3, 4, 5 ];
29 //! const CB_CODES: [u8; 5] = [ 1, 1, 1, 1, 1 ];
30 //!
31 //! let mut tcr = TableCodebookDescReader::new(&CB_CODES, &CB_BITS, map_cb_index);
32 //! let cb = Codebook::new(&mut tcr, CodebookMode::MSB).unwrap();
33 //! ```
34 //!
35 //! Read value using a codebook:
36 //! ```no_run
37 //! use nihav_core::io::bitreader::BitReader;
38 //! use nihav_core::io::codebook::{Codebook, CodebookReader, CodebookMode};
39 //! # use nihav_core::io::codebook::{ShortCodebookDesc, ShortCodebookDescReader, CodebookDescReader, CodebookResult};
40 //!
41 //! # fn foo(br: &mut BitReader) -> CodebookResult<()> {
42 //! # let mut cr = ShortCodebookDescReader::new(vec![ShortCodebookDesc { code: 0b00, bits: 2 }]);
43 //! let cb = Codebook::new(&mut cr, CodebookMode::MSB).unwrap();
44 //! let value = br.read_cb(&cb)?;
45 //! # Ok(())
46 //! # }
47 //! ```
48 //!
49 //! [`MSB`]: ./enum.CodebookMode.html#variant.MSB
50 //! [`LSB`]: ./enum.CodebookMode.html#variant.LSB
51 //! [`CodebookDescReader`]: ./trait.CodebookDescReader.html
52 //! [`TableCodebookDescReader`]: ./struct.TableCodebookDescReader.html
53
54 use std::collections::HashMap;
55 use std::cmp::{max, min};
56 use super::bitreader::BitReader;
57
58 /// A list specifying general codebook operations errors.
59 #[derive(Debug)]
60 pub enum CodebookError {
61 /// Codebook description contains errors.
62 InvalidCodebook,
63 /// Could not allocate memory for codebook.
64 MemoryError,
65 /// Bitstream contains a sequence not present in codebook.
66 InvalidCode,
67 }
68
69 /// Codebook operation modes.
70 #[derive(Debug, Copy, Clone)]
71 pub enum CodebookMode {
72 /// Codes in the codebook should be read most significant bit first.
73 MSB,
74 /// Codes in the codebook should be read least significant bit first.
75 LSB,
76 }
77
78 /// A specialised `Result` type for codebook operations.
79 pub type CodebookResult<T> = Result<T, CodebookError>;
80
81 /// Codebook description for `(code bits, code length, code value)` triplet.
82 ///
83 /// This should be used to create a list of codeword definitions for [`FullCodebookDescReader`].
84 ///
85 /// [`FullCodebookDescReader`]: ./struct.FullCodebookDescReader.html
86 pub struct FullCodebookDesc<S> {
87 /// Codeword bits.
88 pub code: u32,
89 /// Codeword length.
90 pub bits: u8,
91 /// Codeword value (symbol).
92 pub sym: S,
93 }
94
95 /// Codebook description for `(code bits, code length)` pair with array index being used as codeword value.
96 ///
97 /// This should be used to create a list of codeword definitions for [`ShortCodebookDescReader`].
98 ///
99 /// [`ShortCodebookDescReader`]: ./struct.ShortCodebookDescReader.html
100 pub struct ShortCodebookDesc {
101 /// Codeword bits.
102 pub code: u32,
103 /// Codeword length.
104 pub bits: u8,
105 }
106
107 /// The interface for providing a list of codeword definitions to the codebook creator.
108 ///
109 /// The structure implementing this trait should be able to provide the total number of defined codewords and their bits and values. [`ShortCodebookDescReader`] or [`TableCodebookDescReader`] are some examples of such implementation.
110 /// Codeword definitions with zero length are ignored (those may be used to create sparse codebook definitions though).
111 ///
112 /// [`ShortCodebookDescReader`]: ./struct.ShortCodebookDescReader.html
113 /// [`TableCodebookDescReader`]: ./struct.TableCodebookDescReader.html
114 pub trait CodebookDescReader<S> {
115 /// Returns the codeword length for the provided index.
116 fn bits(&mut self, idx: usize) -> u8;
117 /// Returns the codeword bits for the provided index.
118 fn code(&mut self, idx: usize) -> u32;
119 /// Returns the codeword value (aka codeword symbol) for the provided index.
120 fn sym (&mut self, idx: usize) -> S;
121 /// Returns the total number of defined codewords.
122 fn len (&mut self) -> usize;
123 }
124
125 /// The codebook structure for code reading.
126 #[allow(dead_code)]
127 pub struct Codebook<S> {
128 table: Vec<u32>,
129 syms: Vec<S>,
130 lut_bits: u8,
131 }
132
133 /// Trait allowing bitreader to use codebook for decoding bit sequences.
134 pub trait CodebookReader<S> {
135 /// Reads the codeword from the bitstream and returns its value (or [`InvalidCode`] on error).
136 ///
137 /// [`InvalidCode`]: ./enum.CodebookError.html#variant.InvalidCode
138 fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S>;
139 }
140
141 const TABLE_FILL_VALUE: u32 = 0x7F;
142 const MAX_LUT_BITS: u8 = 10;
143
144 fn fill_lut_msb(table: &mut Vec<u32>, off: usize,
145 code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) {
146 if !esc {
147 let fill_len = lut_bits - bits;
148 let fill_size = 1 << fill_len;
149 let fill_code = code << (lut_bits - bits);
150 let lut_value = (symidx << 8) | u32::from(bits);
151 for j in 0..fill_size {
152 let idx = (fill_code + j) as usize;
153 table[idx + off] = lut_value;
154 }
155 } else {
156 let idx = (code as usize) + off;
157 table[idx] = (symidx << 8) | 0x80 | u32::from(bits);
158 }
159 }
160
161 fn fill_lut_lsb(table: &mut Vec<u32>, off: usize,
162 code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) {
163 if !esc {
164 let fill_len = lut_bits - bits;
165 let fill_size = 1 << fill_len;
166 let fill_code = code;
167 let step = lut_bits - fill_len;
168 for j in 0..fill_size {
169 let idx = (fill_code + (j << step)) as usize;
170 table[idx + off] = (symidx << 8) | u32::from(bits);
171 }
172 } else {
173 let idx = (code as usize) + off;
174 table[idx] = (symidx << 8) | 0x80 | u32::from(bits);
175 }
176 }
177
178 fn fill_lut(table: &mut Vec<u32>, mode: CodebookMode,
179 off: usize, code: u32, bits: u8, lut_bits: u8, symidx: u32, esc: bool) -> bool {
180 match mode {
181 CodebookMode::MSB => fill_lut_msb(table, off, code, bits, lut_bits, symidx, esc),
182 CodebookMode::LSB => fill_lut_lsb(table, off, code, bits, lut_bits, symidx, esc),
183 };
184 bits > lut_bits
185 }
186
187 fn resize_table(table: &mut Vec<u32>, bits: u8) -> CodebookResult<u32> {
188 let add_size = (1 << bits) as usize;
189 table.reserve(add_size);
190 let cur_off = table.len() as u32;
191 let new_size = table.len() + add_size;
192 if table.capacity() < new_size { return Err(CodebookError::MemoryError); }
193 table.resize(new_size, TABLE_FILL_VALUE);
194 Ok(cur_off)
195 }
196
197
198 fn extract_lut_part(code: u32, bits: u8, lut_bits: u8, mode: CodebookMode) -> u32 {
199 match mode {
200 CodebookMode::MSB => code >> (bits - lut_bits),
201 CodebookMode::LSB => code & ((1 << lut_bits) - 1),
202 }
203 }
204
205 fn extract_esc_part(code: u32, bits: u8, lut_bits: u8, mode: CodebookMode) -> u32 {
206 match mode {
207 CodebookMode::MSB => code & ((1 << (bits - lut_bits)) - 1),
208 CodebookMode::LSB => code >> lut_bits,
209 }
210 }
211
212 #[derive(Clone,Copy)]
213 struct Code {
214 code: u32,
215 bits: u8,
216 idx: usize,
217 }
218
219 struct CodeBucket {
220 maxlen: u8,
221 offset: usize,
222 codes: Vec<Code>,
223 }
224
225 impl CodeBucket {
226 fn new() -> Self {
227 CodeBucket { maxlen: 0, offset: 0, codes: Vec::new() }
228 }
229 fn add_code(&mut self, c: Code) {
230 if c.bits > self.maxlen { self.maxlen = c.bits; }
231 self.codes.push(c);
232 }
233 }
234
235 type EscapeCodes = HashMap<u32, CodeBucket>;
236
237 fn add_esc_code(cc: &mut EscapeCodes, key: u32, code: u32, bits: u8, idx: usize) {
238 cc.entry(key).or_insert_with(CodeBucket::new);
239 let b = cc.get_mut(&key);
240 if let Some(bucket) = b {
241 bucket.add_code(Code {code, bits, idx });
242 } else { panic!("no bucket when expected!"); }
243 }
244
245 fn build_esc_lut(table: &mut Vec<u32>,
246 mode: CodebookMode,
247 bucket: &CodeBucket) -> CodebookResult<()> {
248 let mut escape_list: EscapeCodes = HashMap::new();
249 let maxlen = if bucket.maxlen > MAX_LUT_BITS { MAX_LUT_BITS } else { bucket.maxlen };
250
251 for code in &bucket.codes {
252 let bits = code.bits;
253 if code.bits <= MAX_LUT_BITS {
254 fill_lut(table, mode, bucket.offset, code.code, bits,
255 maxlen, code.idx as u32, false);
256 } else {
257 let ckey = extract_lut_part(code.code, bits, MAX_LUT_BITS, mode);
258 let cval = extract_esc_part(code.code, bits, MAX_LUT_BITS, mode);
259 add_esc_code(&mut escape_list, ckey, cval, bits - MAX_LUT_BITS, code.idx);
260 }
261 }
262
263 let cur_offset = bucket.offset;
264 for (ckey, sec_bucket) in &mut escape_list {
265 let key = *ckey as u32;
266 let maxlen = min(sec_bucket.maxlen, MAX_LUT_BITS);
267 let new_off = resize_table(table, maxlen)?;
268 fill_lut(table, mode, cur_offset, key, maxlen,
269 MAX_LUT_BITS, new_off, true);
270 sec_bucket.offset = new_off as usize;
271 }
272
273 for sec_bucket in escape_list.values() {
274 build_esc_lut(table, mode, sec_bucket)?;
275 }
276
277 Ok(())
278 }
279
280 impl<S: Copy> Codebook<S> {
281
282 /// Constructs a new `Codebook` instance using provided codebook description and mode.
283 pub fn new(cb: &mut CodebookDescReader<S>, mode: CodebookMode) -> CodebookResult<Self> {
284 let mut maxbits = 0;
285 let mut nnz = 0;
286 let mut escape_list: EscapeCodes = HashMap::new();
287
288 let mut symidx: usize = 0;
289 for i in 0..cb.len() {
290 let bits = cb.bits(i);
291 if bits > 0 {
292 nnz += 1;
293 if cb.code(i) >= (1 << bits) {
294 return Err(CodebookError::InvalidCodebook);
295 }
296 }
297 maxbits = max(bits, maxbits);
298 if bits > MAX_LUT_BITS {
299 let code = cb.code(i);
300 let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode);
301 let cval = extract_esc_part(code, bits, MAX_LUT_BITS, mode);
302 add_esc_code(&mut escape_list, ckey, cval, bits - MAX_LUT_BITS, symidx);
303 }
304 if bits > 0 { symidx += 1; }
305 }
306 if maxbits == 0 { return Err(CodebookError::InvalidCodebook); }
307
308 if maxbits > MAX_LUT_BITS { maxbits = MAX_LUT_BITS; }
309
310 let tab_len = 1 << maxbits;
311 let mut table: Vec<u32> = Vec::with_capacity(tab_len);
312 let mut syms: Vec<S> = Vec::with_capacity(nnz);
313 if table.capacity() < tab_len { return Err(CodebookError::MemoryError); }
314 if syms.capacity() < nnz { return Err(CodebookError::MemoryError); }
315 table.resize(tab_len, TABLE_FILL_VALUE);
316
317 let mut symidx: u32 = 0;
318 for i in 0..cb.len() {
319 let bits = cb.bits(i);
320 let code = cb.code(i);
321 if bits == 0 { continue; }
322 if bits <= MAX_LUT_BITS {
323 fill_lut(&mut table, mode, 0, code, bits, maxbits, symidx, false);
324 } else {
325 let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode) as usize;
326 if table[ckey] == TABLE_FILL_VALUE {
327 let key = ckey as u32;
328 if let Some(bucket) = escape_list.get_mut(&key) {
329 let maxlen = min(bucket.maxlen, MAX_LUT_BITS);
330 let new_off = resize_table(&mut table, maxlen)?;
331 fill_lut(&mut table, mode, 0, key, maxlen, MAX_LUT_BITS, new_off, true);
332 bucket.offset = new_off as usize;
333 }
334 }
335 }
336 symidx += 1;
337 }
338
339 for bucket in escape_list.values() {
340 build_esc_lut(&mut table, mode, &bucket)?;
341 }
342
343 for i in 0..cb.len() {
344 if cb.bits(i) > 0 {
345 syms.push(cb.sym(i));
346 }
347 }
348
349 Ok(Codebook { table, syms, lut_bits: maxbits })
350 }
351 }
352
353 impl<'a, S: Copy> CodebookReader<S> for BitReader<'a> {
354 #[allow(unused_variables)]
355 fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S> {
356 let mut esc = true;
357 let mut idx = 0;
358 let mut lut_bits = cb.lut_bits;
359 while esc {
360 let lut_idx = (self.peek(lut_bits) as usize) + (idx as usize);
361 if cb.table[lut_idx] == TABLE_FILL_VALUE { return Err(CodebookError::InvalidCode); }
362 let bits = cb.table[lut_idx] & 0x7F;
363 esc = (cb.table[lut_idx] & 0x80) != 0;
364 idx = (cb.table[lut_idx] >> 8) as usize;
365 if (bits as isize) > self.left() {
366 return Err(CodebookError::InvalidCode);
367 }
368 let skip_bits = if esc { u32::from(lut_bits) } else { bits };
369 self.skip(skip_bits as u32).unwrap();
370 lut_bits = bits as u8;
371 }
372 Ok(cb.syms[idx])
373 }
374 }
375
376 /// Codebook description that stores a list of codewords and their values.
377 pub struct FullCodebookDescReader<S> {
378 data: Vec<FullCodebookDesc<S>>,
379 }
380
381 impl<S> FullCodebookDescReader<S> {
382 /// Constructs a new `FullCodebookDescReader` instance.
383 pub fn new(data: Vec<FullCodebookDesc<S>>) -> Self {
384 FullCodebookDescReader { data }
385 }
386 }
387
388 impl<S: Copy> CodebookDescReader<S> for FullCodebookDescReader<S> {
389 fn bits(&mut self, idx: usize) -> u8 { self.data[idx].bits }
390 fn code(&mut self, idx: usize) -> u32 { self.data[idx].code }
391 fn sym (&mut self, idx: usize) -> S { self.data[idx].sym }
392 fn len(&mut self) -> usize { self.data.len() }
393 }
394
395 /// Codebook description that stores a list of codewords and their value is equal to the index.
396 pub struct ShortCodebookDescReader {
397 data: Vec<ShortCodebookDesc>,
398 }
399
400 impl ShortCodebookDescReader {
401 /// Constructs a new `ShortCodebookDescReader` instance.
402 pub fn new(data: Vec<ShortCodebookDesc<>>) -> Self {
403 ShortCodebookDescReader { data }
404 }
405 }
406
407 impl CodebookDescReader<u32> for ShortCodebookDescReader {
408 fn bits(&mut self, idx: usize) -> u8 { self.data[idx].bits }
409 fn code(&mut self, idx: usize) -> u32 { self.data[idx].code }
410 fn sym (&mut self, idx: usize) -> u32 { idx as u32 }
411 fn len(&mut self) -> usize { self.data.len() }
412 }
413
414 /// Flexible codebook description that uses two separate arrays for codeword bits and lengths and a function that maps codeword index into its symbol.
415 pub struct TableCodebookDescReader<'a, CodeType:'static, IndexType:'static> {
416 bits: &'a [u8],
417 codes: &'a [CodeType],
418 idx_map: fn(usize) -> IndexType,
419 }
420
421 impl<'a, CodeType, IndexType> TableCodebookDescReader<'a, CodeType, IndexType> {
422 /// Constructs a new `TableCodebookDescReader` instance.
423 pub fn new(codes: &'a [CodeType], bits: &'a [u8], idx_map: fn(usize) -> IndexType) -> Self {
424 Self { bits, codes, idx_map }
425 }
426 }
427 impl<'a, CodeType: Copy+Into<u32>, IndexType> CodebookDescReader<IndexType> for TableCodebookDescReader<'a, CodeType, IndexType>
428 {
429 fn bits(&mut self, idx: usize) -> u8 { self.bits[idx] }
430 fn code(&mut self, idx: usize) -> u32 { self.codes[idx].into() }
431 fn sym (&mut self, idx: usize) -> IndexType { (self.idx_map)(idx) }
432 fn len(&mut self) -> usize { self.bits.len() }
433 }
434
435 #[cfg(test)]
436 mod test {
437 use super::*;
438 use crate::io::bitreader::*;
439
440 #[test]
441 fn test_cb() {
442 const BITS: [u8; 2] = [0b01011011, 0b10111100];
443 let cb_desc: Vec<FullCodebookDesc<i8>> = vec!(
444 FullCodebookDesc { code: 0b0, bits: 1, sym: 16 },
445 FullCodebookDesc { code: 0b10, bits: 2, sym: -3 },
446 FullCodebookDesc { code: 0b110, bits: 3, sym: 42 },
447 FullCodebookDesc { code: 0b1110, bits: 4, sym: -42 }
448 );
449 let buf = &BITS;
450 let mut br = BitReader::new(buf, BitReaderMode::BE);
451 let mut cfr = FullCodebookDescReader::new(cb_desc);
452 let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap();
453 assert_eq!(br.read_cb(&cb).unwrap(), 16);
454 assert_eq!(br.read_cb(&cb).unwrap(), -3);
455 assert_eq!(br.read_cb(&cb).unwrap(), 42);
456 assert_eq!(br.read_cb(&cb).unwrap(), -42);
457 let ret = br.read_cb(&cb);
458 if let Err(e) = ret {
459 assert_eq!(e as i32, CodebookError::InvalidCode as i32);
460 } else {
461 assert_eq!(0, 1);
462 }
463
464 let scb_desc: Vec<ShortCodebookDesc> = vec!(
465 ShortCodebookDesc { code: 0b0, bits: 1 },
466 ShortCodebookDesc { code: 0, bits: 0 },
467 ShortCodebookDesc { code: 0b10, bits: 2 },
468 ShortCodebookDesc { code: 0, bits: 0 },
469 ShortCodebookDesc { code: 0, bits: 0 },
470 ShortCodebookDesc { code: 0b110, bits: 3 },
471 ShortCodebookDesc { code: 0, bits: 0 },
472 ShortCodebookDesc { code: 0b11100, bits: 5 },
473 ShortCodebookDesc { code: 0b11101, bits: 5 },
474 ShortCodebookDesc { code: 0b1111010, bits: 7 },
475 ShortCodebookDesc { code: 0b1111011, bits: 7 },
476 ShortCodebookDesc { code: 0b1111110, bits: 7 },
477 ShortCodebookDesc { code: 0b11111111, bits: 8 }
478 );
479 let mut br2 = BitReader::new(buf, BitReaderMode::BE);
480 let mut cfr = ShortCodebookDescReader::new(scb_desc);
481 let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap();
482 assert_eq!(br2.read_cb(&cb).unwrap(), 0);
483 assert_eq!(br2.read_cb(&cb).unwrap(), 2);
484 assert_eq!(br2.read_cb(&cb).unwrap(), 5);
485 assert_eq!(br2.read_cb(&cb).unwrap(), 8);
486
487 assert_eq!(reverse_bits(0b0000_0101_1011_1011_1101_1111_0111_1111, 32),
488 0b1111_1110_1111_1011_1101_1101_1010_0000);
489
490 const BITS_LE: [u8; 3] = [0b11101111, 0b01110010, 0b01];
491 let buf = &BITS_LE;
492 let scble_desc: Vec<ShortCodebookDesc> = vec!(
493 ShortCodebookDesc { code: 0b00, bits: 2 },
494 ShortCodebookDesc { code: 0, bits: 0 },
495 ShortCodebookDesc { code: 0b01, bits: 2 },
496 ShortCodebookDesc { code: 0, bits: 0 },
497 ShortCodebookDesc { code: 0, bits: 0 },
498 ShortCodebookDesc { code: 0b011, bits: 3 },
499 ShortCodebookDesc { code: 0, bits: 0 },
500 ShortCodebookDesc { code: 0b10111, bits: 5 },
501 ShortCodebookDesc { code: 0b00111, bits: 5 },
502 ShortCodebookDesc { code: 0b0101111, bits: 7 },
503 ShortCodebookDesc { code: 0b0111111, bits: 7 },
504 ShortCodebookDesc { code: 0b1011101111, bits: 10 }
505 );
506 let mut brl = BitReader::new(buf, BitReaderMode::LE);
507 let mut cfr = ShortCodebookDescReader::new(scble_desc);
508 let cb = Codebook::new(&mut cfr, CodebookMode::LSB).unwrap();
509 assert_eq!(brl.read_cb(&cb).unwrap(), 11);
510 assert_eq!(brl.read_cb(&cb).unwrap(), 0);
511 assert_eq!(brl.read_cb(&cb).unwrap(), 7);
512 assert_eq!(brl.read_cb(&cb).unwrap(), 0);
513 }
514 }