av_bitstream/
codebook.rs

1//! Codebook support for bitstream reader.
2//!
3//! Codebook is a set of unique bit strings and values assigned to them.
4
5use std::cmp::{max, min};
6use std::collections::HashMap;
7use std::fmt;
8use std::marker::PhantomData;
9
10use num_traits::AsPrimitive;
11
12use crate::bitread::*;
13
14/// Codebook operations errors.
15#[derive(Debug)]
16pub enum CodebookError {
17    /// The codebook is invalid.
18    InvalidCodebook,
19    /// The analyzed bitstream is not present in the codebook.
20    InvalidCode,
21}
22
23impl std::error::Error for CodebookError {}
24
25impl fmt::Display for CodebookError {
26    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
27        match self {
28            InvalidCodebook => write!(f, "Invalid Codebook"),
29            InvalidCode => write!(f, "Invalid Code"),
30        }
31    }
32}
33
34use self::CodebookError::*;
35
36/// Codebook operation modes.
37#[derive(Debug, Copy, Clone)]
38#[allow(clippy::upper_case_acronyms)]
39pub enum CodebookMode {
40    /// Codes in the codebook should be read most significant bit first.
41    MSB,
42    /// Codes in the codebook should be read least significant bit first.
43    LSB,
44}
45
46/// Codebook description for `(code bits, code length, code value)` triplet.
47pub struct FullCodebookDesc<S> {
48    /// Codeword bits.
49    pub code: u32,
50    /// Codeword length.
51    pub bits: u8,
52    /// Codeword value (symbol).
53    pub sym: S,
54}
55
56/// Codebook description for `(code bits, code length)` pair with array index
57/// being used as codeword value.
58pub struct ShortCodebookDesc {
59    /// Codeword bits.
60    pub code: u32,
61    /// Codeword length.
62    pub bits: u8,
63}
64
65/// This trait defines a series of methods to get some information
66/// from a codebook.
67pub trait CodebookDescReader<S> {
68    /// Returns the codeword length for the provided index.
69    fn bits(&self, idx: usize) -> u8;
70    /// Returns the codeword bits for the provided index.
71    fn code(&self, idx: usize) -> u32;
72    /// Returns the codeword value (codeword symbol) for the provided index.
73    fn sym(&self, idx: usize) -> S;
74
75    /// Returns the total number of defined codewords.
76    fn len(&self) -> usize;
77    /// Tells if the codebook is empty or not.
78    fn is_empty(&self) -> bool;
79}
80
81/// The codebook structure for code reading.
82#[allow(dead_code)]
83pub struct Codebook<S> {
84    table: Vec<u32>,
85    syms: Vec<S>,
86    lut_bits: u8,
87}
88
89/// Adopted by a bitreader to use codebook for decoding bit sequences.
90pub trait CodebookReader<S> {
91    /// Reads the codeword from a bitstream and returns its value.
92    fn read_cb(&mut self, cb: &Codebook<S>) -> Result<S, CodebookError>;
93}
94
95/// Returns the reversed sequence of bits passed as input.
96pub fn reverse_bits(inval: u32) -> u32 {
97    const REV_TAB: [u8; 16] = [
98        0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110, 0b0001, 0b1001, 0b0101,
99        0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
100    ];
101
102    let mut ret = 0;
103    let mut val = inval;
104    for _ in 0..8 {
105        ret = (ret << 4) | u32::from(REV_TAB[(val & 0xF) as usize]);
106        val >>= 4;
107    }
108    ret
109}
110
111const TABLE_FILL_VALUE: u32 = 0x7F;
112const MAX_LUT_BITS: u8 = 10;
113
114fn fill_lut_msb(
115    table: &mut [u32],
116    off: usize,
117    code: u32,
118    bits: u8,
119    lut_bits: u8,
120    symidx: u32,
121    esc: bool,
122) {
123    if !esc {
124        let fill_len = lut_bits - bits;
125        let fill_size = 1 << fill_len;
126        let fill_code = code << (lut_bits - bits);
127        let lut_value = (symidx << 8) | u32::from(bits);
128        for j in 0..fill_size {
129            let idx = (fill_code + j) as usize;
130            table[idx + off] = lut_value;
131        }
132    } else {
133        let idx = (code as usize) + off;
134        table[idx] = (symidx << 8) | 0x80 | u32::from(bits);
135    }
136}
137
138fn fill_lut_lsb(
139    table: &mut [u32],
140    off: usize,
141    code: u32,
142    bits: u8,
143    lut_bits: u8,
144    symidx: u32,
145    esc: bool,
146) {
147    if !esc {
148        let fill_len = lut_bits - bits;
149        let fill_size = 1 << fill_len;
150        let fill_code = code;
151        let step = lut_bits - fill_len;
152        for j in 0..fill_size {
153            let idx = (fill_code + (j << step)) as usize;
154            table[idx + off] = (symidx << 8) | u32::from(bits);
155        }
156    } else {
157        let idx = (code as usize) + off;
158        table[idx] = (symidx << 8) | 0x80 | u32::from(bits);
159    }
160}
161
162fn fill_lut(
163    table: &mut [u32],
164    mode: CodebookMode,
165    off: usize,
166    code: u32,
167    bits: u8,
168    lut_bits: u8,
169    symidx: u32,
170    esc: bool,
171) -> bool {
172    match mode {
173        CodebookMode::MSB => fill_lut_msb(table, off, code, bits, lut_bits, symidx, esc),
174        CodebookMode::LSB => fill_lut_lsb(table, off, code, bits, lut_bits, symidx, esc),
175    };
176    bits > lut_bits
177}
178
179fn resize_table(table: &mut Vec<u32>, bits: u8) -> u32 {
180    let add_size = (1 << bits) as usize;
181    table.reserve(add_size);
182    let cur_off = table.len() as u32;
183    let new_size = table.len() + add_size;
184    table.resize(new_size, TABLE_FILL_VALUE);
185    cur_off
186}
187
188fn extract_lut_part(code: u32, bits: u8, lut_bits: u8, mode: CodebookMode) -> u32 {
189    match mode {
190        CodebookMode::MSB => code >> (bits - lut_bits),
191        CodebookMode::LSB => code & ((1 << lut_bits) - 1),
192    }
193}
194
195fn extract_esc_part(code: u32, bits: u8, lut_bits: u8, mode: CodebookMode) -> u32 {
196    match mode {
197        CodebookMode::MSB => code & ((1 << (bits - lut_bits)) - 1),
198        CodebookMode::LSB => code >> lut_bits,
199    }
200}
201
202#[derive(Clone, Copy)]
203struct Code {
204    code: u32,
205    bits: u8,
206    idx: usize,
207}
208
209struct CodeBucket {
210    maxlen: u8,
211    offset: usize,
212    codes: Vec<Code>,
213}
214
215impl CodeBucket {
216    fn new() -> Self {
217        CodeBucket {
218            maxlen: 0,
219            offset: 0,
220            codes: Vec::new(),
221        }
222    }
223    fn add_code(&mut self, c: Code) {
224        if c.bits > self.maxlen {
225            self.maxlen = c.bits;
226        }
227        self.codes.push(c);
228    }
229}
230
231type EscapeCodes = HashMap<u32, CodeBucket>;
232
233fn add_esc_code(cc: &mut EscapeCodes, key: u32, code: u32, bits: u8, idx: usize) {
234    let bucket = cc.entry(key).or_insert_with(CodeBucket::new);
235    bucket.add_code(Code { code, bits, idx });
236}
237
238fn build_esc_lut(
239    table: &mut Vec<u32>,
240    mode: CodebookMode,
241    bucket: &CodeBucket,
242) -> Result<(), CodebookError> {
243    let mut escape_list: EscapeCodes = HashMap::new();
244    let maxlen = min(bucket.maxlen, MAX_LUT_BITS);
245
246    for code in &bucket.codes {
247        let bits = code.bits;
248        if code.bits <= MAX_LUT_BITS {
249            fill_lut(
250                table,
251                mode,
252                bucket.offset,
253                code.code,
254                bits,
255                maxlen,
256                code.idx as u32,
257                false,
258            );
259        } else {
260            let ckey = extract_lut_part(code.code, bits, MAX_LUT_BITS, mode);
261            let cval = extract_esc_part(code.code, bits, MAX_LUT_BITS, mode);
262            add_esc_code(&mut escape_list, ckey, cval, bits - MAX_LUT_BITS, code.idx);
263        }
264    }
265
266    let cur_offset = bucket.offset;
267    for (ckey, sec_bucket) in &mut escape_list {
268        let key = *ckey;
269        let maxlen = min(sec_bucket.maxlen, MAX_LUT_BITS);
270        let new_off = resize_table(table, maxlen);
271        fill_lut(
272            table,
273            mode,
274            cur_offset,
275            key,
276            maxlen,
277            MAX_LUT_BITS,
278            new_off,
279            true,
280        );
281        sec_bucket.offset = new_off as usize;
282    }
283
284    for sec_bucket in escape_list.values() {
285        build_esc_lut(table, mode, sec_bucket)?;
286    }
287
288    Ok(())
289}
290
291impl<S: Copy> Codebook<S> {
292    /// Constructs a new `Codebook` instance using provided
293    /// codebook description and mode.
294    pub fn new(cb: &dyn CodebookDescReader<S>, mode: CodebookMode) -> Result<Self, CodebookError> {
295        let mut maxbits = 0;
296        let mut nnz = 0;
297        let mut escape_list: EscapeCodes = HashMap::new();
298
299        let mut symidx: usize = 0;
300        for i in 0..cb.len() {
301            let bits = cb.bits(i);
302            if bits > 0 {
303                nnz += 1;
304            }
305            maxbits = max(bits, maxbits);
306            if bits > MAX_LUT_BITS {
307                let code = cb.code(i);
308                let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode);
309                let cval = extract_esc_part(code, bits, MAX_LUT_BITS, mode);
310                add_esc_code(&mut escape_list, ckey, cval, bits - MAX_LUT_BITS, symidx);
311            }
312            if bits > 0 {
313                symidx += 1;
314            }
315        }
316        if maxbits == 0 {
317            return Err(InvalidCodebook);
318        }
319
320        if maxbits > MAX_LUT_BITS {
321            maxbits = MAX_LUT_BITS;
322        }
323
324        let tab_len = 1 << maxbits;
325        let mut table: Vec<u32> = Vec::with_capacity(tab_len);
326        let mut syms: Vec<S> = Vec::with_capacity(nnz);
327        table.resize(tab_len, TABLE_FILL_VALUE);
328
329        let mut symidx: u32 = 0;
330        for i in 0..cb.len() {
331            let bits = cb.bits(i);
332            let code = cb.code(i);
333            if bits == 0 {
334                continue;
335            }
336            if bits <= MAX_LUT_BITS {
337                fill_lut(&mut table, mode, 0, code, bits, maxbits, symidx, false);
338            } else {
339                let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode) as usize;
340                if table[ckey] == TABLE_FILL_VALUE {
341                    let key = ckey as u32;
342                    if let Some(bucket) = escape_list.get_mut(&key) {
343                        let maxlen = min(bucket.maxlen, MAX_LUT_BITS);
344                        let new_off = resize_table(&mut table, maxlen);
345                        fill_lut(
346                            &mut table,
347                            mode,
348                            0,
349                            key,
350                            maxlen,
351                            MAX_LUT_BITS,
352                            new_off,
353                            true,
354                        );
355                        bucket.offset = new_off as usize;
356                    }
357                }
358            }
359            symidx += 1;
360        }
361
362        for bucket in escape_list.values() {
363            build_esc_lut(&mut table, mode, bucket)?;
364        }
365
366        for i in 0..cb.len() {
367            if cb.bits(i) > 0 {
368                syms.push(cb.sym(i));
369            }
370        }
371
372        Ok(Codebook {
373            table,
374            syms,
375            lut_bits: maxbits,
376        })
377    }
378}
379
380impl<'a, S: Copy, B: BitRead<'a>> CodebookReader<S> for B {
381    fn read_cb(&mut self, cb: &Codebook<S>) -> Result<S, CodebookError> {
382        let mut esc = true;
383        let mut idx = 0;
384        let mut lut_bits = cb.lut_bits;
385        while esc {
386            let lut_idx = (self.peek_bits_32(lut_bits as usize) as usize) + idx;
387            if cb.table[lut_idx] == TABLE_FILL_VALUE {
388                return Err(InvalidCode);
389            }
390            let bits = cb.table[lut_idx] & 0x7F;
391            esc = (cb.table[lut_idx] & 0x80) != 0;
392            idx = (cb.table[lut_idx] >> 8) as usize;
393            if bits > self.available() as u32 {
394                return Err(InvalidCode);
395            }
396            let skip_bits = if esc {
397                lut_bits as usize
398            } else {
399                bits as usize
400            };
401            self.skip_bits(skip_bits);
402            lut_bits = bits as u8;
403        }
404        Ok(cb.syms[idx])
405    }
406}
407
408impl<S: Copy> CodebookDescReader<S> for Vec<FullCodebookDesc<S>> {
409    fn bits(&self, idx: usize) -> u8 {
410        self[idx].bits
411    }
412    fn code(&self, idx: usize) -> u32 {
413        self[idx].code
414    }
415    fn sym(&self, idx: usize) -> S {
416        self[idx].sym
417    }
418    fn len(&self) -> usize {
419        Vec::len(self)
420    }
421    fn is_empty(&self) -> bool {
422        Vec::is_empty(self)
423    }
424}
425
426impl CodebookDescReader<u32> for Vec<ShortCodebookDesc> {
427    fn bits(&self, idx: usize) -> u8 {
428        self[idx].bits
429    }
430    fn code(&self, idx: usize) -> u32 {
431        self[idx].code
432    }
433    fn sym(&self, idx: usize) -> u32 {
434        idx as u32
435    }
436    fn len(&self) -> usize {
437        Vec::len(self)
438    }
439    fn is_empty(&self) -> bool {
440        Vec::is_empty(self)
441    }
442}
443
444impl<'a, S: Copy> CodebookDescReader<S> for &'a [FullCodebookDesc<S>] {
445    fn bits(&self, idx: usize) -> u8 {
446        self[idx].bits
447    }
448    fn code(&self, idx: usize) -> u32 {
449        self[idx].code
450    }
451    fn sym(&self, idx: usize) -> S {
452        self[idx].sym
453    }
454    fn len(&self) -> usize {
455        <[FullCodebookDesc<S>]>::len(self)
456    }
457    fn is_empty(&self) -> bool {
458        <[FullCodebookDesc<S>]>::is_empty(self)
459    }
460}
461
462impl<'a> CodebookDescReader<u32> for &'a [ShortCodebookDesc] {
463    fn bits(&self, idx: usize) -> u8 {
464        self[idx].bits
465    }
466    fn code(&self, idx: usize) -> u32 {
467        self[idx].code
468    }
469    fn sym(&self, idx: usize) -> u32 {
470        idx as u32
471    }
472    fn len(&self) -> usize {
473        <[ShortCodebookDesc]>::len(self)
474    }
475    fn is_empty(&self) -> bool {
476        <[ShortCodebookDesc]>::is_empty(self)
477    }
478}
479
480/// Flexible codebook description that uses two separate arrays for
481/// codeword bits and lengths.
482pub struct TableCodebookDescReader<CodeType: 'static, SymType> {
483    bits: &'static [u8],
484    codes: &'static [CodeType],
485    _sym: PhantomData<SymType>,
486}
487
488impl<CodeType, SymType> CodebookDescReader<SymType> for TableCodebookDescReader<CodeType, SymType>
489where
490    CodeType: Copy + Into<u32> + 'static,
491    SymType: Copy + 'static,
492    usize: AsPrimitive<SymType>,
493{
494    fn bits(&self, idx: usize) -> u8 {
495        self.bits[idx]
496    }
497    fn code(&self, idx: usize) -> u32 {
498        self.codes[idx].into()
499    }
500    fn sym(&self, idx: usize) -> SymType {
501        idx.as_()
502    }
503    fn len(&self) -> usize {
504        self.bits.len()
505    }
506    fn is_empty(&self) -> bool {
507        self.bits.is_empty()
508    }
509}
510
511#[cfg(test)]
512mod test {
513    use super::*;
514
515    const BITS: [u8; 8] = [0b01011011, 0b10111100, 0b11111111, 0, 0, 0, 0, 0];
516
517    #[test]
518    fn test_refill_codebook_msb() {
519        // make sure reading codes across 8-byte boundary works
520        let bits = [
521            0b1111_1000,
522            0b0011_1110,
523            0b0000_1111,
524            0b1000_0011,
525            0b1110_0000,
526            0b1111_1000,
527            0b0011_1110,
528            0b0000_1111,
529            0b1000_0000,
530            0,
531            0,
532            0,
533            0,
534            0,
535            0,
536            0,
537        ];
538
539        let cb_desc: Vec<ShortCodebookDesc> = vec![
540            ShortCodebookDesc {
541                code: 0b11111,
542                bits: 5,
543            },
544            ShortCodebookDesc {
545                code: 0b00000,
546                bits: 5,
547            },
548        ];
549        let buf = &bits;
550        let mut br = BitReadBE::new(buf);
551        let cb = Codebook::new(&cb_desc, CodebookMode::MSB).unwrap();
552
553        assert_eq!(br.read_cb(&cb).unwrap(), 0);
554        assert_eq!(br.read_cb(&cb).unwrap(), 1);
555        assert_eq!(br.read_cb(&cb).unwrap(), 0);
556        assert_eq!(br.read_cb(&cb).unwrap(), 1);
557        assert_eq!(br.read_cb(&cb).unwrap(), 0);
558        assert_eq!(br.read_cb(&cb).unwrap(), 1);
559        assert_eq!(br.read_cb(&cb).unwrap(), 0);
560        assert_eq!(br.read_cb(&cb).unwrap(), 1);
561        assert_eq!(br.read_cb(&cb).unwrap(), 0);
562        assert_eq!(br.read_cb(&cb).unwrap(), 1);
563        assert_eq!(br.read_cb(&cb).unwrap(), 0);
564        assert_eq!(br.read_cb(&cb).unwrap(), 1);
565        assert_eq!(br.read_cb(&cb).unwrap(), 0);
566        assert_eq!(br.read_cb(&cb).unwrap(), 1);
567    }
568
569    #[test]
570    fn test_full_codebook_msb() {
571        let cb_desc: Vec<FullCodebookDesc<i8>> = vec![
572            FullCodebookDesc {
573                code: 0b0,
574                bits: 1,
575                sym: 16,
576            },
577            FullCodebookDesc {
578                code: 0b10,
579                bits: 2,
580                sym: -3,
581            },
582            FullCodebookDesc {
583                code: 0b110,
584                bits: 3,
585                sym: 42,
586            },
587            FullCodebookDesc {
588                code: 0b1110,
589                bits: 4,
590                sym: -42,
591            },
592        ];
593        let buf = &BITS;
594        let mut br = BitReadBE::new(buf);
595        let cb = Codebook::new(&cb_desc, CodebookMode::MSB).unwrap();
596
597        assert_eq!(br.read_cb(&cb).unwrap(), 16);
598        assert_eq!(br.read_cb(&cb).unwrap(), -3);
599        assert_eq!(br.read_cb(&cb).unwrap(), 42);
600        assert_eq!(br.read_cb(&cb).unwrap(), -42);
601        let ret = br.read_cb(&cb);
602        if let Err(e) = ret {
603            assert_matches::assert_matches!(e, InvalidCode);
604        } else {
605            assert_eq!(0, 1);
606        }
607    }
608
609    #[test]
610    fn test_short_codebook_msb() {
611        let scb_desc: Vec<ShortCodebookDesc> = vec![
612            ShortCodebookDesc { code: 0b0, bits: 1 },
613            ShortCodebookDesc { code: 0, bits: 0 },
614            ShortCodebookDesc {
615                code: 0b10,
616                bits: 2,
617            },
618            ShortCodebookDesc { code: 0, bits: 0 },
619            ShortCodebookDesc { code: 0, bits: 0 },
620            ShortCodebookDesc {
621                code: 0b110,
622                bits: 3,
623            },
624            ShortCodebookDesc { code: 0, bits: 0 },
625            ShortCodebookDesc {
626                code: 0b11100,
627                bits: 5,
628            },
629            ShortCodebookDesc {
630                code: 0b11101,
631                bits: 5,
632            },
633            ShortCodebookDesc {
634                code: 0b1111010,
635                bits: 7,
636            },
637            ShortCodebookDesc {
638                code: 0b1111011,
639                bits: 7,
640            },
641            ShortCodebookDesc {
642                code: 0b1111110,
643                bits: 7,
644            },
645            ShortCodebookDesc {
646                code: 0b11111111,
647                bits: 8,
648            },
649        ];
650        let buf = &BITS;
651        let mut br2 = BitReadBE::new(buf);
652        let cb = Codebook::new(&scb_desc, CodebookMode::MSB).unwrap();
653        assert_eq!(br2.read_cb(&cb).unwrap(), 0);
654        assert_eq!(br2.read_cb(&cb).unwrap(), 2);
655        assert_eq!(br2.read_cb(&cb).unwrap(), 5);
656        assert_eq!(br2.read_cb(&cb).unwrap(), 8);
657
658        assert_eq!(
659            reverse_bits(0b0000_0101_1011_1011_1101_1111_0111_1111),
660            0b1111_1110_1111_1011_1101_1101_1010_0000
661        );
662    }
663
664    #[test]
665    fn test_short_codebook_lsb() {
666        const BITS_LE: [u8; 8] = [0b11101111, 0b01110010, 0b01, 0, 0, 0, 0, 0];
667        let buf = &BITS_LE;
668        let scble_desc: Vec<ShortCodebookDesc> = vec![
669            ShortCodebookDesc {
670                code: 0b00,
671                bits: 2,
672            },
673            ShortCodebookDesc { code: 0, bits: 0 },
674            ShortCodebookDesc {
675                code: 0b01,
676                bits: 2,
677            },
678            ShortCodebookDesc { code: 0, bits: 0 },
679            ShortCodebookDesc { code: 0, bits: 0 },
680            ShortCodebookDesc {
681                code: 0b011,
682                bits: 3,
683            },
684            ShortCodebookDesc { code: 0, bits: 0 },
685            ShortCodebookDesc {
686                code: 0b10111,
687                bits: 5,
688            },
689            ShortCodebookDesc {
690                code: 0b00111,
691                bits: 5,
692            },
693            ShortCodebookDesc {
694                code: 0b0101111,
695                bits: 7,
696            },
697            ShortCodebookDesc {
698                code: 0b0111111,
699                bits: 7,
700            },
701            ShortCodebookDesc {
702                code: 0b1011101111,
703                bits: 10,
704            },
705        ];
706        let mut brl = BitReadLE::new(buf);
707        let cb = Codebook::new(&scble_desc, CodebookMode::LSB).unwrap();
708        assert_eq!(brl.read_cb(&cb).unwrap(), 11);
709        assert_eq!(brl.read_cb(&cb).unwrap(), 0);
710        assert_eq!(brl.read_cb(&cb).unwrap(), 7);
711        assert_eq!(brl.read_cb(&cb).unwrap(), 0);
712    }
713}
av_bitstream/codebook.rs

av_bitstream/
codebook.rs
