core/compr: make code length limiting in deflate actually work

[nihav.git] / nihav-core / src / compr / deflate.rs

diff --git a/nihav-core/src/compr/deflate.rs b/nihav-core/src/compr/deflate.rs
index 3932473959867d2128ec018b1ac40407c28d7d9b..09e9d8ffb560f2c8dbe504889b5a17f405e10558 100644 (file)
--- a/nihav-core/src/compr/deflate.rs
+++ b/nihav-core/src/compr/deflate.rs
@@ -790,6 +790,7 @@ impl Inflate {
         self.state = InflateState::Start;
     }
 
+    #[allow(clippy::comparison_chain)]
     ///! Decompresses input data into output returning the uncompressed data length.
     pub fn uncompress(src: &[u8], dst: &mut [u8]) -> DecompressResult<usize> {
         let mut csrc = CurrentSource::new(src, BitReaderState::default());
@@ -886,13 +887,11 @@ impl Inflate {
                     }
                     let hdist = csrc.read(5)? as usize + 1;
                     let hclen = csrc.read(4)? as usize + 4;
-                    let mut cur_len_idx = 0;
                     let mut len_lengths = [0; 19];
                     let mut all_lengths = [0; NUM_LITERALS + NUM_DISTS];
 
-                    for _ in 0..hclen {
+                    for cur_len_idx in 0..hclen {
                         len_lengths[LEN_RECODE[cur_len_idx]] = csrc.read(3)? as u8;
-                        cur_len_idx += 1;
                     }
                     let mut len_codes = [ShortCodebookDesc { code: 0, bits: 0 }; 19];
                     lengths_to_codes(&len_lengths, &mut len_codes)?;
@@ -1584,25 +1583,25 @@ fn gen_tree(codes: &mut [u16], lens: &mut [u8], num_codes: &mut usize, stats: &m
         *num_codes = 0;
         return;
     }
-    while tot_w > (1 << max_bits) {
-        for w in stats.iter_mut() {
-            *w = (*w + 1) >> 1;
-        }
-        tot_w = 0;
-        for &w in stats.iter() {
-            tot_w += w;
+    loop {
+        let mut tree = Tree::new();
+        for (sym, &w) in stats.iter().enumerate() {
+            tree.insert(Node{ sym: sym as u16, w: w as u16, idx0: 64000, idx1: 64000 });
         }
-    }
-    let mut tree = Tree::new();
-    for (sym, &w) in stats.iter().enumerate() {
-        tree.insert(Node{ sym: sym as u16, w: w as u16, idx0: 64000, idx1: 64000 });
-    }
-    tree.trim();
-    tree.build();
+        tree.trim();
+        tree.build();
 
-    for n in tree.nodes[..tree.nnodes].iter() {
-        if n.sym != NODE_SYM {
-            lens[n.sym as usize] = n.w as u8;
+        for n in tree.nodes[..tree.nnodes].iter() {
+            if n.sym != NODE_SYM {
+                lens[n.sym as usize] = n.w as u8;
+            }
+        }
+        if !lens.iter().any(|&x| x > max_bits) {
+            break;
+        } else {
+            for w in stats.iter_mut() {
+                *w = (*w + 1) >> 1;
+            }
         }
     }
     lengths_to_codes16(lens, codes);