1 // This was originally based on code from: https://github.com/nwin/lzw
2 // Copyright (c) 2015 nwin
3 // which is under both Apache 2.0 and MIT
4 
5 //! This module provides a bit writer
6 use std::io::{self, Write};
7 
8 #[cfg(target_pointer_width = "64")]
9 #[macro_use]
10 mod arch_dep {
11     /// The data type of the accumulator.
12     /// a 64-bit value allows us to store more before
13     /// each push to the vector, but is sub-optimal
14     /// on 32-bit platforms.
15     pub type AccType = u64;
16     pub const FLUSH_AT: u8 = 48;
17     /// Push pending bits to vector.
18     /// Using a macro here since an inline function.
19     /// didn't optimise properly.
20     macro_rules! push{
21         ($s:ident) => {
22             let len = $s.w.len();
23             $s.w.reserve(6);
24             // Optimization:
25             //
26             // This is basically what `Vec::extend_from_slice` does, but it didn't inline
27             // properly, so we do it manually for now.
28             //
29             // # Unsafe
30             // We reserve enough space right before this, so setting the len manually and doing
31             // unchecked indexing is safe here since we only, and always write to all of the the
32             // uninitialized bytes of the vector.
33             unsafe {
34                 $s.w.set_len(len + 6);
35                 $s.w.get_unchecked_mut(len..).copy_from_slice(&[$s.acc as u8,
36                                                                 ($s.acc >> 8) as u8,
37                                                                 ($s.acc >> 16) as u8,
38                                                                 ($s.acc >> 24) as u8,
39                                                                 ($s.acc >> 32) as u8,
40                                                                 ($s.acc >> 40) as u8
41                 ][..]);
42             }
43 
44         };
45     }
46 }
47 #[cfg(not(target_pointer_width = "64"))]
48 #[macro_use]
49 mod arch_dep {
50     pub type AccType = u32;
51     pub const FLUSH_AT: u8 = 16;
52     macro_rules! push{
53         ($s:ident) => {
54             // Unlike the 64-bit case, using copy_from_slice seemed to worsen performance here.
55             // TODO: Needs benching on a 32-bit system to see what works best.
56             $s.w.push($s.acc as u8);
57             $s.w.push(($s.acc >> 8) as u8);
58         };
59     }
60 }
61 
62 use self::arch_dep::*;
63 
64 /// Writes bits to a byte stream, LSB first.
65 pub struct LsbWriter {
66     // Public for now so it can be replaced after initialization.
67     pub w: Vec<u8>,
68     bits: u8,
69     acc: AccType,
70 }
71 
72 impl LsbWriter {
73     /// Creates a new bit reader
new(writer: Vec<u8>) -> LsbWriter74     pub fn new(writer: Vec<u8>) -> LsbWriter {
75         LsbWriter {
76             w: writer,
77             bits: 0,
78             acc: 0,
79         }
80     }
81 
pending_bits(&self) -> u882     pub fn pending_bits(&self) -> u8 {
83         self.bits
84     }
85 
86     /// Buffer n number of bits, and write them to the vec if there are enough pending bits.
write_bits(&mut self, v: u16, n: u8)87     pub fn write_bits(&mut self, v: u16, n: u8) {
88         // NOTE: This outputs garbage data if n is 0, but v is not 0
89         self.acc |= (v as AccType) << self.bits;
90         self.bits += n;
91         // Waiting until we have FLUSH_AT bits and pushing them all in one batch.
92         while self.bits >= FLUSH_AT {
93             push!(self);
94             self.acc >>= FLUSH_AT;
95             self.bits -= FLUSH_AT;
96         }
97     }
98 
write_bits_finish(&mut self, v: u16, n: u8)99     fn write_bits_finish(&mut self, v: u16, n: u8) {
100         // NOTE: This outputs garbage data if n is 0, but v is not 0
101         self.acc |= (v as AccType) << self.bits;
102         self.bits += n % 8;
103         while self.bits >= 8 {
104             self.w.push(self.acc as u8);
105             self.acc >>= 8;
106             self.bits -= 8;
107         }
108     }
109 
flush_raw(&mut self)110     pub fn flush_raw(&mut self) {
111         let missing = FLUSH_AT - self.bits;
112         // Have to test for self.bits > 0 here,
113         // otherwise flush would output an extra byte when flush was called at a byte boundary
114         if missing > 0 && self.bits > 0 {
115             self.write_bits_finish(0, missing);
116         }
117     }
118 }
119 
120 impl Write for LsbWriter {
write(&mut self, buf: &[u8]) -> io::Result<usize>121     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
122         if self.acc == 0 {
123             self.w.extend_from_slice(buf)
124         } else {
125             for &byte in buf.iter() {
126                 self.write_bits(byte as u16, 8)
127             }
128         }
129         Ok(buf.len())
130     }
131 
flush(&mut self) -> io::Result<()>132     fn flush(&mut self) -> io::Result<()> {
133         self.flush_raw();
134         Ok(())
135     }
136 }
137 
138 #[cfg(test)]
139 mod test {
140     use super::LsbWriter;
141 
142     #[test]
write_bits()143     fn write_bits() {
144         let input = [
145             (3, 3),
146             (10, 8),
147             (88, 7),
148             (0, 2),
149             (0, 5),
150             (0, 0),
151             (238, 8),
152             (126, 8),
153             (161, 8),
154             (10, 8),
155             (238, 8),
156             (174, 8),
157             (126, 8),
158             (174, 8),
159             (65, 8),
160             (142, 8),
161             (62, 8),
162             (10, 8),
163             (1, 8),
164             (161, 8),
165             (78, 8),
166             (62, 8),
167             (158, 8),
168             (206, 8),
169             (10, 8),
170             (64, 7),
171             (0, 0),
172             (24, 5),
173             (0, 0),
174             (174, 8),
175             (126, 8),
176             (193, 8),
177             (174, 8),
178         ];
179         let expected = [
180             83,
181             192,
182             2,
183             220,
184             253,
185             66,
186             21,
187             220,
188             93,
189             253,
190             92,
191             131,
192             28,
193             125,
194             20,
195             2,
196             66,
197             157,
198             124,
199             60,
200             157,
201             21,
202             128,
203             216,
204             213,
205             47,
206             216,
207             21,
208         ];
209         let mut writer = LsbWriter::new(Vec::new());
210         for v in input.iter() {
211             writer.write_bits(v.0, v.1);
212         }
213         writer.flush_raw();
214         assert_eq!(writer.w, expected);
215     }
216 }
217 
218 
219 #[cfg(all(test, feature = "benchmarks"))]
220 mod bench {
221     use test_std::Bencher;
222     use super::LsbWriter;
223     #[bench]
bit_writer(b: &mut Bencher)224     fn bit_writer(b: &mut Bencher) {
225         let input = [
226             (3, 3),
227             (10, 8),
228             (88, 7),
229             (0, 2),
230             (0, 5),
231             (0, 0),
232             (238, 8),
233             (126, 8),
234             (161, 8),
235             (10, 8),
236             (238, 8),
237             (174, 8),
238             (126, 8),
239             (174, 8),
240             (65, 8),
241             (142, 8),
242             (62, 8),
243             (10, 8),
244             (1, 8),
245             (161, 8),
246             (78, 8),
247             (62, 8),
248             (158, 8),
249             (206, 8),
250             (10, 8),
251             (64, 7),
252             (0, 0),
253             (24, 5),
254             (0, 0),
255             (174, 8),
256             (126, 8),
257             (193, 8),
258             (174, 8),
259         ];
260         let mut writer = LsbWriter::new(Vec::with_capacity(100));
261         b.iter(|| for v in input.iter() {
262             let _ = writer.write_bits(v.0, v.1);
263         });
264     }
265 }
266