use crate::deflate::symbol::{self, HuffmanCodec};
use crate::lz77;
use crate::non_blocking::transaction::TransactionalBitReader;
use core::cmp;
use no_std_io2::io::{self, Read};
#[derive(Debug)]
pub struct Decoder<R> {
state: DecoderState,
eos: bool,
bit_reader: TransactionalBitReader<R>,
block_decoder: BlockDecoder,
}
impl<R: Read> Decoder<R> {
pub fn new(inner: R) -> Self {
Decoder {
state: DecoderState::ReadBlockHeader,
eos: false,
bit_reader: TransactionalBitReader::new(inner),
block_decoder: BlockDecoder::new(),
}
}
pub fn as_inner_ref(&self) -> &R {
self.bit_reader.as_inner_ref()
}
pub fn as_inner_mut(&mut self) -> &mut R {
self.bit_reader.as_inner_mut()
}
pub fn into_inner(self) -> R {
self.bit_reader.into_inner()
}
pub(crate) fn bit_reader_mut(&mut self) -> &mut TransactionalBitReader<R> {
&mut self.bit_reader
}
}
impl<R: Read> Read for Decoder<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let mut read_size;
loop {
let next = match self.state {
DecoderState::ReadBlockHeader => {
let (bfinal, btype) = self.bit_reader.transaction(|r| {
let bfinal = r.read_bit()?;
let btype = r.read_bits(2)?;
Ok((bfinal, btype))
})?;
self.eos = bfinal;
self.block_decoder.enter_new_block();
match btype {
0b00 => DecoderState::ReadNonCompressedBlockLen,
0b01 => DecoderState::LoadFixedHuffmanCode,
0b10 => DecoderState::LoadDynamicHuffmanCode,
0b11 => {
return Err(invalid_data_error!(
"btype 0x11 of DEFLATE is reserved(error) value"
));
}
_ => unreachable!(),
}
}
DecoderState::ReadNonCompressedBlockLen => {
let len = self.bit_reader.transaction(|r| {
r.reset();
let mut buf = [0; 2];
r.as_inner_mut().read_exact(&mut buf)?;
let len = u16::from_le_bytes(buf);
r.as_inner_mut().read_exact(&mut buf)?;
let nlen = u16::from_le_bytes(buf);
if !len != nlen {
Err(invalid_data_error!(
"LEN={} is not the one's complement of NLEN={}",
len,
nlen
))
} else {
Ok(len)
}
})?;
DecoderState::ReadNonCompressedBlock { len }
}
DecoderState::ReadNonCompressedBlock { len: 0 } => {
if self.eos {
read_size = 0;
break;
} else {
DecoderState::ReadBlockHeader
}
}
DecoderState::ReadNonCompressedBlock { ref mut len } => {
let buf_len = buf.len();
let buf = &mut buf[..cmp::min(buf_len, *len as usize)];
read_size = self.bit_reader.as_inner_mut().read(buf)?;
self.block_decoder.extend(&buf[..read_size]);
*len -= read_size as u16;
break;
}
DecoderState::LoadFixedHuffmanCode => {
let symbol_decoder = self
.bit_reader
.transaction(|r| symbol::FixedHuffmanCodec.load(r))?;
DecoderState::DecodeBlock(symbol_decoder)
}
DecoderState::LoadDynamicHuffmanCode => {
let symbol_decoder = self
.bit_reader
.transaction(|r| symbol::DynamicHuffmanCodec.load(r))?;
DecoderState::DecodeBlock(symbol_decoder)
}
DecoderState::DecodeBlock(ref mut symbol_decoder) => {
self.block_decoder
.decode(&mut self.bit_reader, symbol_decoder)?;
read_size = self.block_decoder.read(buf)?;
if read_size == 0 && !buf.is_empty() && !self.eos {
DecoderState::ReadBlockHeader
} else {
break;
}
}
};
self.state = next;
}
Ok(read_size)
}
}
#[derive(Debug)]
enum DecoderState {
ReadBlockHeader,
ReadNonCompressedBlockLen,
ReadNonCompressedBlock { len: u16 },
LoadFixedHuffmanCode,
LoadDynamicHuffmanCode,
DecodeBlock(symbol::Decoder),
}
#[derive(Debug)]
struct BlockDecoder {
lz77_decoder: lz77::Lz77Decoder,
eob: bool,
}
impl BlockDecoder {
pub fn new() -> Self {
BlockDecoder {
lz77_decoder: lz77::Lz77Decoder::new(),
eob: false,
}
}
pub fn enter_new_block(&mut self) {
self.eob = false;
}
pub fn decode<R: Read>(
&mut self,
bit_reader: &mut TransactionalBitReader<R>,
symbol_decoder: &mut symbol::Decoder,
) -> io::Result<()> {
if self.eob {
return Ok(());
}
while let Some(s) = self.decode_symbol(bit_reader, symbol_decoder)? {
match s {
symbol::Symbol::Code(code) => {
self.lz77_decoder.decode(code)?;
}
symbol::Symbol::EndOfBlock => {
self.eob = true;
break;
}
}
}
Ok(())
}
fn extend(&mut self, buf: &[u8]) {
self.lz77_decoder.extend_from_slice(buf);
}
fn decode_symbol<R: Read>(
&mut self,
bit_reader: &mut TransactionalBitReader<R>,
symbol_decoder: &mut symbol::Decoder,
) -> io::Result<Option<symbol::Symbol>> {
let result = bit_reader.transaction(|bit_reader| {
let s = symbol_decoder.decode_unchecked(bit_reader);
bit_reader.check_last_error().map(|()| s)
});
match result {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(None),
Err(e) => Err(e),
Ok(s) => Ok(Some(s)),
}
}
}
impl Read for BlockDecoder {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if !self.lz77_decoder.buffer().is_empty() {
self.lz77_decoder.read(buf)
} else if self.eob {
Ok(0)
} else {
Err(io::Error::new(io::ErrorKind::WouldBlock, "Would block"))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::deflate::{EncodeOptions, Encoder};
use crate::util::{nb_read_to_end, WouldBlockReader};
use alloc::{format, string::String, vec::Vec};
use no_std_io2::io::{Read, Write};
#[test]
fn it_works() {
let mut encoder = Encoder::new(Vec::new());
encoder.write_all(b"Hello World!".as_ref()).unwrap();
let encoded_data = encoder.finish().into_result().unwrap();
let mut decoder = Decoder::new(&encoded_data[..]);
let mut decoded_data = Vec::new();
decoder.read_to_end(&mut decoded_data).unwrap();
assert_eq!(decoded_data, b"Hello World!");
}
#[test]
fn non_blocking_io_works() {
let mut encoder = Encoder::new(Vec::new());
encoder.write_all(b"Hello World!".as_ref()).unwrap();
let encoded_data = encoder.finish().into_result().unwrap();
let decoder = Decoder::new(WouldBlockReader::new(&encoded_data[..]));
let decoded_data = nb_read_to_end(decoder).unwrap();
assert_eq!(decoded_data, b"Hello World!");
}
#[test]
fn non_blocking_io_for_large_text_works() {
let text: String = (0..10000)
.into_iter()
.map(|i| format!("test {}", i))
.collect();
let mut encoder = crate::deflate::Encoder::new(Vec::new());
encoder.write_all(text.as_bytes()).unwrap();
let encoded_data = encoder.finish().into_result().unwrap();
let decoder = Decoder::new(WouldBlockReader::new(&encoded_data[..]));
let decoded_data = nb_read_to_end(decoder).unwrap();
assert_eq!(decoded_data, text.as_bytes());
}
#[test]
fn non_compressed_non_blocking_io_works() {
let mut encoder = Encoder::with_options(Vec::new(), EncodeOptions::new().no_compression());
encoder.write_all(b"Hello World!".as_ref()).unwrap();
let encoded_data = encoder.finish().into_result().unwrap();
let decoder = Decoder::new(WouldBlockReader::new(&encoded_data[..]));
let decoded_data = nb_read_to_end(decoder).unwrap();
assert_eq!(decoded_data, b"Hello World!");
}
}