1use super::lossless::BitReader;
2use super::lossless::DecoderError;
3use crate::ImageResult;
4
5/// Rudimentary utility for reading Canonical Huffman Codes.
6/// Based off https://github.com/webmproject/libwebp/blob/7f8472a610b61ec780ef0a8873cd954ac512a505/src/utils/huffman.c
7///
8
9const MAX_ALLOWED_CODE_LENGTH: usize = 15;
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12enum HuffmanTreeNode {
13 Branch(usize), //offset in vector to children
14 Leaf(u16), //symbol stored in leaf
15 Empty,
16}
17
18/// Huffman tree
19#[derive(Clone, Debug, Default)]
20pub(crate) struct HuffmanTree {
21 tree: Vec<HuffmanTreeNode>,
22 max_nodes: usize,
23 num_nodes: usize,
24}
25
26impl HuffmanTree {
27 fn is_full(&self) -> bool {
28 self.num_nodes == self.max_nodes
29 }
30
31 /// Turns a node from empty into a branch and assigns its children
32 fn assign_children(&mut self, node_index: usize) -> usize {
33 let offset_index = self.num_nodes - node_index;
34 self.tree[node_index] = HuffmanTreeNode::Branch(offset_index);
35 self.num_nodes += 2;
36
37 offset_index
38 }
39
40 /// Init a huffman tree
41 fn init(num_leaves: usize) -> ImageResult<HuffmanTree> {
42 if num_leaves == 0 {
43 return Err(DecoderError::HuffmanError.into());
44 }
45
46 let max_nodes = 2 * num_leaves - 1;
47 let tree = vec![HuffmanTreeNode::Empty; max_nodes];
48 let num_nodes = 1;
49
50 let tree = HuffmanTree {
51 tree,
52 max_nodes,
53 num_nodes,
54 };
55
56 Ok(tree)
57 }
58
59 /// Converts code lengths to codes
60 fn code_lengths_to_codes(code_lengths: &[u16]) -> ImageResult<Vec<Option<u16>>> {
61 let max_code_length = *code_lengths
62 .iter()
63 .reduce(|a, b| if a >= b { a } else { b })
64 .unwrap();
65
66 if max_code_length > MAX_ALLOWED_CODE_LENGTH.try_into().unwrap() {
67 return Err(DecoderError::HuffmanError.into());
68 }
69
70 let mut code_length_hist = [0; MAX_ALLOWED_CODE_LENGTH + 1];
71
72 for &length in code_lengths.iter() {
73 code_length_hist[usize::from(length)] += 1;
74 }
75
76 code_length_hist[0] = 0;
77
78 let mut curr_code = 0;
79 let mut next_codes = [None; MAX_ALLOWED_CODE_LENGTH + 1];
80
81 for code_len in 1..=usize::from(max_code_length) {
82 curr_code = (curr_code + code_length_hist[code_len - 1]) << 1;
83 next_codes[code_len] = Some(curr_code);
84 }
85
86 let mut huff_codes = vec![None; code_lengths.len()];
87
88 for (symbol, &length) in code_lengths.iter().enumerate() {
89 let length = usize::from(length);
90 if length > 0 {
91 huff_codes[symbol] = next_codes[length];
92 if let Some(value) = next_codes[length].as_mut() {
93 *value += 1;
94 }
95 } else {
96 huff_codes[symbol] = None;
97 }
98 }
99
100 Ok(huff_codes)
101 }
102
103 /// Adds a symbol to a huffman tree
104 fn add_symbol(&mut self, symbol: u16, code: u16, code_length: u16) -> ImageResult<()> {
105 let mut node_index = 0;
106 let code = usize::from(code);
107
108 for length in (0..code_length).rev() {
109 if node_index >= self.max_nodes {
110 return Err(DecoderError::HuffmanError.into());
111 }
112
113 let node = self.tree[node_index];
114
115 let offset = match node {
116 HuffmanTreeNode::Empty => {
117 if self.is_full() {
118 return Err(DecoderError::HuffmanError.into());
119 }
120 self.assign_children(node_index)
121 }
122 HuffmanTreeNode::Leaf(_) => return Err(DecoderError::HuffmanError.into()),
123 HuffmanTreeNode::Branch(offset) => offset,
124 };
125
126 node_index += offset + ((code >> length) & 1);
127 }
128
129 match self.tree[node_index] {
130 HuffmanTreeNode::Empty => self.tree[node_index] = HuffmanTreeNode::Leaf(symbol),
131 HuffmanTreeNode::Leaf(_) => return Err(DecoderError::HuffmanError.into()),
132 HuffmanTreeNode::Branch(_offset) => return Err(DecoderError::HuffmanError.into()),
133 }
134
135 Ok(())
136 }
137
138 /// Builds a tree implicitly, just from code lengths
139 pub(crate) fn build_implicit(code_lengths: Vec<u16>) -> ImageResult<HuffmanTree> {
140 let mut num_symbols = 0;
141 let mut root_symbol = 0;
142
143 for (symbol, length) in code_lengths.iter().enumerate() {
144 if *length > 0 {
145 num_symbols += 1;
146 root_symbol = symbol.try_into().unwrap();
147 }
148 }
149
150 let mut tree = HuffmanTree::init(num_symbols)?;
151
152 if num_symbols == 1 {
153 tree.add_symbol(root_symbol, 0, 0)?;
154 } else {
155 let codes = HuffmanTree::code_lengths_to_codes(&code_lengths)?;
156
157 for (symbol, &length) in code_lengths.iter().enumerate() {
158 if length > 0 && codes[symbol].is_some() {
159 tree.add_symbol(symbol.try_into().unwrap(), codes[symbol].unwrap(), length)?;
160 }
161 }
162 }
163
164 Ok(tree)
165 }
166
167 /// Builds a tree explicitly from lengths, codes and symbols
168 pub(crate) fn build_explicit(
169 code_lengths: Vec<u16>,
170 codes: Vec<u16>,
171 symbols: Vec<u16>,
172 ) -> ImageResult<HuffmanTree> {
173 let mut tree = HuffmanTree::init(symbols.len())?;
174
175 for i in 0..symbols.len() {
176 tree.add_symbol(symbols[i], codes[i], code_lengths[i])?;
177 }
178
179 Ok(tree)
180 }
181
182 /// Reads a symbol using the bitstream
183 pub(crate) fn read_symbol(&self, bit_reader: &mut BitReader) -> ImageResult<u16> {
184 let mut index = 0;
185 let mut node = self.tree[index];
186
187 while let HuffmanTreeNode::Branch(children_offset) = node {
188 index += children_offset + bit_reader.read_bits::<usize>(1)?;
189 node = self.tree[index];
190 }
191
192 let symbol = match node {
193 HuffmanTreeNode::Branch(_) => unreachable!(),
194 HuffmanTreeNode::Empty => return Err(DecoderError::HuffmanError.into()),
195 HuffmanTreeNode::Leaf(symbol) => symbol,
196 };
197
198 Ok(symbol)
199 }
200}
201