Initial Commit

src/compression.rs

@@ -0,0 +1,166 @@
+use bit_vec::BitVec;
+use rayon::prelude::*;
+use rmp_serde;
+use serde::{Deserialize, Serialize};
+use std::{collections::HashMap, hash::Hash};
+
+use crate::huffman::{self, Tree};
+use Tree::*;
+
+#[derive(serde::Serialize, serde::Deserialize)]
+struct CompressedData<T: Eq + Hash> {
+    encoder: HashMap<T, BitVec>,
+    data: Vec<BitVec>,
+}
+
+pub fn compress<'a, T, FreqsF, TokenExtractor, TokensIter>(
+    lines: &'a [String],
+    get_freqs: FreqsF,
+    line_to_tokens: TokenExtractor,
+) -> Result<Vec<u8>, Box<dyn std::error::Error>>
+where
+    T: Clone + Eq + Hash + Send + Sync + Serialize,
+    FreqsF: Fn(&'a [String]) -> HashMap<T, u64>,
+    TokenExtractor: Fn(&'a str) -> TokensIter + Send + Sync,
+    TokensIter: Iterator<Item = T>,
+{
+    let freqs = get_freqs(lines);
+    let tree = huffman::huffman_tree(&freqs);
+    let encoder = tree.to_encoder();
+
+    let data = lines
+        .par_iter()
+        .map(|line| {
+            let mut bits = BitVec::new();
+            for token in line_to_tokens(line) {
+                bits.extend(encoder.get(&token).unwrap().iter());
+            }
+            bits
+        })
+        .collect();
+
+    let compressed_data = CompressedData { encoder, data };
+    rmp_serde::encode::to_vec(&compressed_data).map_err(|err| err.into())
+}
+
+pub fn extract<'a, T, F>(
+    data: &'a [u8],
+    tokens_to_line: F,
+) -> Result<Vec<String>, Box<dyn std::error::Error>>
+where
+    T: Clone + Eq + Hash + Send + Sync + Deserialize<'a>,
+    F: Fn(Vec<T>) -> String + Send + Sync,
+{
+    let CompressedData { encoder, data }: CompressedData<T> = rmp_serde::decode::from_slice(data)?;
+
+    let decoder = build_decoder(&encoder);
+    let lines = data
+        .par_iter()
+        .map(|line| {
+            let mut tokens = Vec::new();
+            let mut node = &decoder;
+
+            for bit in line.iter() {
+                node = if bit {
+                    node.right.as_ref().expect("invalid Huffman bit sequence")
+                } else {
+                    node.left.as_ref().expect("invalid Huffman bit sequence")
+                };
+
+                if let Some(token) = &node.token {
+                    tokens.push(token.clone());
+                    node = &decoder;
+                }
+            }
+            tokens_to_line(tokens)
+        })
+        .collect();
+
+    Ok(lines)
+}
+
+struct DecoderNode<T> {
+    token: Option<T>,
+    left: Option<Box<DecoderNode<T>>>,
+    right: Option<Box<DecoderNode<T>>>,
+}
+
+impl<T> Default for DecoderNode<T> {
+    fn default() -> Self {
+        DecoderNode {
+            token: None,
+            left: None,
+            right: None,
+        }
+    }
+}
+
+fn build_decoder<T: Clone + Eq + Hash>(encoder: &HashMap<T, BitVec>) -> DecoderNode<T> {
+    let mut root = DecoderNode::default();
+
+    for (token, code) in encoder {
+        let mut node = &mut root;
+        for bit in code.iter() {
+            node = if bit {
+                node.right.get_or_insert_with(|| Box::new(DecoderNode::default()))
+            } else {
+                node.left.get_or_insert_with(|| Box::new(DecoderNode::default()))
+            };
+        }
+        node.token = Some(token.clone());
+    }
+
+    root
+}
+
+impl<T: Eq + Clone + Hash> Tree<T> {
+    pub fn to_encoder(&self) -> HashMap<T, BitVec> {
+        let mut encoder = HashMap::new();
+
+        let mut stack = vec![(self, BitVec::new())];
+        while let Some((node, path)) = stack.pop() {
+            match node {
+                Leaf { token, .. } => {
+                    encoder.insert(token.clone(), path.clone());
+                }
+                Node { left, right, .. } => {
+                    let mut left_path = path.clone();
+                    left_path.push(false);
+                    stack.push((left, left_path));
+
+                    let mut right_path = path.clone();
+                    right_path.push(true);
+                    stack.push((right, right_path));
+                }
+            }
+        }
+
+        encoder
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::freq::{char_frequencies, word_frequencies};
+
+    #[test]
+    fn compress_decompress_test() {
+        let lines = vec![
+            "hey there! nice to meet you.".to_string(),
+            "Serde is a framework for serializing and deserializing Rust data structures"
+                .to_string(),
+        ];
+
+        let data = compress(&lines, char_frequencies, |line| line.chars()).unwrap();
+        let res_lines = extract(&data, |x: Vec<char>| x.into_iter().collect()).unwrap();
+        assert_eq!(&lines, &res_lines);
+
+        let data = compress(&lines, word_frequencies, |line| {
+            line.split_ascii_whitespace().map(|token| token.to_string())
+        })
+        .unwrap();
+        let res_lines = extract(&data, |x: Vec<String>| x.join(" ")).unwrap();
+        assert_eq!(&lines, &res_lines);
+    }
+}

src/freq.rs

@@ -0,0 +1,48 @@
+use rayon::prelude::*;
+use std::collections::HashMap;
+
+pub fn char_frequencies(lines: &[String]) -> HashMap<char, u64> {
+    lines
+        .par_iter()
+        .fold(
+            HashMap::new,
+            |mut freqs: HashMap<_, _>, line: &String| {
+                for ch in line.chars() {
+                    *freqs.entry(ch).or_insert(0) += 1;
+                }
+                freqs
+            },
+        )
+        .reduce(
+            HashMap::new,
+            |mut freqs1, freqs2| {
+                freqs2
+                    .into_iter()
+                    .for_each(|(ch, n)| *freqs1.entry(ch).or_insert(0) += n);
+                freqs1
+            },
+        )
+}
+
+pub fn word_frequencies(lines: &[String]) -> HashMap<String, u64> {
+    lines
+        .par_iter()
+        .fold(
+            HashMap::new,
+            |mut freqs: HashMap<_, _>, line: &String| {
+                for word in line.split_ascii_whitespace() {
+                    *freqs.entry(word.to_string()).or_insert(0) += 1;
+                }
+                freqs
+            },
+        )
+        .reduce(
+            HashMap::new,
+            |mut freqs1, freqs2| {
+                freqs2
+                    .into_iter()
+                    .for_each(|(word, n)| *freqs1.entry(word).or_insert(0) += n);
+                freqs1
+            },
+        )
+}

src/huffman.rs

@@ -0,0 +1,158 @@
+use std::{
+    cmp::Reverse,
+    collections::{BinaryHeap, HashMap},
+};
+use Tree::*;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum Tree<T> {
+    Leaf {
+        freq: u64,
+        token: T,
+    },
+    Node {
+        freq: u64,
+        left: Box<Tree<T>>,
+        right: Box<Tree<T>>,
+    },
+}
+
+#[allow(dead_code)]
+impl<T: Clone> Tree<T> {
+    pub fn freq(&self) -> u64 {
+        match self {
+            Leaf { freq, .. } => *freq,
+            Node { freq, .. } => *freq,
+        }
+    }
+
+    pub fn token(&self) -> Option<T> {
+        match self {
+            Leaf { token, .. } => Some(token.clone()),
+            Node { .. } => None,
+        }
+    }
+
+    pub fn left(&self) -> Option<&Tree<T>> {
+        match self {
+            Node { left, .. } => Some(left),
+            Leaf { .. } => None,
+        }
+    }
+
+    pub fn right(&self) -> Option<&Tree<T>> {
+        match self {
+            Node { right, .. } => Some(right),
+            Leaf { .. } => None,
+        }
+    }
+}
+
+impl<T: Clone + Eq> Ord for Tree<T> {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.freq().cmp(&other.freq())
+    }
+}
+
+impl<T: Clone + Eq> PartialOrd for Tree<T> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+pub fn huffman_tree<T: Eq + Clone>(freqs: &HashMap<T, u64>) -> Tree<T> {
+    let mut heap = BinaryHeap::new();
+    for (token, freq) in freqs {
+        let (freq, token) = (*freq, token.clone());
+        heap.push(Reverse(Leaf { freq, token }))
+    }
+
+    while heap.len() > 1 {
+        let node1 = heap.pop().unwrap().0;
+        let node2 = heap.pop().unwrap().0;
+
+        let merged_node = Node {
+            freq: node1.freq() + node2.freq(),
+            left: Box::new(node1),
+            right: Box::new(node2),
+        };
+        heap.push(Reverse(merged_node));
+    }
+
+    heap.pop().unwrap().0
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::freq::char_frequencies;
+
+    #[test]
+    fn learn_frequencies_test() {
+        let input = vec!["this is an epic sentence".to_string(), "xyz ".to_string()];
+        let freqs = char_frequencies(&input);
+        assert_eq!(freqs[&' '], 5);
+        assert_eq!(freqs[&'t'], 2);
+        assert_eq!(freqs[&'i'], 3);
+        assert_eq!(freqs[&'p'], 1);
+        assert_eq!(freqs[&'z'], 1);
+        assert_eq!(freqs.keys().len(), 13);
+    }
+
+    #[test]
+    fn huffman_tree_test() {
+        let mut freqs = HashMap::new();
+        freqs.insert('a', 40);
+        freqs.insert('b', 35);
+        freqs.insert('c', 20);
+        freqs.insert('d', 5);
+
+        let tree = huffman_tree(&freqs);
+        assert_eq!(tree.freq(), 100);
+
+        // the most frequent character only requires 1 bit
+        assert_eq!(tree.left().and_then(|n| n.token()), Some('a'));
+        assert_eq!(tree.left().map(|n| n.freq()), Some(40));
+
+        // the second most frequent character requires 2 bits
+        assert_eq!(
+            tree.right().and_then(|t| t.right()).and_then(|n| n.token()),
+            Some('b')
+        );
+        assert_eq!(
+            tree.right().and_then(|t| t.right()).map(|n| n.freq()),
+            Some(35)
+        );
+
+        // the least frequent characters require 3 bits
+        assert_eq!(
+            tree.right()
+                .and_then(|t| t.left())
+                .and_then(|t| t.left())
+                .and_then(|n| n.token()),
+            Some('d')
+        );
+        assert_eq!(
+            tree.right()
+                .and_then(|t| t.left())
+                .and_then(|t| t.left())
+                .map(|n| n.freq()),
+            Some(5)
+        );
+
+        assert_eq!(
+            tree.right()
+                .and_then(|t| t.left())
+                .and_then(|t| t.right())
+                .and_then(|n| n.token()),
+            Some('c')
+        );
+        assert_eq!(
+            tree.right()
+                .and_then(|t| t.left())
+                .and_then(|t| t.right())
+                .map(|n| n.freq()),
+            Some(20)
+        );
+    }
+}

src/main.rs

@@ -0,0 +1,106 @@
+mod compression;
+mod freq;
+mod huffman;
+
+use clap::Parser;
+use clap::ValueEnum;
+use std::fs::{self, File};
+use std::io::Write;
+use std::path::PathBuf;
+use std::time;
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+struct Args {
+    #[arg(value_enum)]
+    action: Action,
+    #[arg(value_enum)]
+    mode: Mode,
+    input: PathBuf,
+    output: PathBuf,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
+enum Action {
+    Compress,
+    Extract,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
+enum Mode {
+    Words,
+    Chars,
+}
+
+// cargo run --release -- compress words data/wikisent2.txt data/words.huffman
+// cargo run --release -- extract words data/words.huffman data/extracted.txt
+// (to check if the extraction was correct: `diff data/wikisent2.txt data/extracted.txt`)
+//
+// cargo run --release -- compress chars data/wikisent2.txt data/chars.huffman
+// cargo run --release -- extract chars data/chars.huffman data/extracted.txt
+//
+// to compare with zip:
+// time zip data/test.zip data/wikisent2.txt
+// time unzip data/test.zip -d data/test_zip
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+
+    match args.action {
+        Action::Compress => {
+            let timer = time::Instant::now();
+            let text = fs::read_to_string(args.input)?;
+            let lines: Vec<_> = text.split('\n').map(|x| x.to_string()).collect();
+            let time = timer.elapsed();
+            let lines_count = lines.len();
+            println!("Read the source file with {lines_count} lines in {time:?}");
+
+            let timer = time::Instant::now();
+            let compressed = match args.mode {
+                Mode::Words => compression::compress(&lines, freq::word_frequencies, |line| {
+                    line.split_ascii_whitespace().map(|token| token.to_string())
+                }),
+                Mode::Chars => {
+                    compression::compress(&lines, freq::char_frequencies, |line| line.chars())
+                }
+            }?;
+            let time = timer.elapsed();
+            println!("Compressed as {mode:?} in {time:?}.", mode = args.mode);
+
+            let timer = time::Instant::now();
+            let mut out_f = File::create(&args.output)?;
+            out_f.write_all(&compressed)?;
+            let time = timer.elapsed();
+            println!(
+                "Wrote to {output_path:?} in {time:?}",
+                output_path = args.output
+            );
+        }
+        Action::Extract => {
+            let timer = time::Instant::now();
+            let data = fs::read(&args.input)?;
+            let time = timer.elapsed();
+            println!("Read the compressed file in {time:?}");
+
+            let timer = time::Instant::now();
+            let content = match args.mode {
+                Mode::Words => compression::extract(&data, |tokens: Vec<String>| tokens.join(" "))?,
+                Mode::Chars => {
+                    compression::extract(&data, |tokens: Vec<char>| tokens.into_iter().collect())?
+                }
+            };
+            let time = timer.elapsed();
+            let lines_count = content.len();
+            println!("Extracted file with {lines_count} lines in {time:?}.");
+
+            let timer = time::Instant::now();
+            fs::write(&args.output, content.join("\n"))?;
+            let time = timer.elapsed();
+            println!(
+                "Wrote to {output_path:?} in {time:?}",
+                output_path = args.output
+            );
+        }
+    }
+
+    Ok(())
+}