refactor: improve tests, add error tests

src/merkle/mod.rs

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

src/merkle/partial_mt/mod.rs

@@ -101,17 +101,6 @@ impl PartialMerkleTree {
         self.leaves.contains(&index)
     }
 
-    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();
-            }
-            node_index.move_up()
-        }
-        0
-    }
-
     /// Returns a vector of paths from every leaf to the root.
     pub fn paths(&self) -> Vec<(NodeIndex, ValuePath)> {
         let mut paths = Vec::new();
@@ -247,6 +236,11 @@ impl PartialMerkleTree {
     /// Updates value of the leaf at the specified index returning the old leaf value.
     ///
     /// This also recomputes all hashes between the leaf and the root, updating the root itself.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The depth of the specified node_index is greater than 64 or smaller than 1.
+    /// - The specified node index is not corresponding to the leaf.
     pub fn update_leaf(
         &mut self,
         node_index: NodeIndex,
@@ -282,6 +276,38 @@ impl PartialMerkleTree {
         Ok(old_value)
     }
 
+    // UTILITY FUNCTIONS
+    // --------------------------------------------------------------------------------------------
+
+    /// Utility to visualize a [PartialMerkleTree] in text.
+    pub fn print(&self) -> Result<String, fmt::Error> {
+        let indent = "  ";
+        let mut s = String::new();
+        s.push_str("root: ");
+        s.push_str(&word_to_hex(&self.root())?);
+        s.push('\n');
+        for d in 1..=self.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 = self.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)
+    }
+
     // HELPER METHODS
     // --------------------------------------------------------------------------------------------
 
@@ -301,32 +327,3 @@ 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,6 +1,6 @@
 use super::{
     super::{int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree},
-    Rpo256, ValuePath, Vec, Word,
+    ValuePath, Vec, Word,
 };
 
 // TEST DATA
@@ -13,26 +13,41 @@ 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 NODE30: NodeIndex = NodeIndex::new_unchecked(3, 0);
+const NODE31: NodeIndex = NodeIndex::new_unchecked(3, 1);
 const NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2);
 const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
 
 const VALUES8: [Word; 8] = [
-    int_to_node(1),
-    int_to_node(2),
-    int_to_node(3),
-    int_to_node(4),
-    int_to_node(5),
-    int_to_node(6),
-    int_to_node(7),
-    int_to_node(8),
+    int_to_node(30),
+    int_to_node(31),
+    int_to_node(32),
+    int_to_node(33),
+    int_to_node(34),
+    int_to_node(35),
+    int_to_node(36),
+    int_to_node(37),
 ];
 
 // TESTS
 // ================================================================================================
 
-// with_paths CONSTRUCTOR TESTS
-// ------------------------------------------------------------------------------------------------
-
+// For the Partial Merkle Tree tests we will use parts of the Merkle Tree which full form is
+// illustrated below:
+//
+//              __________ root __________
+//             /                          \
+//       ____ 10 ____                ____ 11 ____
+//      /            \              /            \
+//     20            21            22            23
+//   /    \        /    \        /    \        /    \
+// (30)  (31)    (32)  (33)    (34)  (35)    (36)  (37)
+//
+// Where node number is a concatenation of its depth and index. For example, node with
+// NodeIndex(3, 5) will be labled as `35`. Leaves of the tree are shown as nodes with parenthesis
+// (33).
+
+/// Checks that root returned by `root()` function is equal to the expected one.
 #[test]
 fn get_root() {
     let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
@@ -47,6 +62,9 @@ fn get_root() {
     assert_eq!(pmt.root(), expected_root.into());
 }
 
+/// This test checks correctness of the `add_path()` and `get_path()` functions. First it creates a
+/// PMT using `add_path()` by adding Merkle Paths from node 33 and node 22 to the empty PMT. Then
+/// it checks that paths returned by `get_path()` function are equal to the expected ones.
 #[test]
 fn add_and_get_paths() {
     let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
@@ -72,6 +90,7 @@ fn add_and_get_paths() {
     assert_eq!(expected_root, *actual_root);
 }
 
+/// Checks that function `get_node` used on nodes 10 and 32 returns expected values.
 #[test]
 fn get_node() {
     let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
@@ -87,67 +106,36 @@ fn get_node() {
     assert_eq!(ms.get_node(expected_root, NODE10).unwrap(), *pmt.get_node(NODE10).unwrap());
 }
 
+/// Updates leaves of the PMT using `update_leaf()` function and checks that new root of the tree
+/// is equal to the expected one.
 #[test]
 fn update_leaf() {
-    let mut mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
+    let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
     let root = mt.root();
 
-    let ms = MerkleStore::from(&mt);
+    let mut ms = MerkleStore::from(&mt);
     let path33 = ms.get_path(root, NODE33).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, new_value32).unwrap();
-    let expected_root = mt.root();
+    let expected_root = ms.set_node(root, NODE32, new_value32).unwrap().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();
+    let new_value20 = int_to_node(120);
+    let expected_root = ms.set_node(expected_root, NODE20, new_value20).unwrap().root;
 
-    pmt.update_leaf(NODE20, new_value20).unwrap();
+    pmt.update_leaf(NODE20, new_value20.into()).unwrap();
     let actual_root = pmt.root();
 
     assert_eq!(expected_root, *actual_root);
 }
 
-#[test]
-fn check_leaf_depth() {
-    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 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);
-}
-
+/// Checks that paths of the PMT returned by `paths()` function are equal to the expected ones.
 #[test]
 fn get_paths() {
     let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
@@ -161,6 +149,19 @@ fn get_paths() {
     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();
+    // After PMT creation with path33 (33; 32, 20, 11) and path22 (22; 23, 10) we will have this
+    // tree:
+    //
+    //           ______root______
+    //          /                \
+    //      ___10___           ___11___
+    //     /        \         /        \
+    //   (20)       21      (22)      (23)
+    //            /    \
+    //          (32)  (33)
+    //
+    // Which have leaf nodes 20, 22, 23, 32 and 33. Hence overall we will have 5 paths -- one path
+    // for each leaf.
 
     let leaves = vec![NODE20, NODE22, NODE23, NODE32, NODE33];
     let expected_paths: Vec<(NodeIndex, ValuePath)> = leaves
@@ -181,6 +182,7 @@ fn get_paths() {
     assert_eq!(expected_paths, actual_paths);
 }
 
+// Checks correctness of leaves determination when using the `leaves()` function.
 #[test]
 fn leaves() {
     let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
@@ -192,6 +194,17 @@ fn leaves() {
     let path22 = ms.get_path(expected_root, NODE22).unwrap();
 
     let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
+    // After PMT creation with path33 (33; 32, 20, 11) we will have this tree:
+    //
+    //           ______root______
+    //          /                \
+    //      ___10___            (11)
+    //     /         \
+    //   (20)        21
+    //             /    \
+    //           (32)  (33)
+    //
+    // Which have leaf nodes 11, 20, 32 and 33.
 
     let value11 = mt.get_node(NODE11).unwrap().into();
     let value20 = mt.get_node(NODE20).unwrap().into();
@@ -204,6 +217,17 @@ fn leaves() {
     assert!(expected_leaves.eq(pmt.leaves()));
 
     pmt.add_path(2, path22.value.into(), path22.path).unwrap();
+    // After adding the path22 (22; 23, 10) to the existing PMT we will have this tree:
+    //
+    //           ______root______
+    //          /                \
+    //      ___10___           ___11___
+    //     /        \         /        \
+    //   (20)       21      (22)      (23)
+    //            /    \
+    //          (32)  (33)
+    //
+    // Which have leaf nodes 20, 22, 23, 32 and 33.
 
     let value20 = mt.get_node(NODE20).unwrap().into();
     let value22 = mt.get_node(NODE22).unwrap().into();
@@ -222,3 +246,68 @@ fn leaves() {
     let expected_leaves = leaves.iter().map(|&tuple| tuple);
     assert!(expected_leaves.eq(pmt.leaves()));
 }
+
+/// Checks that addition of the path with different root will cause an error.
+#[test]
+fn err_add_path() {
+    let path33 = vec![int_to_node(1), int_to_node(2), int_to_node(3)].into();
+    let path22 = vec![int_to_node(4), int_to_node(5)].into();
+
+    let mut pmt = PartialMerkleTree::new();
+    pmt.add_path(3, int_to_node(6).into(), path33).unwrap();
+
+    assert!(pmt.add_path(2, int_to_node(7).into(), path22).is_err());
+}
+
+/// Checks that the request of the node which is not in the PMT will cause an error.
+#[test]
+fn err_get_node() {
+    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 pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
+
+    assert!(pmt.get_node(NODE22).is_err());
+    assert!(pmt.get_node(NODE23).is_err());
+    assert!(pmt.get_node(NODE30).is_err());
+    assert!(pmt.get_node(NODE31).is_err());
+}
+
+/// Checks that the request of the path from the leaf which is not in the PMT will cause an error.
+#[test]
+fn err_get_path() {
+    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 pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
+
+    assert!(pmt.get_path(NODE22).is_err());
+    assert!(pmt.get_path(NODE23).is_err());
+    assert!(pmt.get_path(NODE30).is_err());
+    assert!(pmt.get_path(NODE31).is_err());
+}
+
+#[test]
+fn err_update_leaf() {
+    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 mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
+
+    assert!(pmt.update_leaf(NODE22, int_to_node(22).into()).is_err());
+    assert!(pmt.update_leaf(NODE23, int_to_node(23).into()).is_err());
+    assert!(pmt.update_leaf(NODE30, int_to_node(30).into()).is_err());
+    assert!(pmt.update_leaf(NODE31, int_to_node(31).into()).is_err());
+}