bit_io/
writer.rs

1use std::io::{self, Write};
2
3const DEFAULT_BUFFER_SIZE: usize = 1023;
4
5// bit writer (MSB to LSB)
6// TODO: make it abstract over BitStorage
7pub struct Writer<W> {
8    buffer: Box<[u8]>,
9    size: usize, // in bits
10    writer: W
11}
12
13impl<W> Writer<W> where W: Write {
14    pub fn new(writer: W) -> Self {
15        Self::with_capacity(DEFAULT_BUFFER_SIZE, writer)
16    }
17
18    pub fn with_capacity(capacity: usize, writer: W) -> Self {
19        Writer {
20            buffer: vec![0; capacity + 1].into_boxed_slice(),
21            size: 0,
22            writer
23        }
24    }
25
26    pub fn write_bit(&mut self, bit: bool) -> io::Result<()> {
27        let free_bits = self.buffer.len() * 8 - self.size;
28        if free_bits == 0 {
29            self.flush()?;
30        }
31        
32        if bit {
33            let byte_index = self.size / 8;
34            let bit_index = 7 - self.size % 8;
35            self.buffer[byte_index] |= 1 << bit_index;
36        }
37
38        self.size += 1; 
39        Ok(())
40    }
41
42    // write |count| bits of |bits| from MSB to LSB
43    // TODO: optimize
44    // TODO: abstract over |bits|
45    // TODO: use byteorder and just write in BE order into 
46    //       buffer when size_of::<Bits>() > 1 and buffer is aligned
47    pub fn write_bits(&mut self, bits: u8, count: usize) -> io::Result<()> {
48        for i in 0..count {
49            let bit = bits & (1 << (7 - i)) != 0;
50            self.write_bit(bit)?
51        }
52
53        Ok(())
54    }
55
56    // flush writer (byte-level)
57    // the remaining bits of last byte will be 0
58    pub fn flush(&mut self) -> io::Result<()> {
59        let bytes = (self.size + 7) / 8;
60        self.writer.write_all(&self.buffer[..bytes])?;
61        self.reset(bytes);
62        Ok(())
63    }
64
65    fn reset(&mut self, bytes: usize) {
66        // reset buffer
67        unsafe {
68            ::std::ptr::write_bytes(self.buffer.as_mut_ptr(), 0, bytes);
69        }
70        self.size = 0;
71    }
72
73    pub fn finish(mut self) -> io::Result<W> {
74        self.flush()?;
75        Ok(self.writer)
76    }
77}
78
79
80#[cfg(test)]
81mod tests {
82    use super::Writer;
83
84    #[test]
85    fn within_byte() {
86        let mut writer = Writer::new(Vec::new());
87        writer.write_bits(0b10100000, 4).unwrap();
88        writer.write_bits(0b01010000, 4).unwrap();
89        let buffer = writer.finish().unwrap();
90        assert_eq!(&buffer, &[0b10100101]);
91    }
92
93    #[test]
94    fn overlapping() {
95        let mut writer = Writer::new(Vec::new());
96        writer.write_bits(0b10101000, 6).unwrap();
97        writer.write_bits(0b10101000, 6).unwrap();
98        let buffer = writer.finish().unwrap();
99        assert_eq!(&buffer, &[0b10101010, 0b10100000]);
100    }
101
102    #[test]
103    fn buffer_overflow() {
104        let mut writer = Writer::with_capacity(0, Vec::new());
105        writer.write_bits(0b10101000, 6).unwrap();
106        writer.write_bits(0b10101000, 6).unwrap();
107        let buffer = writer.finish().unwrap();
108        assert_eq!(&buffer, &[0b10101010, 0b10100000]);
109    }
110}
bit_io/writer.rs

bit_io/
writer.rs
