1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use crate::bitstream::LsbWriter;
use crate::huffman_table::HuffmanTable;
use crate::lzvalue::LZType;
#[cfg(test)]
use std::mem;

// The first bits of each block, which describe the type of the block
// `-TTF` - TT = type, 00 = stored, 01 = fixed, 10 = dynamic, 11 = reserved, F - 1 if final block
// `0000`;
const FIXED_FIRST_BYTE: u16 = 0b010;
const FIXED_FIRST_BYTE_FINAL: u16 = 0b011;
const DYNAMIC_FIRST_BYTE: u16 = 0b100;
const DYNAMIC_FIRST_BYTE_FINAL: u16 = 0b101;

#[allow(dead_code)]
pub enum BType {
    NoCompression = 0b00,
    FixedHuffman = 0b01,
    DynamicHuffman = 0b10, // Reserved = 0b11, //Error
}

/// A struct wrapping a writer that writes data compressed using the provided Huffman table
pub struct EncoderState {
    pub huffman_table: HuffmanTable,
    pub writer: LsbWriter,
}

impl EncoderState {
    /// Creates a new encoder state using the provided Huffman table and writer
    pub fn new(writer: Vec<u8>) -> EncoderState {
        EncoderState {
            huffman_table: HuffmanTable::empty(),
            writer: LsbWriter::new(writer),
        }
    }

    #[cfg(test)]
    /// Creates a new encoder state using the fixed Huffman table
    pub fn fixed(writer: Vec<u8>) -> EncoderState {
        EncoderState {
            huffman_table: HuffmanTable::fixed_table(),
            writer: LsbWriter::new(writer),
        }
    }

    pub fn inner_vec(&mut self) -> &mut Vec<u8> {
        &mut self.writer.w
    }

    /// Encodes a literal value to the writer
    fn write_literal(&mut self, value: u8) {
        let code = self.huffman_table.get_literal(value);
        debug_assert!(code.length > 0);
        self.writer.write_bits(code.code, code.length);
    }

    /// Write a LZvalue to the contained writer, returning Err if the write operation fails
    pub fn write_lzvalue(&mut self, value: LZType) {
        match value {
            LZType::Literal(l) => self.write_literal(l),
            LZType::StoredLengthDistance(l, d) => {
                let (code, extra_bits_code) = self.huffman_table.get_length_huffman(l);
                debug_assert!(
                    code.length != 0,
                    format!("Code: {:?}, Value: {:?}", code, value)
                );
                self.writer.write_bits(code.code, code.length);
                self.writer
                    .write_bits(extra_bits_code.code, extra_bits_code.length);

                let (code, extra_bits_code) = self.huffman_table.get_distance_huffman(d);
                debug_assert!(
                    code.length != 0,
                    format!("Code: {:?}, Value: {:?}", code, value)
                );

                self.writer.write_bits(code.code, code.length);
                self.writer
                    .write_bits(extra_bits_code.code, extra_bits_code.length)
            }
        };
    }

    /// Write the start of a block, returning Err if the write operation fails.
    pub fn write_start_of_block(&mut self, fixed: bool, final_block: bool) {
        if final_block {
            // The final block has one bit flipped to indicate it's
            // the final one
            if fixed {
                self.writer.write_bits(FIXED_FIRST_BYTE_FINAL, 3)
            } else {
                self.writer.write_bits(DYNAMIC_FIRST_BYTE_FINAL, 3)
            }
        } else if fixed {
            self.writer.write_bits(FIXED_FIRST_BYTE, 3)
        } else {
            self.writer.write_bits(DYNAMIC_FIRST_BYTE, 3)
        }
    }

    /// Write the end of block code
    pub fn write_end_of_block(&mut self) {
        let code = self.huffman_table.get_end_of_block();
        self.writer.write_bits(code.code, code.length)
    }

    /// Flush the contained writer and it's bitstream wrapper.
    pub fn flush(&mut self) {
        self.writer.flush_raw()
    }

    pub fn set_huffman_to_fixed(&mut self) {
        self.huffman_table.set_to_fixed()
    }

    /// Reset the encoder state with a new writer, returning the old one if flushing
    /// succeeds.
    #[cfg(test)]
    pub fn reset(&mut self, writer: Vec<u8>) -> Vec<u8> {
        // Make sure the writer is flushed
        // Ideally this should be done before this function is called, but we
        // do it here just in case.
        self.flush();
        // Reset the huffman table
        // This probably isn't needed, but again, we do it just in case to avoid leaking any data
        // If this turns out to be a performance issue, it can probably be ignored later.
        self.huffman_table = HuffmanTable::empty();
        mem::replace(&mut self.writer.w, writer)
    }
}