feat: add leaves() iterator to SimpleSmt

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<Word> {
+    pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
         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

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<Item = (u64, &Word)> {
+        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<Item = InnerNodeInfo> + '_ {
         self.branches.values().map(|e| InnerNodeInfo {

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<I, R>(&self, roots: I) -> MerkleStore
+    where
+        I: Iterator<Item = R>,
+        R: Borrow<Word>,
+    {
         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);
             }
         }
     }

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();
+    // 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);
+
+    // --- extract subtrees 1 and 3 -------------------------------------------
+    // this should give the same result as above as subtree2 is nested withing subtree1
+
+    let substore = store.subset([subtree1.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 that all subtrees are in the store
+    check_mstore_subtree(&substore, &subtree1);
+    check_mstore_subtree(&substore, &subtree2);
+    check_mstore_subtree(&substore, &subtree3);
+}
+
+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 = subtree1.get_path(index).unwrap();
-        assert_eq!(path1.path, path2);
-    }
-
-    // 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 path2 = subtree2.get_path(index).unwrap();
-        assert_eq!(path1.path, path2);
-    }
-
-    // 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();
-        assert_eq!(&path1.value, value);
-
-        let path2 = subtree3.get_path(index).unwrap();
+        let path2 = subtree.get_path(index).unwrap();
         assert_eq!(path1.path, path2);
     }
 }