Rust 实现哈夫曼编码解码压缩算法

哈夫曼编码（Huffman Coding）是一种树形的编码格式，经常用于数据压缩。网上很多的有很多种 C / C++ / Java 实现，本文以 Rust 作为视角来看看如何利用哈夫曼编码实现文件的压缩与解压。

代码已同步至 github：github.com/oloshe/rust…

压缩

计算字符权重 CharWeightMap

计算每个字符出现的次数作为权重

type Weight = u64;
/// 字符权重
pub struct CharWeightMap {
    pub inner: HashMap<char, Weight>
}

impl CharWeightMap {
    pub fn build(input: &String) -> Self {
        let mut map = HashMap::new();
        for (_, c) in input.char_indices() {
            map.entry(c).or_insert(0).add_assign(1);
        }
        Self { inner: map }
    }
    pub fn len(&self) -> usize {
        self.inner.len()
    }
    pub fn iter(&self) -> Iter<char, Weight> {
        self.inner.iter()  
    }
}

构建哈夫曼树 HuffmanTree

type RefHuffmanTree = Rc<RefCell<HuffmanTree>>;
type Weight = u64;

/// 哈夫曼树
pub struct HuffmanTree {
    pub value: Option<char>,
    pub weight: Weight,
    pub parent: Option<RefHuffmanTree>,
    pub left: Option<RefHuffmanTree>,
    pub right: Option<RefHuffmanTree>,
}

impl HuffmanTree {
    pub fn new() -> Self {
        Self {
            value: None,
            weight: 0,
            parent: None,
            left: None,
            right: None,
        }
    }

    pub fn build(char_weight: CharWeightMap) -> RefHuffmanTree
    {
        // 原始结点数量
        let n = char_weight.len();
        // 构建完整哈夫曼树总共需要的结点数量
        let total = 2 * n - 1;
        // 初始化所有结点
        let vec = (0..total)
            .map(|_| Rc::new(RefCell::new(Self::new())))
            .collect::<Vec<Rc<RefCell<HuffmanTree>>>>();

        // 字符结点赋值
        char_weight.iter()
            .enumerate()
            .into_iter()
            .for_each(|(index, (ch, weight))| {
                // println!("{}: {} ({})", index, &weight, ch);
                vec[index].borrow_mut().value = Some(*ch);
                vec[index].borrow_mut().weight = *weight;
            });

        for index in n..total {
            // 找到 [0, index-1] 中权重最小的结点
            let m1 = Self::find_min(&vec[..index]).unwrap();
            // 标记父结点为 index 上的结点，下次就不会找到这个
            m1.borrow_mut().parent = Some(vec[index].clone());
            // 找到 [0, index-1] 中权重第二小的结点
            let m2 = Self::find_min(&vec[..index]).unwrap();
            // 标记该结点的父结点为 index 上的结点。
            m2.borrow_mut().parent = Some(vec[index].clone());

            let w1 = m1.as_ref().borrow().weight;
            let w2 = m2.as_ref().borrow().weight;
            let weight = w1 + w2;

            vec[index].borrow_mut().weight = weight;
            vec[index].borrow_mut().left = Some(m1.clone());
            vec[index].borrow_mut().right = Some(m2.clone());
        }
        // 最后一个结点即为构建好的完整哈夫曼树
        vec.last().unwrap().clone()
    }

    /// 获取最小的值
    fn find_min(tree_slice: &[Rc<RefCell<HuffmanTree>>]) -> Option<Rc<RefCell<HuffmanTree>>> {
        let mut min = Weight::MAX;
        let mut result = None;
        for tree in tree_slice {
            let tree_cell = tree.as_ref();
            if tree_cell.borrow().parent.is_none() && tree_cell.borrow().weight < min {
                min = tree_cell.borrow().weight;
                result = Some(tree.clone());
            }
        }
        result
    }
}

二进制映射 HuffmanBinaryMap

用 Vec<bool> 来表示二进制。tree_dfs 为深度优先遍历哈夫曼树。

/// 字符二进制映射，表示字符对应的二进制位，可用 bitvec 替代
pub struct HuffmanBinaryMap {
    pub inner: HashMap<char, Vec<bool>>
}

impl HuffmanBinaryMap {
    pub fn build(huffman_tree: RefHuffmanTree) -> Self {
        let mut map = HashMap::new();
        Self::tree_dfs(&Some(huffman_tree), &mut map, &mut vec![]);
        Self { inner: map }
    }
    fn tree_dfs(
        tree: &Option<RefHuffmanTree>, 
        map: &mut HashMap<char, Vec<bool>>,
        vec: &mut Vec<bool>
    ) {
        if let Some(tree) = tree {
            let tree = tree.as_ref().borrow();
            if let Some(ch) = tree.value {
                map.insert(ch, vec.clone());
            }
            vec.push(false);
            Self::tree_dfs(&tree.left, map, vec);
            let last = vec.last_mut().unwrap();
            *last = true;
            Self::tree_dfs(&tree.right, map, vec);
            vec.pop();
        }
    }
}

/// 用于写入配置文件
impl Display for HuffmanBinaryMap {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut buf = String::new();
        self.inner.iter()
            .for_each(|(c, vec)| {
                let mut bit_str = String::new();
                vec.iter().for_each(|b| {
                    bit_str += if *b { "1" } else { "0" }
                });
                buf += format!("{}:{}\n", *c as u32, bit_str).as_str();
            });
        f.write_str(buf.as_str())
    }
}

space 为末尾补齐八位二进制数的数量；capacity 用于解码时优化内存分配。

余下部分冒号左边为字符对应的 u32码 ，右边为它对应的为二进制数。

（为什么用u32 ？如果用字符表示会有特殊字符问题，比如回车和冒号）

bit_map 用了上面实现的 Display trait 。

编码 HuffmanCodec::encode

编码返回了一个元组，分别为压缩后的字节数组和配置文件内容。配置文件描述了如何解压。

pub struct HuffmanCodec;

impl HuffmanCodec {
    /// 哈夫曼编码
    pub fn encode(source: &String) -> (Vec<u8>, String) {
        // 构建字符权重映射
        let weight_map = CharWeightMap::build(&source);
        // 构建哈夫曼树
        let tree = HuffmanTree::build(weight_map);
        // 哈夫曼二进制映射表
        let bit_map = HuffmanBinaryMap::build(tree);
        // println!("{}", bit_map);
        let mut result: Vec<u8> = vec![];
        let (mut buf, mut count) = (0, 0);
        for (_, ch) in source.char_indices() {
            let vec = bit_map.inner.get(&ch).unwrap();
            vec.iter().for_each(|b| {
                buf <<= 1;
                if *b { buf |= 1 }
                count += 1;
                if count >= 8 {
                    result.push(buf);
                    buf = 0;
                    count = 0;
                }
            })
        }
        // 末尾补位数量
        let mut space = 0u8;
        if count != 0 {
            space = 8 - count;
            buf <<= space;
            result.push(buf);
        }
        // 返回的结果
        (
            result, // 压缩后的字节数组
            format!("space:{}\ncapacity:{}\n{}", space, source.len(), bit_map), // 配置文件内容
        )
    }
}

可以思考一下为什么要补齐 8 位

压缩文件

fn hfm_compress(file: &str) {
    let filename = {
        let mut arr = file.split(".").collect::<Vec<&str>>();
        arr.pop();
        arr.join(".")
    };
    // 打开文件并读取字符串到内存中
    let mut input_file = File::open(file).expect("未找到该文件");
    let mut source_buf = String::new();
    input_file.read_to_string(&mut source_buf).expect("读取文件失败");

    // 哈夫曼编码
    let (u8_arr, config) = HuffmanCodec::encode(&source_buf);

    // 创建压缩文件
    let output_file_name = format!("{}.hfm", filename);
    let mut output_file = File::create(&output_file_name).unwrap();
    output_file.write(&u8_arr).unwrap();

    // 创建压缩配置文件
    let output_cfg_file_name = format!("{}.hfm.config", filename);
    let mut output_cfg_file = File::create(&output_cfg_file_name).unwrap();
    output_cfg_file.write(config.as_bytes()).unwrap();
    println!("\n压缩成功！\n文件保存为: {}\n配置文件: {}", output_file_name, output_cfg_file_name);
    let size_before = source_buf.as_bytes().len();
    let size_after = u8_arr.len();
    println!("压缩前大小：{} 字节", size_before);
    println!("压缩后大小：{} 字节", size_after);
    println!("压缩比率： {:.2}%", ((size_after as f64 / size_before as f64) * 100.0));
}

解压

解压配置 DecodeConfig

配置文件的读取跟前面写入配置文件内容相关，可以根据自己的结构去定义。

/// 配置文件的配置
pub struct DecodeConfig {
    pub inner: HashMap<String, char>,
    pub space: u8,
    pub capacity: usize,
}
impl DecodeConfig {
    /// 从配置文件读取 
    pub fn build(source: &String) -> Self {
        let mut map = HashMap::default();
        let (mut space, mut capacity) = (0u8, 0usize);
        let arr = source.split("\n");
        for s in arr {
            let pair: Vec<&str> = s.split(":").collect();
            if pair.len() != 2 {
                continue;
            }
            let (ch, bit) = (pair[0], pair[1]);
            match ch {
                "space" => {
                    space = u8::from_str_radix(bit, 10).unwrap();
                    continue;
                },
                "capacity" => {
                    capacity = usize::from_str_radix(bit, 10).unwrap();
                    continue;
                },
                _ => (),
            }
            map.insert(bit.to_owned(), char::from_u32(u32::from_str_radix(ch, 10).unwrap()).unwrap());
        };
        Self { inner: map, space, capacity }
    }
    pub fn get(&self, k: &String) -> Option<&char> {
        self.inner.get(k)
    }
}

解码 HuffmanCodec::decode

下面的 format!("{u8:>0width$b}", u8=num, width=8) 表示 以二进制形式输出并以 0 来补齐8位 。

pub struct HuffmanCodec;

impl HuffmanCodec {
    pub fn decode(source: &[u8], decode_map: &DecodeConfig) -> String {
        // 防止内存频繁分配，直接定义容量
        let mut result = String::with_capacity(decode_map.capacity);
        let bit_str = source.iter()
            .map(|num| {
                format!("{u8:>0width$b}", u8=num, width=8)
            })
            .collect::<Vec<String>>()
            .join("");
        // println!("二进制序列：{}", bit_str);

        let mut tmp_str = String::with_capacity(20);
        let last_idx = bit_str.len() - decode_map.space as usize;
        for (i, ch) in bit_str.char_indices() {
            if i >= last_idx {
                break;
            }
            tmp_str.push(ch);
            if let Some(mch) = decode_map.get(&tmp_str) {
                result.push(*mch);
                tmp_str.clear();
            }
        }
        result
    }
}

解压文件

fn hfm_decompress(file: &str, config_file: &str, save_file: &str) {
    // 压缩文件
    let mut encodede_file = File::open(file).expect(format!("未找到文件:{}",file).as_str());
    let mut buf = vec![];
    encodede_file.read_to_end(&mut buf).unwrap();

    // 读取配置文件
    let mut config = File::open(config_file).expect(format!("未找到配置文件:{}", config_file).as_str());
    let mut buf2 = String::new();
    config.read_to_string(&mut buf2).unwrap();
    
    // 构建配置
    let char_map = DecodeConfig::build(&buf2);

    // 解码
    let result = HuffmanCodec::decode(&buf, &char_map);

    let mut savef = File::create(save_file).unwrap();
    savef.write(result.as_bytes()).unwrap();

    println!("\n解压成功！\n文件已保存至：{}", save_file);
}

总结

可以看到 Rust 实现相比其他语言实现还是有明显的差距。总体来说逻辑和结构更加清晰，这也是我认为 Rust 相比于 C 的更优的地方之一。

本篇文章完整地过了一遍哈夫曼压缩解压的过程，不足之处还望多多指教和交流。具体到细节有很多实现方法是可以按照自己的想法去实现的，本文仅作参考。

代码已经同步推送至 github，感兴趣的同学欢迎 star & 点赞。

github.com/oloshe/rust…