From 23f448fb3362eae6df030d2613d8ff818c22f33c Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Tue, 30 May 2023 20:23:36 +0300 Subject: [PATCH 1/9] feat: partial Merkle tree --- src/merkle/mod.rs | 3 + src/merkle/partial_mt/mod.rs | 270 +++++++++++++++++++++++++++++++++ src/merkle/partial_mt/tests.rs | 232 ++++++++++++++++++++++++++++ 3 files changed, 505 insertions(+) create mode 100644 src/merkle/partial_mt/mod.rs create mode 100644 src/merkle/partial_mt/tests.rs diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 631b960..6f7d2ab 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -39,6 +39,9 @@ pub use store::MerkleStore; mod node; pub use node::InnerNodeInfo; +mod partial_mt; +pub use partial_mt::PartialMerkleTree; + // ERRORS // ================================================================================================ diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs new file mode 100644 index 0000000..ea5fa9b --- /dev/null +++ b/src/merkle/partial_mt/mod.rs @@ -0,0 +1,270 @@ +use super::{ + BTreeMap, BTreeSet, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, + Word, EMPTY_WORD, +}; + +#[cfg(test)] +mod tests; + +// PARTIAL MERKLE TREE +// ================================================================================================ + +/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values. +/// +/// The root of the tree is recomputed on each new leaf update. +pub struct PartialMerkleTree { + root: RpoDigest, + max_depth: u8, + nodes: BTreeMap, + leaves: BTreeSet, +} + +impl Default for PartialMerkleTree { + fn default() -> Self { + Self::new() + } +} + +impl PartialMerkleTree { + // CONSTANTS + // -------------------------------------------------------------------------------------------- + + /// An RpoDigest consisting of 4 ZERO elements. + pub const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD); + + /// Minimum supported depth. + pub const MIN_DEPTH: u8 = 1; + + /// Maximum supported depth. + pub const MAX_DEPTH: u8 = 64; + + // 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. + /// + /// # 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 + where + R: IntoIterator, + I: Iterator + ExactSizeIterator, + { + // 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())); + } + + 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) + } + + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns the root of this Merkle tree. + pub fn root(&self) -> Word { + self.root.into() + } + + /// Returns the depth of this Merkle tree. + // TODO: maybe it's better to rename it to the `max_depth` + pub fn depth(&self) -> u8 { + self.max_depth + } + + /// Returns a node at the specified NodeIndex. + /// + /// # 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) + } + + /// Returns a value of the leaf at the specified NodeIndex. + /// + /// # Errors + /// Returns an error if the NodeIndex is not contained in the leaves set. + pub fn get_leaf(&self, index: NodeIndex) -> Result { + if !self.leaves.contains(&index) { + // This error not really suitable in this situation, should I create a new error? + Err(MerkleError::InvalidIndex { + depth: index.depth(), + value: index.value(), + }) + } else { + self.nodes + .get(&index) + .ok_or(MerkleError::NodeNotInSet(index)) + .map(|hash| **hash) + } + } + + /// Returns a map of the all + pub fn paths(&self) -> Result, MerkleError> { + let mut paths = BTreeMap::new(); + for leaf_index in self.leaves.iter() { + let index = *leaf_index; + paths.insert(leaf_index, self.get_path(index)?); + } + Ok(paths) + } + + /// Returns a Merkle path from the node at the specified index to the root. + /// + /// The node itself is not included in the path. + /// + /// # Errors + /// Returns an error if: + /// - the specified index has depth set to 0 or the depth is greater than the depth of this + /// Merkle tree. + /// - the specified index is not contained in the nodes map. + 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() { + return Err(MerkleError::DepthTooBig(index.depth() as u64)); + } + + if !self.nodes.contains_key(&index) { + return Err(MerkleError::NodeNotInSet(index)); + } + + 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)? + }; + index.move_up(); + let sibling_hash = + self.nodes.get(&sibling_index).cloned().unwrap_or(Self::EMPTY_DIGEST); + path.push(Word::from(sibling_hash)); + } + Ok(MerklePath::new(path)) + } + + // ITERATORS + // -------------------------------------------------------------------------------------------- + + /// Returns an iterator over the leaves of this [PartialMerkleTree]. + pub fn leaves(&self) -> impl Iterator { + self.nodes + .iter() + .filter(|(index, _)| self.leaves.contains(index)) + .map(|(index, hash)| (*index, &(**hash))) + } + + /// Returns an iterator over the inner nodes of this Merkle tree. + pub fn inner_nodes(&self) -> impl Iterator + '_ { + let inner_nodes = self.nodes.iter().filter(|(index, _)| !self.leaves.contains(index)); + inner_nodes.map(|(index, digest)| { + let left_index = NodeIndex::new(index.depth() + 1, index.value() * 2) + .expect("Failure to get left child index"); + let right_index = NodeIndex::new(index.depth() + 1, index.value() * 2 + 1) + .expect("Failure to get right child index"); + let left_hash = self.nodes.get(&left_index).cloned().unwrap_or(Self::EMPTY_DIGEST); + let right_hash = self.nodes.get(&right_index).cloned().unwrap_or(Self::EMPTY_DIGEST); + InnerNodeInfo { + value: **digest, + left: *left_hash, + right: *right_hash, + } + }) + } + + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// 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. + pub fn update_leaf( + &mut self, + node_index: NodeIndex, + value: RpoDigest, + ) -> Result { + // check correctness of the depth and update it + Self::check_depth(node_index.depth())?; + self.update_depth(node_index.depth()); + + // insert NodeIndex to the leaves Set + self.leaves.insert(node_index); + + // add node value to the nodes Map + let old_value = self.nodes.insert(node_index, value).unwrap_or(Self::EMPTY_DIGEST); + + // if the old value and new value are the same, there is nothing to update + if value == old_value { + return Ok(value); + } + + let mut node_index = node_index; + let mut value = value; + for _ in 0..node_index.depth() { + 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) + } 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)) + }; + node_index.move_up(); + value = Rpo256::merge(&[left, right]); + self.nodes.insert(node_index, value); + } + + self.root = value; + Ok(old_value) + } + + // HELPER METHODS + // -------------------------------------------------------------------------------------------- + + /// Updates depth value with the maximum of current and provided depth. + fn update_depth(&mut self, new_depth: u8) { + self.max_depth = new_depth.max(self.max_depth); + } + + /// Returns an error if the depth is 0 or is greater than 64. + fn check_depth(depth: u8) -> Result<(), MerkleError> { + // validate the range of the depth. + if depth < Self::MIN_DEPTH { + return Err(MerkleError::DepthTooSmall(depth)); + } else if Self::MAX_DEPTH < depth { + return Err(MerkleError::DepthTooBig(depth as u64)); + } + Ok(()) + } +} diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs new file mode 100644 index 0000000..8cf6c1f --- /dev/null +++ b/src/merkle/partial_mt/tests.rs @@ -0,0 +1,232 @@ +use super::{ + super::{int_to_node, MerkleTree, NodeIndex, RpoDigest}, + BTreeMap, InnerNodeInfo, MerkleError, PartialMerkleTree, Rpo256, Vec, Word, EMPTY_WORD, +}; + +// TEST DATA +// ================================================================================================ + +const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0); +const NODE11: NodeIndex = NodeIndex::new_unchecked(1, 1); + +const NODE20: NodeIndex = NodeIndex::new_unchecked(2, 0); +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 KEYS4: [NodeIndex; 4] = [NODE20, NODE21, NODE22, NODE23]; + +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)), +]; + +const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8]; + +// TESTS +// ================================================================================================ + +#[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); +} + +#[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()); +} + +#[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 expected = vec![ + InnerNodeInfo { + value: root, + left: l1n0, + right: l1n1, + }, + InnerNodeInfo { + value: l1n0, + left: l2n0, + right: l2n1, + }, + InnerNodeInfo { + value: l1n1, + left: l2n2, + right: l2n3, + }, + ]; + assert_eq!(nodes, expected); + + Ok(()) +} + +#[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()); +} + +// 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()) +} From b4324475b6eed4f12af7e8cbca5c1797de78cb85 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Thu, 1 Jun 2023 17:29:01 +0300 Subject: [PATCH 2/9] 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() + } } From ebf71c2dc7c1981bf1524aa11566c946136b586a Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Fri, 2 Jun 2023 21:57:33 +0300 Subject: [PATCH 3/9] refactor: optimize code, remove not momentarily necessary functions --- src/merkle/mod.rs | 5 + src/merkle/partial_mt/mod.rs | 197 +++++++++++++-------------------- src/merkle/partial_mt/tests.rs | 145 +++++++----------------- 3 files changed, 118 insertions(+), 229 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 6f7d2ab..c3cbdf6 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -101,3 +101,8 @@ impl std::error::Error for MerkleError {} const fn int_to_node(value: u64) -> Word { [Felt::new(value), ZERO, ZERO, ZERO] } + +#[cfg(test)] +const fn int_to_digest(value: u64) -> RpoDigest { + RpoDigest::new([Felt::new(value), ZERO, ZERO, ZERO]) +} diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index 8b793b6..be46802 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -1,15 +1,25 @@ use super::{ - BTreeMap, BTreeSet, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, + BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec, Word, EMPTY_WORD, }; #[cfg(test)] mod tests; +// CONSTANTS +// ================================================================================================ + +/// Index of the root node. +const ROOT_INDEX: NodeIndex = NodeIndex::root(); + +/// An RpoDigest consisting of 4 ZERO elements. +const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD); + // PARTIAL MERKLE TREE // ================================================================================================ -/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values. +/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values. Partial Merkle +/// Tree allows to create Merkle Tree by providing Merkle paths of different lengths. /// /// The root of the tree is recomputed on each new leaf update. pub struct PartialMerkleTree { @@ -28,17 +38,12 @@ impl PartialMerkleTree { // CONSTANTS // -------------------------------------------------------------------------------------------- - /// An RpoDigest consisting of 4 ZERO elements. - pub const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD); - /// Minimum supported depth. pub const MIN_DEPTH: u8 = 1; /// Maximum supported depth. pub const MAX_DEPTH: u8 = 64; - pub const ROOT_INDEX: NodeIndex = NodeIndex::new_unchecked(0, 0); - // CONSTRUCTORS // -------------------------------------------------------------------------------------------- @@ -56,7 +61,7 @@ impl PartialMerkleTree { /// Analogous to [Self::add_path]. pub fn with_paths(paths: I) -> Result where - I: IntoIterator, + I: IntoIterator, { // create an empty tree let tree = PartialMerkleTree::new(); @@ -71,8 +76,8 @@ impl PartialMerkleTree { // -------------------------------------------------------------------------------------------- /// Returns the root of this Merkle tree. - pub fn root(&self) -> Word { - *self.nodes.get(&Self::ROOT_INDEX).cloned().unwrap_or(Self::EMPTY_DIGEST) + pub fn root(&self) -> RpoDigest { + self.nodes.get(&ROOT_INDEX).cloned().unwrap_or(EMPTY_DIGEST) } /// Returns the depth of this Merkle tree. @@ -101,38 +106,22 @@ impl PartialMerkleTree { } 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. - /// - /// # Errors - /// Returns an error if the NodeIndex is not contained in the leaves set. - pub fn get_leaf(&self, index: NodeIndex) -> Result { - if !self.leaves.contains(&index) { - // This error not really suitable in this situation, should I create a new error? - Err(MerkleError::InvalidIndex { - depth: index.depth(), - value: index.value(), - }) - } else { - self.nodes - .get(&index) - .ok_or(MerkleError::NodeNotInSet(index)) - .map(|hash| **hash) - } + Ok(0_u8) } - /// 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() { - let index = *leaf_index; - paths.insert(leaf_index, self.get_path(index)?); - } - Ok(paths) + /// Returns a vector of paths from every leaf to the root. + pub fn paths(&self) -> Vec<(NodeIndex, ValuePath)> { + let mut paths = Vec::new(); + self.leaves.iter().for_each(|leaf| { + paths.push(( + *leaf, + ValuePath { + value: *self.get_node(*leaf).expect("Failed to get leaf node"), + path: self.get_path(*leaf).expect("Failed to get path"), + }, + )); + }); + paths } /// Returns a Merkle path from the node at the specified index to the root. @@ -157,11 +146,11 @@ impl PartialMerkleTree { let mut path = Vec::new(); for _ in 0..index.depth() { - let sibling_index = Self::get_sibling_index(&index)?; + let sibling_index = index.sibling(); index.move_up(); - let sibling_hash = - self.nodes.get(&sibling_index).cloned().unwrap_or(Self::EMPTY_DIGEST); - path.push(Word::from(sibling_hash)); + let sibling = + self.nodes.get(&sibling_index).cloned().expect("Sibling node not in the map"); + path.push(Word::from(sibling)); } Ok(MerklePath::new(path)) } @@ -170,28 +159,18 @@ impl PartialMerkleTree { // -------------------------------------------------------------------------------------------- /// Returns an iterator over the leaves of this [PartialMerkleTree]. - pub fn leaves(&self) -> impl Iterator { - self.nodes - .iter() - .filter(|(index, _)| self.leaves.contains(index)) - .map(|(index, hash)| (*index, &(**hash))) - } - - /// Returns an iterator over the inner nodes of this Merkle tree. - pub fn inner_nodes(&self) -> impl Iterator + '_ { - let inner_nodes = self.nodes.iter().filter(|(index, _)| !self.leaves.contains(index)); - inner_nodes.map(|(index, digest)| { - let left_index = NodeIndex::new(index.depth() + 1, index.value() * 2) - .expect("Failure to get left child index"); - let right_index = NodeIndex::new(index.depth() + 1, index.value() * 2 + 1) - .expect("Failure to get right child index"); - let left_hash = self.nodes.get(&left_index).cloned().unwrap_or(Self::EMPTY_DIGEST); - let right_hash = self.nodes.get(&right_index).cloned().unwrap_or(Self::EMPTY_DIGEST); - InnerNodeInfo { - value: **digest, - left: *left_hash, - right: *right_hash, - } + pub fn leaves(&self) -> impl Iterator + '_ { + 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() + ) + }), + ) }) } @@ -208,55 +187,60 @@ impl PartialMerkleTree { /// different root). pub fn add_path( &mut self, - index_value: NodeIndex, - value: Word, - mut path: MerklePath, + index_value: u64, + value: RpoDigest, + path: MerklePath, ) -> Result<(), MerkleError> { + let index_value = NodeIndex::new(path.len() as u8, index_value)?; + Self::check_depth(index_value.depth())?; self.update_depth(index_value.depth()); - // add node index to the leaves set + // add provided node and its sibling to the leaves set self.leaves.insert(index_value); - let sibling_node_index = Self::get_sibling_index(&index_value)?; + let sibling_node_index = index_value.sibling(); self.leaves.insert(sibling_node_index); - // add first two nodes to the nodes map - self.nodes.insert(index_value, value.into()); + // add provided node and its sibling to the nodes map + self.nodes.insert(index_value, value); 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| { + let node = Rpo256::merge(&index_value.build_node(value, path[0].into())); + let root = path.iter().skip(1).copied().fold(node, |node, hash| { index_value.move_up(); // insert calculated node to the nodes map - self.nodes.insert(index_value, root); + self.nodes.insert(index_value, node); - 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); + let sibling_node = index_value.sibling(); + // node became a leaf only if it is a new node (it wasn't in nodes map) + if !self.nodes.contains_key(&sibling_node) { + self.leaves.insert(sibling_node); + } + + // node stops being a leaf if the path contains a node which is a child of this leaf + let mut parent = index_value; + parent.move_up(); + if self.leaves.contains(&parent) { + self.leaves.remove(&parent); + } // 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())) + Rpo256::merge(&index_value.build_node(node, 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())); + if self.root() == EMPTY_DIGEST { + self.nodes.insert(ROOT_INDEX, root); + } else if self.root() != root { + return Err(MerkleError::ConflictingRoots([*self.root(), *root].to_vec())); } - self.update_leaves()?; + // self.update_leaves()?; Ok(()) } @@ -277,7 +261,7 @@ impl PartialMerkleTree { self.leaves.insert(node_index); // add node value to the nodes Map - let old_value = self.nodes.insert(node_index, value).unwrap_or(Self::EMPTY_DIGEST); + let old_value = self.nodes.insert(node_index, value).unwrap_or(EMPTY_DIGEST); // if the old value and new value are the same, there is nothing to update if value == old_value { @@ -333,33 +317,4 @@ 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 2a57ec8..1efbb6a 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -1,22 +1,14 @@ use crate::hash::rpo::RpoDigest; use super::{ - super::{int_to_node, NodeIndex}, - InnerNodeInfo, MerkleError, PartialMerkleTree, Rpo256, Vec, Word, + super::{int_to_digest, int_to_node, NodeIndex}, + PartialMerkleTree, Rpo256, }; // 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); - -const NODE20: NodeIndex = NodeIndex::new_unchecked(2, 0); -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 NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2); const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3); @@ -29,29 +21,29 @@ const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3); #[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); + let leaf0 = int_to_digest(0); + let leaf1 = int_to_digest(1); + let leaf2 = int_to_digest(2); + let leaf3 = int_to_digest(3); let parent0 = calculate_parent_hash(leaf0, 0, leaf1); let parent1 = calculate_parent_hash(leaf2, 2, leaf3); let root_exp = calculate_parent_hash(parent0, 0, parent1); - let set = super::PartialMerkleTree::with_paths([(NODE20, leaf0, vec![leaf1, parent1].into())]) - .unwrap(); + let set = + super::PartialMerkleTree::with_paths([(0, leaf0, vec![*leaf1, *parent1].into())]).unwrap(); assert_eq!(set.root(), root_exp); } #[test] 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 value32 = int_to_digest(32); + let value33 = int_to_digest(33); + let value20 = int_to_digest(20); + let value22 = int_to_digest(22); + let value23 = int_to_digest(23); let value21 = Rpo256::merge(&[value32, value33]); let value10 = Rpo256::merge(&[value20, value21]); @@ -62,8 +54,8 @@ fn add_and_get_paths() { let path_22 = vec![*value23, *value10]; let pmt = PartialMerkleTree::with_paths([ - (NODE33, *value33, path_33.clone().into()), - (NODE22, *value22, path_22.clone().into()), + (3, value33, path_33.clone().into()), + (2, value22, path_22.clone().into()), ]) .unwrap(); let stored_path_33 = pmt.get_path(NODE33).unwrap(); @@ -76,20 +68,20 @@ fn add_and_get_paths() { #[test] 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 hash_6 = int_to_digest(6); let index = NodeIndex::make(3, 6); - let pmt = PartialMerkleTree::with_paths([(index, hash_6, path_6.into())]).unwrap(); + let pmt = PartialMerkleTree::with_paths([(index.value(), hash_6, path_6.into())]).unwrap(); - assert_eq!(int_to_node(6u64), *pmt.get_node(index).unwrap()); + assert_eq!(int_to_digest(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 value32 = int_to_digest(32); + let value33 = int_to_digest(33); + let value20 = int_to_digest(20); + let value22 = int_to_digest(22); + let value23 = int_to_digest(23); let value21 = Rpo256::merge(&[value32, value33]); let value10 = Rpo256::merge(&[value20, value21]); @@ -99,13 +91,11 @@ fn update_leaf() { let path_22 = vec![*value23, *value10]; - let mut pmt = PartialMerkleTree::with_paths([ - (NODE33, *value33, path_33.into()), - (NODE22, *value22, path_22.into()), - ]) - .unwrap(); + let mut pmt = + PartialMerkleTree::with_paths([(3, value33, path_33.into()), (2, value22, path_22.into())]) + .unwrap(); - let new_value32 = int_to_node(132).into(); + let new_value32 = int_to_digest(132); let new_value21 = Rpo256::merge(&[new_value32, value33]); let new_value10 = Rpo256::merge(&[value20, new_value21]); let expected_root = Rpo256::merge(&[new_value10, value11]); @@ -116,83 +106,22 @@ fn update_leaf() { let new_root = pmt.root(); - assert_eq!(new_root, *expected_root); -} - -#[test] -fn test_inner_node_iterator() -> Result<(), MerkleError> { - 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: *value10, - right: *value11, - }, - InnerNodeInfo { - value: *value10, - left: *value20, - right: *value21, - }, - InnerNodeInfo { - value: *value11, - left: *value22, - right: *value23, - }, - InnerNodeInfo { - value: *value21, - left: *value32, - right: *value33, - }, - ]; - assert_eq!(nodes, expected); - - Ok(()) + assert_eq!(new_root, expected_root); } #[test] 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 value32 = int_to_digest(32); + let value33 = int_to_digest(33); + let value20 = int_to_digest(20); + let value22 = int_to_digest(22); + let value23 = int_to_digest(23); let value11 = Rpo256::merge(&[value22, value23]); let path_33 = vec![*value32, *value20, *value11]; - let pmt = PartialMerkleTree::with_paths([(NODE33, *value33, path_33.into())]).unwrap(); + let pmt = PartialMerkleTree::with_paths([(3, value33, path_33.into())]).unwrap(); assert_eq!(pmt.get_leaf_depth(0).unwrap(), 2); assert_eq!(pmt.get_leaf_depth(1).unwrap(), 2); @@ -211,11 +140,11 @@ fn check_leaf_depth() { /// - 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 { +fn calculate_parent_hash(node: RpoDigest, node_pos: u64, sibling: RpoDigest) -> RpoDigest { let parity = node_pos & 1; if parity == 0 { - Rpo256::merge(&[node.into(), sibling.into()]).into() + Rpo256::merge(&[node, sibling]) } else { - Rpo256::merge(&[sibling.into(), node.into()]).into() + Rpo256::merge(&[sibling, node]) } } From 55cc71dadfc5ce801b083e4f274196523cced61c Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Sat, 3 Jun 2023 15:38:24 +0300 Subject: [PATCH 4/9] fix: fix add_path func leaf determination --- src/merkle/partial_mt/mod.rs | 12 +++++------- src/merkle/partial_mt/tests.rs | 3 +++ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index be46802..26eb8a3 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -213,19 +213,17 @@ impl PartialMerkleTree { // insert calculated node to the nodes map self.nodes.insert(index_value, node); + // if the calculated node was a leaf, remove it from leaves set. + if self.leaves.contains(&index_value) { + self.leaves.remove(&index_value); + } + let sibling_node = index_value.sibling(); // node became a leaf only if it is a new node (it wasn't in nodes map) if !self.nodes.contains_key(&sibling_node) { self.leaves.insert(sibling_node); } - // node stops being a leaf if the path contains a node which is a child of this leaf - let mut parent = index_value; - parent.move_up(); - if self.leaves.contains(&parent) { - self.leaves.remove(&parent); - } - // insert node from Merkle path to the nodes map self.nodes.insert(sibling_node, hash.into()); diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index 1efbb6a..35d41ef 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -133,6 +133,9 @@ fn check_leaf_depth() { assert_eq!(pmt.get_leaf_depth(7).unwrap(), 1); } +// TODO: add test for add_path function and check correctness of leaf determination (requires +// inner_nodes iter) + // HELPER FUNCTIONS // -------------------------------------------------------------------------------------------- From 43f1a4cb64ff7103b46bf758f1e5479c2d457988 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Sat, 3 Jun 2023 15:47:27 +0300 Subject: [PATCH 5/9] refactor: MerkleStore clippy fix --- src/merkle/store/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index d4c1ffb..8198435 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -404,7 +404,7 @@ impl MerkleStore { if let Some(node) = source.nodes.get(&root) { // if the node has already been inserted, no need to process it further as all of its // descendants should be already cloned from the source store - if matches!(self.nodes.insert(root, *node), None) { + if self.nodes.insert(root, *node).is_none() { self.clone_tree_from(node.left, source); self.clone_tree_from(node.right, source); } From 2708a2364901e68f27a0eba5baa1a21fa7d0b557 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Mon, 5 Jun 2023 18:02:16 +0300 Subject: [PATCH 6/9] refactor: optimize code, fix bugs --- src/merkle/mod.rs | 5 -- src/merkle/partial_mt/mod.rs | 70 +++++++--------- src/merkle/partial_mt/tests.rs | 149 ++++++++++++++------------------- 3 files changed, 93 insertions(+), 131 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index c3cbdf6..6f7d2ab 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -101,8 +101,3 @@ impl std::error::Error for MerkleError {} const fn int_to_node(value: u64) -> Word { [Felt::new(value), ZERO, ZERO, ZERO] } - -#[cfg(test)] -const fn int_to_digest(value: u64) -> RpoDigest { - RpoDigest::new([Felt::new(value), ZERO, ZERO, ZERO]) -} diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index 26eb8a3..7115557 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -22,6 +22,7 @@ const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD); /// Tree allows to create Merkle Tree by providing Merkle paths of different lengths. /// /// The root of the tree is recomputed on each new leaf update. +#[derive(Debug, Clone, PartialEq, Eq)] pub struct PartialMerkleTree { max_depth: u8, nodes: BTreeMap, @@ -112,12 +113,12 @@ impl PartialMerkleTree { /// Returns a vector of paths from every leaf to the root. pub fn paths(&self) -> Vec<(NodeIndex, ValuePath)> { let mut paths = Vec::new(); - self.leaves.iter().for_each(|leaf| { + self.leaves.iter().for_each(|&leaf| { paths.push(( - *leaf, + leaf, ValuePath { - value: *self.get_node(*leaf).expect("Failed to get leaf node"), - path: self.get_path(*leaf).expect("Failed to get path"), + value: *self.get_node(leaf).expect("Failed to get leaf node"), + path: self.get_path(leaf).expect("Failed to get path"), }, )); }); @@ -160,10 +161,10 @@ impl PartialMerkleTree { /// Returns an iterator over the leaves of this [PartialMerkleTree]. pub fn leaves(&self) -> impl Iterator + '_ { - self.leaves.iter().map(|leaf| { + self.leaves.iter().map(|&leaf| { ( - *leaf, - self.get_node(*leaf).unwrap_or_else(|_| { + leaf, + self.get_node(leaf).unwrap_or_else(|_| { panic!( "Leaf with node index ({}, {}) is not in the nodes map", leaf.depth(), @@ -214,19 +215,25 @@ impl PartialMerkleTree { self.nodes.insert(index_value, node); // if the calculated node was a leaf, remove it from leaves set. - if self.leaves.contains(&index_value) { - self.leaves.remove(&index_value); - } + self.leaves.remove(&index_value); let sibling_node = index_value.sibling(); - // node became a leaf only if it is a new node (it wasn't in nodes map) - if !self.nodes.contains_key(&sibling_node) { + + // 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. + // - 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. + // - 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. + if self.nodes.insert(sibling_node, hash.into()).is_none() { 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(node, hash.into())) }); @@ -238,8 +245,6 @@ impl PartialMerkleTree { return Err(MerkleError::ConflictingRoots([*self.root(), *root].to_vec())); } - // self.update_leaves()?; - Ok(()) } @@ -250,7 +255,7 @@ impl PartialMerkleTree { &mut self, node_index: NodeIndex, value: RpoDigest, - ) -> Result { + ) -> Result, MerkleError> { // check correctness of the depth and update it Self::check_depth(node_index.depth())?; self.update_depth(node_index.depth()); @@ -259,38 +264,19 @@ impl PartialMerkleTree { self.leaves.insert(node_index); // add node value to the nodes Map - let old_value = self.nodes.insert(node_index, value).unwrap_or(EMPTY_DIGEST); + let old_value = self.nodes.insert(node_index, value); // if the old value and new value are the same, there is nothing to update - if value == old_value { - return Ok(value); + if old_value.is_some() && value == old_value.unwrap() { + return Ok(old_value); } let mut node_index = node_index; let mut value = value; for _ in 0..node_index.depth() { - 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() - .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() - .ok_or(MerkleError::NodeNotInSet(right_index))?, - ) - }; + let sibling = self.nodes.get(&node_index.sibling()).expect("sibling should exist"); + value = Rpo256::merge(&node_index.build_node(value, *sibling)); node_index.move_up(); - value = Rpo256::merge(&[left, right]); self.nodes.insert(node_index, value); } diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index 35d41ef..c612f66 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -1,18 +1,29 @@ -use crate::hash::rpo::RpoDigest; - use super::{ - super::{int_to_digest, int_to_node, NodeIndex}, - PartialMerkleTree, Rpo256, + super::{int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree}, + Word, }; // TEST DATA // ================================================================================================ +const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0); + const NODE22: NodeIndex = NodeIndex::new_unchecked(2, 2); 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), +]; + // TESTS // ================================================================================================ @@ -21,107 +32,92 @@ const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3); #[test] fn get_root() { - let leaf0 = int_to_digest(0); - let leaf1 = int_to_digest(1); - let leaf2 = int_to_digest(2); - let leaf3 = int_to_digest(3); + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); - let parent0 = calculate_parent_hash(leaf0, 0, leaf1); - let parent1 = calculate_parent_hash(leaf2, 2, leaf3); + let mut store = MerkleStore::new(); + let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); - let root_exp = calculate_parent_hash(parent0, 0, parent1); + let path33 = ms.get_path(expected_root, NODE33).unwrap(); - let set = - super::PartialMerkleTree::with_paths([(0, leaf0, vec![*leaf1, *parent1].into())]).unwrap(); + let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap(); - assert_eq!(set.root(), root_exp); + assert_eq!(pmt.root(), expected_root.into()); } #[test] fn add_and_get_paths() { - let value32 = int_to_digest(32); - let value33 = int_to_digest(33); - let value20 = int_to_digest(20); - let value22 = int_to_digest(22); - let value23 = int_to_digest(23); - - let value21 = Rpo256::merge(&[value32, value33]); - let value10 = Rpo256::merge(&[value20, value21]); - let value11 = Rpo256::merge(&[value22, value23]); + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); - let path_33 = vec![*value32, *value20, *value11]; + let mut store = MerkleStore::new(); + let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); - let path_22 = vec![*value23, *value10]; + let expected_path33 = ms.get_path(expected_root, NODE33).unwrap(); + let expected_path22 = ms.get_path(expected_root, NODE22).unwrap(); let pmt = PartialMerkleTree::with_paths([ - (3, value33, path_33.clone().into()), - (2, value22, path_22.clone().into()), + (3_u64, expected_path33.value.into(), expected_path33.path.clone()), + (2, expected_path22.value.into(), expected_path22.path.clone()), ]) .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); + let path33 = pmt.get_path(NODE33).unwrap(); + let path22 = pmt.get_path(NODE22).unwrap(); + + assert_eq!(expected_path33.path, path33); + assert_eq!(expected_path22.path, path22); } #[test] fn get_node() { - let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)]; - let hash_6 = int_to_digest(6); - let index = NodeIndex::make(3, 6); - let pmt = PartialMerkleTree::with_paths([(index.value(), hash_6, path_6.into())]).unwrap(); + 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 path33 = ms.get_path(expected_root, NODE33).unwrap(); - assert_eq!(int_to_digest(6u64), pmt.get_node(index).unwrap()); + let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap(); + + assert_eq!(ms.get_node(expected_root, NODE32).unwrap(), *pmt.get_node(NODE32).unwrap()); + assert_eq!(ms.get_node(expected_root, NODE10).unwrap(), *pmt.get_node(NODE10).unwrap()); } #[test] fn update_leaf() { - let value32 = int_to_digest(32); - let value33 = int_to_digest(33); - let value20 = int_to_digest(20); - let value22 = int_to_digest(22); - let value23 = int_to_digest(23); - - let value21 = Rpo256::merge(&[value32, value33]); - let value10 = Rpo256::merge(&[value20, value21]); - let value11 = Rpo256::merge(&[value22, value23]); + let mut mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let root = mt.root(); - let path_33 = vec![*value32, *value20, *value11]; - - let path_22 = vec![*value23, *value10]; + let mut store = MerkleStore::new(); + let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); + let path33 = ms.get_path(root, NODE33).unwrap(); let mut pmt = - PartialMerkleTree::with_paths([(3, value33, path_33.into()), (2, value22, path_22.into())]) - .unwrap(); - - let new_value32 = int_to_digest(132); - let new_value21 = Rpo256::merge(&[new_value32, value33]); - let new_value10 = Rpo256::merge(&[value20, new_value21]); - let expected_root = Rpo256::merge(&[new_value10, value11]); + PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap(); - let old_leaf = pmt.update_leaf(NODE32, new_value32).unwrap(); + let new_value32 = int_to_node(132); + mt.update_leaf(2_u64, new_value32).unwrap(); + let expected_root = mt.root(); - assert_eq!(value32, old_leaf); + pmt.update_leaf(NODE32, new_value32.into()).unwrap(); + let actual_root = pmt.root(); - let new_root = pmt.root(); - - assert_eq!(new_root, expected_root); + assert_eq!(expected_root, *actual_root); } #[test] fn check_leaf_depth() { - let value32 = int_to_digest(32); - let value33 = int_to_digest(33); - let value20 = int_to_digest(20); - let value22 = int_to_digest(22); - let value23 = int_to_digest(23); + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); - let value11 = Rpo256::merge(&[value22, value23]); + let mut store = MerkleStore::new(); + let ms = MerkleStore::extend(&mut store, mt.inner_nodes()); - let path_33 = vec![*value32, *value20, *value11]; + let path33 = ms.get_path(expected_root, NODE33).unwrap(); - let pmt = PartialMerkleTree::with_paths([(3, value33, path_33.into())]).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); @@ -131,23 +127,8 @@ fn check_leaf_depth() { 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()); } // TODO: add test for add_path function and check correctness of leaf determination (requires // inner_nodes iter) - -// HELPER FUNCTIONS -// -------------------------------------------------------------------------------------------- - -/// 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: RpoDigest, node_pos: u64, sibling: RpoDigest) -> RpoDigest { - let parity = node_pos & 1; - if parity == 0 { - Rpo256::merge(&[node, sibling]) - } else { - Rpo256::merge(&[sibling, node]) - } -} From 218a64b5c780807df17698aa84f893489709670f Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Tue, 6 Jun 2023 01:30:14 +0300 Subject: [PATCH 7/9] refactor: small fixes --- src/merkle/partial_mt/mod.rs | 34 ++++++++++++++++++---------------- src/merkle/partial_mt/tests.rs | 32 +++++++++++++------------------- 2 files changed, 31 insertions(+), 35 deletions(-) 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 From 766702e37af3ff65a147d0786ec19b0cf760e6c5 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Fri, 9 Jun 2023 13:33:56 +0300 Subject: [PATCH 8/9] refactor: improve tests, small fixes --- src/merkle/mod.rs | 2 +- src/merkle/partial_mt/mod.rs | 40 ++++++++-- src/merkle/partial_mt/tests.rs | 140 +++++++++++++++++++++++++++------ src/utils.rs | 6 ++ 4 files changed, 158 insertions(+), 30 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 6f7d2ab..73f149e 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::PartialMerkleTree; +pub use partial_mt::{pmt_to_text, PartialMerkleTree}; // ERRORS // ================================================================================================ diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index ab7b644..fd9f1c6 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -2,9 +2,8 @@ use super::{ BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec, Word, EMPTY_WORD, }; - -extern crate alloc; -use alloc::format; +use crate::utils::{format, string::String, word_to_hex}; +use core::fmt; #[cfg(test)] mod tests; @@ -102,10 +101,8 @@ impl PartialMerkleTree { self.leaves.contains(&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."); + 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(); @@ -304,3 +301,32 @@ 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 6151421..85d3036 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -1,14 +1,17 @@ use super::{ super::{int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree}, - Word, + Rpo256, ValuePath, Vec, Word, }; // TEST DATA // ================================================================================================ const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0); +const NODE11: NodeIndex = NodeIndex::new_unchecked(1, 1); +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 NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2); const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3); @@ -39,7 +42,7 @@ fn get_root() { let path33 = ms.get_path(expected_root, NODE33).unwrap(); - let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap(); + let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); assert_eq!(pmt.root(), expected_root.into()); } @@ -54,17 +57,19 @@ fn add_and_get_paths() { let expected_path33 = ms.get_path(expected_root, NODE33).unwrap(); let expected_path22 = ms.get_path(expected_root, NODE22).unwrap(); - let pmt = PartialMerkleTree::with_paths([ - (3_u64, expected_path33.value.into(), expected_path33.path.clone()), - (2, expected_path22.value.into(), expected_path22.path.clone()), - ]) - .unwrap(); + let mut pmt = PartialMerkleTree::new(); + pmt.add_path(3, expected_path33.value.into(), expected_path33.path.clone()) + .unwrap(); + pmt.add_path(2, expected_path22.value.into(), expected_path22.path.clone()) + .unwrap(); let path33 = pmt.get_path(NODE33).unwrap(); let path22 = pmt.get_path(NODE22).unwrap(); + let actual_root = pmt.root(); assert_eq!(expected_path33.path, path33); assert_eq!(expected_path22.path, path22); + assert_eq!(expected_root, *actual_root); } #[test] @@ -76,7 +81,7 @@ fn get_node() { let path33 = ms.get_path(expected_root, NODE33).unwrap(); - let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap(); + let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); assert_eq!(ms.get_node(expected_root, NODE32).unwrap(), *pmt.get_node(NODE32).unwrap()); assert_eq!(ms.get_node(expected_root, NODE10).unwrap(), *pmt.get_node(NODE10).unwrap()); @@ -90,17 +95,28 @@ fn update_leaf() { let ms = MerkleStore::from(&mt); let path33 = ms.get_path(root, NODE33).unwrap(); - let mut pmt = - PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).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_u64, new_value32).unwrap(); + mt.update_leaf(2, new_value32).unwrap(); let expected_root = mt.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(); + + pmt.update_leaf(NODE20, new_value20).unwrap(); + let actual_root = pmt.root(); + + assert_eq!(expected_root, *actual_root); } #[test] @@ -112,17 +128,97 @@ fn check_leaf_depth() { let path33 = ms.get_path(expected_root, NODE33).unwrap(); - let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).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); +} + +#[test] +fn get_paths() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); - 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); + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + let path22 = ms.get_path(expected_root, NODE22).unwrap(); + + 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(); + + let leaves = vec![NODE20, NODE22, NODE23, NODE32, NODE33]; + let expected_paths: Vec<(NodeIndex, ValuePath)> = leaves + .iter() + .map(|&leaf| { + ( + leaf, + ValuePath { + value: mt.get_node(leaf).unwrap().into(), + path: mt.get_path(leaf).unwrap(), + }, + ) + }) + .collect(); + + let actual_paths = pmt.paths(); + + assert_eq!(expected_paths, actual_paths); } -// TODO: add test for add_path function and check correctness of leaf determination (requires -// inner_nodes iter) +#[test] +fn leaves() { + 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 path22 = ms.get_path(expected_root, NODE22).unwrap(); + + let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + + let value11 = mt.get_node(NODE11).unwrap().into(); + let value20 = mt.get_node(NODE20).unwrap().into(); + let value32 = mt.get_node(NODE32).unwrap().into(); + let value33 = mt.get_node(NODE33).unwrap().into(); + + let leaves = vec![(NODE11, value11), (NODE20, value20), (NODE32, value32), (NODE33, value33)]; + + let expected_leaves = leaves.iter().map(|&tuple| tuple); + assert!(expected_leaves.eq(pmt.leaves())); + + pmt.add_path(2, path22.value.into(), path22.path).unwrap(); + + let value20 = mt.get_node(NODE20).unwrap().into(); + let value22 = mt.get_node(NODE22).unwrap().into(); + let value23 = mt.get_node(NODE23).unwrap().into(); + let value32 = mt.get_node(NODE32).unwrap().into(); + let value33 = mt.get_node(NODE33).unwrap().into(); + + let leaves = vec![ + (NODE20, value20), + (NODE22, value22), + (NODE23, value23), + (NODE32, value32), + (NODE33, value33), + ]; + + let expected_leaves = leaves.iter().map(|&tuple| tuple); + assert!(expected_leaves.eq(pmt.leaves())); +} diff --git a/src/utils.rs b/src/utils.rs index 522d9ea..b9b7849 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -2,6 +2,12 @@ use super::Word; use crate::utils::string::String; use core::fmt::{self, Write}; +#[cfg(not(feature = "std"))] +pub use alloc::format; + +#[cfg(feature = "std")] +pub use std::format; + // RE-EXPORTS // ================================================================================================ pub use winter_utils::{ From b9def61e28c52ec7d956d5a276d2fa4590431dc5 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Tue, 13 Jun 2023 16:14:07 +0300 Subject: [PATCH 9/9] 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()); +}