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