I'm looking for a general review with emphasis on idiomaticness and error handling. extern crate byteorder; extern crate rand; use std::collections::hash_map::HashMap; use std::collections::BinaryHeap; use byteorder::{ReadBytesExt, WriteBytesExt, NativeEndian}; use std::{fmt, error, result, io, path, cmp, fs}; use std::io::Read; use std::io::Write; type Result<T> = result::Result<T, HuffmanError>; const BITS: usize = 8; #[derive(Debug)] pub enum HuffmanError { Io(io::Error), ParseTree, AlphabetMismatch, Empty, } impl fmt::Display for HuffmanError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { use HuffmanError::*; use std::error::Error; match *self { Io(ref err) => write!(f, "IO error: {}", err), ParseTree => write!(f, "Parse tree error: {}", self.description()), Empty => write!(f, "Empty error: {}", self.description()), AlphabetMismatch => write!(f, "Alphabet mismatch error: {}", self.description()), } } } impl error::Error for HuffmanError { fn description(&self) -> &str { use HuffmanError::*; match *self { Io(ref err) => err.description(), ParseTree => "invalid encoded huffman tree", AlphabetMismatch => "alphabet doesn't match parsed tree", Empty => "empty stream", } } fn cause(&self) -> Option<&error::Error> { match *self { HuffmanError::Io(ref err) => Some(err), _ => None, } } } impl From<io::Error> for HuffmanError { fn from(err: io::Error) -> HuffmanError { HuffmanError::Io(err) } } #[derive(Eq, Debug)] enum HuffmanTree { Inner { frecuency: u32, left: Box<HuffmanTree>, right: Box<HuffmanTree>, }, Leaf { frecuency: u32, character: char }, } fn characters_mode(file: &str) -> HashMap<char, u32> { let mut table = HashMap::new(); for c in file.chars() { *(table.entry(c).or_insert(0)) += 1; } table } fn create_priority_queue(input: &str) -> BinaryHeap<HuffmanTree> { characters_mode(input) .iter() .map(|(c, f)| { HuffmanTree::Leaf { character: *c, frecuency: *f, } }) .collect() } impl HuffmanTree { fn new(input: &str) -> Result<HuffmanTree> { let mut pq = create_priority_queue(input); for _ in 1..pq.len() { let min1 = pq.pop().unwrap(); let min2 = pq.pop().unwrap(); pq.push(min1.join(min2)); } pq.pop().ok_or(HuffmanError::Empty) } fn frec(&self) -> u32 { use HuffmanTree::*; match self { &Inner { frecuency, .. } | &Leaf { frecuency, .. } => frecuency, } } fn join(self, other: HuffmanTree) -> HuffmanTree { HuffmanTree::Inner { frecuency: self.frec() + other.frec(), left: Box::new(self), right: Box::new(other), } } fn create_char_mapper_recur(&self, bit_vec: &mut BitVector, map: &mut HashMap<char, BitVector>) { use HuffmanTree::*; match self { &Inner { ref left, ref right, .. } => { bit_vec.push(true); left.create_char_mapper_recur(bit_vec, map); bit_vec.push(false); right.create_char_mapper_recur(bit_vec, map); } &Leaf { character, .. } => { map.insert(character, bit_vec.clone()); } } bit_vec.pop(); } fn create_char_mapper(&self) -> HashMap<char, BitVector> { let mut bit_vec = BitVector::new(); let mut map = HashMap::new(); if let &HuffmanTree::Leaf { .. } = self { bit_vec.push(true); } self.create_char_mapper_recur(&mut bit_vec, &mut map); map } fn decode<I, T>(encoded_walk: &mut I, mut chars: &mut T) -> Result<HuffmanTree> where I: Iterator<Item = bool>, T: Iterator<Item = char> { match encoded_walk.next() { Some(x) if x => { let left = Self::decode(encoded_walk, chars)?; let right = Self::decode(encoded_walk, chars)?; Ok(left.join(right)) } Some(_) => { let c = chars.next() .ok_or(HuffmanError::AlphabetMismatch)?; Ok(HuffmanTree::Leaf { frecuency: 0, character: c, }) } None => Err(HuffmanError::ParseTree), } } fn serialize<W: std::io::Write>(&self, writer: &mut W) -> Result<()> { let (encoded_walk, alphabet) = self.encode(); let walk_bit_len = encoded_walk.len() as u64; writer.write_u64::<NativeEndian>(walk_bit_len)?; writer.write(&encoded_walk.bits)?; let encoded_alphabet = alphabet.as_bytes(); let alphabet_byte_len = encoded_alphabet.len() as u64; writer.write_u64::<NativeEndian>(alphabet_byte_len)?; writer.write_all(encoded_alphabet)?; Ok(()) } fn de_serialize<R: std::io::Read>(reader: &mut R) -> Result<HuffmanTree> { use std::io::Read; let walk_len = reader.read_u64::<NativeEndian>()?; let mut walk_bytes = Vec::new(); reader.take((walk_len + BITS as u64 - 1) / BITS as u64) .read_to_end(&mut walk_bytes)?; let bit_vec = BitVector { bits: walk_bytes, size: walk_len as usize, }; let chars_len = reader.read_u64::<NativeEndian>()?; let mut chars = String::new(); reader.take(chars_len) .read_to_string(&mut chars)?; let bit_iter = &mut bit_vec.iter(); Self::decode(bit_iter, &mut chars.chars()) } fn encode_recur(&self, bit_vec: &mut BitVector, alphabet: &mut String) { use HuffmanTree::*; match self { &Inner { ref left, ref right, .. } => { bit_vec.push(true); left.encode_recur(bit_vec, alphabet); right.encode_recur(bit_vec, alphabet); } &Leaf { character, .. } => { bit_vec.push(false); alphabet.push(character); } } } fn encode(&self) -> (BitVector, String) { let mut bit_vector = BitVector::new(); let mut alphabet = String::new(); self.encode_recur(&mut bit_vector, &mut alphabet); (bit_vector, alphabet) } fn encode_string(&self, file_str: &str) -> BitVector { let char_map = self.create_char_mapper(); let itr = file_str.chars() .map(|c| char_map.get(&c).expect("Assertion error")); let mut bit_vec = BitVector::new(); for code in itr { bit_vec.append(&code); } bit_vec } fn decode_string<I>(&self, bit_iter: I) -> String where I: Iterator<Item = bool> { use HuffmanTree::*; let mut node = self; let mut output = String::new(); for bit in bit_iter { if let &Inner { ref left, ref right, .. } = node { node = if bit { left } else { right } } if let &Leaf { character, .. } = node { output.push(character); node = self; } } output } } impl Ord for HuffmanTree { fn cmp(&self, other: &HuffmanTree) -> cmp::Ordering { self.frec().cmp(&other.frec()).reverse() } } impl PartialOrd for HuffmanTree { fn partial_cmp(&self, other: &HuffmanTree) -> Option<cmp::Ordering> { Some(self.cmp(other)) } } impl PartialEq for HuffmanTree { fn eq(&self, other: &HuffmanTree) -> bool { self.frec() == other.frec() } } #[derive(Clone, Debug)] struct BitVector { bits: Vec<u8>, size: usize, } impl BitVector { fn new() -> BitVector { BitVector { bits: Vec::new(), size: 0, } } fn push(&mut self, bit: bool) { let leftover = self.size % BITS; if leftover == 0 { self.bits.push(0); } let last_byte = self.bits .last_mut() .expect("Assertion error"); *last_byte |= (bit as u8) << leftover; self.size += 1; } fn pop(&mut self) { if self.len() == 0 { return; } let len = self.size - 1; self.put(len, false); self.size = len; if self.len() % BITS == 0 { self.bits.pop(); } } #[allow(dead_code)] fn push_all(&mut self, bits: &[bool]) { for bit in bits { self.push(*bit); } } fn append(&mut self, other: &BitVector) { let leftover = self.size % BITS; let empty_bits = BITS - leftover; let len = self.bits.len(); self.bits.extend(other.bits.iter().cloned()); self.size += other.size; if leftover == 0 { return; } for i in len..self.bits.len() { self.move_bits(empty_bits, i); } if (self.size - 1) / BITS != self.bits.len() - 1 { self.bits.pop(); } } fn check(&self, i: usize) { assert!(i < self.size, format!("Index out of bounds. Index: {} >= len: {}", i, self.size)); } #[allow(dead_code)] fn get(&self, i: usize) -> bool { self.check(i); (1 & (self.bits[i / BITS] >> (i % BITS))) != 0 } fn put(&mut self, i: usize, bit: bool) { self.check(i); let (byte, bit_pos) = (i / BITS, i % BITS); self.bits[byte] &= !(1 << bit_pos); self.bits[byte] |= (bit as u8) << bit_pos; } fn move_bits(&mut self, bits: usize, i: usize) { let leftover = BITS - bits; let (from, to) = (self.bits[i], self.bits[i - 1]); self.bits[i - 1] = to ^ (from << leftover); self.bits[i] = from >> bits; } fn iter<'a>(&'a self) -> BitVectorIter<'a> { let itr = self.bits .iter() .flat_map(|byte| (0..BITS).map(move |i| ((byte >> i) & 1) != 0)) .take(self.size); BitVectorIter { iter: Box::new(itr) } } fn len(&self) -> usize { self.size } #[allow(dead_code)] fn byte_len(&self) -> usize { self.bits.len() } } struct BitVectorIter<'a> { iter: Box<std::iter::Iterator<Item = bool> + 'a>, } impl<'a> IntoIterator for &'a BitVector { type Item = bool; type IntoIter = BitVectorIter<'a>; fn into_iter(self) -> Self::IntoIter { self.iter() } } impl<'a> Iterator for BitVectorIter<'a> { type Item = bool; fn next(&mut self) -> Option<Self::Item> { self.iter.next() } } pub trait HuffmanCompress { fn compress<T: std::io::Write>(&mut self, &mut T) -> Result<()>; } pub trait HuffmanDeCompress { fn de_compress<T: std::io::Write>(&mut self, writer: &mut T) -> Result<()> { let string = self.de_compress_to_string()?; writer.write_all(&string.as_bytes())?; Ok(()) } fn de_compress_to_string(&mut self) -> Result<String>; } pub trait HuffmanCodes: HuffmanDeCompress + HuffmanCompress {} impl<T: std::io::Read> HuffmanCodes for T {} impl<T: std::io::Read> HuffmanCompress for T { fn compress<W: std::io::Write>(&mut self, writer: &mut W) -> Result<()> { let mut file_str = String::new(); self.read_to_string(&mut file_str)?; let tree = HuffmanTree::new(&file_str)?; let encoded_file = tree.encode_string(&file_str); let junk = BITS - (encoded_file.len() % BITS); let mask = (((junk != BITS) as i8) << (BITS - 1)) >> (BITS - 1); tree.serialize(writer)?; writer.write_i8(junk as i8 & mask)?; writer.write(&encoded_file.bits)?; Ok(()) } } impl<T: std::io::Read> HuffmanDeCompress for T { fn de_compress_to_string(&mut self) -> Result<String> { let tree = HuffmanTree::de_serialize(self)?; let junk = self.read_i8()?; let mut bytes = Vec::new(); self.read_to_end(&mut bytes)?; let bit_vec = BitVector { size: BITS * bytes.len() - junk as usize, bits: bytes, }; let string = tree.decode_string(bit_vec.iter()); Ok(string) } } pub fn encode<P, T>(path: P, target: T) -> Result<()> where P: AsRef<path::Path>, T: AsRef<path::Path> { let f1 = fs::File::open(path)?; let mut reader = io::BufReader::new(f1); let f2 = fs::File::create(target)?; let mut writer = io::BufWriter::new(f2); reader.compress(&mut writer)?; Ok(()) } pub fn decode<P>(path: P) -> Result<String> where P: AsRef<path::Path> { let f = fs::File::open(path)?; let mut reader = &mut io::BufReader::new(f); Ok(reader.de_compress_to_string()?) } fn read_file<P: AsRef<path::Path>>(path: P) -> std::io::Result<String> { let f = fs::File::open(path)?; let mut reader = io::BufReader::new(f); let mut result = String::new(); reader.read_to_string(&mut result)?; Ok(result) } fn rand_string(len: usize) -> String { use rand::Rng; rand::thread_rng() .gen_iter::<char>() .take(len) .collect() } #[test] fn it_works() { for _ in 1..100 { let string = rand_string(5000); if let Err(err) = test(&string) { panic!(format!("{}", err)) } } } fn test(text: &str) -> Result<()> { const FILE: &'static str = "test"; const CFILE: &'static str = "compress"; { let f = fs::File::create(FILE)?; let mut writer = io::BufWriter::new(f); writer.write_all(&text.as_bytes())?; } encode(FILE, CFILE)?; let x = decode(CFILE)?; let y = read_file(FILE)?; assert_eq!(x, y); assert_eq!(x, text); Ok(()) }