1#![cfg_attr(all(coverage_nightly, test), feature(coverage_attribute))]
42
43use std::io;
44
45use scuffle_bytes_util::{BitReader, BitWriter};
46
47pub trait BitReaderExpGolombExt {
53 fn read_exp_golomb(&mut self) -> io::Result<u64>;
55
56 fn read_signed_exp_golomb(&mut self) -> io::Result<i64> {
58 let exp_glob = self.read_exp_golomb()?;
59
60 if exp_glob % 2 == 0 {
61 Ok(-((exp_glob / 2) as i64))
62 } else {
63 Ok((exp_glob / 2) as i64 + 1)
64 }
65 }
66}
67
68impl<R: io::Read> BitReaderExpGolombExt for BitReader<R> {
69 fn read_exp_golomb(&mut self) -> io::Result<u64> {
70 let mut leading_zeros = 0;
71 while !self.read_bit()? {
72 leading_zeros += 1;
73 }
74
75 let mut result = 1;
76 for _ in 0..leading_zeros {
77 result <<= 1;
78 result |= self.read_bit()? as u64;
79 }
80
81 Ok(result - 1)
82 }
83}
84
85pub trait BitWriterExpGolombExt {
91 fn write_exp_golomb(&mut self, input: u64) -> io::Result<()>;
93
94 fn write_signed_exp_golomb(&mut self, number: i64) -> io::Result<()> {
96 let number = if number <= 0 {
97 -number as u64 * 2
98 } else {
99 number as u64 * 2 - 1
100 };
101
102 self.write_exp_golomb(number)
103 }
104}
105
106impl<W: io::Write> BitWriterExpGolombExt for BitWriter<W> {
107 fn write_exp_golomb(&mut self, input: u64) -> io::Result<()> {
108 let mut number = input + 1;
109 let mut leading_zeros = 0;
110 while number > 1 {
111 number >>= 1;
112 leading_zeros += 1;
113 }
114
115 for _ in 0..leading_zeros {
116 self.write_bit(false)?;
117 }
118
119 self.write_bits(input + 1, leading_zeros + 1)?;
120
121 Ok(())
122 }
123}
124
125#[cfg(test)]
126#[cfg_attr(all(test, coverage_nightly), coverage(off))]
127mod tests {
128 use bytes::Buf;
129 use scuffle_bytes_util::{BitReader, BitWriter};
130
131 use crate::{BitReaderExpGolombExt, BitWriterExpGolombExt};
132
133 pub fn get_remaining_bits(reader: &BitReader<std::io::Cursor<Vec<u8>>>) -> usize {
134 let remaining = reader.get_ref().remaining();
135
136 if reader.is_aligned() {
137 remaining * 8
138 } else {
139 remaining * 8 + (8 - reader.bit_pos() as usize)
140 }
141 }
142
143 #[test]
144 fn test_exp_glob_decode() {
145 let mut bit_writer = BitWriter::<Vec<u8>>::default();
146
147 bit_writer.write_bits(0b1, 1).unwrap(); bit_writer.write_bits(0b010, 3).unwrap(); bit_writer.write_bits(0b011, 3).unwrap(); bit_writer.write_bits(0b00100, 5).unwrap(); bit_writer.write_bits(0b00101, 5).unwrap(); bit_writer.write_bits(0b00110, 5).unwrap(); bit_writer.write_bits(0b00111, 5).unwrap(); let data = bit_writer.finish().unwrap();
156
157 let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
158
159 let remaining_bits = get_remaining_bits(&bit_reader);
160
161 let result = bit_reader.read_exp_golomb().unwrap();
162 assert_eq!(result, 0);
163 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
164
165 let result = bit_reader.read_exp_golomb().unwrap();
166 assert_eq!(result, 1);
167 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
168
169 let result = bit_reader.read_exp_golomb().unwrap();
170 assert_eq!(result, 2);
171 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
172
173 let result = bit_reader.read_exp_golomb().unwrap();
174 assert_eq!(result, 3);
175 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
176
177 let result = bit_reader.read_exp_golomb().unwrap();
178 assert_eq!(result, 4);
179 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
180
181 let result = bit_reader.read_exp_golomb().unwrap();
182 assert_eq!(result, 5);
183 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
184
185 let result = bit_reader.read_exp_golomb().unwrap();
186 assert_eq!(result, 6);
187 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
188 }
189
190 #[test]
191 fn test_signed_exp_glob_decode() {
192 let mut bit_writer = BitWriter::<Vec<u8>>::default();
193
194 bit_writer.write_bits(0b1, 1).unwrap(); bit_writer.write_bits(0b010, 3).unwrap(); bit_writer.write_bits(0b011, 3).unwrap(); bit_writer.write_bits(0b00100, 5).unwrap(); bit_writer.write_bits(0b00101, 5).unwrap(); bit_writer.write_bits(0b00110, 5).unwrap(); bit_writer.write_bits(0b00111, 5).unwrap(); let data = bit_writer.finish().unwrap();
203
204 let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
205
206 let remaining_bits = get_remaining_bits(&bit_reader);
207
208 let result = bit_reader.read_signed_exp_golomb().unwrap();
209 assert_eq!(result, 0);
210 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
211
212 let result = bit_reader.read_signed_exp_golomb().unwrap();
213 assert_eq!(result, 1);
214 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
215
216 let result = bit_reader.read_signed_exp_golomb().unwrap();
217 assert_eq!(result, -1);
218 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
219
220 let result = bit_reader.read_signed_exp_golomb().unwrap();
221 assert_eq!(result, 2);
222 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
223
224 let result = bit_reader.read_signed_exp_golomb().unwrap();
225 assert_eq!(result, -2);
226 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
227
228 let result = bit_reader.read_signed_exp_golomb().unwrap();
229 assert_eq!(result, 3);
230 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
231
232 let result = bit_reader.read_signed_exp_golomb().unwrap();
233 assert_eq!(result, -3);
234 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
235 }
236
237 #[test]
238 fn test_exp_glob_encode() {
239 let mut bit_writer = BitWriter::<Vec<u8>>::default();
240
241 bit_writer.write_exp_golomb(0).unwrap();
242 bit_writer.write_exp_golomb(1).unwrap();
243 bit_writer.write_exp_golomb(2).unwrap();
244 bit_writer.write_exp_golomb(3).unwrap();
245 bit_writer.write_exp_golomb(4).unwrap();
246 bit_writer.write_exp_golomb(5).unwrap();
247 bit_writer.write_exp_golomb(6).unwrap();
248 bit_writer.write_exp_golomb(u64::MAX - 1).unwrap();
249
250 let data = bit_writer.finish().unwrap();
251
252 let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
253
254 let remaining_bits = get_remaining_bits(&bit_reader);
255
256 let result = bit_reader.read_exp_golomb().unwrap();
257 assert_eq!(result, 0);
258 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
259
260 let result = bit_reader.read_exp_golomb().unwrap();
261 assert_eq!(result, 1);
262 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
263
264 let result = bit_reader.read_exp_golomb().unwrap();
265 assert_eq!(result, 2);
266 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
267
268 let result = bit_reader.read_exp_golomb().unwrap();
269 assert_eq!(result, 3);
270 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
271
272 let result = bit_reader.read_exp_golomb().unwrap();
273 assert_eq!(result, 4);
274 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
275
276 let result = bit_reader.read_exp_golomb().unwrap();
277 assert_eq!(result, 5);
278 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
279
280 let result = bit_reader.read_exp_golomb().unwrap();
281 assert_eq!(result, 6);
282 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
283
284 let result = bit_reader.read_exp_golomb().unwrap();
285 assert_eq!(result, u64::MAX - 1);
286 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 154);
287 }
288
289 #[test]
290 fn test_signed_exp_glob_encode() {
291 let mut bit_writer = BitWriter::<Vec<u8>>::default();
292
293 bit_writer.write_signed_exp_golomb(0).unwrap();
294 bit_writer.write_signed_exp_golomb(1).unwrap();
295 bit_writer.write_signed_exp_golomb(-1).unwrap();
296 bit_writer.write_signed_exp_golomb(2).unwrap();
297 bit_writer.write_signed_exp_golomb(-2).unwrap();
298 bit_writer.write_signed_exp_golomb(3).unwrap();
299 bit_writer.write_signed_exp_golomb(-3).unwrap();
300 bit_writer.write_signed_exp_golomb(i64::MAX).unwrap();
301
302 let data = bit_writer.finish().unwrap();
303
304 let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
305
306 let remaining_bits = get_remaining_bits(&bit_reader);
307
308 let result = bit_reader.read_signed_exp_golomb().unwrap();
309 assert_eq!(result, 0);
310 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
311
312 let result = bit_reader.read_signed_exp_golomb().unwrap();
313 assert_eq!(result, 1);
314 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
315
316 let result = bit_reader.read_signed_exp_golomb().unwrap();
317 assert_eq!(result, -1);
318 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
319
320 let result = bit_reader.read_signed_exp_golomb().unwrap();
321 assert_eq!(result, 2);
322 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
323
324 let result = bit_reader.read_signed_exp_golomb().unwrap();
325 assert_eq!(result, -2);
326 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
327
328 let result = bit_reader.read_signed_exp_golomb().unwrap();
329 assert_eq!(result, 3);
330 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
331
332 let result = bit_reader.read_signed_exp_golomb().unwrap();
333 assert_eq!(result, -3);
334 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
335
336 let result = bit_reader.read_signed_exp_golomb().unwrap();
337 assert_eq!(result, i64::MAX);
338 assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 154);
339 }
340}