feat: change constructor from with_leaves to with_paths

src/merkle/partial_mt/mod.rs

@@ -13,7 +13,6 @@ mod tests;
 ///
 /// The root of the tree is recomputed on each new leaf update.
 pub struct PartialMerkleTree {
-    root: RpoDigest,
     max_depth: u8,
     nodes: BTreeMap<NodeIndex, RpoDigest>,
     leaves: BTreeSet<NodeIndex>,
@@ -38,50 +37,34 @@ impl PartialMerkleTree {
     /// Maximum supported depth.
     pub const MAX_DEPTH: u8 = 64;
 
+    pub const ROOT_INDEX: NodeIndex = NodeIndex::new_unchecked(0, 0);
+
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
 
     /// Returns a new emply [PartialMerkleTree].
     pub fn new() -> Self {
         PartialMerkleTree {
-            root: Self::EMPTY_DIGEST,
             max_depth: 0,
             nodes: BTreeMap::new(),
             leaves: BTreeSet::new(),
         }
     }
 
-    /// Returns a new [PartialMerkleTree] instantiated with leaves set as specified by the provided
-    /// entries.
+    /// Appends the provided paths iterator into the set.
     ///
-    /// # Errors
-    /// Returns an error if:
-    /// - If the depth is 0 or is greater than 64.
-    /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
-    /// - The provided entries contain multiple values for the same key.
-    pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
+    /// Analogous to [Self::add_path].
+    pub fn with_paths<I>(paths: I) -> Result<Self, MerkleError>
     where
-        R: IntoIterator<IntoIter = I>,
-        I: Iterator<Item = (NodeIndex, RpoDigest)> + ExactSizeIterator,
+        I: IntoIterator<Item = (NodeIndex, Word, MerklePath)>,
     {
         // create an empty tree
-        let mut tree = PartialMerkleTree::new();
-
-        // check if the number of leaves can be accommodated by the tree's depth; we use a min
-        // depth of 63 because we consider passing in a vector of size 2^64 infeasible.
-        let entries = entries.into_iter();
-        let max = (1_u64 << 63) as usize;
-        if entries.len() > max {
-            return Err(MerkleError::InvalidNumEntries(max, entries.len()));
-        }
+        let tree = PartialMerkleTree::new();
 
-        for (node_index, rpo_digest) in entries {
-            let old_value = tree.update_leaf(node_index, rpo_digest)?;
-            if old_value != Self::EMPTY_DIGEST {
-                return Err(MerkleError::DuplicateValuesForIndex(node_index.value()));
-            }
-        }
-        Ok(tree)
+        paths.into_iter().try_fold(tree, |mut tree, (index, value, path)| {
+            tree.add_path(index, value, path)?;
+            Ok(tree)
+        })
     }
 
     // PUBLIC ACCESSORS
@@ -89,12 +72,11 @@ impl PartialMerkleTree {
 
     /// Returns the root of this Merkle tree.
     pub fn root(&self) -> Word {
-        self.root.into()
+        *self.nodes.get(&Self::ROOT_INDEX).cloned().unwrap_or(Self::EMPTY_DIGEST)
     }
 
     /// Returns the depth of this Merkle tree.
-    // TODO: maybe it's better to rename it to the `max_depth`
-    pub fn depth(&self) -> u8 {
+    pub fn max_depth(&self) -> u8 {
         self.max_depth
     }
 
@@ -102,11 +84,26 @@ impl PartialMerkleTree {
     ///
     /// # Errors
     /// Returns an error if the specified NodeIndex is not contained in the nodes map.
-    pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
-        self.nodes
-            .get(&index)
-            .ok_or(MerkleError::NodeNotInSet(index))
-            .map(|hash| **hash)
+    pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
+        self.nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index)).map(|hash| *hash)
+    }
+
+    /// Returns true if provided index contains in the leaves set, false otherwise.
+    pub fn is_leaf(&self, index: NodeIndex) -> bool {
+        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)?;
+        for _ in 0..node_index.depth() {
+            if self.leaves.contains(&node_index) {
+                return Ok(node_index.depth());
+            }
+            node_index.move_up()
+        }
+        // we don't have an error for this case, maybe it makes sense to create a new error, something like
+        // NoLeafForIndex("There is no leaf for provided index"). But it will be used almost never.
+        Err(MerkleError::NodeNotInSet(node_index))
     }
 
     /// Returns a value of the leaf at the specified NodeIndex.
@@ -128,7 +125,7 @@ impl PartialMerkleTree {
         }
     }
 
-    /// Returns a map of the all
+    /// Returns a map of paths from every leaf to the root.
     pub fn paths(&self) -> Result<BTreeMap<&NodeIndex, MerklePath>, MerkleError> {
         let mut paths = BTreeMap::new();
         for leaf_index in self.leaves.iter() {
@@ -150,7 +147,7 @@ impl PartialMerkleTree {
     pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
         if index.is_root() {
             return Err(MerkleError::DepthTooSmall(index.depth()));
-        } else if index.depth() > self.depth() {
+        } else if index.depth() > self.max_depth() {
             return Err(MerkleError::DepthTooBig(index.depth() as u64));
         }
 
@@ -160,12 +157,7 @@ impl PartialMerkleTree {
 
         let mut path = Vec::new();
         for _ in 0..index.depth() {
-            let is_right = index.is_value_odd();
-            let sibling_index = if is_right {
-                NodeIndex::new(index.depth(), index.value() - 1)?
-            } else {
-                NodeIndex::new(index.depth(), index.value() + 1)?
-            };
+            let sibling_index = Self::get_sibling_index(&index)?;
             index.move_up();
             let sibling_hash =
                 self.nodes.get(&sibling_index).cloned().unwrap_or(Self::EMPTY_DIGEST);
@@ -206,6 +198,69 @@ impl PartialMerkleTree {
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
+    /// Adds the nodes of the specified Merkle path to this [PartialMerkleTree]. The `index_value`
+    /// and `value` parameters specify the leaf node at which the path starts.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The depth of the specified node_index is greater than 64 or smaller than 1.
+    /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
+    ///   different root).
+    pub fn add_path(
+        &mut self,
+        index_value: NodeIndex,
+        value: Word,
+        mut path: MerklePath,
+    ) -> Result<(), MerkleError> {
+        Self::check_depth(index_value.depth())?;
+        self.update_depth(index_value.depth());
+
+        // add node index to the leaves set
+        self.leaves.insert(index_value);
+        let sibling_node_index = Self::get_sibling_index(&index_value)?;
+        self.leaves.insert(sibling_node_index);
+
+        // add first two nodes to the nodes map
+        self.nodes.insert(index_value, value.into());
+        self.nodes.insert(sibling_node_index, path[0].into());
+
+        // update the current path
+        let parity = index_value.value() & 1;
+        path.insert(parity as usize, value);
+
+        // traverse to the root, updating the nodes
+        let mut index_value = index_value;
+        let root = Rpo256::merge(&[path[0].into(), path[1].into()]);
+        let root = path.iter().skip(2).copied().fold(root, |root, hash| {
+            index_value.move_up();
+            // insert calculated node to the nodes map
+            self.nodes.insert(index_value, root);
+
+            let sibling_node = Self::get_sibling_index_unchecked(&index_value);
+            // assume for now that all path nodes are leaves and add them to the leaves set
+            self.leaves.insert(sibling_node);
+
+            // insert node from Merkle path to the nodes map
+            self.nodes.insert(sibling_node, hash.into());
+
+            Rpo256::merge(&index_value.build_node(root, hash.into()))
+        });
+
+        let old_root = self.nodes.get(&Self::ROOT_INDEX).cloned().unwrap_or(Self::EMPTY_DIGEST);
+
+        // if the path set is empty (the root is all ZEROs), set the root to the root of the added
+        // path; otherwise, the root of the added path must be identical to the current root
+        if old_root == Self::EMPTY_DIGEST {
+            self.nodes.insert(Self::ROOT_INDEX, root);
+        } else if old_root != root {
+            return Err(MerkleError::ConflictingRoots([*old_root, *root].to_vec()));
+        }
+
+        self.update_leaves()?;
+
+        Ok(())
+    }
+
     /// 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.
@@ -235,17 +290,28 @@ impl PartialMerkleTree {
             let is_right = node_index.is_value_odd();
             let (left, right) = if is_right {
                 let left_index = NodeIndex::new(node_index.depth(), node_index.value() - 1)?;
-                (self.nodes.get(&left_index).cloned().unwrap_or(Self::EMPTY_DIGEST), value)
+                (
+                    self.nodes
+                        .get(&left_index)
+                        .cloned()
+                        .ok_or(MerkleError::NodeNotInSet(left_index))?,
+                    value,
+                )
             } else {
                 let right_index = NodeIndex::new(node_index.depth(), node_index.value() + 1)?;
-                (value, self.nodes.get(&right_index).cloned().unwrap_or(Self::EMPTY_DIGEST))
+                (
+                    value,
+                    self.nodes
+                        .get(&right_index)
+                        .cloned()
+                        .ok_or(MerkleError::NodeNotInSet(right_index))?,
+                )
             };
             node_index.move_up();
             value = Rpo256::merge(&[left, right]);
             self.nodes.insert(node_index, value);
         }
 
-        self.root = value;
         Ok(old_value)
     }
 
@@ -267,4 +333,33 @@ impl PartialMerkleTree {
         }
         Ok(())
     }
+
+    fn get_sibling_index(node_index: &NodeIndex) -> Result<NodeIndex, MerkleError> {
+        if node_index.is_value_odd() {
+            NodeIndex::new(node_index.depth(), node_index.value() - 1)
+        } else {
+            NodeIndex::new(node_index.depth(), node_index.value() + 1)
+        }
+    }
+
+    fn get_sibling_index_unchecked(node_index: &NodeIndex) -> NodeIndex {
+        if node_index.is_value_odd() {
+            NodeIndex::new_unchecked(node_index.depth(), node_index.value() - 1)
+        } else {
+            NodeIndex::new_unchecked(node_index.depth(), node_index.value() + 1)
+        }
+    }
+
+    // Removes from the leaves set indexes of nodes which have descendants.
+    fn update_leaves(&mut self) -> Result<(), MerkleError> {
+        for leaf_node in self.leaves.clone().iter() {
+            let left_child = NodeIndex::new(leaf_node.depth() + 1, leaf_node.value() * 2)?;
+            let right_child = NodeIndex::new(leaf_node.depth() + 1, leaf_node.value() * 2 + 1)?;
+            if self.nodes.contains_key(&left_child) || self.nodes.contains_key(&right_child) {
+                self.leaves.remove(leaf_node);
+            }
+        }
+
+        Ok(())
+    }
 }

src/merkle/partial_mt/tests.rs

@@ -1,11 +1,15 @@
+use crate::hash::rpo::RpoDigest;
+
 use super::{
-    super::{int_to_node, MerkleTree, NodeIndex, RpoDigest},
-    BTreeMap, InnerNodeInfo, MerkleError, PartialMerkleTree, Rpo256, Vec, Word, EMPTY_WORD,
+    super::{int_to_node, NodeIndex},
+    InnerNodeInfo, MerkleError, PartialMerkleTree, Rpo256, Vec, Word,
 };
 
 // TEST DATA
 // ================================================================================================
 
+const ROOT_NODE: NodeIndex = NodeIndex::new_unchecked(0, 0);
+
 const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0);
 const NODE11: NodeIndex = NodeIndex::new_unchecked(1, 1);
 
@@ -14,125 +18,161 @@ const NODE21: NodeIndex = NodeIndex::new_unchecked(2, 1);
 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 NODE34: NodeIndex = NodeIndex::new_unchecked(3, 4);
-const NODE35: NodeIndex = NodeIndex::new_unchecked(3, 5);
-const NODE36: NodeIndex = NodeIndex::new_unchecked(3, 6);
-const NODE37: NodeIndex = NodeIndex::new_unchecked(3, 7);
+const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
+
+// TESTS
+// ================================================================================================
 
-const KEYS4: [NodeIndex; 4] = [NODE20, NODE21, NODE22, NODE23];
+// with_paths CONSTRUCTOR TESTS
+// ------------------------------------------------------------------------------------------------
 
-const WVALUES4: [Word; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)];
-const DVALUES4: [RpoDigest; 4] = [
-    RpoDigest::new(int_to_node(1)),
-    RpoDigest::new(int_to_node(2)),
-    RpoDigest::new(int_to_node(3)),
-    RpoDigest::new(int_to_node(4)),
-];
+#[test]
+fn get_root() {
+    let leaf0 = int_to_node(0);
+    let leaf1 = int_to_node(1);
+    let leaf2 = int_to_node(2);
+    let leaf3 = int_to_node(3);
 
-const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8];
+    let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
+    let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
 
-// TESTS
-// ================================================================================================
+    let root_exp = calculate_parent_hash(parent0, 0, parent1);
+
+    let set = super::PartialMerkleTree::with_paths([(NODE20, leaf0, vec![leaf1, parent1].into())])
+        .unwrap();
+
+    assert_eq!(set.root(), root_exp);
+}
 
 #[test]
-fn build_partial_tree() {
-    // insert single value
-    let mut pmt = PartialMerkleTree::new();
-
-    let mut values = ZERO_VALUES8.to_vec();
-    let key = NODE36;
-    let new_node = int_to_node(7);
-    values[key.value() as usize] = new_node;
-
-    let hash0 = Rpo256::merge(&[int_to_node(0).into(), int_to_node(0).into()]);
-    let hash00 = Rpo256::merge(&[hash0, hash0]);
-
-    pmt.update_leaf(NODE10, hash00).expect("Failed to update leaf");
-    pmt.update_leaf(NODE22, hash0).expect("Failed to update leaf");
-    let old_value = pmt.update_leaf(key, new_node.into()).expect("Failed to update leaf");
-
-    let mt2 = MerkleTree::new(values.clone()).unwrap();
-    assert_eq!(mt2.root(), pmt.root());
-    assert_eq!(mt2.get_path(NODE36).unwrap(), pmt.get_path(NODE36).unwrap());
-    assert_eq!(*old_value, EMPTY_WORD);
-
-    // insert second value at distinct leaf branch
-    let key = NODE32;
-    let new_node = int_to_node(3);
-    values[key.value() as usize] = new_node;
-    pmt.update_leaf(NODE20, hash0).expect("Failed to update leaf");
-    let old_value = pmt.update_leaf(key, new_node.into()).expect("Failed to update leaf");
-    let mt3 = MerkleTree::new(values).unwrap();
-    assert_eq!(mt3.root(), pmt.root());
-    assert_eq!(mt3.get_path(NODE32).unwrap(), pmt.get_path(NODE32).unwrap());
-    assert_eq!(*old_value, EMPTY_WORD);
+fn add_and_get_paths() {
+    let value32 = int_to_node(32).into();
+    let value33 = int_to_node(33).into();
+    let value20 = int_to_node(20).into();
+    let value22 = int_to_node(22).into();
+    let value23 = int_to_node(23).into();
+
+    let value21 = Rpo256::merge(&[value32, value33]);
+    let value10 = Rpo256::merge(&[value20, value21]);
+    let value11 = Rpo256::merge(&[value22, value23]);
+
+    let path_33 = vec![*value32, *value20, *value11];
+
+    let path_22 = vec![*value23, *value10];
+
+    let pmt = PartialMerkleTree::with_paths([
+        (NODE33, *value33, path_33.clone().into()),
+        (NODE22, *value22, path_22.clone().into()),
+    ])
+    .unwrap();
+    let stored_path_33 = pmt.get_path(NODE33).unwrap();
+    let stored_path_22 = pmt.get_path(NODE22).unwrap();
+
+    assert_eq!(path_33, *stored_path_33);
+    assert_eq!(path_22, *stored_path_22);
 }
 
 #[test]
-fn test_depth2_tree() {
-    let tree = PartialMerkleTree::with_leaves(KEYS4.into_iter().zip(DVALUES4.into_iter())).unwrap();
-
-    // check internal structure
-    let (root, node2, node3) = compute_internal_nodes();
-    assert_eq!(root, tree.root());
-    assert_eq!(node2, tree.get_node(NODE10).unwrap());
-    assert_eq!(node3, tree.get_node(NODE11).unwrap());
-
-    // check get_node()
-    assert_eq!(WVALUES4[0], tree.get_node(NODE20).unwrap());
-    assert_eq!(WVALUES4[1], tree.get_node(NODE21).unwrap());
-    assert_eq!(WVALUES4[2], tree.get_node(NODE22).unwrap());
-    assert_eq!(WVALUES4[3], tree.get_node(NODE23).unwrap());
-
-    // check get_path(): depth 2
-    assert_eq!(vec![WVALUES4[1], node3], *tree.get_path(NODE20).unwrap());
-    assert_eq!(vec![WVALUES4[0], node3], *tree.get_path(NODE21).unwrap());
-    assert_eq!(vec![WVALUES4[3], node2], *tree.get_path(NODE22).unwrap());
-    assert_eq!(vec![WVALUES4[2], node2], *tree.get_path(NODE23).unwrap());
-
-    // check get_path(): depth 1
-    assert_eq!(vec![node3], *tree.get_path(NODE10).unwrap());
-    assert_eq!(vec![node2], *tree.get_path(NODE11).unwrap());
+fn get_node() {
+    let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
+    let hash_6 = int_to_node(6);
+    let index = NodeIndex::make(3, 6);
+    let pmt = PartialMerkleTree::with_paths([(index, hash_6, path_6.into())]).unwrap();
+
+    assert_eq!(int_to_node(6u64), *pmt.get_node(index).unwrap());
+}
+
+#[test]
+fn update_leaf() {
+    let value32 = int_to_node(32).into();
+    let value33 = int_to_node(33).into();
+    let value20 = int_to_node(20).into();
+    let value22 = int_to_node(22).into();
+    let value23 = int_to_node(23).into();
+
+    let value21 = Rpo256::merge(&[value32, value33]);
+    let value10 = Rpo256::merge(&[value20, value21]);
+    let value11 = Rpo256::merge(&[value22, value23]);
+
+    let path_33 = vec![*value32, *value20, *value11];
+
+    let path_22 = vec![*value23, *value10];
+
+    let mut pmt = PartialMerkleTree::with_paths([
+        (NODE33, *value33, path_33.into()),
+        (NODE22, *value22, path_22.into()),
+    ])
+    .unwrap();
+
+    let new_value32 = int_to_node(132).into();
+    let new_value21 = Rpo256::merge(&[new_value32, value33]);
+    let new_value10 = Rpo256::merge(&[value20, new_value21]);
+    let expected_root = Rpo256::merge(&[new_value10, value11]);
+
+    let old_leaf = pmt.update_leaf(NODE32, new_value32).unwrap();
+
+    assert_eq!(value32, old_leaf);
+
+    let new_root = pmt.root();
+
+    assert_eq!(new_root, *expected_root);
 }
 
 #[test]
 fn test_inner_node_iterator() -> Result<(), MerkleError> {
-    let tree = PartialMerkleTree::with_leaves(KEYS4.into_iter().zip(DVALUES4.into_iter())).unwrap();
-
-    // check depth 2
-    assert_eq!(WVALUES4[0], tree.get_node(NODE20).unwrap());
-    assert_eq!(WVALUES4[1], tree.get_node(NODE21).unwrap());
-    assert_eq!(WVALUES4[2], tree.get_node(NODE22).unwrap());
-    assert_eq!(WVALUES4[3], tree.get_node(NODE23).unwrap());
-
-    // get parent nodes
-    let root = tree.root();
-    let l1n0 = tree.get_node(NODE10)?;
-    let l1n1 = tree.get_node(NODE11)?;
-    let l2n0 = tree.get_node(NODE20)?;
-    let l2n1 = tree.get_node(NODE21)?;
-    let l2n2 = tree.get_node(NODE22)?;
-    let l2n3 = tree.get_node(NODE23)?;
-
-    let nodes: Vec<InnerNodeInfo> = tree.inner_nodes().collect();
+    let value32 = int_to_node(32).into();
+    let value33 = int_to_node(33).into();
+    let value20 = int_to_node(20).into();
+    let value22 = int_to_node(22).into();
+    let value23 = int_to_node(23).into();
+
+    let value21 = Rpo256::merge(&[value32, value33]);
+    let value10 = Rpo256::merge(&[value20, value21]);
+    let value11 = Rpo256::merge(&[value22, value23]);
+    let root = Rpo256::merge(&[value10, value11]);
+
+    let path_33 = vec![*value32, *value20, *value11];
+
+    let path_22 = vec![*value23, *value10];
+
+    let pmt = PartialMerkleTree::with_paths([
+        (NODE33, *value33, path_33.into()),
+        (NODE22, *value22, path_22.into()),
+    ])
+    .unwrap();
+
+    assert_eq!(root, pmt.get_node(ROOT_NODE).unwrap());
+    assert_eq!(value10, pmt.get_node(NODE10).unwrap());
+    assert_eq!(value11, pmt.get_node(NODE11).unwrap());
+    assert_eq!(value20, pmt.get_node(NODE20).unwrap());
+    assert_eq!(value21, pmt.get_node(NODE21).unwrap());
+    assert_eq!(value22, pmt.get_node(NODE22).unwrap());
+    assert_eq!(value23, pmt.get_node(NODE23).unwrap());
+    assert_eq!(value32, pmt.get_node(NODE32).unwrap());
+    assert_eq!(value33, pmt.get_node(NODE33).unwrap());
+
+    let nodes: Vec<InnerNodeInfo> = pmt.inner_nodes().collect();
     let expected = vec![
         InnerNodeInfo {
-            value: root,
-            left: l1n0,
-            right: l1n1,
+            value: *root,
+            left: *value10,
+            right: *value11,
+        },
+        InnerNodeInfo {
+            value: *value10,
+            left: *value20,
+            right: *value21,
         },
         InnerNodeInfo {
-            value: l1n0,
-            left: l2n0,
-            right: l2n1,
+            value: *value11,
+            left: *value22,
+            right: *value23,
         },
         InnerNodeInfo {
-            value: l1n1,
-            left: l2n2,
-            right: l2n3,
+            value: *value21,
+            left: *value32,
+            right: *value33,
         },
     ];
     assert_eq!(nodes, expected);
@@ -141,92 +181,41 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> {
 }
 
 #[test]
-fn small_tree_opening_is_consistent() {
-    //        ____k____
-    //       /         \
-    //     _i_         _j_
-    //    /   \       /   \
-    //   e     f     g     h
-    //  / \   / \   / \   / \
-    // a   b 0   0 c   0 0   d
-
-    let z = Word::from(RpoDigest::default());
-
-    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
-    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
-    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
-    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
-
-    let e = Word::from(Rpo256::merge(&[a.into(), b.into()]));
-    let f = Word::from(Rpo256::merge(&[z.into(), z.into()]));
-    let g = Word::from(Rpo256::merge(&[c.into(), z.into()]));
-    let h = Word::from(Rpo256::merge(&[z.into(), d.into()]));
-
-    let i = Word::from(Rpo256::merge(&[e.into(), f.into()]));
-    let j = Word::from(Rpo256::merge(&[g.into(), h.into()]));
-
-    let k = Word::from(Rpo256::merge(&[i.into(), j.into()]));
-
-    // let depth = 3;
-    // let entries = vec![(0, a), (1, b), (4, c), (7, d)];
-    // let tree = SimpleSmt::with_leaves(depth, entries).unwrap();
-    let entries = BTreeMap::from([
-        (NODE30, a.into()),
-        (NODE31, b.into()),
-        (NODE34, c.into()),
-        (NODE37, d.into()),
-        (NODE21, f.into()),
-    ]);
-
-    let tree = PartialMerkleTree::with_leaves(entries).unwrap();
-
-    assert_eq!(tree.root(), k);
-
-    let cases: Vec<(NodeIndex, Vec<Word>)> = vec![
-        (NODE30, vec![b, f, j]),
-        (NODE31, vec![a, f, j]),
-        (NODE34, vec![z, h, i]),
-        (NODE37, vec![z, g, i]),
-        (NODE20, vec![f, j]),
-        (NODE21, vec![e, j]),
-        (NODE22, vec![h, i]),
-        (NODE23, vec![g, i]),
-        (NODE10, vec![j]),
-        (NODE11, vec![i]),
-    ];
-
-    for (index, path) in cases {
-        let opening = tree.get_path(index).unwrap();
-
-        assert_eq!(path, *opening);
-    }
-}
-
-#[test]
-fn fail_on_duplicates() {
-    let entries = [
-        (NODE31, int_to_node(1).into()),
-        (NODE35, int_to_node(2).into()),
-        (NODE31, int_to_node(3).into()),
-    ];
-    let smt = PartialMerkleTree::with_leaves(entries);
-    assert!(smt.is_err());
-}
-
-#[test]
-fn with_no_duplicates_empty_node() {
-    let entries = [(NODE31, int_to_node(0).into()), (NODE35, int_to_node(2).into())];
-    let smt = PartialMerkleTree::with_leaves(entries);
-    assert!(smt.is_ok());
+fn check_leaf_depth() {
+    let value32: RpoDigest = int_to_node(32).into();
+    let value33: RpoDigest = int_to_node(33).into();
+    let value20: RpoDigest = int_to_node(20).into();
+    let value22 = int_to_node(22).into();
+    let value23 = int_to_node(23).into();
+
+    let value11 = Rpo256::merge(&[value22, value23]);
+
+    let path_33 = vec![*value32, *value20, *value11];
+
+    let pmt = PartialMerkleTree::with_paths([(NODE33, *value33, path_33.into())]).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);
 }
 
 // HELPER FUNCTIONS
 // --------------------------------------------------------------------------------------------
 
-fn compute_internal_nodes() -> (Word, Word, Word) {
-    let node2 = Rpo256::hash_elements(&[WVALUES4[0], WVALUES4[1]].concat());
-    let node3 = Rpo256::hash_elements(&[WVALUES4[2], WVALUES4[3]].concat());
-    let root = Rpo256::merge(&[node2, node3]);
-
-    (root.into(), node2.into(), node3.into())
+/// Calculates the hash of the parent node by two sibling ones
+/// - node — current node
+/// - node_pos — position of the current node
+/// - sibling — neighboring vertex in the tree
+fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
+    let parity = node_pos & 1;
+    if parity == 0 {
+        Rpo256::merge(&[node.into(), sibling.into()]).into()
+    } else {
+        Rpo256::merge(&[sibling.into(), node.into()]).into()
+    }
 }