Further refactor lossless encoder (#177)

Author: fintelia

src/encoder.rs

@@ -9,10 +9,9 @@
 #[cfg(all(test, feature = "_benchmarks"))]
 extern crate test;
 
-pub use self::decoder::{
-    DecodingError, LoopCount, UpsamplingMethod, WebPDecodeOptions, WebPDecoder,
-};
-pub use self::encoder::{ColorType, EncoderParams, EncodingError, WebPEncoder};
+pub use decoder::{DecodingError, LoopCount, UpsamplingMethod, WebPDecodeOptions, WebPDecoder};
+pub use encoder::{ColorType, EncodingError, WebPEncoder};
+pub use lossless::EncoderParams;
 
 mod alpha_blending;
 mod decoder;

src/lib.rs

@@ -1,6 +1,6 @@
 use std::io::BufRead;
 
-use super::decoder::BitReader;
+use super::BitReader;
 use crate::decoder::DecodingError;
 
 const MAX_ALLOWED_CODE_LENGTH: usize = 15;

src/lossless/huffman.rs src/lossless/decoder/huffman.rssrc/lossless/huffman.rs renamed to src/lossless/decoder/huffman.rs

@@ -2,39 +2,19 @@
 //!
 //! [Lossless spec](https://developers.google.com/speed/webp/docs/webp_lossless_bitstream_specification)
 
+mod huffman;
+mod reverse_transform;
+
 use std::io::BufRead;
 use std::mem;
 
-use super::huffman::HuffmanTree;
-use super::reverse_transform::{
+use super::{CODE_LENGTH_CODES, CODE_LENGTH_CODE_ORDER, DISTANCE_MAP};
+use crate::decoder::DecodingError;
+use huffman::HuffmanTree;
+use reverse_transform::{
     apply_color_indexing_transform, apply_color_transform, apply_predictor_transform,
     apply_subtract_green_transform, TransformType,
 };
-use crate::decoder::DecodingError;
-
-const CODE_LENGTH_CODES: usize = 19;
-const CODE_LENGTH_CODE_ORDER: [usize; CODE_LENGTH_CODES] = [
-    17, 18, 0, 1, 2, 3, 4, 5, 16, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
-];
-
-#[rustfmt::skip]
-const DISTANCE_MAP: [(i8, i8); 120] = [
-    (0, 1),  (1, 0),  (1, 1),  (-1, 1), (0, 2),  (2, 0),  (1, 2),  (-1, 2),
-    (2, 1),  (-2, 1), (2, 2),  (-2, 2), (0, 3),  (3, 0),  (1, 3),  (-1, 3),
-    (3, 1),  (-3, 1), (2, 3),  (-2, 3), (3, 2),  (-3, 2), (0, 4),  (4, 0),
-    (1, 4),  (-1, 4), (4, 1),  (-4, 1), (3, 3),  (-3, 3), (2, 4),  (-2, 4),
-    (4, 2),  (-4, 2), (0, 5),  (3, 4),  (-3, 4), (4, 3),  (-4, 3), (5, 0),
-    (1, 5),  (-1, 5), (5, 1),  (-5, 1), (2, 5),  (-2, 5), (5, 2),  (-5, 2),
-    (4, 4),  (-4, 4), (3, 5),  (-3, 5), (5, 3),  (-5, 3), (0, 6),  (6, 0),
-    (1, 6),  (-1, 6), (6, 1),  (-6, 1), (2, 6),  (-2, 6), (6, 2),  (-6, 2),
-    (4, 5),  (-4, 5), (5, 4),  (-5, 4), (3, 6),  (-3, 6), (6, 3),  (-6, 3),
-    (0, 7),  (7, 0),  (1, 7),  (-1, 7), (5, 5),  (-5, 5), (7, 1),  (-7, 1),
-    (4, 6),  (-4, 6), (6, 4),  (-6, 4), (2, 7),  (-2, 7), (7, 2),  (-7, 2),
-    (3, 7),  (-3, 7), (7, 3),  (-7, 3), (5, 6),  (-5, 6), (6, 5),  (-6, 5),
-    (8, 0),  (4, 7),  (-4, 7), (7, 4),  (-7, 4), (8, 1),  (8, 2),  (6, 6),
-    (-6, 6), (8, 3),  (5, 7),  (-5, 7), (7, 5),  (-7, 5), (8, 4),  (6, 7),
-    (-6, 7), (7, 6),  (-7, 6), (8, 5),  (7, 7),  (-7, 7), (8, 6),  (8, 7)
-];
 
 const GREEN: usize = 0;
 const RED: usize = 1;

src/lossless/decoder.rs src/lossless/decoder/mod.rssrc/lossless/decoder.rs renamed to src/lossless/decoder/mod.rs

@@ -2,7 +2,7 @@ use std::ops::Range;
 
 use crate::decoder::DecodingError;
 
-use super::decoder::subsample_size;
+use super::subsample_size;
 
 #[derive(Debug, Clone)]
 pub(crate) enum TransformType {

src/lossless/reverse_transform.rs src/lossless/decoder/reverse_transform.rssrc/lossless/reverse_transform.rs renamed to src/lossless/decoder/reverse_transform.rs

@@ -0,0 +1,210 @@
+use std::collections::BinaryHeap;
+use std::io::{self, Write};
+
+use super::BitWriter;
+use crate::lossless::CODE_LENGTH_CODE_ORDER;
+
+pub(crate) fn write_single_entry_huffman_tree<W: Write>(
+    w: &mut BitWriter<W>,
+    symbol: u8,
+) -> io::Result<()> {
+    w.write_bits(1, 2)?;
+    if symbol <= 1 {
+        w.write_bits(0, 1)?;
+        w.write_bits(u64::from(symbol), 1)?;
+    } else {
+        w.write_bits(1, 1)?;
+        w.write_bits(u64::from(symbol), 8)?;
+    }
+    Ok(())
+}
+
+pub(crate) fn write_huffman_tree<W: Write>(
+    w: &mut BitWriter<W>,
+    frequencies: &[u32],
+    lengths: &mut [u8],
+    codes: &mut [u16],
+) -> io::Result<()> {
+    if !build_huffman_tree(frequencies, lengths, codes, 15) {
+        let symbol = frequencies
+            .iter()
+            .position(|&frequency| frequency > 0)
+            .unwrap_or(0);
+        return write_single_entry_huffman_tree(w, symbol as u8);
+    }
+
+    let mut code_length_lengths = [0u8; 16];
+    let mut code_length_codes = [0u16; 16];
+    let mut code_length_frequencies = [0u32; 16];
+    for &length in lengths.iter() {
+        code_length_frequencies[length as usize] += 1;
+    }
+    let single_code_length_length = !build_huffman_tree(
+        &code_length_frequencies,
+        &mut code_length_lengths,
+        &mut code_length_codes,
+        7,
+    );
+
+    // Write the huffman tree
+    w.write_bits(0, 1)?; // normal huffman tree
+    w.write_bits(19 - 4, 4)?; // num_code_lengths - 4
+
+    for i in CODE_LENGTH_CODE_ORDER {
+        if i > 15 || code_length_frequencies[i] == 0 {
+            w.write_bits(0, 3)?;
+        } else if single_code_length_length {
+            w.write_bits(1, 3)?;
+        } else {
+            w.write_bits(u64::from(code_length_lengths[i]), 3)?;
+        }
+    }
+
+    match lengths.len() {
+        256 => {
+            w.write_bits(1, 1)?; // max_symbol is stored
+            w.write_bits(3, 3)?; // max_symbol_nbits / 2 - 2
+            w.write_bits(254, 8)?; // max_symbol - 2
+        }
+        280 => w.write_bits(0, 1)?,
+        _ => unreachable!(),
+    }
+
+    // Write the huffman codes
+    if !single_code_length_length {
+        for &len in lengths.iter() {
+            w.write_bits(
+                u64::from(code_length_codes[len as usize]),
+                code_length_lengths[len as usize],
+            )?;
+        }
+    }
+
+    Ok(())
+}
+
+fn build_huffman_tree(
+    frequencies: &[u32],
+    lengths: &mut [u8],
+    codes: &mut [u16],
+    length_limit: u8,
+) -> bool {
+    assert_eq!(frequencies.len(), lengths.len());
+    assert_eq!(frequencies.len(), codes.len());
+
+    if frequencies.iter().filter(|&&f| f > 0).count() <= 1 {
+        lengths.fill(0);
+        codes.fill(0);
+        return false;
+    }
+
+    #[derive(Eq, PartialEq, Copy, Clone, Debug)]
+    struct Item(u32, u16);
+    impl Ord for Item {
+        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+            other.0.cmp(&self.0)
+        }
+    }
+    impl PartialOrd for Item {
+        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+            Some(self.cmp(other))
+        }
+    }
+
+    // Build a huffman tree
+    let mut internal_nodes = Vec::new();
+    let mut nodes = BinaryHeap::from_iter(
+        frequencies
+            .iter()
+            .enumerate()
+            .filter(|(_, &frequency)| frequency > 0)
+            .map(|(i, &frequency)| Item(frequency, i as u16)),
+    );
+    while nodes.len() > 1 {
+        let Item(frequency1, index1) = nodes.pop().unwrap();
+        let mut root = nodes.peek_mut().unwrap();
+        internal_nodes.push((index1, root.1));
+        *root = Item(
+            frequency1 + root.0,
+            internal_nodes.len() as u16 + frequencies.len() as u16 - 1,
+        );
+    }
+
+    // Walk the tree to assign code lengths
+    lengths.fill(0);
+    let mut stack = Vec::new();
+    stack.push((nodes.pop().unwrap().1, 0));
+    while let Some((node, depth)) = stack.pop() {
+        let node = node as usize;
+        if node < frequencies.len() {
+            lengths[node] = depth as u8;
+        } else {
+            let (left, right) = internal_nodes[node - frequencies.len()];
+            stack.push((left, depth + 1));
+            stack.push((right, depth + 1));
+        }
+    }
+
+    // Limit the codes to length length_limit
+    let mut max_length = 0;
+    for &length in lengths.iter() {
+        max_length = max_length.max(length);
+    }
+    if max_length > length_limit {
+        let mut counts = [0u32; 16];
+        for &length in lengths.iter() {
+            counts[length.min(length_limit) as usize] += 1;
+        }
+
+        let mut total = 0;
+        for (i, count) in counts
+            .iter()
+            .enumerate()
+            .skip(1)
+            .take(length_limit as usize)
+        {
+            total += count << (length_limit as usize - i);
+        }
+
+        while total > 1u32 << length_limit {
+            let mut i = length_limit as usize - 1;
+            while counts[i] == 0 {
+                i -= 1;
+            }
+            counts[i] -= 1;
+            counts[length_limit as usize] -= 1;
+            counts[i + 1] += 2;
+            total -= 1;
+        }
+
+        // assign new lengths
+        let mut len = length_limit;
+        let mut indexes = frequencies.iter().copied().enumerate().collect::<Vec<_>>();
+        indexes.sort_unstable_by_key(|&(_, frequency)| frequency);
+        for &(i, frequency) in &indexes {
+            if frequency > 0 {
+                while counts[len as usize] == 0 {
+                    len -= 1;
+                }
+                lengths[i] = len;
+                counts[len as usize] -= 1;
+            }
+        }
+    }
+
+    // Assign codes
+    codes.fill(0);
+    let mut code = 0u32;
+    for len in 1..=length_limit {
+        for (i, &length) in lengths.iter().enumerate() {
+            if length == len {
+                codes[i] = (code as u16).reverse_bits() >> (16 - len);
+                code += 1;
+            }
+        }
+        code <<= 1;
+    }
+    assert_eq!(code, 2 << length_limit);
+
+    true
+}