Skip to content

Commit

Permalink
Ensure huffman tree is valid (#48)
Browse files Browse the repository at this point in the history
  • Loading branch information
fintelia authored Jan 2, 2024
1 parent 9aa4bbc commit b27b798
Showing 1 changed file with 11 additions and 6 deletions.
17 changes: 11 additions & 6 deletions src/huffman.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,29 +64,34 @@ impl HuffmanTree {
.iter()
.reduce(|a, b| if a >= b { a } else { b })
.unwrap();

if max_code_length > MAX_ALLOWED_CODE_LENGTH.try_into().unwrap() {
return Err(DecodingError::HuffmanError);
}

// Build histogram
let mut code_length_hist = [0; MAX_ALLOWED_CODE_LENGTH + 1];

for &length in code_lengths.iter() {
code_length_hist[usize::from(length)] += 1;
}

code_length_hist[0] = 0;

// Ensure code lengths produce a valid huffman tree
let mut total = 0;
for (code_len, &count) in code_length_hist.iter().enumerate() {
total += (count as u32) << (MAX_ALLOWED_CODE_LENGTH - code_len);
}
if total != 1 << MAX_ALLOWED_CODE_LENGTH {
return Err(DecodingError::HuffmanError);
}

// Assign codes
let mut curr_code = 0;
let mut next_codes = [None; MAX_ALLOWED_CODE_LENGTH + 1];

for code_len in 1..=usize::from(max_code_length) {
curr_code = (curr_code + code_length_hist[code_len - 1]) << 1;
next_codes[code_len] = Some(curr_code);
}

let mut huff_codes = vec![None; code_lengths.len()];

for (symbol, &length) in code_lengths.iter().enumerate() {
let length = usize::from(length);
if length > 0 {
Expand Down

0 comments on commit b27b798

Please sign in to comment.