diff --git a/src/deflate/deflate_token.rs b/src/deflate/deflate_token.rs index ec7a990..fc78971 100644 --- a/src/deflate/deflate_token.rs +++ b/src/deflate/deflate_token.rs @@ -18,6 +18,15 @@ pub enum DeflateToken { Reference(DeflateTokenReference), } +impl DeflateToken { + pub fn new_ref(len: u32, dist: u32, irregular258: bool) -> DeflateToken { + DeflateToken::Reference(DeflateTokenReference::new(len, dist, irregular258)) + } + pub fn new_lit(lit: u8) -> DeflateToken { + DeflateToken::Literal(lit) + } +} + /// In the case of a distance and length, the length is the number of bytes to copy from the /// previous bytes, and the distance is the number of bytes back to start copying from. /// @@ -70,7 +79,7 @@ pub enum DeflateHuffmanType { }, } -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub enum DeflateTokenBlock { Huffman { tokens: Vec, diff --git a/src/deflate/deflate_writer.rs b/src/deflate/deflate_writer.rs index 3712127..d510ff6 100644 --- a/src/deflate/deflate_writer.rs +++ b/src/deflate/deflate_writer.rs @@ -66,7 +66,7 @@ impl DeflateWriter { DeflateHuffmanType::Static { .. } => { self.bitwriter.write(1, 2, &mut self.output); let huffman_writer = HuffmanWriter::start_fixed_huffman_table(); - self.encode_block_with_decoder(tokens, &huffman_writer); + self.encode_huffman(tokens, &huffman_writer); } DeflateHuffmanType::Dynamic { huffman_encoding, .. @@ -77,7 +77,7 @@ impl DeflateWriter { &mut self.output, )?; - self.encode_block_with_decoder(tokens, &huffman_writer); + self.encode_huffman(tokens, &huffman_writer); } }, } @@ -90,11 +90,7 @@ impl DeflateWriter { self.bitwriter.flush_whole_bytes(&mut self.output); } - fn encode_block_with_decoder( - &mut self, - tokens: &Vec, - huffman_writer: &HuffmanWriter, - ) { + fn encode_huffman(&mut self, tokens: &Vec, huffman_writer: &HuffmanWriter) { for token in tokens { match token { DeflateToken::Literal(lit) => { @@ -111,7 +107,7 @@ impl DeflateWriter { &mut self.output, LITLEN_CODE_COUNT as u16 - 2, ); - self.bitwriter.write(5, 31, &mut self.output); + self.bitwriter.write(31, 5, &mut self.output); } else { let lencode = quantize_length(reference.len()); huffman_writer.write_literal( @@ -128,22 +124,22 @@ impl DeflateWriter { &mut self.output, ); } + } - let distcode = quantize_distance(reference.dist()); - huffman_writer.write_distance( - &mut self.bitwriter, + let distcode = quantize_distance(reference.dist()); + huffman_writer.write_distance( + &mut self.bitwriter, + &mut self.output, + distcode as u16, + ); + + let distextra = DIST_EXTRA_TABLE[distcode]; + if distextra > 0 { + self.bitwriter.write( + reference.dist() - 1 - DIST_BASE_TABLE[distcode] as u32, + distextra.into(), &mut self.output, - distcode as u16, ); - - let distextra = DIST_EXTRA_TABLE[distcode]; - if distextra > 0 { - self.bitwriter.write( - reference.dist() - 1 - DIST_BASE_TABLE[distcode] as u32, - distextra.into(), - &mut self.output, - ); - } } } } @@ -152,3 +148,56 @@ impl DeflateWriter { huffman_writer.write_literal(&mut self.bitwriter, &mut self.output, 256); } } + +/// Create a set of blocks and read them back to see if they are identical +#[test] +fn roundtrip_deflate_writer() { + use super::deflate_reader::DeflateReader; + use std::io::Cursor; + + let mut w = DeflateWriter::new(); + + let blocks = [ + DeflateTokenBlock::Huffman { + tokens: vec![DeflateToken::Literal(0), DeflateToken::Literal(1)], + huffman_type: DeflateHuffmanType::Static { incomplete: false }, + }, + DeflateTokenBlock::Stored { + uncompressed: vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + padding_bits: 0b101, // there are 3 bits of padding + }, + DeflateTokenBlock::Huffman { + tokens: vec![ + DeflateToken::Literal(0), + DeflateToken::Literal(1), + DeflateToken::new_ref(100, 1, false), + DeflateToken::new_ref(258, 1, true), + DeflateToken::Literal(3), + ], + huffman_type: DeflateHuffmanType::Static { incomplete: false }, + }, + ]; + + for i in 0..blocks.len() { + w.encode_block(&blocks[i], i == blocks.len() - 1).unwrap(); + } + w.flush_with_padding(0); + + let output = w.detach_output(); + + let mut r = DeflateReader::new(Cursor::new(output)); + + let mut newcontent = Vec::new(); + loop { + let mut last = false; + newcontent.push(r.read_block(&mut last).unwrap()); + if last { + break; + } + } + + assert_eq!(blocks.len(), newcontent.len()); + for i in 0..blocks.len() { + assert_eq!(blocks[i], newcontent[i], "block {}", i); + } +}