From b9def61e28c52ec7d956d5a276d2fa4590431dc5 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Tue, 13 Jun 2023 16:14:07 +0300 Subject: [PATCH] refactor: improve tests, add error tests --- src/merkle/mod.rs | 2 +- src/merkle/partial_mt/mod.rs | 77 +++++++------ src/merkle/partial_mt/tests.rs | 193 ++++++++++++++++++++++++--------- 3 files changed, 179 insertions(+), 93 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 73f149e..6f7d2ab 100644 --- a/src/merkle/mod.rs +++ b/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 // ================================================================================================ diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index fd9f1c6..3807de3 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/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 { + 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 { - 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) -} diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index 85d3036..ba62896 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/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()); +}