]>
Commit | Line | Data |
---|---|---|
1 | use std::collections::HashMap; | |
2 | use std::cmp::{max, min}; | |
3 | use io::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 { nnz = nnz + 1; } | |
199 | maxbits = max(bits, maxbits); | |
200 | if bits > MAX_LUT_BITS { | |
201 | let code = cb.code(i); | |
202 | let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode); | |
203 | let cval = extract_esc_part(code, bits, MAX_LUT_BITS, mode); | |
204 | add_esc_code(&mut escape_list, ckey, cval, bits - MAX_LUT_BITS, symidx); | |
205 | } | |
206 | if bits > 0 { symidx = symidx + 1; } | |
207 | } | |
208 | if maxbits == 0 { return Err(CodebookError::InvalidCodebook); } | |
209 | ||
210 | if maxbits > MAX_LUT_BITS { maxbits = MAX_LUT_BITS; } | |
211 | ||
212 | let tab_len = 1 << maxbits; | |
213 | let mut table: Vec<u32> = Vec::with_capacity(tab_len); | |
214 | let mut syms: Vec<S> = Vec::with_capacity(nnz); | |
215 | if table.capacity() < tab_len { return Err(CodebookError::MemoryError); } | |
216 | if syms.capacity() < nnz { return Err(CodebookError::MemoryError); } | |
217 | table.resize(tab_len, TABLE_FILL_VALUE); | |
218 | ||
219 | let mut symidx: u32 = 0; | |
220 | for i in 0..cb.len() { | |
221 | let bits = cb.bits(i); | |
222 | let code = cb.code(i); | |
223 | if bits == 0 { continue; } | |
224 | if bits <= MAX_LUT_BITS { | |
225 | fill_lut(&mut table, mode, 0, code, bits, maxbits, symidx, false); | |
226 | } else { | |
227 | let ckey = extract_lut_part(code, bits, MAX_LUT_BITS, mode) as usize; | |
228 | if table[ckey] == TABLE_FILL_VALUE { | |
229 | let key = ckey as u32; | |
230 | if let Some(bucket) = escape_list.get_mut(&key) { | |
231 | let maxlen = min(bucket.maxlen, MAX_LUT_BITS); | |
232 | let new_off = resize_table(&mut table, maxlen)?; | |
233 | fill_lut(&mut table, mode, 0, key, maxlen, MAX_LUT_BITS, new_off, true); | |
234 | bucket.offset = new_off as usize; | |
235 | } | |
236 | } | |
237 | } | |
238 | symidx = symidx + 1; | |
239 | } | |
240 | ||
241 | for (_, bucket) in &escape_list { | |
242 | build_esc_lut(&mut table, mode, &bucket)?; | |
243 | } | |
244 | ||
245 | for i in 0..cb.len() { | |
246 | if cb.bits(i) > 0 { | |
247 | syms.push(cb.sym(i)); | |
248 | } | |
249 | } | |
250 | ||
251 | Ok(Codebook { table: table, syms: syms, lut_bits: maxbits }) | |
252 | } | |
253 | } | |
254 | ||
255 | impl<'a, S: Copy> CodebookReader<S> for BitReader<'a> { | |
256 | #[allow(unused_variables)] | |
257 | fn read_cb(&mut self, cb: &Codebook<S>) -> CodebookResult<S> { | |
258 | let mut esc = true; | |
259 | let mut idx = 0; | |
260 | let mut lut_bits = cb.lut_bits; | |
261 | while esc { | |
262 | let lut_idx = (self.peek(lut_bits) as usize) + (idx as usize); | |
263 | if cb.table[lut_idx] == TABLE_FILL_VALUE { return Err(CodebookError::InvalidCode); } | |
264 | let bits = cb.table[lut_idx] & 0x7F; | |
265 | esc = (cb.table[lut_idx] & 0x80) != 0; | |
266 | idx = (cb.table[lut_idx] >> 8) as usize; | |
267 | if (bits as isize) > self.left() { | |
268 | return Err(CodebookError::InvalidCode); | |
269 | } | |
270 | let skip_bits = if esc { lut_bits as u32 } else { bits }; | |
271 | if let Err(_) = self.skip(skip_bits as u32) {} | |
272 | lut_bits = bits as u8; | |
273 | } | |
274 | Ok(cb.syms[idx]) | |
275 | } | |
276 | } | |
277 | ||
278 | pub struct FullCodebookDescReader<S> { | |
279 | data: Vec<FullCodebookDesc<S>>, | |
280 | } | |
281 | ||
282 | impl<S> FullCodebookDescReader<S> { | |
283 | pub fn new(data: Vec<FullCodebookDesc<S>>) -> Self { | |
284 | FullCodebookDescReader { data: data } | |
285 | } | |
286 | } | |
287 | ||
288 | impl<S: Copy> CodebookDescReader<S> for FullCodebookDescReader<S> { | |
289 | fn bits(&mut self, idx: usize) -> u8 { self.data[idx].bits } | |
290 | fn code(&mut self, idx: usize) -> u32 { self.data[idx].code } | |
291 | fn sym (&mut self, idx: usize) -> S { self.data[idx].sym } | |
292 | fn len(&mut self) -> usize { self.data.len() } | |
293 | } | |
294 | ||
295 | pub struct ShortCodebookDescReader { | |
296 | data: Vec<ShortCodebookDesc>, | |
297 | } | |
298 | ||
299 | impl ShortCodebookDescReader { | |
300 | pub fn new(data: Vec<ShortCodebookDesc<>>) -> Self { | |
301 | ShortCodebookDescReader { data: data } | |
302 | } | |
303 | } | |
304 | ||
305 | impl CodebookDescReader<u32> for ShortCodebookDescReader { | |
306 | fn bits(&mut self, idx: usize) -> u8 { self.data[idx].bits } | |
307 | fn code(&mut self, idx: usize) -> u32 { self.data[idx].code } | |
308 | fn sym (&mut self, idx: usize) -> u32 { idx as u32 } | |
309 | fn len(&mut self) -> usize { self.data.len() } | |
310 | } | |
311 | ||
312 | #[cfg(test)] | |
313 | mod test { | |
314 | use super::*; | |
315 | use io::bitreader::*; | |
316 | ||
317 | #[test] | |
318 | fn test_cb() { | |
319 | const BITS: [u8; 2] = [0b01011011, 0b10111100]; | |
320 | let cb_desc: Vec<FullCodebookDesc<i8>> = vec!( | |
321 | FullCodebookDesc { code: 0b0, bits: 1, sym: 16 }, | |
322 | FullCodebookDesc { code: 0b10, bits: 2, sym: -3 }, | |
323 | FullCodebookDesc { code: 0b110, bits: 3, sym: 42 }, | |
324 | FullCodebookDesc { code: 0b1110, bits: 4, sym: -42 } | |
325 | ); | |
326 | let buf = &BITS; | |
327 | let mut br = BitReader::new(buf, buf.len(), BitReaderMode::BE); | |
328 | let mut cfr = FullCodebookDescReader::new(cb_desc); | |
329 | let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap(); | |
330 | assert_eq!(br.read_cb(&cb).unwrap(), 16); | |
331 | assert_eq!(br.read_cb(&cb).unwrap(), -3); | |
332 | assert_eq!(br.read_cb(&cb).unwrap(), 42); | |
333 | assert_eq!(br.read_cb(&cb).unwrap(), -42); | |
334 | let ret = br.read_cb(&cb); | |
335 | if let Err(e) = ret { | |
336 | assert_eq!(e as i32, CodebookError::InvalidCode as i32); | |
337 | } else { | |
338 | assert_eq!(0, 1); | |
339 | } | |
340 | ||
341 | let scb_desc: Vec<ShortCodebookDesc> = vec!( | |
342 | ShortCodebookDesc { code: 0b0, bits: 1 }, | |
343 | ShortCodebookDesc { code: 0, bits: 0 }, | |
344 | ShortCodebookDesc { code: 0b10, bits: 2 }, | |
345 | ShortCodebookDesc { code: 0, bits: 0 }, | |
346 | ShortCodebookDesc { code: 0, bits: 0 }, | |
347 | ShortCodebookDesc { code: 0b110, bits: 3 }, | |
348 | ShortCodebookDesc { code: 0, bits: 0 }, | |
349 | ShortCodebookDesc { code: 0b11100, bits: 5 }, | |
350 | ShortCodebookDesc { code: 0b11101, bits: 5 }, | |
351 | ShortCodebookDesc { code: 0b1111010, bits: 7 }, | |
352 | ShortCodebookDesc { code: 0b1111011, bits: 7 }, | |
353 | ShortCodebookDesc { code: 0b1111110, bits: 7 }, | |
354 | ShortCodebookDesc { code: 0b11111111, bits: 8 } | |
355 | ); | |
356 | let mut br2 = BitReader::new(buf, buf.len(), BitReaderMode::BE); | |
357 | let mut cfr = ShortCodebookDescReader::new(scb_desc); | |
358 | let cb = Codebook::new(&mut cfr, CodebookMode::MSB).unwrap(); | |
359 | assert_eq!(br2.read_cb(&cb).unwrap(), 0); | |
360 | assert_eq!(br2.read_cb(&cb).unwrap(), 2); | |
361 | assert_eq!(br2.read_cb(&cb).unwrap(), 5); | |
362 | assert_eq!(br2.read_cb(&cb).unwrap(), 8); | |
363 | ||
364 | assert_eq!(reverse_bits(0b0000_0101_1011_1011_1101_1111_0111_1111, 32), | |
365 | 0b1111_1110_1111_1011_1101_1101_1010_0000); | |
366 | ||
367 | const BITS_LE: [u8; 3] = [0b11101111, 0b01110010, 0b01]; | |
368 | let buf = &BITS_LE; | |
369 | let scble_desc: Vec<ShortCodebookDesc> = vec!( | |
370 | ShortCodebookDesc { code: 0b00, bits: 2 }, | |
371 | ShortCodebookDesc { code: 0, bits: 0 }, | |
372 | ShortCodebookDesc { code: 0b01, bits: 2 }, | |
373 | ShortCodebookDesc { code: 0, bits: 0 }, | |
374 | ShortCodebookDesc { code: 0, bits: 0 }, | |
375 | ShortCodebookDesc { code: 0b011, bits: 3 }, | |
376 | ShortCodebookDesc { code: 0, bits: 0 }, | |
377 | ShortCodebookDesc { code: 0b10111, bits: 5 }, | |
378 | ShortCodebookDesc { code: 0b00111, bits: 5 }, | |
379 | ShortCodebookDesc { code: 0b0101111, bits: 7 }, | |
380 | ShortCodebookDesc { code: 0b0111111, bits: 7 }, | |
381 | ShortCodebookDesc { code: 0b1011101111, bits: 10 } | |
382 | ); | |
383 | let mut brl = BitReader::new(buf, buf.len(), BitReaderMode::LE); | |
384 | let mut cfr = ShortCodebookDescReader::new(scble_desc); | |
385 | let cb = Codebook::new(&mut cfr, CodebookMode::LSB).unwrap(); | |
386 | assert_eq!(brl.read_cb(&cb).unwrap(), 11); | |
387 | assert_eq!(brl.read_cb(&cb).unwrap(), 0); | |
388 | assert_eq!(brl.read_cb(&cb).unwrap(), 7); | |
389 | assert_eq!(brl.read_cb(&cb).unwrap(), 0); | |
390 | } | |
391 | } |