deflate/
deflate_state.rs

1use std::io::Write;
2use std::{cmp, io, mem};
3
4use crate::compress::Flush;
5use crate::compression_options::{CompressionOptions, MAX_HASH_CHECKS};
6use crate::encoder_state::EncoderState;
7pub use crate::huffman_table::MAX_MATCH;
8use crate::huffman_table::NUM_LITERALS_AND_LENGTHS;
9use crate::input_buffer::InputBuffer;
10use crate::length_encode::{EncodedLength, LeafVec};
11use crate::lz77::LZ77State;
12use crate::output_writer::DynamicWriter;
13
14/// A counter used for checking values in debug mode.
15/// Does nothing when debug assertions are disabled.
16#[derive(Default)]
17pub struct DebugCounter {
18    #[cfg(debug_assertions)]
19    count: u64,
20}
21
22impl DebugCounter {
23    #[cfg(debug_assertions)]
24    pub const fn get(&self) -> u64 {
25        self.count
26    }
27
28    #[cfg(not(debug_assertions))]
29    pub const fn get(&self) -> u64 {
30        0
31    }
32
33    #[cfg(debug_assertions)]
34    pub fn reset(&mut self) {
35        self.count = 0;
36    }
37
38    #[cfg(not(debug_assertions))]
39    pub fn reset(&self) {}
40
41    #[cfg(debug_assertions)]
42    pub fn add(&mut self, val: u64) {
43        self.count += val;
44    }
45
46    #[cfg(not(debug_assertions))]
47    pub fn add(&self, _: u64) {}
48}
49
50pub struct LengthBuffers {
51    pub leaf_buf: LeafVec,
52    pub length_buf: Vec<EncodedLength>,
53}
54
55impl LengthBuffers {
56    #[inline]
57    fn new() -> LengthBuffers {
58        LengthBuffers {
59            leaf_buf: Vec::with_capacity(NUM_LITERALS_AND_LENGTHS),
60            length_buf: Vec::with_capacity(19),
61        }
62    }
63}
64
65/// A struct containing all the stored state used for the encoder.
66pub struct DeflateState<W: Write> {
67    /// State of lz77 compression.
68    pub lz77_state: LZ77State,
69    pub input_buffer: InputBuffer,
70    pub compression_options: CompressionOptions,
71    /// State the Huffman part of the compression and the output buffer.
72    pub encoder_state: EncoderState,
73    /// The buffer containing the raw output of the lz77-encoding.
74    pub lz77_writer: DynamicWriter,
75    /// Buffers used when generating Huffman code lengths.
76    pub length_buffers: LengthBuffers,
77    /// Total number of bytes consumed/written to the input buffer.
78    pub bytes_written: u64,
79    /// Wrapped writer.
80    /// Option is used to allow us to implement `Drop` and `finish()` at the same time for the
81    /// writer structs.
82    pub inner: Option<W>,
83    /// The position in the output buffer where data should be flushed from, to keep track of
84    /// what data has been output in case not all data is output when writing to the wrapped
85    /// writer.
86    pub output_buf_pos: usize,
87    pub flush_mode: Flush,
88    /// Whether we need to flush everything before continuing.
89    /// Currently only used after having output a sync flush.
90    /// This is implemented in a somewhat clunky manner at the moment,
91    /// ideally it should be done in a more fail-safe way to avoid
92    /// further bugs.
93    pub needs_flush: bool,
94    /// Number of bytes written as calculated by sum of block input lengths.
95    /// Used to check that they are correct when `debug_assertions` are enabled.
96    pub bytes_written_control: DebugCounter,
97}
98
99impl<W: Write> DeflateState<W> {
100    pub fn new(compression_options: CompressionOptions, writer: W) -> DeflateState<W> {
101        DeflateState {
102            input_buffer: InputBuffer::empty(),
103            lz77_state: LZ77State::new(
104                compression_options.max_hash_checks,
105                cmp::min(compression_options.lazy_if_less_than, MAX_HASH_CHECKS),
106                compression_options.matching_type,
107            ),
108            encoder_state: EncoderState::new(Vec::with_capacity(1024 * 32)),
109            lz77_writer: DynamicWriter::new(),
110            length_buffers: LengthBuffers::new(),
111            compression_options,
112            bytes_written: 0,
113            inner: Some(writer),
114            output_buf_pos: 0,
115            flush_mode: Flush::None,
116            needs_flush: false,
117            bytes_written_control: DebugCounter::default(),
118        }
119    }
120
121    #[inline]
122    pub fn output_buf(&mut self) -> &mut Vec<u8> {
123        self.encoder_state.inner_vec()
124    }
125
126    /// Resets the status of the decoder, leaving the compression options intact
127    ///
128    /// If flushing the current writer succeeds, it is replaced with the provided one,
129    /// buffers and status (except compression options) is reset and the old writer
130    /// is returned.
131    ///
132    /// If flushing fails, the rest of the writer is not cleared.
133    pub fn reset(&mut self, writer: W) -> io::Result<W> {
134        self.encoder_state.flush();
135        self.inner
136            .as_mut()
137            .expect("Missing writer!")
138            .write_all(self.encoder_state.inner_vec())?;
139        self.encoder_state.inner_vec().clear();
140        self.input_buffer = InputBuffer::empty();
141        self.lz77_writer.clear();
142        self.lz77_state.reset();
143        self.bytes_written = 0;
144        self.output_buf_pos = 0;
145        self.flush_mode = Flush::None;
146        self.needs_flush = false;
147        if cfg!(debug_assertions) {
148            self.bytes_written_control.reset();
149        }
150        mem::replace(&mut self.inner, Some(writer))
151            .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "Missing writer"))
152    }
153}