bit_io/
reader.rs

1use std::io::{self, Read};
2
3const DEFAULT_BUFFER_SIZE: usize = 1023;
4
5// bit reader (MSB to LSB)
6// TODO: make it abstract over BitStorage
7pub struct Reader<R> {
8    buffer: Box<[u8]>,
9    position: usize,  // in bits
10    size: usize, // also in bits
11    reader: R
12}
13
14impl<R> Reader<R> where R: Read {
15    pub fn new(reader: R) -> Self {
16        Self::with_capacity(DEFAULT_BUFFER_SIZE, reader)
17    }
18
19    pub fn with_capacity(capacity: usize, reader: R) -> Self {
20        Reader {
21            buffer: vec![0; capacity + 1].into_boxed_slice(),
22            position: 0,
23            size: 0,
24            reader
25        }
26    }
27
28    pub fn read_bit(&mut self) -> io::Result<bool> {
29        if self.size == self.position {
30            self.fill_buffer()?;
31        }
32
33        let byte_index = self.position / 8;
34        let bit_index = 7 - self.position % 8;
35        let bit = self.buffer[byte_index] & (1 << bit_index) != 0;
36
37        println!("Reading at {} -> {}", self.position, bit);
38
39        self.position += 1;
40        Ok(bit)
41    }
42
43    // TODO: abstract over |bits|
44    pub fn read_bits(&mut self, bits: &mut u8, count: usize) -> io::Result<()> {
45        for i in 0..count {
46            let bit = self.read_bit()?;
47            if bit {
48                *bits |= 1 << (7 - i);
49            }
50        }
51
52        Ok(())
53    }
54
55    fn fill_buffer(&mut self) -> io::Result<()> {
56        
57        let n = self.reader.read(&mut self.buffer)?;
58        println!("Refill: {}", n);
59        if n == 0 {
60            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Unexpected EOF"));
61        }
62
63        self.size = n * 8;
64        self.position = 0;
65        Ok(())
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use std::io::Cursor;
72    use super::*;
73
74    #[test]
75    fn within_byte() {
76        let buffer = vec![0b10100101];
77        let cursor = Cursor::new(buffer);
78        let mut reader = Reader::new(cursor);
79        let mut bits = 0u8;
80
81        reader.read_bits(&mut bits, 4).unwrap();
82        assert_eq!(bits, 0b10100000);
83
84        bits = 0;
85        reader.read_bits(&mut bits, 4).unwrap();
86        assert_eq!(bits, 0b01010000);
87
88        let e = reader.read_bit();
89        assert!(e.is_err());
90        assert_eq!(e.unwrap_err().kind(), io::ErrorKind::UnexpectedEof);
91    }
92
93    #[test]
94    fn overlapping() {
95        let buffer = vec![0b10101010, 0b10100000];
96        let cursor = Cursor::new(buffer);
97        let mut reader = Reader::new(cursor);
98        let mut bits = 0u8;
99
100        reader.read_bits(&mut bits, 6).unwrap();
101        assert_eq!(bits, 0b10101000);
102        bits = 0;
103
104        reader.read_bits(&mut bits, 6).unwrap();
105        assert_eq!(bits, 0b10101000);
106        bits = 0;
107
108        reader.read_bits(&mut bits, 4).unwrap();        
109        assert_eq!(bits, 0);
110
111        let e = reader.read_bit();
112        assert!(e.is_err());
113        assert_eq!(e.unwrap_err().kind(), io::ErrorKind::UnexpectedEof);
114    }
115
116    #[test]
117    fn buffer_overflow() {
118        let buffer = vec![0b10101010, 0b10100000];
119        let cursor = Cursor::new(buffer);
120        let mut reader = Reader::with_capacity(0, cursor);
121        let mut bits = 0u8;
122
123        reader.read_bits(&mut bits, 6).unwrap();
124        assert_eq!(bits, 0b10101000);
125        bits = 0;
126
127        reader.read_bits(&mut bits, 6).unwrap();
128        assert_eq!(bits, 0b10101000);
129        bits = 0;
130
131        reader.read_bits(&mut bits, 4).unwrap();        
132        assert_eq!(bits, 0);
133
134        let e = reader.read_bit();
135        assert!(e.is_err());
136        assert_eq!(e.unwrap_err().kind(), io::ErrorKind::UnexpectedEof);
137    }
138}