1 #[cfg(test)]
2 use std::mem;
3 use huffman_table::HuffmanTable;
4 use bitstream::LsbWriter;
5 use 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`;
10 const FIXED_FIRST_BYTE: u16 = 0b010;
11 const FIXED_FIRST_BYTE_FINAL: u16 = 0b011;
12 const DYNAMIC_FIRST_BYTE: u16 = 0b100;
13 const DYNAMIC_FIRST_BYTE_FINAL: u16 = 0b101;
14 
15 #[allow(dead_code)]
16 pub 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
23 pub struct EncoderState {
24     pub huffman_table: HuffmanTable,
25     pub writer: LsbWriter,
26 }
27 
28 impl EncoderState {
29     /// Creates a new encoder state using the provided huffman table and writer
new(writer: Vec<u8>) -> EncoderState30     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
fixed(writer: Vec<u8>) -> EncoderState39     pub fn fixed(writer: Vec<u8>) -> EncoderState {
40         EncoderState {
41             huffman_table: HuffmanTable::fixed_table(),
42             writer: LsbWriter::new(writer),
43         }
44     }
45 
inner_vec(&mut self) -> &mut Vec<u8>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
write_literal(&mut self, value: u8)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
write_lzvalue(&mut self, value: LZType)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.
write_start_of_block(&mut self, fixed: bool, final_block: bool)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
write_end_of_block(&mut self)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.
flush(&mut self)108     pub fn flush(&mut self) {
109         self.writer.flush_raw()
110     }
111 
set_huffman_to_fixed(&mut self)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)]
reset(&mut self, writer: Vec<u8>) -> Vec<u8>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 }
131