blob: 6aba35d4c1bc27a2c2729e0f4c58d126257a1d66 [file] [log] [blame]
// 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);
});
}
}