av_bitstream/
bitread.rs

1//! Bitstream reader functionality.
2
3use crate::byteread::*;
4
5/// Used to interact with a sequence of 64 bits, taking into account the
6/// relative endianness.
7///
8/// This trait also offers some methods to read and remove bits from an internal
9/// 64-bit cache.
10pub trait BitReadEndian {
11    /// Peeks n bits from the cache.
12    fn peek_val(&mut self, n: usize) -> u64;
13    /// Merges two sequences of bits together.
14    fn merge_val(msp: u64, lsp: u64, msb: usize, lsb: usize) -> u64;
15    /// Builds a new cache.
16    fn build_cache(cache: u64, refill: u64, cache_size: usize) -> u64;
17    /// Removes n bits from the cache.
18    fn skip_rem(&mut self, n: usize);
19}
20
21/// Used to extract a sequence of bits from an internal buffer.
22pub trait BitReadFill {
23    /// Tells if it is still possible to read bits from an internal buffer.
24    fn can_refill(&self) -> bool;
25    /// Gets a 32-bits sequence from an internal buffer.
26    fn fill32(&self) -> u64;
27    /// Gets a 64-bits sequence from an internal buffer.
28    fn fill64(&self) -> u64;
29}
30
31/// Used to interact with an internal buffer.
32pub trait BitReadInternal: BitReadEndian + BitReadFill {
33    /// Gets the number of bits left in an internal buffer.
34    fn left(&self) -> usize;
35    /// Extracts a 32-bit sequence from an internal buffer and saves it
36    /// within an internal cache.
37    fn refill32(&mut self);
38    /// Extracts a 64-bit sequence from an internal buffer and saves it
39    /// within an internal cache.
40    fn refill64(&mut self);
41
42    /// Returns n bits from an internal buffer.
43    fn get_val(&mut self, n: usize) -> u64 {
44        let ret = self.peek_val(n);
45
46        self.skip_rem(n);
47
48        ret
49    }
50}
51
52/// Used to define a bitreader.
53pub trait BitRead<'a>: BitReadInternal + Copy {
54    /// Creates a new bitreader with an internal buffer associated to it.
55    fn new(_: &'a [u8]) -> Self;
56    /// Tells the number of bits read from the internal buffer.
57    fn consumed(&self) -> usize;
58    /// Tells the number of bits still available in the internal buffer.
59    fn available(&self) -> usize;
60
61    /// Discard a certain number of bits from the internal buffer.
62    fn skip_bits(&mut self, size: usize);
63
64    /// Returns a single bit from the internal buffer.
65    #[inline]
66    fn get_bit(&mut self) -> bool {
67        if self.left() == 0 {
68            self.refill64();
69        }
70
71        self.get_val(1) != 0
72    }
73
74    /// Returns n bits from the internal buffer as a 64-bit sequence.
75    #[inline]
76    fn get_bits_64(&mut self, mut n: usize) -> u64 {
77        let mut left = 0;
78        let mut ret = 0;
79
80        if n == 0 {
81            return 0;
82        }
83
84        if self.left() < n {
85            n -= self.left();
86            left = self.left();
87            ret = self.get_val(left);
88            self.refill64();
89        }
90
91        Self::merge_val(self.get_val(n), ret, left, n)
92    }
93
94    /// Returns n bits from the internal buffer as a 32-bit sequence.
95    #[inline]
96    fn get_bits_32(&mut self, n: usize) -> u32 {
97        if n == 0 {
98            return 0;
99        }
100
101        if self.left() <= n {
102            self.refill32();
103        }
104
105        self.get_val(n) as u32
106    }
107
108    /// Peeks the next bit present in the internal buffer.
109    #[inline]
110    fn peek_bit(&mut self) -> bool {
111        let mut tmp = *self;
112
113        tmp.get_bit()
114    }
115
116    /// Peeks the next 32-bit sequence present in the internal buffer.
117    #[inline]
118    fn peek_bits_32(&mut self, n: usize) -> u32 {
119        let mut tmp = *self;
120
121        tmp.get_bits_32(n)
122    }
123
124    /// Peeks the next 64-bit sequence present in the internal buffer.
125    #[inline]
126    fn peek_bits_64(&self, n: usize) -> u64 {
127        let mut tmp = *self;
128
129        tmp.get_bits_64(n)
130    }
131
132    /// Aligns the bits present in the internal buffer.
133    #[inline]
134    fn align_bits(&mut self) {
135        let left = self.left() & 7;
136
137        self.skip_bits(left);
138    }
139}
140
141#[doc(hidden)]
142#[macro_export]
143macro_rules! endian_reader {
144    {$name: ident, $docname: expr} => {
145        #[doc = "A "]
146        #[doc = $docname]
147        #[doc = " reader."]
148        #[doc = ""]
149        #[doc = "It contains the data structures necessary to create a "]
150        #[doc = $docname]
151        #[doc = " reader."]
152        #[derive(Debug, Clone, Copy)]
153        #[allow(clippy::upper_case_acronyms)]
154        pub struct $name<'a> {
155            buffer : &'a[u8], /// read buffer, 8-bytes padded
156            index : usize,
157            cache : u64,
158            left : usize,
159        }
160        impl <'a> BitReadInternal for $name<'a> {
161            #[inline]
162            fn left(&self) -> usize {
163                self.left
164            }
165            #[inline]
166            fn refill32(&mut self) -> () {
167                if !self.can_refill() {
168                    return;
169                }
170                let val = self.fill32();
171
172                self.cache  = Self::build_cache(self.cache, val, self.left);
173
174                self.index += 4;
175                self.left  += 32;
176            }
177            #[inline]
178            fn refill64(&mut self) -> () {
179                if !self.can_refill() {
180                    return;
181                }
182
183                self.cache  = self.fill64();
184                self.index += 8;
185                self.left   = 64;
186            }
187        }
188
189        impl <'a> BitRead<'a> for $name<'a> {
190            fn new(buffer: &'a[u8]) -> $name<'a> {
191                let mut reader = $name {
192                    buffer,
193                    index: 0,
194                    cache: 0,
195                    left: 0
196                };
197
198                reader.refill64();
199                return reader;
200            }
201            #[inline]
202            fn consumed(&self) -> usize {
203                self.index * 8 - self.left
204            }
205
206            #[inline]
207            fn available(&self) -> usize {
208                (self.buffer.len() - self.index) * 8 + self.left
209            }
210
211            #[inline]
212            fn skip_bits(&mut self, mut n:usize) -> () {
213                if self.left < n {
214                    n -= self.left;
215                    if n > 64 {
216                        let skip = n / 8;
217
218                        n -= skip * 8;
219                        self.index += skip;
220                    }
221                    self.skip_rem(n);
222                    self.refill64();
223                }
224
225                self.skip_rem(n);
226            }
227
228        }
229    }
230}
231
232#[doc(hidden)]
233#[macro_export]
234macro_rules! little_endian_reader {
235    {$name: ident} => {
236        endian_reader!{ $name, "little-endian" }
237
238        impl <'a> BitReadEndian for $name<'a> {
239            #[inline]
240            fn peek_val(&mut self, n:usize) -> u64 {
241                let v = self.cache & ((1u64.checked_shl(n as u32).unwrap_or(0)).overflowing_sub(1).0);
242
243                return v;
244            }
245            #[inline]
246            fn skip_rem(&mut self, n:usize) -> () {
247                self.cache = self.cache.checked_shr(n as u32).unwrap_or(0);
248                self.left = self.left.saturating_sub(n);
249            }
250            #[inline]
251            fn merge_val(msp:u64, lsp:u64, msb:usize, _:usize) -> u64 {
252                msp << msb | lsp
253            }
254            #[inline]
255            fn build_cache(cache:u64, refill:u64, cache_size:usize) -> u64 {
256                cache | refill << cache_size
257            }
258        }
259    }
260}
261
262little_endian_reader! { BitReadLE }
263
264impl<'a> BitReadFill for BitReadLE<'a> {
265    #[inline]
266    fn can_refill(&self) -> bool {
267        self.index + 8 <= self.buffer.len()
268    }
269    #[inline(always)]
270    fn fill32(&self) -> u64 {
271        u64::from(get_u32l(&self.buffer[self.index..]))
272    }
273    #[inline(always)]
274    fn fill64(&self) -> u64 {
275        get_u64l(&self.buffer[self.index..])
276    }
277}
278
279#[doc(hidden)]
280#[macro_export]
281macro_rules! big_endian_reader {
282    {$name: ident} => {
283        endian_reader!{ $name, "big-endian" }
284
285        impl <'a> BitReadEndian for $name<'a> {
286            #[inline]
287            fn peek_val(&mut self, n:usize) -> u64 {
288                self.cache.checked_shr(64 - n as u32).unwrap_or(0)
289            }
290            #[inline]
291            fn skip_rem(&mut self, n:usize) -> () {
292                self.cache = self.cache.checked_shl(n as u32).unwrap_or(0);
293                self.left = self.left.saturating_sub(n);
294            }
295            #[inline]
296            fn merge_val(msp:u64, lsp:u64, _:usize, lsb:usize) -> u64 {
297                msp | lsp.checked_shl(lsb as u32).unwrap_or(0)
298            }
299            #[inline]
300            fn build_cache(cache:u64, refill:u64, cache_size:usize) -> u64 {
301                cache | refill << (32 - cache_size)
302            }
303        }
304    }
305}
306
307big_endian_reader! { BitReadBE }
308
309impl<'a> BitReadFill for BitReadBE<'a> {
310    #[inline]
311    fn can_refill(&self) -> bool {
312        self.index + 8 <= self.buffer.len()
313    }
314    #[inline(always)]
315    fn fill32(&self) -> u64 {
316        u64::from(get_u32b(&self.buffer[self.index..]))
317    }
318    #[inline(always)]
319    fn fill64(&self) -> u64 {
320        get_u64b(&self.buffer[self.index..])
321    }
322}
323
324#[cfg(test)]
325mod test {
326    pub const CHECKBOARD0101: [u8; 128] = [0b01010101; 128];
327    pub const CHECKBOARD0011: [u8; 128] = [0b00110011; 128];
328
329    mod le {
330        use super::super::*;
331        use super::*;
332
333        #[test]
334        fn get_bit() {
335            let b = &CHECKBOARD0101;
336            let mut reader = BitReadLE::new(b);
337
338            assert!(reader.get_bit());
339            assert!(!reader.get_bit());
340        }
341
342        #[test]
343        fn get_bits_64() {
344            let b = &CHECKBOARD0101;
345            let mut reader = BitReadLE::new(b);
346
347            assert_eq!(reader.get_bits_64(1), 1);
348            assert_eq!(reader.get_bits_64(2), 2);
349            assert_eq!(reader.get_bits_64(4), 10);
350            assert_eq!(reader.get_bits_64(1), 0);
351            assert_eq!(reader.get_bits_64(8), 85);
352            assert_eq!(reader.get_bits_64(32), 1431655765);
353            assert_eq!(reader.get_bits_64(64), 6148914691236517205);
354        }
355
356        #[test]
357        fn get_bits_64_double() {
358            let b = &CHECKBOARD0101;
359
360            let mut reader = BitReadLE::new(b);
361
362            assert_eq!(reader.get_bits_64(64), 6148914691236517205);
363            assert_eq!(reader.get_bits_64(64), 6148914691236517205);
364
365            let mut reader = BitReadLE::new(b);
366
367            assert_eq!(reader.get_bits_64(32), 1431655765);
368            assert_eq!(reader.get_bits_64(32), 1431655765);
369            assert_eq!(reader.get_bits_64(32), 1431655765);
370            assert_eq!(reader.get_bits_64(32), 1431655765);
371        }
372
373        #[test]
374        fn peek_bits_64() {
375            let reader = BitReadLE {
376                buffer: &CHECKBOARD0101,
377                index: 0,
378                cache: 0,
379                left: 0,
380            };
381
382            assert_eq!(reader.peek_bits_64(1), 1);
383            assert_eq!(reader.peek_bits_64(1), 1);
384            assert_eq!(reader.peek_bits_64(2), 1);
385            assert_eq!(reader.peek_bits_64(2), 1);
386            assert_eq!(reader.peek_bits_64(32), 1431655765);
387            assert_eq!(reader.peek_bits_64(64), 6148914691236517205);
388        }
389
390        #[test]
391        fn get_bits_32() {
392            let mut reader = BitReadLE {
393                buffer: &CHECKBOARD0101,
394                index: 0,
395                cache: 0,
396                left: 0,
397            };
398
399            assert!(reader.get_bits_32(1) == 1);
400            assert!(reader.get_bits_32(2) == 2);
401            assert!(reader.get_bits_32(4) == 10);
402            assert!(reader.get_bits_32(1) == 0);
403            assert!(reader.get_bits_32(8) == 85);
404
405            assert_eq!(reader.get_bits_32(1), 1);
406            for _ in 0..8 {
407                assert_eq!(reader.get_bits_32(8), 170);
408            }
409        }
410
411        #[test]
412        fn peek_bits_32() {
413            let b = &CHECKBOARD0101;
414            let mut reader = BitReadLE::new(b);
415
416            assert!(reader.peek_bits_32(1) == 1);
417            assert!(reader.peek_bits_32(1) == 1);
418            assert!(reader.peek_bits_32(2) == 1);
419            assert!(reader.peek_bits_32(2) == 1);
420        }
421
422        #[test]
423        fn skip_bits() {
424            let b = &CHECKBOARD0101;
425            let mut reader = BitReadLE::new(b);
426
427            reader.skip_bits(0);
428            assert!(reader.peek_bits_32(1) == 1);
429            reader.skip_bits(2);
430            assert!(reader.peek_bits_32(1) == 1);
431            reader.skip_bits(2);
432            assert!(reader.peek_bits_32(1) == 1);
433        }
434
435        #[test]
436        fn align_bits() {
437            let b = &CHECKBOARD0011;
438            let mut reader = BitReadLE::new(b);
439
440            reader.align_bits();
441            assert!(reader.get_bits_64(3) == 3);
442            reader.align_bits();
443            assert!(reader.get_bits_64(4) == 3);
444            reader.skip_bits(1);
445            reader.align_bits();
446            assert!(reader.get_bits_64(4) == 3);
447        }
448
449        #[test]
450        fn overread() {
451            let b = &CHECKBOARD0011;
452            let mut reader = BitReadLE::new(b);
453
454            reader.skip_bits(128 * 8 + 2);
455            reader.get_bits_64(6);
456        }
457    }
458    mod be {
459        use super::super::*;
460        use super::*;
461
462        #[test]
463        fn get_bit() {
464            let b = &CHECKBOARD0101;
465            let mut reader = BitReadBE::new(b);
466
467            assert!(!reader.get_bit());
468            assert!(reader.get_bit());
469        }
470
471        #[test]
472        fn get_bits_64() {
473            let b = &CHECKBOARD0101;
474            let mut reader = BitReadBE::new(b);
475
476            assert_eq!(reader.get_bits_64(1), 0);
477            assert_eq!(reader.get_bits_64(2), 2);
478            assert_eq!(reader.get_bits_64(4), 10);
479            assert_eq!(reader.get_bits_64(1), 1);
480            assert_eq!(reader.get_bits_64(8), 85);
481            assert_eq!(reader.get_bits_64(32), 1431655765);
482            assert_eq!(reader.get_bits_64(64), 6148914691236517205);
483        }
484
485        #[test]
486        fn get_bits_64_double() {
487            let b = &CHECKBOARD0101;
488
489            let mut reader = BitReadBE::new(b);
490
491            assert_eq!(reader.get_bits_64(64), 6148914691236517205);
492            assert_eq!(reader.get_bits_64(64), 6148914691236517205);
493
494            let mut reader = BitReadBE::new(b);
495
496            assert_eq!(reader.get_bits_64(32), 1431655765);
497            assert_eq!(reader.get_bits_64(32), 1431655765);
498            assert_eq!(reader.get_bits_64(32), 1431655765);
499            assert_eq!(reader.get_bits_64(32), 1431655765);
500        }
501
502        #[test]
503        fn peek_bits_64() {
504            let b = &CHECKBOARD0101;
505            let reader = BitReadBE::new(b);
506
507            assert_eq!(reader.peek_bits_64(1), 0);
508            assert_eq!(reader.peek_bits_64(1), 0);
509            assert_eq!(reader.peek_bits_64(2), 1);
510            assert_eq!(reader.peek_bits_64(2), 1);
511            assert_eq!(reader.peek_bits_64(32), 1431655765);
512            assert_eq!(reader.peek_bits_64(64), 6148914691236517205);
513        }
514
515        #[test]
516        fn get_bits_32() {
517            let b = &CHECKBOARD0101;
518            let mut reader = BitReadBE::new(b);
519
520            assert!(reader.get_bits_32(1) == 0);
521            assert!(reader.get_bits_32(2) == 2);
522            assert!(reader.get_bits_32(4) == 10);
523            assert!(reader.get_bits_32(1) == 1);
524            assert!(reader.get_bits_32(8) == 85);
525
526            assert_eq!(reader.get_bits_32(1), 0);
527            for _ in 0..8 {
528                assert_eq!(reader.get_bits_32(8), 170);
529            }
530        }
531
532        #[test]
533        fn peek_bits_32() {
534            let b = &CHECKBOARD0101;
535            let mut reader = BitReadBE::new(b);
536
537            assert!(reader.peek_bits_32(1) == 0);
538            assert!(reader.peek_bits_32(1) == 0);
539            assert!(reader.peek_bits_32(2) == 1);
540            assert!(reader.peek_bits_32(2) == 1);
541        }
542
543        #[test]
544        fn skip_bits() {
545            let b = &CHECKBOARD0101;
546            let mut reader = BitReadBE::new(b);
547
548            reader.skip_bits(0);
549            assert!(reader.peek_bits_32(1) == 0);
550            reader.skip_bits(2);
551            assert!(reader.peek_bits_32(1) == 0);
552            reader.skip_bits(2);
553            assert!(reader.peek_bits_32(1) == 0);
554        }
555
556        #[test]
557        fn align_bits() {
558            let b = &CHECKBOARD0011;
559            let mut reader = BitReadBE::new(b);
560
561            reader.align_bits();
562            assert!(reader.get_bits_64(3) == 1);
563            reader.align_bits();
564            assert!(reader.get_bits_64(4) == 3);
565            reader.skip_bits(1);
566            reader.align_bits();
567            assert!(reader.get_bits_64(4) == 3);
568        }
569
570        #[test]
571        fn overread() {
572            let b = &CHECKBOARD0011;
573            let mut reader = BitReadBE::new(b);
574
575            reader.skip_bits(128 * 8 + 2);
576            reader.get_bits_64(6);
577        }
578    }
579}
av_bitstream/bitread.rs

av_bitstream/
bitread.rs
