From b4324475b6eed4f12af7e8cbca5c1797de78cb85 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Thu, 1 Jun 2023 17:29:01 +0300 Subject: [PATCH] feat: change constructor from with_leaves to with_paths --- src/merkle/partial_mt/mod.rs | 189 ++++++++++++----- src/merkle/partial_mt/tests.rs | 357 ++++++++++++++++----------------- 2 files changed, 315 insertions(+), 231 deletions(-) diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index ea5fa9b..8b793b6 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/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, leaves: BTreeSet, @@ -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(entries: R) -> Result + /// Analogous to [Self::add_path]. + pub fn with_paths(paths: I) -> Result where - R: IntoIterator, - I: Iterator + ExactSizeIterator, + I: IntoIterator, { // 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 { - self.nodes - .get(&index) - .ok_or(MerkleError::NodeNotInSet(index)) - .map(|hash| **hash) + pub fn get_node(&self, index: NodeIndex) -> Result { + 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 { + 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, 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 { 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 { + 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(()) + } } diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index 8cf6c1f..2a57ec8 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/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 = 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 = 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)> = 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() + } }