Proper decode.

1 files changed, 52 insertions, 6 deletions

diff --git a/rust/huffman/src/main.rs b/rust/huffman/src/main.rs
index d2fc8a1..9b61918 100644
--- a/rust/huffman/src/main.rs
+++ b/rust/huffman/src/main.rs
@@ -1,14 +1,16 @@
 extern crate itertools;
 use itertools::Itertools;
+use itertools::Unfold;
 use std::cmp::Eq;
 use std::cmp::Ordering;
+use std::fmt::Debug;
 use std::collections::binary_heap::BinaryHeap;
 
 static MESSAGE: &'static str = r#"
   Hello! What's going on? This is a program.
   The brown fox jumped over the lazy dog.
   xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
-  \o/
+  æøå
 "#;
 
 #[derive(Debug, Eq, PartialEq)]
@@ -20,6 +22,15 @@ enum Tree<T> {
 #[derive(Debug, Eq, PartialEq)]
 enum Direction { Left, Right }
 
+impl ToString for Direction {
+    fn to_string(&self) -> String {
+        match *self {
+            Direction::Left => "0".to_string(),
+            Direction::Right => "1".to_string()
+        }
+    }
+}
+
 fn cost<T>(t: &Tree<T>) -> i32 {
     match *t {
         Tree::Leaf(cost, _) => cost,
@@ -57,6 +68,30 @@ fn make_path<T: Eq>(tree: &Tree<T>, elem: &T) -> Vec<Direction> {
     result
 }
 
+fn find_char2<'a, T: Debug+Eq+Copy, I>(tree: &Tree<T>, path: &mut I) -> Option<T>
+    where I: Iterator<Item=&'a Direction> {
+    match (tree) {
+        &Tree::Leaf(_, ref c) => Some(*c),
+        &Tree::Branch(_, ref left, ref right) =>
+            match (path.next()) {
+                Some(&Direction::Left) =>
+                    find_char2(left.as_ref(), path),
+                Some(&Direction::Right) =>
+                    find_char2(right.as_ref(), path),
+                _ => None
+            },
+    }
+}
+
+fn decode<'a, T: Debug+Eq+Copy, I>(tree: &Tree<T>, path: &'a mut I)
+    -> Vec<T>
+    where I: Iterator<Item=&'a Direction> {
+    Unfold::new((tree, path),
+                |&mut (ref tree, ref mut path)| {
+                    find_char2(tree, path)
+                }).collect::<Vec<_>>()
+}
+
 fn find_char<T: Eq+Copy>(tree: &Tree<T>, path: &[Direction]) -> Option<T> {
     match (tree, path.split_first()) {
         (&Tree::Branch(_, ref left, _), Some((&Direction::Left, ds))) =>
@@ -77,8 +112,7 @@ fn main() {
         .chars()
         .sorted().into_iter()
         .group_by(|x| *x)
-        .map(|(elem, xs)| (elem, xs.len() as i32))
-        .map(|(elem, cost)| Box::new(Tree::Leaf(cost, elem)))
+        .map(|(elem, xs)| Box::new(Tree::Leaf(xs.len() as i32, elem)))
         .collect::<BinaryHeap<_>>();
 
     let mut left;
@@ -94,12 +128,24 @@ fn main() {
                          left, right)));
     }
 
-    let restored_message = MESSAGE
+    let code = MESSAGE
         .chars()
         .map(|c| make_path(left.as_ref(), &c))
-        .flat_map(|p| find_char(left.as_ref(), &p))
+        .collect::<Vec<_>>();
+
+    println!("{}", code.iter()
+             .flatten().map(|d| d.to_string())
+             .collect::<String>());
+
+    let restored_message0 =
+        decode(left.as_ref(), &mut code.iter().flatten())
+         .into_iter().collect::<String>();
+
+    let restored_message = code.iter()
+        .flat_map(|p| find_char2(left.as_ref(), &mut p.iter()))
         .collect::<String>();
 
     println!("{}", MESSAGE);
-    println!("{}", restored_message);
+    println!("Restored message");
+    println!("{}", restored_message0);
 }