Mercurial > touhou

comparison utils/src/bitstream.rs @ 757:21b186be2590

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