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