diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index 7115557..ab7b644 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -3,6 +3,9 @@ use super::{ Word, EMPTY_WORD, }; +extern crate alloc; +use alloc::format; + #[cfg(test)] mod tests; @@ -99,15 +102,17 @@ impl PartialMerkleTree { self.leaves.contains(&index) } - pub fn get_leaf_depth(&self, index: u64) -> Result { - let mut node_index = NodeIndex::new(self.max_depth(), 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."); for _ in 0..node_index.depth() { if self.leaves.contains(&node_index) { - return Ok(node_index.depth()); + return node_index.depth(); } node_index.move_up() } - Ok(0_u8) + 0 } /// Returns a vector of paths from every leaf to the root. @@ -164,13 +169,7 @@ impl PartialMerkleTree { self.leaves.iter().map(|&leaf| { ( leaf, - self.get_node(leaf).unwrap_or_else(|_| { - panic!( - "Leaf with node index ({}, {}) is not in the nodes map", - leaf.depth(), - leaf.value() - ) - }), + self.get_node(leaf).expect(&format!("Leaf with {leaf} is not in the nodes map")), ) }) } @@ -221,12 +220,12 @@ impl PartialMerkleTree { // Insert node from Merkle path to the nodes map. This sibling node becomes a leaf only // if it is a new node (it wasn't in nodes map). - // Node can be in 3 states: internal node, leaf of the tree and not a node at all. + // Node can be in 3 states: internal node, leaf of the tree and not a tree node at all. // - Internal node can only stay in this state -- addition of a new path can't make it // a leaf or remove it from the tree. // - Leaf node can stay in the same state (remain a leaf) or can become an internal // node. In the first case we don't need to do anything, and the second case is handled - // in the line 219. + // by the call of `self.leaves.remove(&index_value);` // - New node can be a calculated node or a "sibling" node from a Merkle Path: // --- Calculated node, obviously, never can be a leaf. // --- Sibling node can be only a leaf, because otherwise it is not a new node. @@ -255,7 +254,7 @@ impl PartialMerkleTree { &mut self, node_index: NodeIndex, value: RpoDigest, - ) -> Result, MerkleError> { + ) -> Result { // check correctness of the depth and update it Self::check_depth(node_index.depth())?; self.update_depth(node_index.depth()); @@ -264,10 +263,13 @@ impl PartialMerkleTree { self.leaves.insert(node_index); // add node value to the nodes Map - let old_value = self.nodes.insert(node_index, value); + let old_value = self + .nodes + .insert(node_index, value) + .ok_or(MerkleError::NodeNotInSet(node_index))?; // if the old value and new value are the same, there is nothing to update - if old_value.is_some() && value == old_value.unwrap() { + if value == old_value { return Ok(old_value); } diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index c612f66..6151421 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -35,8 +35,7 @@ fn get_root() { let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); let expected_root = mt.root(); - let mut store = MerkleStore::new(); - let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); + let ms = MerkleStore::from(&mt); let path33 = ms.get_path(expected_root, NODE33).unwrap(); @@ -50,8 +49,7 @@ fn add_and_get_paths() { let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); let expected_root = mt.root(); - let mut store = MerkleStore::new(); - let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); + let ms = MerkleStore::from(&mt); let expected_path33 = ms.get_path(expected_root, NODE33).unwrap(); let expected_path22 = ms.get_path(expected_root, NODE22).unwrap(); @@ -74,8 +72,7 @@ fn get_node() { let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); let expected_root = mt.root(); - let mut store = MerkleStore::new(); - let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); + let ms = MerkleStore::from(&mt); let path33 = ms.get_path(expected_root, NODE33).unwrap(); @@ -90,8 +87,7 @@ fn update_leaf() { let mut mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); let root = mt.root(); - let mut store = MerkleStore::new(); - let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); + let ms = MerkleStore::from(&mt); let path33 = ms.get_path(root, NODE33).unwrap(); let mut pmt = @@ -112,22 +108,20 @@ fn check_leaf_depth() { let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); let expected_root = mt.root(); - let mut store = MerkleStore::new(); - let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); + let ms = MerkleStore::from(&mt); let path33 = ms.get_path(expected_root, NODE33).unwrap(); let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap(); - assert_eq!(pmt.get_leaf_depth(0).unwrap(), 2); - assert_eq!(pmt.get_leaf_depth(1).unwrap(), 2); - assert_eq!(pmt.get_leaf_depth(2).unwrap(), 3); - assert_eq!(pmt.get_leaf_depth(3).unwrap(), 3); - assert_eq!(pmt.get_leaf_depth(4).unwrap(), 1); - assert_eq!(pmt.get_leaf_depth(5).unwrap(), 1); - assert_eq!(pmt.get_leaf_depth(6).unwrap(), 1); - assert_eq!(pmt.get_leaf_depth(7).unwrap(), 1); - assert!(pmt.get_leaf_depth(8).is_err()); + 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); } // TODO: add test for add_path function and check correctness of leaf determination (requires