Merge pull request #117 from 0xPolygonMiden/hacka-simplify-consuming-merkle-tree

feat: add node iterator to MerkleTree

src/merkle/merkle_tree.rs

@@ -1,4 +1,6 @@
-use super::{Felt, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word};
+use super::{
+    Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word,
+};
 use crate::{
     utils::{string::String, uninit_vector, word_to_hex},
     FieldElement,
@@ -12,7 +14,7 @@ use winter_math::log2;
 /// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two).
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct MerkleTree {
-    pub(crate) nodes: Vec<Word>,
+    nodes: Vec<Word>,
 }
 
 impl MerkleTree {
@@ -158,6 +160,47 @@ impl MerkleTree {
 
         Ok(())
     }
+
+    /// An iterator over every inner node in the tree. The iterator order is unspecified.
+    pub fn inner_nodes(&self) -> MerkleTreeNodes<'_> {
+        MerkleTreeNodes {
+            nodes: &self.nodes,
+            index: 1, // index 0 is just padding, start at 1
+        }
+    }
+}
+
+// ITERATORS
+// ================================================================================================
+
+/// An iterator over every inner node of the [MerkleTree].
+///
+/// Use this to extract the data of the tree, there is no guarantee on the order of the elements.
+pub struct MerkleTreeNodes<'a> {
+    nodes: &'a Vec<Word>,
+    index: usize,
+}
+
+impl<'a> Iterator for MerkleTreeNodes<'a> {
+    type Item = InnerNodeInfo;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index < self.nodes.len() / 2 {
+            let value = self.index;
+            let left = self.index * 2;
+            let right = left + 1;
+
+            self.index += 1;
+
+            Some(InnerNodeInfo {
+                value: self.nodes[value],
+                left: self.nodes[left],
+                right: self.nodes[right],
+            })
+        } else {
+            None
+        }
+    }
 }
 
 /// Utility to vizualize a [MerkleTree] in text.
@@ -212,7 +255,7 @@ pub fn path_to_text(path: &MerklePath) -> Result<String, fmt::Error> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::merkle::int_to_node;
+    use crate::merkle::{int_to_node, InnerNodeInfo};
     use core::mem::size_of;
     use proptest::prelude::*;
 
@@ -323,6 +366,40 @@ mod tests {
         assert_eq!(expected_tree.nodes, tree.nodes);
     }
 
+    #[test]
+    fn nodes() -> Result<(), MerkleError> {
+        let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
+        let root = tree.root();
+        let l1n0 = tree.get_node(NodeIndex::new(1, 0))?;
+        let l1n1 = tree.get_node(NodeIndex::new(1, 1))?;
+        let l2n0 = tree.get_node(NodeIndex::new(2, 0))?;
+        let l2n1 = tree.get_node(NodeIndex::new(2, 1))?;
+        let l2n2 = tree.get_node(NodeIndex::new(2, 2))?;
+        let l2n3 = tree.get_node(NodeIndex::new(2, 3))?;
+
+        let nodes: Vec<InnerNodeInfo> = tree.inner_nodes().collect();
+        let expected = vec![
+            InnerNodeInfo {
+                value: root,
+                left: l1n0,
+                right: l1n1,
+            },
+            InnerNodeInfo {
+                value: l1n0,
+                left: l2n0,
+                right: l2n1,
+            },
+            InnerNodeInfo {
+                value: l1n1,
+                left: l2n2,
+                right: l2n3,
+            },
+        ];
+        assert_eq!(nodes, expected);
+
+        Ok(())
+    }
+
     proptest! {
         #[test]
         fn arbitrary_word_can_be_represented_as_digest(

src/merkle/mod.rs

@@ -32,6 +32,9 @@ pub use mmr::{Mmr, MmrPeaks};
 mod store;
 pub use store::MerkleStore;
 
+mod node;
+pub use node::InnerNodeInfo;
+
 // ERRORS
 // ================================================================================================
 

src/merkle/node.rs

@@ -0,0 +1,9 @@
+use super::Word;
+
+/// Representation of a node with two children used for iterating over containers.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct InnerNodeInfo {
+    pub value: Word,
+    pub left: Word,
+    pub right: Word,
+}

src/merkle/store/mod.rs

@@ -249,33 +249,18 @@ impl MerkleStore {
             return Err(MerkleError::DepthTooSmall(leaves.len() as u8));
         }
 
-        let layers = leaves.len().ilog2();
         let tree = MerkleTree::new(leaves)?;
-
-        let mut depth = 0;
-        let mut parent_offset = 1;
-        let mut child_offset = 2;
-        while depth < layers {
-            let layer_size = 1usize << depth;
-            for _ in 0..layer_size {
-                // merkle tree is using level form representation, so left and right siblings are
-                // next to each other
-                let left = tree.nodes[child_offset];
-                let right = tree.nodes[child_offset + 1];
-                self.nodes.insert(
-                    tree.nodes[parent_offset].into(),
-                    Node {
-                        left: left.into(),
-                        right: right.into(),
-                    },
-                );
-                parent_offset += 1;
-                child_offset += 2;
-            }
-            depth += 1;
+        for node in tree.inner_nodes() {
+            self.nodes.insert(
+                node.value.into(),
+                Node {
+                    left: node.left.into(),
+                    right: node.right.into(),
+                },
+            );
         }
 
-        Ok(tree.nodes[1])
+        Ok(tree.root())
     }
 
     /// Adds a Sparse Merkle tree defined by the specified `entries` to the store, and returns the