From a1968efbfca521a6629a355b303764234ef5356e Mon Sep 17 00:00:00 2001 From: Romain Ruetschi Date: Fri, 2 Dec 2016 14:54:31 +0100 Subject: [PATCH] Implement `Iterator` and `IntoIterator` for `MerkleTree`. --- benches/lib.rs | 19 +++++++ src/lib.rs | 1 + src/merkletree.rs | 46 +++++++++++++--- src/tests.rs | 49 +++++++++++++++++ src/tree.rs | 131 +++++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 236 insertions(+), 10 deletions(-) diff --git a/benches/lib.rs b/benches/lib.rs index 80ded07..cd70425 100644 --- a/benches/lib.rs +++ b/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); + } + }); +} + diff --git a/src/lib.rs b/src/lib.rs index d9c817d..72ef3fb 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/merkletree.rs b/src/merkletree.rs index 52d1c8b..3aff7ad 100644 --- a/src/merkletree.rs +++ b/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 { @@ -28,8 +27,8 @@ pub struct MerkleTree { impl MerkleTree where T: Into> + 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) -> Option { + /// Returns `None` if `values` is empty. + pub fn from_vec(algorithm: &'static Algorithm, values: Vec) -> Option { if values.is_empty() { return None } @@ -39,7 +38,7 @@ impl MerkleTree where T: Into> + 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 MerkleTree where T: Into> + 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 MerkleTree where T: Into> + Clone { let root = cur.remove(0); Some(MerkleTree { - algorithm: algo, + algorithm: algorithm, root: root, height: height, count: count @@ -111,4 +110,35 @@ impl MerkleTree where T: Into> + Clone { ) } + /// Creates an `Iterator` over the values contained in this Merkle tree. + pub fn iter(&self) -> LeavesIterator { + self.root.iter() + } + } + +impl IntoIterator for MerkleTree { + + type Item = T; + type IntoIter = LeavesIntoIterator; + + /// 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 { + + 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() + } + +} + diff --git a/src/tests.rs b/src/tests.rs index a9b0c6d..d7d267d 100644 --- a/src/tests.rs +++ b/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::>(); + let tree = MerkleTree::from_vec(digest, values.clone()).unwrap(); + let iter = tree.iter().map(|x| x.clone()).collect::>(); + + assert_eq!(values, iter); +} + +#[test] +fn test_tree_into_iter() { + let values = (1..10).map(|x| vec![x]).collect::>(); + let tree = MerkleTree::from_vec(digest, values.clone()).unwrap(); + let iter = tree.into_iter().map(|x| x.clone()).collect::>(); + + assert_eq!(values, iter); +} + +#[test] +fn test_tree_into_iter_loop() { + let values = (1..10).map(|x| vec![x]).collect::>(); + 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::>(); + 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::>(); + + assert_eq!(refs, collected); +} + diff --git a/src/tree.rs b/src/tree.rs index eb5c903..68f9702 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -24,7 +24,7 @@ pub enum Tree { } } -impl Tree where T: Into> + Clone { +impl Tree { /// Create a new tree pub fn new(hash: Vec, value: T) -> Self { @@ -35,7 +35,9 @@ impl Tree where T: Into> + Clone { } /// Create a new leaf - pub fn make_leaf(algo: &'static Algorithm, value: T) -> Tree { + pub fn make_leaf(algo: &'static Algorithm, value: T) -> Tree + where T: Into> + Clone { + let hash = algo.hash_bytes(&value.clone().into()); Tree::new(hash, value) } @@ -48,5 +50,130 @@ impl Tree where T: Into> + Clone { } } + /// Returns a borrowing iterator over the leaves of the tree. + pub fn iter(&self) -> LeavesIterator { + 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> +} + +impl <'a, T> LeavesIterator<'a, T> { + + fn new(root: &'a Tree) -> 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) { + 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 { + current_value: Option, + right_nodes: Vec> +} + +impl LeavesIntoIterator { + + fn new(root: Tree) -> 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) { + loop { + match tree { + Tree::Node { left, right, .. } => { + self.right_nodes.push(*right); + tree = *left; + }, + + Tree::Leaf { value, .. } => { + self.current_value = Some(value); + break; + } + } + } + } + +} + +impl Iterator for LeavesIntoIterator { + + type Item = T; + + fn next(&mut self) -> Option { + let result = self.current_value.take(); + + if let Some(rest) = self.right_nodes.pop() { + self.add_left(rest); + } + + result + } + +} + +impl IntoIterator for Tree { + + type Item = T; + type IntoIter = LeavesIntoIterator; + + fn into_iter(self) -> Self::IntoIter { + LeavesIntoIterator::new(self) + } + }