deflate/
encoder_state.rs

1#[cfg(test)]
2use std::mem;
3use huffman_table::HuffmanTable;
4use bitstream::LsbWriter;
5use lzvalue::LZType;
6
7// The first bits of each block, which describe the type of the block
8// `-TTF` - TT = type, 00 = stored, 01 = fixed, 10 = dynamic, 11 = reserved, F - 1 if final block
9// `0000`;
10const FIXED_FIRST_BYTE: u16 = 0b010;
11const FIXED_FIRST_BYTE_FINAL: u16 = 0b011;
12const DYNAMIC_FIRST_BYTE: u16 = 0b100;
13const DYNAMIC_FIRST_BYTE_FINAL: u16 = 0b101;
14
15#[allow(dead_code)]
16pub enum BType {
17    NoCompression = 0b00,
18    FixedHuffman = 0b01,
19    DynamicHuffman = 0b10, // Reserved = 0b11, //Error
20}
21
22/// A struct wrapping a writer that writes data compressed using the provided huffman table
23pub struct EncoderState {
24    pub huffman_table: HuffmanTable,
25    pub writer: LsbWriter,
26}
27
28impl EncoderState {
29    /// Creates a new encoder state using the provided huffman table and writer
30    pub fn new(writer: Vec<u8>) -> EncoderState {
31        EncoderState {
32            huffman_table: HuffmanTable::empty(),
33            writer: LsbWriter::new(writer),
34        }
35    }
36
37    #[cfg(test)]
38    /// Creates a new encoder state using the fixed huffman table
39    pub fn fixed(writer: Vec<u8>) -> EncoderState {
40        EncoderState {
41            huffman_table: HuffmanTable::fixed_table(),
42            writer: LsbWriter::new(writer),
43        }
44    }
45
46    pub fn inner_vec(&mut self) -> &mut Vec<u8> {
47        &mut self.writer.w
48    }
49
50    /// Encodes a literal value to the writer
51    fn write_literal(&mut self, value: u8) {
52        let code = self.huffman_table.get_literal(value);
53        debug_assert!(code.length > 0);
54        self.writer.write_bits(code.code, code.length);
55    }
56
57    /// Write a LZvalue to the contained writer, returning Err if the write operation fails
58    pub fn write_lzvalue(&mut self, value: LZType) {
59        match value {
60            LZType::Literal(l) => self.write_literal(l),
61            LZType::StoredLengthDistance(l, d) => {
62                let (code, extra_bits_code) = self.huffman_table.get_length_huffman(l);
63                debug_assert!(
64                    code.length != 0,
65                    format!("Code: {:?}, Value: {:?}", code, value)
66                );
67                self.writer.write_bits(code.code, code.length);
68                self.writer
69                    .write_bits(extra_bits_code.code, extra_bits_code.length);
70
71                let (code, extra_bits_code) = self.huffman_table.get_distance_huffman(d);
72                debug_assert!(
73                    code.length != 0,
74                    format!("Code: {:?}, Value: {:?}", code, value)
75                );
76
77                self.writer.write_bits(code.code, code.length);
78                self.writer
79                    .write_bits(extra_bits_code.code, extra_bits_code.length)
80            }
81        };
82    }
83
84    /// Write the start of a block, returning Err if the write operation fails.
85    pub fn write_start_of_block(&mut self, fixed: bool, final_block: bool) {
86        if final_block {
87            // The final block has one bit flipped to indicate it's
88            // the final one
89            if fixed {
90                self.writer.write_bits(FIXED_FIRST_BYTE_FINAL, 3)
91            } else {
92                self.writer.write_bits(DYNAMIC_FIRST_BYTE_FINAL, 3)
93            }
94        } else if fixed {
95            self.writer.write_bits(FIXED_FIRST_BYTE, 3)
96        } else {
97            self.writer.write_bits(DYNAMIC_FIRST_BYTE, 3)
98        }
99    }
100
101    /// Write the end of block code
102    pub fn write_end_of_block(&mut self) {
103        let code = self.huffman_table.get_end_of_block();
104        self.writer.write_bits(code.code, code.length)
105    }
106
107    /// Flush the contained writer and it's bitstream wrapper.
108    pub fn flush(&mut self) {
109        self.writer.flush_raw()
110    }
111
112    pub fn set_huffman_to_fixed(&mut self) {
113        self.huffman_table.set_to_fixed()
114    }
115
116    /// Reset the encoder state with a new writer, returning the old one if flushing
117    /// succeeds.
118    #[cfg(test)]
119    pub fn reset(&mut self, writer: Vec<u8>) -> Vec<u8> {
120        // Make sure the writer is flushed
121        // Ideally this should be done before this function is called, but we
122        // do it here just in case.
123        self.flush();
124        // Reset the huffman table
125        // This probably isn't needed, but again, we do it just in case to avoid leaking any data
126        // If this turns out to be a performance issue, it can probably be ignored later.
127        self.huffman_table = HuffmanTable::empty();
128        mem::replace(&mut self.writer.w, writer)
129    }
130}