use std::io;
use super::{Numeric, SignedNumeric, BitQueue, BitQueueBE, BitQueueLE};
use huffman::WriteHuffmanTree;
pub trait BitWrite {
fn write_bit(&mut self, bit: bool) -> Result<(), io::Error>;
fn write<U>(&mut self, bits: u32, value: U) -> Result<(), io::Error>
where U: Numeric;
fn write_signed<S>(&mut self, bits: u32, value: S) -> Result<(), io::Error>
where S: SignedNumeric;
fn write_bytes(&mut self, buf: &[u8]) -> Result<(), io::Error>;
fn write_huffman<T>(&mut self,
tree: &WriteHuffmanTree<T>,
symbol: T) ->
Result<(), io::Error> where T: Ord + Copy;
fn write_unary0(&mut self, mut value: u32) -> Result<(), io::Error> {
while value > 8 {
self.write(8, 0xFFu8)?;
value -= 8;
}
if value > 0 {
self.write(value, (1 << value) - 1)?;
}
self.write_bit(false)
}
fn write_unary1(&mut self, mut value: u32) -> Result<(), io::Error> {
while value > 8 {
self.write(8, 0u8)?;
value -= 8;
}
if value > 0 {
self.write(value, 0)?;
}
self.write_bit(true)
}
fn byte_aligned(&self) -> bool;
fn byte_align(&mut self) -> Result<(), io::Error> {
while !self.byte_aligned() {
self.write_bit(false)?;
}
Ok(())
}
}
macro_rules! define_write_bit {
() => {
#[inline(always)]
fn write_bit(&mut self, bit: bool) -> Result<(), io::Error> {
self.bitqueue.push(1, if bit {1} else {0});
if self.bitqueue.len() == 8 {
write_byte(self.writer, self.bitqueue.pop(8))
} else {
Ok(())
}
}
}
}
macro_rules! define_write {
($bitqueue:ident) => {
fn write<U>(&mut self, bits: u32, value: U) -> Result<(), io::Error>
where U: Numeric {
let mut acc = $bitqueue::from_value(value, bits);
write_unaligned(&mut self.writer, &mut acc, &mut self.bitqueue)
.and_then(|()|
write_aligned(&mut self.writer, &mut self.byte_buf, &mut acc))
.and_then(|()|
Ok(self.bitqueue.push(acc.len(), acc.value().to_u8())))
}
}
}
macro_rules! define_write_bytes {
() => {
fn write_bytes(&mut self, buf: &[u8]) -> Result<(), io::Error> {
if self.byte_aligned() {
self.writer.write_all(buf)
} else {
for b in buf {
self.write(8, *b)?;
}
Ok(())
}
}
}
}
pub struct BitWriterBE<'a> {
writer: &'a mut io::Write,
bitqueue: BitQueueBE<u8>,
byte_buf: Vec<u8>
}
impl<'a> BitWriterBE<'a> {
pub fn new(writer: &mut io::Write) -> BitWriterBE {
BitWriterBE{writer: writer,
bitqueue: BitQueueBE::new(),
byte_buf: Vec::new()}
}
}
impl<'a> BitWrite for BitWriterBE<'a> {
define_write_bit!();
define_write!(BitQueueBE);
define_write_bytes!();
fn write_signed<S>(&mut self, bits: u32, value: S) -> Result<(), io::Error>
where S: SignedNumeric {
if value.is_negative() {
self.write_bit(true)
.and_then(|()| self.write(bits - 1, value.as_unsigned(bits)))
} else {
self.write_bit(false)
.and_then(|()| self.write(bits - 1, value))
}
}
#[inline]
fn write_huffman<T>(&mut self,
tree: &WriteHuffmanTree<T>,
symbol: T) ->
Result<(), io::Error> where T: Ord + Copy {
let (bits, value) = tree.get_be(symbol);
self.write(bits, value)
}
#[inline]
fn byte_aligned(&self) -> bool {self.bitqueue.is_empty()}
}
pub struct BitWriterLE<'a> {
writer: &'a mut io::Write,
bitqueue: BitQueueLE<u8>,
byte_buf: Vec<u8>
}
impl<'a> BitWriterLE<'a> {
pub fn new(writer: &mut io::Write) -> BitWriterLE {
BitWriterLE{writer: writer,
bitqueue: BitQueueLE::new(),
byte_buf: Vec::new()}
}
}
impl<'a> BitWrite for BitWriterLE<'a> {
define_write_bit!();
define_write!(BitQueueLE);
define_write_bytes!();
fn write_signed<S>(&mut self, bits: u32, value: S) -> Result<(), io::Error>
where S: SignedNumeric {
if value.is_negative() {
self.write(bits - 1, value.as_unsigned(bits))
.and_then(|()| self.write_bit(true))
} else {
self.write(bits - 1, value)
.and_then(|()| self.write_bit(false))
}
}
#[inline]
fn write_huffman<T>(&mut self,
tree: &WriteHuffmanTree<T>,
symbol: T) ->
Result<(), io::Error> where T: Ord + Copy {
let (bits, value) = tree.get_le(symbol);
self.write(bits, value)
}
#[inline]
fn byte_aligned(&self) -> bool {self.bitqueue.is_empty()}
}
#[inline]
fn write_byte(writer: &mut io::Write, byte: u8) -> Result<(),io::Error> {
let buf = [byte];
writer.write_all(&buf)
}
fn write_unaligned<N>(writer: &mut io::Write,
acc: &mut BitQueue<N>,
rem: &mut BitQueue<u8>) -> Result<(), io::Error>
where N: Numeric {
if rem.is_empty() {
Ok(())
} else {
use std::cmp::min;
let bits_to_transfer = min(8 - rem.len(), acc.len());
rem.push(bits_to_transfer, acc.pop(bits_to_transfer).to_u8());
if rem.len() == 8 {
write_byte(writer, rem.pop(8))
} else {
Ok(())
}
}
}
fn write_aligned<N>(writer: &mut io::Write,
byte_buf: &mut Vec<u8>,
acc: &mut BitQueue<N>) -> Result<(), io::Error>
where N: Numeric {
let bytes_to_write = (acc.len() / 8) as usize;
if bytes_to_write > 0 {
byte_buf.clear();
byte_buf.extend((0..bytes_to_write).map(|_| acc.pop(8).to_u8()));
writer.write_all(&byte_buf)
} else {
Ok(())
}
}