refactor: small fixes

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<u8, MerkleError> {
-        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<Option<RpoDigest>, MerkleError> {
+    ) -> Result<RpoDigest, MerkleError> {
         // 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);
         }
 

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