diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 763a415..4548739 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -118,15 +118,15 @@ impl MerkleTree { // -------------------------------------------------------------------------------------------- /// Returns an iterator over the leaves of this [MerkleTree]. - pub fn leaves(&self) -> core::slice::Iter { + pub fn leaves(&self) -> impl Iterator { let leaves_start = self.nodes.len() / 2; - self.nodes[leaves_start..].iter() + self.nodes.iter().skip(leaves_start).enumerate().map(|(i, v)| (i as u64, v)) } /// Returns n iterator over every inner node of this [MerkleTree]. /// /// The iterator order is unspecified. - pub fn inner_nodes(&self) -> InnerNodeIterator<'_> { + pub fn inner_nodes(&self) -> InnerNodeIterator { InnerNodeIterator { nodes: &self.nodes, index: 1, // index 0 is just padding, start at 1 diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 6b0a6c2..d56f294 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -182,6 +182,14 @@ impl SimpleSmt { self.get_path(index) } + // ITERATORS + // -------------------------------------------------------------------------------------------- + + /// Returns an iterator over the leaves of this [SimpleSmt]. + pub fn leaves(&self) -> impl Iterator { + self.leaves.iter().map(|(i, w)| (*i, w)) + } + /// Returns an iterator over the inner nodes of this Merkle tree. pub fn inner_nodes(&self) -> impl Iterator + '_ { self.branches.values().map(|e| InnerNodeInfo { diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 560a406..40eaabd 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,9 +1,9 @@ -use super::mmr::Mmr; use super::{ - BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, MerkleTree, - NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word, + mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, + MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word, }; use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; +use core::borrow::Borrow; #[cfg(test)] mod tests; @@ -264,11 +264,15 @@ impl MerkleStore { /// nodes which are descendants of the specified roots. /// /// The roots for which no descendants exist in this Merkle store are ignored. - pub fn subset(&self, roots: &[Word]) -> MerkleStore { + pub fn subset(&self, roots: I) -> MerkleStore + where + I: Iterator, + R: Borrow, + { let mut store = MerkleStore::new(); for root in roots { - let root = RpoDigest::from(*root); - self.clone_tree_into(root, &mut store); + let root = RpoDigest::from(*root.borrow()); + store.clone_tree_from(root, self); } store } @@ -393,17 +397,17 @@ impl MerkleStore { // HELPER METHODS // -------------------------------------------------------------------------------------------- - /// Recursively clones a tree starting at the specified root into the specified target. + /// Recursively clones a tree with the specified root from the specified source into self. /// - /// If this Merkle store does not contain a tree with the specified root, this is a noop. - fn clone_tree_into(&self, root: RpoDigest, target: &mut Self) { - // process the node only if it is in this store - if let Some(node) = self.nodes.get(&root) { + /// If the source store does not contain a tree with the specified root, this is a noop. + fn clone_tree_from(&mut self, root: RpoDigest, source: &Self) { + // process the node only if it is in the source + if let Some(node) = source.nodes.get(&root) { // if the node has already been inserted, no need to process it further as all of its - // descendants should be already in the target store - if matches!(target.nodes.insert(root, *node), None) { - self.clone_tree_into(node.left, target); - self.clone_tree_into(node.right, target); + // descendants should be already cloned from the source store + if matches!(self.nodes.insert(root, *node), None) { + self.clone_tree_from(node.left, source); + self.clone_tree_from(node.right, source); } } } diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index beda857..0bdfa4f 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -750,38 +750,39 @@ fn mstore_subset() { let subtree2 = MerkleTree::new(VALUES8[2..4].to_vec()).unwrap(); let subtree3 = MerkleTree::new(VALUES8[6..].to_vec()).unwrap(); - let substore = store.subset(&[subtree1.root(), subtree2.root(), subtree3.root()]); + // --- extract all 3 subtrees --------------------------------------------- + + let substore = store.subset([subtree1.root(), subtree2.root(), subtree3.root()].iter()); // number of nodes should increase by 4: 3 nodes form subtree1 and 1 node from subtree3 assert_eq!(substore.nodes.len(), empty_store_num_nodes + 4); - // make sure paths for subtree1 are correct - for (i, value) in subtree1.leaves().enumerate() { - let index = NodeIndex::new(2, i as u64).unwrap(); - let path1 = substore.get_path(subtree1.root(), index).unwrap(); - assert_eq!(&path1.value, value); + // make sure paths that all subtrees are in the store + check_mstore_subtree(&substore, &subtree1); + check_mstore_subtree(&substore, &subtree2); + check_mstore_subtree(&substore, &subtree3); - let path2 = subtree1.get_path(index).unwrap(); - assert_eq!(path1.path, path2); - } + // --- extract subtrees 1 and 3 ------------------------------------------- + // this should give the same result as above as subtree2 is nested withing subtree1 - // make sure paths for subtree2 are correct - for (i, value) in subtree2.leaves().enumerate() { - let index = NodeIndex::new(1, i as u64).unwrap(); - let path1 = substore.get_path(subtree2.root(), index).unwrap(); - assert_eq!(&path1.value, value); + let substore = store.subset([subtree1.root(), subtree3.root()].iter()); - let path2 = subtree2.get_path(index).unwrap(); - assert_eq!(path1.path, path2); - } + // number of nodes should increase by 4: 3 nodes form subtree1 and 1 node from subtree3 + assert_eq!(substore.nodes.len(), empty_store_num_nodes + 4); + + // make sure paths that all subtrees are in the store + check_mstore_subtree(&substore, &subtree1); + check_mstore_subtree(&substore, &subtree2); + check_mstore_subtree(&substore, &subtree3); +} - // make sure paths for subtree3 are correct - for (i, value) in subtree3.leaves().enumerate() { - let index = NodeIndex::new(1, i as u64).unwrap(); - let path1 = substore.get_path(subtree3.root(), index).unwrap(); +fn check_mstore_subtree(store: &MerkleStore, subtree: &MerkleTree) { + for (i, value) in subtree.leaves() { + let index = NodeIndex::new(subtree.depth(), i).unwrap(); + let path1 = store.get_path(subtree.root(), index).unwrap(); assert_eq!(&path1.value, value); - let path2 = subtree3.get_path(index).unwrap(); + let path2 = subtree.get_path(index).unwrap(); assert_eq!(path1.path, path2); } }