637
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
1 //! Bitstream module.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
2
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
3 use std::io;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
4
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
5 /// Wrapper around any `Read` trait, to allow bit operations.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
6 pub struct BitStream<R: io::Read + io::Seek> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
7 io: R,
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
8 remaining_bits: usize,
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
9 byte: u8,
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
10 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
11
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
12 impl<R: io::Read + io::Seek> BitStream<R> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
13 /// Create a new bitstream.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
14 pub fn new(io: R) -> BitStream<R> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
15 BitStream {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
16 io,
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
17 remaining_bits: 0,
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
18 byte: 0,
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
19 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
20 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
21
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
22 /// Seek inside the bitstream, ditching any unused data read.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
23 pub fn seek(&mut self, seek_from: io::SeekFrom) -> io::Result<u64> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
24 self.remaining_bits = 0;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
25 self.byte = 0;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
26 self.io.seek(seek_from)
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
27 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
28
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
29 fn fill_byte(&mut self) -> io::Result<()> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
30 assert!(self.remaining_bits == 0);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
31
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
32 let mut buf = [0u8; 1];
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
33 self.io.read_exact(&mut buf)?;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
34 self.byte = buf[0];
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
35 self.remaining_bits = 8;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
36 Ok(())
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
37 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
38
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
39 /// Read only one bit from the stream.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
40 pub fn read_bit(&mut self) -> io::Result<bool> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
41 if self.remaining_bits == 0 {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
42 self.fill_byte()?;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
43 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
44 self.remaining_bits -= 1;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
45 Ok((self.byte >> self.remaining_bits) & 0x01 != 0)
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
46 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
47
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
48 /// Read `nb_bits` bits from the stream.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
49 pub fn read(&mut self, nb_bits: usize) -> io::Result<usize> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
50 let mut nb_bits2 = nb_bits;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
51 let mut value: usize = 0;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
52 while nb_bits2 > 0 {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
53 if self.remaining_bits == 0 {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
54 self.fill_byte()?;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
55 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
56 let read = if nb_bits2 > self.remaining_bits { self.remaining_bits } else { nb_bits2 };
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
57 nb_bits2 -= read;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
58 self.remaining_bits -= read;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
59 value |= (self.byte as usize >> self.remaining_bits) << nb_bits2;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
60 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
61 Ok(value & ((1 << nb_bits) - 1))
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
62 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
63
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
64 /// Read a given amount of bytes.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
65 pub fn read_bytes(&mut self, nb_bytes: usize) -> io::Result<Vec<u8>> {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
66 let mut buf = vec![0u8; nb_bytes];
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
67 self.io.read_exact(&mut buf)?;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
68 Ok(buf)
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
69 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
70 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
71
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
72 #[cfg(test)]
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
73 mod tests {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
74 use super::*;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
75 use crate::util::SeekableSlice;
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
76
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
77 #[test]
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
78 fn bit_by_bit() {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
79 let data = SeekableSlice::new(&[1, 2, 3]);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
80 let mut bitstream = BitStream::new(data);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
81
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
82 // 1
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
83 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
84 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
85 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
86 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
87 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
88 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
89 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
90 assert_eq!(bitstream.read_bit().unwrap(), true);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
91
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
92 // 2
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
93 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
94 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
95 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
96 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
97 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
98 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
99 assert_eq!(bitstream.read_bit().unwrap(), true);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
100 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
101
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
102 // 3
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
103 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
104 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
105 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
106 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
107 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
108 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
109 assert_eq!(bitstream.read_bit().unwrap(), true);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
110 assert_eq!(bitstream.read_bit().unwrap(), true);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
111
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
112 // Can’t read after the end.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
113 bitstream.read_bit().unwrap_err();
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
114 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
115
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
116 #[test]
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
117 fn byte_by_byte() {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
118 let data = SeekableSlice::new(&[1, 2, 3]);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
119 let mut bitstream = BitStream::new(data);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
120
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
121 assert_eq!(bitstream.read(8).unwrap(), 1);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
122 assert_eq!(bitstream.read(8).unwrap(), 2);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
123 assert_eq!(bitstream.read(8).unwrap(), 3);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
124
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
125 // Can’t read after the end.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
126 bitstream.read(1).unwrap_err();
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
127 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
128
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
129 #[test]
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
130 fn unaligned_bytes() {
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
131 let data = SeekableSlice::new(&[0, 129, 1, 128]);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
132 let mut bitstream = BitStream::new(data);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
133
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
134 assert_eq!(bitstream.read_bit().unwrap(), false);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
135 assert_eq!(bitstream.read(8).unwrap(), 1);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
136 assert_eq!(bitstream.read(8).unwrap(), 2);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
137 assert_eq!(bitstream.read(8).unwrap(), 3);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
138 assert_eq!(bitstream.read(7).unwrap(), 0);
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
139
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
140 // Can’t read after the end.
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
141 bitstream.read(1).unwrap_err();
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
142 }
|
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
diff
changeset
|
143 }
|