| // This was originally based on code from: https://github.com/nwin/lzw |
| // Copyright (c) 2015 nwin |
| // which is under both Apache 2.0 and MIT |
| |
| //! This module provides a bit writer |
| use std::io::{self, Write}; |
| |
| #[cfg(target_pointer_width = "64")] |
| #[macro_use] |
| mod arch_dep { |
| /// The data type of the accumulator. |
| /// a 64-bit value allows us to store more before |
| /// each push to the vector, but is sub-optimal |
| /// on 32-bit platforms. |
| pub type AccType = u64; |
| pub const FLUSH_AT: u8 = 48; |
| /// Push pending bits to vector. |
| /// Using a macro here since an inline function. |
| /// didn't optimise properly. |
| macro_rules! push{ |
| ($s:ident) => { |
| let len = $s.w.len(); |
| $s.w.reserve(6); |
| // Optimization: |
| // |
| // This is basically what `Vec::extend_from_slice` does, but it didn't inline |
| // properly, so we do it manually for now. |
| // |
| // # Unsafe |
| // We reserve enough space right before this, so setting the len manually and doing |
| // unchecked indexing is safe here since we only, and always write to all of the the |
| // uninitialized bytes of the vector. |
| unsafe { |
| $s.w.set_len(len + 6); |
| $s.w.get_unchecked_mut(len..).copy_from_slice(&[$s.acc as u8, |
| ($s.acc >> 8) as u8, |
| ($s.acc >> 16) as u8, |
| ($s.acc >> 24) as u8, |
| ($s.acc >> 32) as u8, |
| ($s.acc >> 40) as u8 |
| ][..]); |
| } |
| |
| }; |
| } |
| } |
| #[cfg(not(target_pointer_width = "64"))] |
| #[macro_use] |
| mod arch_dep { |
| pub type AccType = u32; |
| pub const FLUSH_AT: u8 = 16; |
| macro_rules! push{ |
| ($s:ident) => { |
| // Unlike the 64-bit case, using copy_from_slice seemed to worsen performance here. |
| // TODO: Needs benching on a 32-bit system to see what works best. |
| $s.w.push($s.acc as u8); |
| $s.w.push(($s.acc >> 8) as u8); |
| }; |
| } |
| } |
| |
| use self::arch_dep::*; |
| |
| /// Writes bits to a byte stream, LSB first. |
| pub struct LsbWriter { |
| // Public for now so it can be replaced after initialization. |
| pub w: Vec<u8>, |
| bits: u8, |
| acc: AccType, |
| } |
| |
| impl LsbWriter { |
| /// Creates a new bit reader |
| pub fn new(writer: Vec<u8>) -> LsbWriter { |
| LsbWriter { |
| w: writer, |
| bits: 0, |
| acc: 0, |
| } |
| } |
| |
| pub fn pending_bits(&self) -> u8 { |
| self.bits |
| } |
| |
| /// Buffer n number of bits, and write them to the vec if there are enough pending bits. |
| pub fn write_bits(&mut self, v: u16, n: u8) { |
| // NOTE: This outputs garbage data if n is 0, but v is not 0 |
| self.acc |= (v as AccType) << self.bits; |
| self.bits += n; |
| // Waiting until we have FLUSH_AT bits and pushing them all in one batch. |
| while self.bits >= FLUSH_AT { |
| push!(self); |
| self.acc >>= FLUSH_AT; |
| self.bits -= FLUSH_AT; |
| } |
| } |
| |
| fn write_bits_finish(&mut self, v: u16, n: u8) { |
| // NOTE: This outputs garbage data if n is 0, but v is not 0 |
| self.acc |= (v as AccType) << self.bits; |
| self.bits += n % 8; |
| while self.bits >= 8 { |
| self.w.push(self.acc as u8); |
| self.acc >>= 8; |
| self.bits -= 8; |
| } |
| } |
| |
| pub fn flush_raw(&mut self) { |
| let missing = FLUSH_AT - self.bits; |
| // Have to test for self.bits > 0 here, |
| // otherwise flush would output an extra byte when flush was called at a byte boundary |
| if missing > 0 && self.bits > 0 { |
| self.write_bits_finish(0, missing); |
| } |
| } |
| } |
| |
| impl Write for LsbWriter { |
| fn write(&mut self, buf: &[u8]) -> io::Result<usize> { |
| if self.acc == 0 { |
| self.w.extend_from_slice(buf) |
| } else { |
| for &byte in buf.iter() { |
| self.write_bits(byte as u16, 8) |
| } |
| } |
| Ok(buf.len()) |
| } |
| |
| fn flush(&mut self) -> io::Result<()> { |
| self.flush_raw(); |
| Ok(()) |
| } |
| } |
| |
| #[cfg(test)] |
| mod test { |
| use super::LsbWriter; |
| |
| #[test] |
| fn write_bits() { |
| let input = [ |
| (3, 3), |
| (10, 8), |
| (88, 7), |
| (0, 2), |
| (0, 5), |
| (0, 0), |
| (238, 8), |
| (126, 8), |
| (161, 8), |
| (10, 8), |
| (238, 8), |
| (174, 8), |
| (126, 8), |
| (174, 8), |
| (65, 8), |
| (142, 8), |
| (62, 8), |
| (10, 8), |
| (1, 8), |
| (161, 8), |
| (78, 8), |
| (62, 8), |
| (158, 8), |
| (206, 8), |
| (10, 8), |
| (64, 7), |
| (0, 0), |
| (24, 5), |
| (0, 0), |
| (174, 8), |
| (126, 8), |
| (193, 8), |
| (174, 8), |
| ]; |
| let expected = [ |
| 83, |
| 192, |
| 2, |
| 220, |
| 253, |
| 66, |
| 21, |
| 220, |
| 93, |
| 253, |
| 92, |
| 131, |
| 28, |
| 125, |
| 20, |
| 2, |
| 66, |
| 157, |
| 124, |
| 60, |
| 157, |
| 21, |
| 128, |
| 216, |
| 213, |
| 47, |
| 216, |
| 21, |
| ]; |
| let mut writer = LsbWriter::new(Vec::new()); |
| for v in input.iter() { |
| writer.write_bits(v.0, v.1); |
| } |
| writer.flush_raw(); |
| assert_eq!(writer.w, expected); |
| } |
| } |
| |
| |
| #[cfg(all(test, feature = "benchmarks"))] |
| mod bench { |
| use test_std::Bencher; |
| use super::LsbWriter; |
| #[bench] |
| fn bit_writer(b: &mut Bencher) { |
| let input = [ |
| (3, 3), |
| (10, 8), |
| (88, 7), |
| (0, 2), |
| (0, 5), |
| (0, 0), |
| (238, 8), |
| (126, 8), |
| (161, 8), |
| (10, 8), |
| (238, 8), |
| (174, 8), |
| (126, 8), |
| (174, 8), |
| (65, 8), |
| (142, 8), |
| (62, 8), |
| (10, 8), |
| (1, 8), |
| (161, 8), |
| (78, 8), |
| (62, 8), |
| (158, 8), |
| (206, 8), |
| (10, 8), |
| (64, 7), |
| (0, 0), |
| (24, 5), |
| (0, 0), |
| (174, 8), |
| (126, 8), |
| (193, 8), |
| (174, 8), |
| ]; |
| let mut writer = LsbWriter::new(Vec::with_capacity(100)); |
| b.iter(|| for v in input.iter() { |
| let _ = writer.write_bits(v.0, v.1); |
| }); |
| } |
| } |