refactor: improve tests, small fixes

src/merkle/mod.rs

@@ -40,7 +40,7 @@ mod node;
 pub use node::InnerNodeInfo;
 
 mod partial_mt;
-pub use partial_mt::PartialMerkleTree;
+pub use partial_mt::{pmt_to_text, PartialMerkleTree};
 
 // ERRORS
 // ================================================================================================

src/merkle/partial_mt/mod.rs

@@ -2,9 +2,8 @@ use super::{
     BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec,
     Word, EMPTY_WORD,
 };
-
-extern crate alloc;
-use alloc::format;
+use crate::utils::{format, string::String, word_to_hex};
+use core::fmt;
 
 #[cfg(test)]
 mod tests;
@@ -102,10 +101,8 @@ impl PartialMerkleTree {
         self.leaves.contains(&index)
     }
 
-    pub fn get_leaf_depth(&self, index: u64) -> u8 {
-        let trunc_value = index.min(2_u64.pow(self.max_depth() as u32));
-        let mut node_index =
-            NodeIndex::new(self.max_depth(), trunc_value).expect("Truncated value is not valid.");
+    pub fn get_leaf_depth(&self, node_index: NodeIndex) -> u8 {
+        let mut node_index = node_index;
         for _ in 0..node_index.depth() {
             if self.leaves.contains(&node_index) {
                 return node_index.depth();
@@ -304,3 +301,32 @@ impl PartialMerkleTree {
         Ok(())
     }
 }
+
+/// Utility to visualize a [PartialMerkleTree] in text.
+pub fn pmt_to_text(tree: &PartialMerkleTree) -> Result<String, fmt::Error> {
+    let indent = "  ";
+    let mut s = String::new();
+    s.push_str("root: ");
+    s.push_str(&word_to_hex(&tree.root())?);
+    s.push('\n');
+    for d in 1..=tree.max_depth() {
+        let entries = 2u64.pow(d.into());
+        for i in 0..entries {
+            let index = NodeIndex::new(d, i).expect("The index must always be valid");
+            let node = tree.get_node(index);
+            let node = match node {
+                Err(_) => continue,
+                Ok(node) => node,
+            };
+
+            for _ in 0..d {
+                s.push_str(indent);
+            }
+            s.push_str(&format!("({}, {}): ", index.depth(), index.value()));
+            s.push_str(&word_to_hex(&node)?);
+            s.push('\n');
+        }
+    }
+
+    Ok(s)
+}

src/merkle/partial_mt/tests.rs

@@ -1,14 +1,17 @@
 use super::{
     super::{int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree},
-    Word,
+    Rpo256, ValuePath, Vec, Word,
 };
 
 // TEST DATA
 // ================================================================================================
 
 const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0);
+const NODE11: NodeIndex = NodeIndex::new_unchecked(1, 1);
 
+const NODE20: NodeIndex = NodeIndex::new_unchecked(2, 0);
 const NODE22: NodeIndex = NodeIndex::new_unchecked(2, 2);
+const NODE23: NodeIndex = NodeIndex::new_unchecked(2, 3);
 
 const NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2);
 const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
@@ -39,7 +42,7 @@ fn get_root() {
 
     let path33 = ms.get_path(expected_root, NODE33).unwrap();
 
-    let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
+    let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
 
     assert_eq!(pmt.root(), expected_root.into());
 }
@@ -54,17 +57,19 @@ fn add_and_get_paths() {
     let expected_path33 = ms.get_path(expected_root, NODE33).unwrap();
     let expected_path22 = ms.get_path(expected_root, NODE22).unwrap();
 
-    let pmt = PartialMerkleTree::with_paths([
-        (3_u64, expected_path33.value.into(), expected_path33.path.clone()),
-        (2, expected_path22.value.into(), expected_path22.path.clone()),
-    ])
-    .unwrap();
+    let mut pmt = PartialMerkleTree::new();
+    pmt.add_path(3, expected_path33.value.into(), expected_path33.path.clone())
+        .unwrap();
+    pmt.add_path(2, expected_path22.value.into(), expected_path22.path.clone())
+        .unwrap();
 
     let path33 = pmt.get_path(NODE33).unwrap();
     let path22 = pmt.get_path(NODE22).unwrap();
+    let actual_root = pmt.root();
 
     assert_eq!(expected_path33.path, path33);
     assert_eq!(expected_path22.path, path22);
+    assert_eq!(expected_root, *actual_root);
 }
 
 #[test]
@@ -76,7 +81,7 @@ fn get_node() {
 
     let path33 = ms.get_path(expected_root, NODE33).unwrap();
 
-    let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
+    let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
 
     assert_eq!(ms.get_node(expected_root, NODE32).unwrap(), *pmt.get_node(NODE32).unwrap());
     assert_eq!(ms.get_node(expected_root, NODE10).unwrap(), *pmt.get_node(NODE10).unwrap());
@@ -90,17 +95,28 @@ fn update_leaf() {
     let ms = MerkleStore::from(&mt);
     let path33 = ms.get_path(root, NODE33).unwrap();
 
-    let mut pmt =
-        PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
+    let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
 
     let new_value32 = int_to_node(132);
-    mt.update_leaf(2_u64, new_value32).unwrap();
+    mt.update_leaf(2, new_value32).unwrap();
     let expected_root = mt.root();
 
     pmt.update_leaf(NODE32, new_value32.into()).unwrap();
     let actual_root = pmt.root();
 
     assert_eq!(expected_root, *actual_root);
+
+    let mut new_vals = VALUES8.clone();
+    new_vals[1] = int_to_node(131);
+    new_vals[2] = int_to_node(132);
+    let new_value20 = Rpo256::merge(&[new_vals[0].into(), new_vals[1].into()]);
+    let mt = MerkleTree::new(new_vals.to_vec()).unwrap();
+    let expected_root = mt.root();
+
+    pmt.update_leaf(NODE20, new_value20).unwrap();
+    let actual_root = pmt.root();
+
+    assert_eq!(expected_root, *actual_root);
 }
 
 #[test]
@@ -112,17 +128,97 @@ fn check_leaf_depth() {
 
     let path33 = ms.get_path(expected_root, NODE33).unwrap();
 
-    let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
+    let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
+
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(4, 1)), 2);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(4, 6)), 3);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(4, 10)), 1);
+
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 1)), 2);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 2)), 3);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 5)), 1);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 7)), 1);
+
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 0)), 2);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 1)), 0);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 2)), 1);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 3)), 1);
+
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(1, 0)), 0);
+    assert_eq!(pmt.get_leaf_depth(NodeIndex::make(1, 1)), 1);
+}
+
+#[test]
+fn get_paths() {
+    let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
+    let expected_root = mt.root();
+
+    let ms = MerkleStore::from(&mt);
 
-    assert_eq!(pmt.get_leaf_depth(0), 2);
-    assert_eq!(pmt.get_leaf_depth(1), 2);
-    assert_eq!(pmt.get_leaf_depth(2), 3);
-    assert_eq!(pmt.get_leaf_depth(3), 3);
-    assert_eq!(pmt.get_leaf_depth(4), 1);
-    assert_eq!(pmt.get_leaf_depth(5), 1);
-    assert_eq!(pmt.get_leaf_depth(6), 1);
-    assert_eq!(pmt.get_leaf_depth(7), 1);
+    let path33 = ms.get_path(expected_root, NODE33).unwrap();
+    let path22 = ms.get_path(expected_root, NODE22).unwrap();
+
+    let mut pmt = PartialMerkleTree::new();
+    pmt.add_path(3, path33.value.into(), path33.path.clone()).unwrap();
+    pmt.add_path(2, path22.value.into(), path22.path.clone()).unwrap();
+
+    let leaves = vec![NODE20, NODE22, NODE23, NODE32, NODE33];
+    let expected_paths: Vec<(NodeIndex, ValuePath)> = leaves
+        .iter()
+        .map(|&leaf| {
+            (
+                leaf,
+                ValuePath {
+                    value: mt.get_node(leaf).unwrap().into(),
+                    path: mt.get_path(leaf).unwrap(),
+                },
+            )
+        })
+        .collect();
+
+    let actual_paths = pmt.paths();
+
+    assert_eq!(expected_paths, actual_paths);
 }
 
-// TODO: add test for add_path function and check correctness of leaf determination (requires
-// inner_nodes iter)
+#[test]
+fn leaves() {
+    let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
+    let expected_root = mt.root();
+
+    let ms = MerkleStore::from(&mt);
+
+    let path33 = ms.get_path(expected_root, NODE33).unwrap();
+    let path22 = ms.get_path(expected_root, NODE22).unwrap();
+
+    let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
+
+    let value11 = mt.get_node(NODE11).unwrap().into();
+    let value20 = mt.get_node(NODE20).unwrap().into();
+    let value32 = mt.get_node(NODE32).unwrap().into();
+    let value33 = mt.get_node(NODE33).unwrap().into();
+
+    let leaves = vec![(NODE11, value11), (NODE20, value20), (NODE32, value32), (NODE33, value33)];
+
+    let expected_leaves = leaves.iter().map(|&tuple| tuple);
+    assert!(expected_leaves.eq(pmt.leaves()));
+
+    pmt.add_path(2, path22.value.into(), path22.path).unwrap();
+
+    let value20 = mt.get_node(NODE20).unwrap().into();
+    let value22 = mt.get_node(NODE22).unwrap().into();
+    let value23 = mt.get_node(NODE23).unwrap().into();
+    let value32 = mt.get_node(NODE32).unwrap().into();
+    let value33 = mt.get_node(NODE33).unwrap().into();
+
+    let leaves = vec![
+        (NODE20, value20),
+        (NODE22, value22),
+        (NODE23, value23),
+        (NODE32, value32),
+        (NODE33, value33),
+    ];
+
+    let expected_leaves = leaves.iter().map(|&tuple| tuple);
+    assert!(expected_leaves.eq(pmt.leaves()));
+}

src/utils.rs

@@ -2,6 +2,12 @@ use super::Word;
 use crate::utils::string::String;
 use core::fmt::{self, Write};
 
+#[cfg(not(feature = "std"))]
+pub use alloc::format;
+
+#[cfg(feature = "std")]
+pub use std::format;
+
 // RE-EXPORTS
 // ================================================================================================
 pub use winter_utils::{