]>
Commit | Line | Data |
---|---|---|
1 | #[derive(Debug)] | |
2 | pub enum BitReaderMode { | |
3 | BE, | |
4 | LE16, | |
5 | LE32, | |
6 | } | |
7 | ||
8 | #[derive(Debug)] | |
9 | pub enum BitReaderError { | |
10 | BitstreamEnd, | |
11 | TooManyBitsRequested, | |
12 | } | |
13 | ||
14 | use self::BitReaderError::*; | |
15 | ||
16 | type BitReaderResult<T> = Result<T, BitReaderError>; | |
17 | ||
18 | #[derive(Debug)] | |
19 | pub struct BitReader<'a> { | |
20 | cache: u64, | |
21 | bits: u8, | |
22 | pos: usize, | |
23 | end: usize, | |
24 | src: &'a [u8], | |
25 | mode: BitReaderMode, | |
26 | } | |
27 | ||
28 | impl<'a> BitReader<'a> { | |
29 | ||
30 | pub fn new(src: &'a [u8], size: usize, mode: BitReaderMode) -> Self { | |
31 | if src.len() < size { panic!("size is less than needed"); } | |
32 | BitReader{ cache: 0, pos: 0, bits: 0, end: size, src: src, mode: mode } | |
33 | } | |
34 | ||
35 | pub fn tell(&self) -> usize { | |
36 | self.pos * 8 - (self.bits as usize) | |
37 | } | |
38 | ||
39 | pub fn left(&self) -> isize { | |
40 | ((self.end as isize) - (self.pos as isize)) * 8 + (self.bits as isize) | |
41 | } | |
42 | ||
43 | fn fill32be(&mut self, src: &[u8]) { | |
44 | let nw = (((src[0] as u32) << 24) | | |
45 | ((src[1] as u32) << 16) | | |
46 | ((src[2] as u32) << 8) | | |
47 | ((src[3] as u32) << 0)) as u64; | |
48 | self.cache |= nw << (32 - self.bits); | |
49 | } | |
50 | ||
51 | fn fill32le16(&mut self, src: &[u8], realbits: u8) { | |
52 | let mut nw = (((src[1] as u32) << 24) | | |
53 | ((src[0] as u32) << 16) | | |
54 | ((src[3] as u32) << 8) | | |
55 | ((src[2] as u32) << 0)) as u64; | |
56 | if realbits <= 16 { nw >>= 16; } | |
57 | self.cache |= nw << self.bits; | |
58 | } | |
59 | ||
60 | fn fill32le32(&mut self, src: &[u8]) { | |
61 | let nw = (((src[3] as u32) << 24) | | |
62 | ((src[2] as u32) << 16) | | |
63 | ((src[1] as u32) << 8) | | |
64 | ((src[0] as u32) << 0)) as u64; | |
65 | self.cache |= nw << self.bits; | |
66 | } | |
67 | ||
68 | fn refill(&mut self) -> BitReaderResult<()> { | |
69 | if self.pos >= self.end { return Err(BitstreamEnd) } | |
70 | while self.bits <= 32 { | |
71 | if self.pos + 4 <= self.end { | |
72 | let buf = &self.src[self.pos..]; | |
73 | match self.mode { | |
74 | BitReaderMode::BE => self.fill32be (buf), | |
75 | BitReaderMode::LE16 => self.fill32le16(buf, 32), | |
76 | BitReaderMode::LE32 => self.fill32le32(buf), | |
77 | } | |
78 | self.pos += 4; | |
79 | self.bits += 32; | |
80 | } else { | |
81 | let mut buf: [u8; 4] = [0, 0, 0, 0]; | |
82 | let mut newbits: u8 = 0; | |
83 | for i in 0..3 { | |
84 | if self.pos < self.end { | |
85 | buf[i] = self.src[self.pos]; | |
86 | self.pos = self.pos + 1; | |
87 | newbits += 8; | |
88 | } | |
89 | } | |
90 | if newbits == 0 { break; } | |
91 | match self.mode { | |
92 | BitReaderMode::BE => self.fill32be (&buf), | |
93 | BitReaderMode::LE16 => self.fill32le16(&buf, newbits), | |
94 | BitReaderMode::LE32 => self.fill32le32(&buf), | |
95 | } | |
96 | self.bits += newbits; | |
97 | } | |
98 | } | |
99 | Ok(()) | |
100 | } | |
101 | ||
102 | fn read_cache(&mut self, nbits: u8) -> u32 { | |
103 | let res = match self.mode { | |
104 | BitReaderMode::BE => (self.cache as u64) >> (64 - nbits), | |
105 | _ => ((1u64 << nbits) - 1) & self.cache, | |
106 | }; | |
107 | res as u32 | |
108 | } | |
109 | ||
110 | fn skip_cache(&mut self, nbits: u8) { | |
111 | match self.mode { | |
112 | BitReaderMode::BE => self.cache <<= nbits, | |
113 | _ => self.cache >>= nbits, | |
114 | }; | |
115 | self.bits -= nbits; | |
116 | } | |
117 | ||
118 | fn reset_cache(&mut self) { | |
119 | self.bits = 0; | |
120 | self.cache = 0; | |
121 | } | |
122 | ||
123 | pub fn read(&mut self, nbits: u8) -> BitReaderResult<u32> { | |
124 | if nbits > 32 { return Err(TooManyBitsRequested) } | |
125 | if self.bits < nbits { | |
126 | if let Err(err) = self.refill() { return Err(err) } | |
127 | if self.bits < nbits { return Err(BitstreamEnd) } | |
128 | } | |
129 | let res = self.read_cache(nbits); | |
130 | self.skip_cache(nbits); | |
131 | Ok(res as u32) | |
132 | } | |
133 | ||
134 | pub fn peek(&mut self, nbits: u8) -> u32 { | |
135 | if nbits > 32 { return 0 } | |
136 | if self.bits < nbits { let _ = self.refill(); } | |
137 | self.read_cache(nbits) | |
138 | } | |
139 | ||
140 | pub fn skip(&mut self, nbits: u32) -> BitReaderResult<()> { | |
141 | if self.bits as u32 >= nbits { | |
142 | self.skip_cache(nbits as u8); | |
143 | return Ok(()); | |
144 | } | |
145 | let mut skip_bits = nbits - (self.bits as u32); | |
146 | self.reset_cache(); | |
147 | self.pos += ((skip_bits / 32) * 4) as usize; | |
148 | skip_bits = skip_bits & 0x1F; | |
149 | self.refill()?; | |
150 | if skip_bits > 0 { | |
151 | self.skip_cache(skip_bits as u8); | |
152 | } | |
153 | Ok(()) | |
154 | } | |
155 | ||
156 | pub fn seek(&mut self, nbits: u32) -> BitReaderResult<()> { | |
157 | if ((nbits + 7) >> 3) as usize > self.end { return Err(TooManyBitsRequested); } | |
158 | self.reset_cache(); | |
159 | self.pos = ((nbits / 32) * 4) as usize; | |
160 | self.skip(nbits & 0x1F) | |
161 | } | |
162 | } | |
163 | ||
164 | #[cfg(test)] | |
165 | mod test { | |
166 | use super::*; | |
167 | ||
168 | #[test] | |
169 | fn br_works() { | |
170 | const DATA: [u8; 18] = [0b00011011; 18]; | |
171 | let src = &DATA; | |
172 | let mut br = BitReader::new(src, src.len(), BitReaderMode::LE16); | |
173 | ||
174 | for _ in 0..8 { | |
175 | assert_eq!(br.read(16).unwrap(), 0x1B1B); | |
176 | } | |
177 | } | |
178 | } |