From b27b798bc668c7a80acbdbe428871aa1d4d9a040 Mon Sep 17 00:00:00 2001
From: Jonathan Behrens <fintelia@gmail.com>
Date: Mon, 1 Jan 2024 23:27:58 -0500
Subject: [PATCH] Ensure huffman tree is valid (#48)

---
 src/huffman.rs | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/huffman.rs b/src/huffman.rs
index 3f59b76..0fe82d5 100644
--- a/src/huffman.rs
+++ b/src/huffman.rs
@@ -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 {