1use super::BlockType;
2use super::symbol;
3use crate::bit;
4use crate::finish::{Complete, Finish};
5use crate::lz77;
6use alloc::vec::Vec;
7use core::cmp;
8use no_std_io2::io;
9
10pub const DEFAULT_BLOCK_SIZE: usize = 1024 * 1024;
12
13const MAX_NON_COMPRESSED_BLOCK_SIZE: usize = 0xFFFF;
14
15#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub struct EncodeOptions<E = lz77::DefaultLz77Encoder> {
18 block_size: usize,
19 dynamic_huffman: bool,
20 lz77: Option<E>,
21}
22impl Default for EncodeOptions<lz77::DefaultLz77Encoder> {
23 fn default() -> Self {
24 Self::new()
25 }
26}
27impl EncodeOptions<lz77::DefaultLz77Encoder> {
28 pub fn new() -> Self {
38 EncodeOptions {
39 block_size: DEFAULT_BLOCK_SIZE,
40 dynamic_huffman: true,
41 lz77: Some(lz77::DefaultLz77Encoder::new()),
42 }
43 }
44}
45impl<E> EncodeOptions<E>
46where
47 E: lz77::Lz77Encode,
48{
49 pub fn with_lz77(lz77: E) -> Self {
60 EncodeOptions {
61 block_size: DEFAULT_BLOCK_SIZE,
62 dynamic_huffman: true,
63 lz77: Some(lz77),
64 }
65 }
66
67 pub fn no_compression(mut self) -> Self {
78 self.lz77 = None;
79 self
80 }
81
82 pub fn block_size(mut self, size: usize) -> Self {
94 self.block_size = size;
95 self
96 }
97
98 pub fn fixed_huffman_codes(mut self) -> Self {
108 self.dynamic_huffman = false;
109 self
110 }
111
112 fn get_block_type(&self) -> BlockType {
113 if self.lz77.is_none() {
114 BlockType::Raw
115 } else if self.dynamic_huffman {
116 BlockType::Dynamic
117 } else {
118 BlockType::Fixed
119 }
120 }
121 fn get_block_size(&self) -> usize {
122 if self.lz77.is_none() {
123 cmp::min(self.block_size, MAX_NON_COMPRESSED_BLOCK_SIZE)
124 } else {
125 self.block_size
126 }
127 }
128}
129
130#[derive(Debug)]
132pub struct Encoder<W, E = lz77::DefaultLz77Encoder> {
133 writer: bit::BitWriter<W>,
134 block: Block<E>,
135}
136impl<W> Encoder<W, lz77::DefaultLz77Encoder>
137where
138 W: io::Write,
139{
140 pub fn new(inner: W) -> Self {
157 Self::with_options(inner, EncodeOptions::default())
158 }
159}
160impl<W, E> Encoder<W, E>
161where
162 W: io::Write,
163 E: lz77::Lz77Encode,
164{
165 pub fn with_options(inner: W, options: EncodeOptions<E>) -> Self {
183 Encoder {
184 writer: bit::BitWriter::new(inner),
185 block: Block::new(options),
186 }
187 }
188
189 pub fn finish(mut self) -> Finish<W, io::Error> {
204 match self.block.finish(&mut self.writer) {
205 Ok(_) => Finish::new(self.writer.into_inner(), None),
206 Err(e) => Finish::new(self.writer.into_inner(), Some(e)),
207 }
208 }
209
210 pub fn as_inner_ref(&self) -> &W {
212 self.writer.as_inner_ref()
213 }
214
215 pub fn as_inner_mut(&mut self) -> &mut W {
217 self.writer.as_inner_mut()
218 }
219
220 pub fn into_inner(self) -> W {
222 self.writer.into_inner()
223 }
224
225 pub(crate) fn zlib_sync_flush(&mut self) -> io::Result<()> {
226 self.block.flush(&mut self.writer, false)?;
227
228 self.writer.write_bit(false)?;
229 self.writer.write_bits(2, BlockType::Raw as u16)?;
230 self.writer.flush()?;
231 self.writer.as_inner_mut().write_all(&[0, 0, 255, 255])?;
232
233 self.writer.as_inner_mut().flush()
234 }
235}
236impl<W, E> io::Write for Encoder<W, E>
237where
238 W: io::Write,
239 E: lz77::Lz77Encode,
240{
241 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
242 self.block.write(&mut self.writer, buf)?;
243 Ok(buf.len())
244 }
245 fn flush(&mut self) -> io::Result<()> {
246 self.block.flush(&mut self.writer, false)?;
247 self.writer.as_inner_mut().flush()
248 }
249}
250impl<W, E> Complete for Encoder<W, E>
251where
252 W: io::Write,
253 E: lz77::Lz77Encode,
254{
255 fn complete(self) -> io::Result<()> {
256 self.finish().into_result().map(|_| ())
257 }
258}
259
260#[derive(Debug)]
261struct Block<E> {
262 block_type: BlockType,
263 block_size: usize,
264 block_buf: BlockBuf<E>,
265}
266impl<E> Block<E>
267where
268 E: lz77::Lz77Encode,
269{
270 fn new(options: EncodeOptions<E>) -> Self {
271 Block {
272 block_type: options.get_block_type(),
273 block_size: options.get_block_size(),
274 block_buf: BlockBuf::new(options.lz77, options.dynamic_huffman),
275 }
276 }
277 fn write<W>(&mut self, writer: &mut bit::BitWriter<W>, buf: &[u8]) -> io::Result<()>
278 where
279 W: io::Write,
280 {
281 self.block_buf.append(buf);
282 while self.block_buf.len() >= self.block_size {
283 self.flush(writer, false)?;
284 }
285 Ok(())
286 }
287 fn flush<W>(&mut self, writer: &mut bit::BitWriter<W>, is_final: bool) -> io::Result<()>
288 where
289 W: io::Write,
290 {
291 writer.write_bit(is_final)?;
292 writer.write_bits(2, self.block_type as u16)?;
293 self.block_buf.flush(writer)?;
294 Ok(())
295 }
296 fn finish<W>(mut self, writer: &mut bit::BitWriter<W>) -> io::Result<()>
297 where
298 W: io::Write,
299 {
300 self.flush(writer, true)?;
301 writer.flush()?;
302 Ok(())
303 }
304}
305
306#[derive(Debug)]
307enum BlockBuf<E> {
308 Raw(RawBuf),
309 Fixed(CompressBuf<symbol::FixedHuffmanCodec, E>),
310 Dynamic(CompressBuf<symbol::DynamicHuffmanCodec, E>),
311}
312impl<E> BlockBuf<E>
313where
314 E: lz77::Lz77Encode,
315{
316 fn new(lz77: Option<E>, dynamic: bool) -> Self {
317 if let Some(lz77) = lz77 {
318 if dynamic {
319 BlockBuf::Dynamic(CompressBuf::new(symbol::DynamicHuffmanCodec, lz77))
320 } else {
321 BlockBuf::Fixed(CompressBuf::new(symbol::FixedHuffmanCodec, lz77))
322 }
323 } else {
324 BlockBuf::Raw(RawBuf::new())
325 }
326 }
327 fn append(&mut self, buf: &[u8]) {
328 match *self {
329 BlockBuf::Raw(ref mut b) => b.append(buf),
330 BlockBuf::Fixed(ref mut b) => b.append(buf),
331 BlockBuf::Dynamic(ref mut b) => b.append(buf),
332 }
333 }
334 fn len(&self) -> usize {
335 match *self {
336 BlockBuf::Raw(ref b) => b.len(),
337 BlockBuf::Fixed(ref b) => b.len(),
338 BlockBuf::Dynamic(ref b) => b.len(),
339 }
340 }
341 fn flush<W>(&mut self, writer: &mut bit::BitWriter<W>) -> io::Result<()>
342 where
343 W: io::Write,
344 {
345 match *self {
346 BlockBuf::Raw(ref mut b) => b.flush(writer),
347 BlockBuf::Fixed(ref mut b) => b.flush(writer),
348 BlockBuf::Dynamic(ref mut b) => b.flush(writer),
349 }
350 }
351}
352
353#[derive(Debug)]
354struct RawBuf {
355 buf: Vec<u8>,
356}
357impl RawBuf {
358 fn new() -> Self {
359 RawBuf { buf: Vec::new() }
360 }
361 fn append(&mut self, buf: &[u8]) {
362 self.buf.extend_from_slice(buf);
363 }
364 fn len(&self) -> usize {
365 self.buf.len()
366 }
367 fn flush<W>(&mut self, writer: &mut bit::BitWriter<W>) -> io::Result<()>
368 where
369 W: io::Write,
370 {
371 let size = cmp::min(self.buf.len(), MAX_NON_COMPRESSED_BLOCK_SIZE);
372 writer.flush()?;
373 writer
374 .as_inner_mut()
375 .write_all(&(size as u16).to_le_bytes())?;
376 writer
377 .as_inner_mut()
378 .write_all(&(!size as u16).to_le_bytes())?;
379 writer.as_inner_mut().write_all(&self.buf[..size])?;
380 self.buf.drain(0..size);
381 Ok(())
382 }
383}
384
385#[derive(Debug)]
386struct CompressBuf<H, E> {
387 huffman: H,
388 lz77: E,
389 buf: Vec<symbol::Symbol>,
390 original_size: usize,
391}
392impl<H, E> CompressBuf<H, E>
393where
394 H: symbol::HuffmanCodec,
395 E: lz77::Lz77Encode,
396{
397 fn new(huffman: H, lz77: E) -> Self {
398 CompressBuf {
399 huffman,
400 lz77,
401 buf: Vec::new(),
402 original_size: 0,
403 }
404 }
405 fn append(&mut self, buf: &[u8]) {
406 self.original_size += buf.len();
407 self.lz77.encode(buf, &mut self.buf);
408 }
409 fn len(&self) -> usize {
410 self.original_size
411 }
412 fn flush<W>(&mut self, writer: &mut bit::BitWriter<W>) -> io::Result<()>
413 where
414 W: io::Write,
415 {
416 self.lz77.flush(&mut self.buf);
417 self.buf.push(symbol::Symbol::EndOfBlock);
418 let symbol_encoder = self.huffman.build(&self.buf)?;
419 self.huffman.save(writer, &symbol_encoder)?;
420 for s in self.buf.drain(..) {
421 symbol_encoder.encode(writer, &s)?;
422 }
423 self.original_size = 0;
424 Ok(())
425 }
426}
427
428#[cfg(test)]
429mod tests {
430 use super::super::Decoder;
431 use super::*;
432 use no_std_io2::io::{Read as _, Write as _};
433
434 #[test]
435 fn test_issues_52() {
436 let input = crate::deflate::test_data::ISSUE_52_INPUT;
438
439 const LIMIT_1: usize = 16_031;
440 const LIMIT_2: usize = LIMIT_1 + 1;
441
442 let mut encoder = Encoder::new(Vec::new());
445 encoder.write_all(&input[0..LIMIT_1]).unwrap();
446 let compressed: Vec<u8> = encoder.finish().into_result().unwrap();
447
448 assert!(LIMIT_1 > compressed.len());
449
450 let mut encoder = Encoder::new(Vec::new());
453 encoder.write_all(&input[0..LIMIT_2]).unwrap();
454 let compressed: Vec<u8> = encoder.finish().into_result().unwrap();
455
456 assert!(LIMIT_2 > compressed.len());
457 }
458
459 #[test]
460 fn test_issue_27() {
461 let writes = ["fooooooooooooooooo", "bar", "baz"];
464
465 let mut encoder = Encoder::new(Vec::new());
466 for _ in 0..2 {
467 for string in &writes {
468 encoder.write(string.as_bytes()).expect("Write failed");
469 }
470 encoder.flush().expect("Flush failed");
471 }
472 let finished = encoder.finish().unwrap();
473 #[cfg(feature = "std")]
474 println!("{:?}", finished.0);
475
476 let mut output = Vec::new();
477 Decoder::new(&finished.0[..])
478 .read_to_end(&mut output)
479 .unwrap();
480 assert_eq!(
481 output,
482 "fooooooooooooooooobarbazfooooooooooooooooobarbaz".as_bytes()
483 );
484 }
485}