1 use std::io::Write;
2 use std::{cmp, io, mem};
3 
4 use crate::compress::Flush;
5 use crate::compression_options::{CompressionOptions, MAX_HASH_CHECKS};
6 use crate::encoder_state::EncoderState;
7 pub use crate::huffman_table::MAX_MATCH;
8 use crate::huffman_table::NUM_LITERALS_AND_LENGTHS;
9 use crate::input_buffer::InputBuffer;
10 use crate::length_encode::{EncodedLength, LeafVec};
11 use crate::lz77::LZ77State;
12 use 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)]
17 pub struct DebugCounter {
18     #[cfg(debug_assertions)]
19     count: u64,
20 }
21 
22 impl DebugCounter {
23     #[cfg(debug_assertions)]
get(&self) -> u6424     pub fn get(&self) -> u64 {
25         self.count
26     }
27 
28     #[cfg(not(debug_assertions))]
get(&self) -> u6429     pub fn get(&self) -> u64 {
30         0
31     }
32 
33     #[cfg(debug_assertions)]
reset(&mut self)34     pub fn reset(&mut self) {
35         self.count = 0;
36     }
37 
38     #[cfg(not(debug_assertions))]
reset(&self)39     pub fn reset(&self) {}
40 
41     #[cfg(debug_assertions)]
add(&mut self, val: u64)42     pub fn add(&mut self, val: u64) {
43         self.count += val;
44     }
45 
46     #[cfg(not(debug_assertions))]
add(&self, _: u64)47     pub fn add(&self, _: u64) {}
48 }
49 
50 pub struct LengthBuffers {
51     pub leaf_buf: LeafVec,
52     pub length_buf: Vec<EncodedLength>,
53 }
54 
55 impl LengthBuffers {
56     #[inline]
new() -> LengthBuffers57     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.
66 pub 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 
99 impl<W: Write> DeflateState<W> {
new(compression_options: CompressionOptions, writer: W) -> 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]
output_buf(&mut self) -> &mut Vec<u8>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.
reset(&mut self, writer: W) -> io::Result<W>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 }
154