split NihAV into subcrates
[nihav.git] / nihav-core / src / io / codebook.rs
CommitLineData
0ca65ffd
KS
1use std::collections::HashMap;
2use std::cmp::{max, min};
aca89041 3use super::bitreader::BitReader;
4667915a
KS
4
5#[derive(Debug)]
6pub enum CodebookError {
7 InvalidCodebook,
8 MemoryError,
9 InvalidCode,
10}
11
0ca65ffd
KS
12#[derive(Debug, Copy, Clone)]
13pub enum CodebookMode {
14 MSB,
15 LSB,
16}
17
4667915a
KS
18type CodebookResult<T> = Result<T, CodebookError>;
19
20pub struct FullCodebookDesc<S> {
21 code: u32,
22 bits: u8,
23 sym: S,
24}
25
26pub struct ShortCodebookDesc {
27 code: u32,
28 bits: u8,
29}
30
31pub 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)]
39pub struct Codebook<S> {
40 table: Vec<u32>,
41 syms: Vec<S>,
42 lut_bits: u8,
43}
44
45pub trait CodebookReader<S> {
46 fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S>;
47}
48
0ca65ffd
KS
49const TABLE_FILL_VALUE: u32 = 0x7F;
50const MAX_LUT_BITS: u8 = 10;
51
52fn 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
69fn 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
86fn 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
95fn 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
106fn 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
113fn 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)]
121struct Code {
122 code: u32,
123 bits: u8,
124 idx: usize,
125}
126
127struct CodeBucket {
128 maxlen: u8,
129 offset: usize,
130 codes: Vec<Code>,
131}
132
133impl 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
143type EscapeCodes = HashMap<u32, CodeBucket>;
144
145fn 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
153fn 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;
26af5ca8 161 if code.bits <= MAX_LUT_BITS {
0ca65ffd
KS
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 }
1a151e53 184
0ca65ffd
KS
185 Ok(())
186}
187
4667915a 188impl<S: Copy> Codebook<S> {
0ca65ffd
KS
189
190 pub fn new(cb: &mut CodebookDescReader<S>, mode: CodebookMode) -> CodebookResult<Self> {
4667915a
KS
191 let mut maxbits = 0;
192 let mut nnz = 0;
0ca65ffd
KS
193 let mut escape_list: EscapeCodes = HashMap::new();
194
195 let mut symidx: usize = 0;
4667915a
KS
196 for i in 0..cb.len() {
197 let bits = cb.bits(i);
26488721
KS
198 if bits > 0 {
199 nnz = nnz + 1;
200 if cb.code(i) >= (1 << bits) {
201 return Err(CodebookError::InvalidCodebook);
202 }
203 }
0ca65ffd
KS
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);
4667915a 210 }
0ca65ffd 211 if bits > 0 { symidx = symidx + 1; }
4667915a
KS
212 }
213 if maxbits == 0 { return Err(CodebookError::InvalidCodebook); }
214
0ca65ffd
KS
215 if maxbits > MAX_LUT_BITS { maxbits = MAX_LUT_BITS; }
216
4667915a 217 let tab_len = 1 << maxbits;
0ca65ffd
KS
218 let mut table: Vec<u32> = Vec::with_capacity(tab_len);
219 let mut syms: Vec<S> = Vec::with_capacity(nnz);
4667915a 220 if table.capacity() < tab_len { return Err(CodebookError::MemoryError); }
0ca65ffd
KS
221 if syms.capacity() < nnz { return Err(CodebookError::MemoryError); }
222 table.resize(tab_len, TABLE_FILL_VALUE);
4667915a
KS
223
224 let mut symidx: u32 = 0;
225 for i in 0..cb.len() {
226 let bits = cb.bits(i);
0ca65ffd 227 let code = cb.code(i);
4667915a 228 if bits == 0 { continue; }
0ca65ffd
KS
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 }
4667915a
KS
242 }
243 symidx = symidx + 1;
244 }
245
0ca65ffd
KS
246 for (_, bucket) in &escape_list {
247 build_esc_lut(&mut table, mode, &bucket)?;
248 }
249
4667915a
KS
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
260impl<'a, S: Copy> CodebookReader<S> for BitReader<'a> {
261 #[allow(unused_variables)]
262 fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S> {
0ca65ffd
KS
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;
4667915a 278 }
0ca65ffd 279 Ok(cb.syms[idx])
4667915a
KS
280 }
281}
282
283pub struct FullCodebookDescReader<S> {
284 data: Vec<FullCodebookDesc<S>>,
285}
286
287impl<S> FullCodebookDescReader<S> {
288 pub fn new(data: Vec<FullCodebookDesc<S>>) -> Self {
289 FullCodebookDescReader { data: data }
290 }
291}
292
293impl<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
300pub struct ShortCodebookDescReader {
301 data: Vec<ShortCodebookDesc>,
302}
303
304impl ShortCodebookDescReader {
305 pub fn new(data: Vec<ShortCodebookDesc<>>) -> Self {
306 ShortCodebookDescReader { data: data }
307 }
308}
309
310impl 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
6465a946
KS
317pub struct TableCodebookDescReader<CodeType:'static, IndexType:'static> {
318 bits: &'static [u8],
319 codes: &'static [CodeType],
320 idx_map: fn(usize) -> IndexType,
321}
322
323impl<'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}
328impl<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
4667915a
KS
336#[cfg(test)]
337mod test {
338 use super::*;
aca89041 339 use crate::io::bitreader::*;
4667915a
KS
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);
0ca65ffd 353 let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap();
4667915a
KS
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 },
0ca65ffd
KS
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 }
4667915a
KS
379 );
380 let mut br2 = BitReader::new(buf, buf.len(), BitReaderMode::BE);
381 let mut cfr = ShortCodebookDescReader::new(scb_desc);
0ca65ffd 382 let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap();
4667915a
KS
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);
0ca65ffd
KS
386 assert_eq!(br2.read_cb(&cb).unwrap(), 8);
387
06fd8c88 388 assert_eq!(reverse_bits(0b0000_0101_1011_1011_1101_1111_0111_1111, 32),
0ca65ffd
KS
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);
4667915a
KS
414 }
415}