Implement Iterator and IntoIterator for MerkleTree.

benches/lib.rs

@@ -104,3 +104,22 @@ fn bench_big_rnd_proof_check(b: &mut Bencher) {
         }
     });
 }
+
+#[bench]
+fn bench_big_rnd_iter(b: &mut Bencher) {
+    let mut values = vec![vec![0u8; 256]; 160];
+    let mut rng    = rand::IsaacRng::new_unseeded();
+
+    for mut v in &mut values {
+        rng.fill_bytes(&mut v);
+    }
+
+    let tree = MerkleTree::from_vec(digest, values.clone()).unwrap();
+
+    b.iter(|| {
+        for value in &tree {
+            test::black_box(value);
+        }
+    });
+}
+

src/lib.rs

@@ -22,6 +22,7 @@ pub use proof::Proof;
 mod hashutils;
 
 mod tree;
+pub use tree::{ LeavesIterator, LeavesIntoIterator };
 
 #[cfg(feature = "serialization-protobuf")]
 mod proto;

src/merkletree.rs

@@ -1,14 +1,13 @@
 
 use ring::digest::Algorithm;
 
-use tree::Tree;
+use tree::{ Tree, LeavesIterator, LeavesIntoIterator };
 use hashutils::HashUtils;
 
 use proof::{ Proof, Lemma };
 
 /// A Merkle tree is a binary tree, with values of type `T` at the leafs,
-/// and where every node holds the hash of the concatenation of the hashes of
-/// its children nodes.
+/// and where every internal node holds the hash of the concatenation of the hashes of its children nodes.
 #[derive(Clone, Debug)]
 pub struct MerkleTree<T> {
 
@@ -28,8 +27,8 @@ pub struct MerkleTree<T> {
 impl <T> MerkleTree<T> where T: Into<Vec<u8>> + Clone {
 
     /// Constructs a Merkle Tree from a vector of data blocks.
-    /// Returns None if `values` is empty.
-    pub fn from_vec(algo: &'static Algorithm, values: Vec<T>) -> Option<Self> {
+    /// Returns `None` if `values` is empty.
+    pub fn from_vec(algorithm: &'static Algorithm, values: Vec<T>) -> Option<Self> {
         if values.is_empty() {
             return None
         }
@@ -39,7 +38,7 @@ impl <T> MerkleTree<T> where T: Into<Vec<u8>> + Clone {
         let mut cur    = Vec::with_capacity(count);
 
         for v in values {
-            let leaf = Tree::make_leaf(algo, v);
+            let leaf = Tree::make_leaf(algorithm, v);
             cur.push(leaf);
         }
 
@@ -53,7 +52,7 @@ impl <T> MerkleTree<T> where T: Into<Vec<u8>> + Clone {
                     let left  = cur.remove(0);
                     let right = cur.remove(0);
 
-                    let combined_hash = algo.combine_hashes(
+                    let combined_hash = algorithm.combine_hashes(
                         left.hash(),
                         right.hash()
                     );
@@ -78,7 +77,7 @@ impl <T> MerkleTree<T> where T: Into<Vec<u8>> + Clone {
         let root = cur.remove(0);
 
         Some(MerkleTree {
-            algorithm: algo,
+            algorithm: algorithm,
             root: root,
             height: height,
             count: count
@@ -111,4 +110,35 @@ impl <T> MerkleTree<T> where T: Into<Vec<u8>> + Clone {
         )
     }
 
+    /// Creates an `Iterator` over the values contained in this Merkle tree.
+    pub fn iter(&self) -> LeavesIterator<T> {
+        self.root.iter()
+    }
+
 }
+
+impl <T> IntoIterator for MerkleTree<T> {
+
+    type Item     = T;
+    type IntoIter = LeavesIntoIterator<T>;
+
+    /// Creates a consuming iterator, that is, one that moves each value out of the Merkle tree.
+    /// The tree cannot be used after calling this.
+    fn into_iter(self) -> Self::IntoIter {
+        self.root.into_iter()
+    }
+
+}
+
+impl <'a, T> IntoIterator for &'a MerkleTree<T> {
+
+    type Item     = &'a T;
+    type IntoIter = LeavesIterator<'a, T>;
+
+    /// Creates a borrowing `Iterator` over the values contained in this Merkle tree.
+    fn into_iter(self) -> Self::IntoIter {
+        self.root.iter()
+    }
+
+}
+

src/tests.rs

@@ -173,3 +173,52 @@ fn test_mutate_proof_first_lemma() {
         i += 1;
     }
 }
+
+#[test]
+fn test_tree_iter() {
+    let values  = (1..10).map(|x| vec![x]).collect::<Vec<_>>();
+    let tree    = MerkleTree::from_vec(digest, values.clone()).unwrap();
+    let iter    = tree.iter().map(|x| x.clone()).collect::<Vec<_>>();
+
+    assert_eq!(values, iter);
+}
+
+#[test]
+fn test_tree_into_iter() {
+    let values  = (1..10).map(|x| vec![x]).collect::<Vec<_>>();
+    let tree    = MerkleTree::from_vec(digest, values.clone()).unwrap();
+    let iter    = tree.into_iter().map(|x| x.clone()).collect::<Vec<_>>();
+
+    assert_eq!(values, iter);
+}
+
+#[test]
+fn test_tree_into_iter_loop() {
+    let values  = (1..10).map(|x| vec![x]).collect::<Vec<_>>();
+    let tree    = MerkleTree::from_vec(digest, values.clone()).unwrap();
+
+    let mut collected = Vec::new();
+
+    for value in tree {
+        collected.push(value);
+    }
+
+    assert_eq!(values, collected);
+}
+
+#[test]
+fn test_tree_into_iter_loop_borrow() {
+    let values  = (1..10).map(|x| vec![x]).collect::<Vec<_>>();
+    let tree    = MerkleTree::from_vec(digest, values.clone()).unwrap();
+
+    let mut collected = Vec::new();
+
+    for value in &tree {
+        collected.push(value);
+    }
+
+    let refs = values.iter().collect::<Vec<_>>();
+
+    assert_eq!(refs, collected);
+}
+

src/tree.rs

@@ -24,7 +24,7 @@ pub enum Tree<T> {
     }
 }
 
-impl <T> Tree<T> where T: Into<Vec<u8>> + Clone {
+impl <T> Tree<T> {
 
     /// Create a new tree
     pub fn new(hash: Vec<u8>, value: T) -> Self {
@@ -35,7 +35,9 @@ impl <T> Tree<T> where T: Into<Vec<u8>> + Clone {
     }
 
     /// Create a new leaf
-    pub fn make_leaf(algo: &'static Algorithm, value: T) -> Tree<T> {
+    pub fn make_leaf(algo: &'static Algorithm, value: T) -> Tree<T>
+            where T: Into<Vec<u8>> + Clone {
+
         let hash = algo.hash_bytes(&value.clone().into());
         Tree::new(hash, value)
     }
@@ -48,5 +50,130 @@ impl <T> Tree<T> where T: Into<Vec<u8>> + Clone {
         }
     }
 
+    /// Returns a borrowing iterator over the leaves of the tree.
+    pub fn iter(&self) -> LeavesIterator<T> {
+        LeavesIterator::new(self)
+    }
+
+}
+
+/// An borrowing iterator over the leaves of a `Tree`.  
+/// Adapted from http://codereview.stackexchange.com/q/110283.
+#[allow(missing_debug_implementations)]
+pub struct LeavesIterator<'a, T> where T: 'a {
+    current_value: Option<&'a T>,
+    right_nodes: Vec<&'a Tree<T>>
+}
+
+impl <'a, T> LeavesIterator<'a, T> {
+
+    fn new(root: &'a Tree<T>) -> Self {
+        let mut iter = LeavesIterator {
+            current_value: None,
+            right_nodes: Vec::new()
+        };
+
+        iter.add_left(root);
+
+        iter
+    }
+
+    fn add_left(&mut self, mut tree: &'a Tree<T>) {
+        loop {
+            match *tree {
+                Tree::Node { ref left, ref right, .. } => {
+                    self.right_nodes.push(right);
+                    tree = left;
+                },
+
+                Tree::Leaf { ref value, .. } => {
+                    self.current_value = Some(value);
+                    break;
+                }
+            }
+        }
+    }
+
+}
+
+impl <'a, T> Iterator for LeavesIterator<'a, T> {
+
+    type Item = &'a T;
+
+    fn next(&mut self) -> Option<&'a T> {
+        let result = self.current_value.take();
+
+        if let Some(rest) = self.right_nodes.pop() {
+            self.add_left(rest);
+        }
+
+        result
+    }
+
+}
+
+/// An iterator over the leaves of a `Tree`.
+#[allow(missing_debug_implementations)]
+pub struct LeavesIntoIterator<T> {
+    current_value: Option<T>,
+    right_nodes: Vec<Tree<T>>
+}
+
+impl <T> LeavesIntoIterator<T> {
+
+    fn new(root: Tree<T>) -> Self {
+        let mut iter = LeavesIntoIterator {
+            current_value: None,
+            right_nodes: Vec::new()
+        };
+
+        iter.add_left(root);
+
+        iter
+    }
+
+    fn add_left(&mut self, mut tree: Tree<T>) {
+        loop {
+            match tree {
+                Tree::Node { left, right, .. } => {
+                    self.right_nodes.push(*right);
+                    tree = *left;
+                },
+
+                Tree::Leaf { value, .. } => {
+                    self.current_value = Some(value);
+                    break;
+                }
+            }
+        }
+    }
+
+}
+
+impl <T> Iterator for LeavesIntoIterator<T> {
+
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        let result = self.current_value.take();
+
+        if let Some(rest) = self.right_nodes.pop() {
+            self.add_left(rest);
+        }
+
+        result
+    }
+
+}
+
+impl <T> IntoIterator for Tree<T> {
+
+    type Item     = T;
+    type IntoIter = LeavesIntoIterator<T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        LeavesIntoIterator::new(self)
+    }
+
 }