Skip to main content

libflate/deflate/
encode.rs

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
10/// The default size of a DEFLATE block.
11pub const DEFAULT_BLOCK_SIZE: usize = 1024 * 1024;
12
13const MAX_NON_COMPRESSED_BLOCK_SIZE: usize = 0xFFFF;
14
15/// Options for a DEFLATE encoder.
16#[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    /// Makes a default instance.
29    ///
30    /// # Examples
31    /// ```
32    /// use libflate::deflate::{Encoder, EncodeOptions};
33    ///
34    /// let options = EncodeOptions::new();
35    /// let encoder = Encoder::with_options(Vec::new(), options);
36    /// ```
37    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    /// Specifies the LZ77 encoder used to compress input data.
50    ///
51    /// # Example
52    /// ```
53    /// use libflate::lz77::DefaultLz77Encoder;
54    /// use libflate::deflate::{Encoder, EncodeOptions};
55    ///
56    /// let options = EncodeOptions::with_lz77(DefaultLz77Encoder::new());
57    /// let encoder = Encoder::with_options(Vec::new(), options);
58    /// ```
59    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    /// Disables LZ77 compression.
68    ///
69    /// # Example
70    /// ```
71    /// use libflate::lz77::DefaultLz77Encoder;
72    /// use libflate::deflate::{Encoder, EncodeOptions};
73    ///
74    /// let options = EncodeOptions::new().no_compression();
75    /// let encoder = Encoder::with_options(Vec::new(), options);
76    /// ```
77    pub fn no_compression(mut self) -> Self {
78        self.lz77 = None;
79        self
80    }
81
82    /// Specifies the hint of the size of a DEFLATE block.
83    ///
84    /// The default value is `DEFAULT_BLOCK_SIZE`.
85    ///
86    /// # Example
87    /// ```
88    /// use libflate::deflate::{Encoder, EncodeOptions};
89    ///
90    /// let options = EncodeOptions::new().block_size(512 * 1024);
91    /// let encoder = Encoder::with_options(Vec::new(), options);
92    /// ```
93    pub fn block_size(mut self, size: usize) -> Self {
94        self.block_size = size;
95        self
96    }
97
98    /// Specifies to compress with fixed huffman codes.
99    ///
100    /// # Example
101    /// ```
102    /// use libflate::deflate::{Encoder, EncodeOptions};
103    ///
104    /// let options = EncodeOptions::new().fixed_huffman_codes();
105    /// let encoder = Encoder::with_options(Vec::new(), options);
106    /// ```
107    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/// DEFLATE encoder.
131#[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    /// Makes a new encoder instance.
141    ///
142    /// Encoded DEFLATE stream is written to `inner`.
143    ///
144    /// # Examples
145    /// ```
146    /// use no_std_io2::io::Write;
147    /// use libflate::deflate::Encoder;
148    ///
149    /// let mut encoder = Encoder::new(Vec::new());
150    /// encoder.write_all(b"Hello World!".as_ref()).unwrap();
151    ///
152    /// assert_eq!(encoder.finish().into_result().unwrap(),
153    ///            [5, 192, 49, 13, 0, 0, 8, 3, 65, 43, 224, 6, 7, 24, 128, 237,
154    ///            147, 38, 245, 63, 244, 230, 65, 181, 50, 215, 1]);
155    /// ```
156    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    /// Makes a new encoder instance with specified options.
166    ///
167    /// Encoded DEFLATE stream is written to `inner`.
168    ///
169    /// # Examples
170    /// ```
171    /// use no_std_io2::io::Write;
172    /// use libflate::deflate::{Encoder, EncodeOptions};
173    ///
174    /// let options = EncodeOptions::new().no_compression();
175    /// let mut encoder = Encoder::with_options(Vec::new(), options);
176    /// encoder.write_all(b"Hello World!".as_ref()).unwrap();
177    ///
178    /// assert_eq!(encoder.finish().into_result().unwrap(),
179    ///            [1, 12, 0, 243, 255, 72, 101, 108, 108, 111, 32, 87, 111,
180    ///             114, 108, 100, 33]);
181    /// ```
182    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    /// Flushes internal buffer and returns the inner stream.
190    ///
191    /// # Examples
192    /// ```
193    /// use no_std_io2::io::Write;
194    /// use libflate::deflate::Encoder;
195    ///
196    /// let mut encoder = Encoder::new(Vec::new());
197    /// encoder.write_all(b"Hello World!".as_ref()).unwrap();
198    ///
199    /// assert_eq!(encoder.finish().into_result().unwrap(),
200    ///            [5, 192, 49, 13, 0, 0, 8, 3, 65, 43, 224, 6, 7, 24, 128, 237,
201    ///            147, 38, 245, 63, 244, 230, 65, 181, 50, 215, 1]);
202    /// ```
203    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    /// Returns the immutable reference to the inner stream.
211    pub fn as_inner_ref(&self) -> &W {
212        self.writer.as_inner_ref()
213    }
214
215    /// Returns the mutable reference to the inner stream.
216    pub fn as_inner_mut(&mut self) -> &mut W {
217        self.writer.as_inner_mut()
218    }
219
220    /// Unwraps the `Encoder`, returning the inner stream.
221    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        // see: https://github.com/sile/libflate/issues/52
437        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        // Attempt 1 (should succeed)
443        //
444        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        // Attempt 2 (will fail without patch)
451        //
452        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        // See: https://github.com/sile/libflate/issues/27
462
463        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}