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
generated()22 impl DebugCounter {
23 #[cfg(debug_assertions)]
24 pub fn get(&self) -> u64 {
25 self.count
26 }
27
28 #[cfg(not(debug_assertions))]
29 pub 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
50 pub struct LengthBuffers {
51 pub leaf_buf: LeafVec,
52 pub length_buf: Vec<EncodedLength>,
53 }
54
55 impl 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.
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> {
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 }
154