From 59f772322113c0e0d7624868a8e8f3821ffdf069 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Fri, 26 May 2023 14:49:58 -0700 Subject: [PATCH 01/19] chore: update crete version to v0.6.0 --- CHANGELOG.md | 2 ++ Cargo.toml | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c1fc15..6f61b43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,5 @@ +## 0.6.0 (TBD) + ## 0.5.0 (2023-05-26) * Implemented `TieredSmt` (#152, #153). diff --git a/Cargo.toml b/Cargo.toml index f29a1bb..8c9a3b5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,12 +1,12 @@ [package] name = "miden-crypto" -version = "0.5.0" +version = "0.6.0" description = "Miden Cryptographic primitives" authors = ["miden contributors"] readme = "README.md" license = "MIT" repository = "https://github.com/0xPolygonMiden/crypto" -documentation = "https://docs.rs/miden-crypto/0.5.0" +documentation = "https://docs.rs/miden-crypto/0.6.0" categories = ["cryptography", "no-std"] keywords = ["miden", "crypto", "hash", "merkle"] edition = "2021" From 23f448fb3362eae6df030d2613d8ff818c22f33c Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Tue, 30 May 2023 20:23:36 +0300 Subject: [PATCH 02/19] 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 03/19] 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 04/19] 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 05/19] 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 06/19] 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 07/19] 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 08/19] 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 09/19] 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 fe9aa8c28c1b37b873627ba2f0061020a0d67719 Mon Sep 17 00:00:00 2001 From: tohrnii <100405913+tohrnii@users.noreply.github.com> Date: Fri, 9 Jun 2023 21:18:13 +0100 Subject: [PATCH 10/19] refactor: refactor crypto APIs to use RpoDigest instead of Word --- benches/store.rs | 4 +- src/hash/rpo/tests.rs | 7 +- src/merkle/merkle_tree.rs | 99 +++++++++++--------- src/merkle/mmr/accumulator.rs | 19 +++- src/merkle/mmr/full.rs | 20 ++-- src/merkle/mmr/tests.rs | 160 ++++++++++++++++++++++++++------ src/merkle/mod.rs | 14 ++- src/merkle/node.rs | 8 +- src/merkle/path.rs | 44 +++++---- src/merkle/path_set.rs | 77 ++++++++------- src/merkle/simple_smt/mod.rs | 35 +++---- src/merkle/simple_smt/tests.rs | 122 +++++++++++++----------- src/merkle/store/mod.rs | 98 +++++++++++--------- src/merkle/store/tests.rs | 165 +++++++++++++++++---------------- src/merkle/tiered_smt/mod.rs | 14 +-- src/merkle/tiered_smt/tests.rs | 82 +++++++++++----- 16 files changed, 591 insertions(+), 377 deletions(-) diff --git a/benches/store.rs b/benches/store.rs index aaa9806..f3d9cfa 100644 --- a/benches/store.rs +++ b/benches/store.rs @@ -409,7 +409,7 @@ fn update_leaf_merkletree(c: &mut Criterion) { // The MerkleTree automatically updates its internal root, the Store maintains // the old root and adds the new one. Here we update the root to have a fair // comparison - store_root = store.set_node(root, index, value).unwrap().root; + store_root = store.set_node(root, index, value.into()).unwrap().root; black_box(store_root) }, BatchSize::SmallInput, @@ -455,7 +455,7 @@ fn update_leaf_simplesmt(c: &mut Criterion) { // The MerkleTree automatically updates its internal root, the Store maintains // the old root and adds the new one. Here we update the root to have a fair // comparison - store_root = store.set_node(root, index, value).unwrap().root; + store_root = store.set_node(root, index, value.into()).unwrap().root; black_box(store_root) }, BatchSize::SmallInput, diff --git a/src/hash/rpo/tests.rs b/src/hash/rpo/tests.rs index d6379bb..8b2a258 100644 --- a/src/hash/rpo/tests.rs +++ b/src/hash/rpo/tests.rs @@ -2,7 +2,10 @@ use super::{ Felt, FieldElement, Hasher, Rpo256, RpoDigest, StarkField, ALPHA, INV_ALPHA, ONE, STATE_WIDTH, ZERO, }; -use crate::utils::collections::{BTreeSet, Vec}; +use crate::{ + utils::collections::{BTreeSet, Vec}, + Word, +}; use core::convert::TryInto; use proptest::prelude::*; use rand_utils::rand_value; @@ -232,7 +235,7 @@ proptest! { } } -const EXPECTED: [[Felt; 4]; 19] = [ +const EXPECTED: [Word; 19] = [ [ Felt::new(1502364727743950833), Felt::new(5880949717274681448), diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 4548739..0fc99a7 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -1,10 +1,5 @@ -use super::{ - Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word, -}; -use crate::{ - utils::{string::String, uninit_vector, word_to_hex}, - FieldElement, -}; +use super::{InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word}; +use crate::utils::{string::String, uninit_vector, word_to_hex}; use core::{fmt, slice}; use winter_math::log2; @@ -14,7 +9,7 @@ use winter_math::log2; /// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two). #[derive(Debug, Clone, PartialEq, Eq)] pub struct MerkleTree { - nodes: Vec, + nodes: Vec, } impl MerkleTree { @@ -34,10 +29,12 @@ impl MerkleTree { // create un-initialized vector to hold all tree nodes let mut nodes = unsafe { uninit_vector(2 * n) }; - nodes[0] = [Felt::ZERO; 4]; + nodes[0] = RpoDigest::default(); // copy leaves into the second part of the nodes vector - nodes[n..].copy_from_slice(&leaves); + nodes[n..].iter_mut().zip(leaves).for_each(|(node, leaf)| { + *node = RpoDigest::from(leaf); + }); // re-interpret nodes as an array of two nodes fused together // Safety: `nodes` will never move here as it is not bound to an external lifetime (i.e. @@ -47,7 +44,7 @@ impl MerkleTree { // calculate all internal tree nodes for i in (1..n).rev() { - nodes[i] = Rpo256::merge(&pairs[i]).into(); + nodes[i] = Rpo256::merge(&pairs[i]); } Ok(Self { nodes }) @@ -57,7 +54,7 @@ impl MerkleTree { // -------------------------------------------------------------------------------------------- /// Returns the root of this Merkle tree. - pub fn root(&self) -> Word { + pub fn root(&self) -> RpoDigest { self.nodes[1] } @@ -74,7 +71,7 @@ impl MerkleTree { /// Returns an error if: /// * The specified depth is greater than the depth of the tree. /// * The specified index is not valid for the specified depth. - pub fn get_node(&self, index: NodeIndex) -> Result { + pub fn get_node(&self, index: NodeIndex) -> Result { if index.is_root() { return Err(MerkleError::DepthTooSmall(index.depth())); } else if index.depth() > self.depth() { @@ -120,7 +117,7 @@ impl MerkleTree { /// Returns an iterator over the leaves of this [MerkleTree]. pub fn leaves(&self) -> impl Iterator { let leaves_start = self.nodes.len() / 2; - self.nodes.iter().skip(leaves_start).enumerate().map(|(i, v)| (i as u64, v)) + self.nodes.iter().skip(leaves_start).enumerate().map(|(i, v)| (i as u64, &**v)) } /// Returns n iterator over every inner node of this [MerkleTree]. @@ -159,13 +156,13 @@ impl MerkleTree { // update the current node let pos = index.to_scalar_index() as usize; - self.nodes[pos] = value; + self.nodes[pos] = value.into(); // traverse to the root, updating each node with the merged values of its parents for _ in 0..index.depth() { index.move_up(); let pos = index.to_scalar_index() as usize; - let value = Rpo256::merge(&pairs[pos]).into(); + let value = Rpo256::merge(&pairs[pos]); self.nodes[pos] = value; } @@ -180,7 +177,7 @@ impl MerkleTree { /// /// Use this to extract the data of the tree, there is no guarantee on the order of the elements. pub struct InnerNodeIterator<'a> { - nodes: &'a Vec, + nodes: &'a Vec, index: usize, } @@ -258,21 +255,25 @@ pub fn path_to_text(path: &MerklePath) -> Result { #[cfg(test)] mod tests { use super::*; - use crate::merkle::{int_to_node, InnerNodeInfo}; + use crate::{ + merkle::{int_to_leaf, InnerNodeInfo}, + Felt, Word, WORD_SIZE, + }; use core::mem::size_of; use proptest::prelude::*; - const LEAVES4: [Word; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; + const LEAVES4: [Word; WORD_SIZE] = + [int_to_leaf(1), int_to_leaf(2), int_to_leaf(3), int_to_leaf(4)]; const LEAVES8: [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_leaf(1), + int_to_leaf(2), + int_to_leaf(3), + int_to_leaf(4), + int_to_leaf(5), + int_to_leaf(6), + int_to_leaf(7), + int_to_leaf(8), ]; #[test] @@ -282,7 +283,7 @@ mod tests { // leaves were copied correctly for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) { - assert_eq!(a, b); + assert_eq!(*a, RpoDigest::from(*b)); } let (root, node2, node3) = compute_internal_nodes(); @@ -299,10 +300,10 @@ mod tests { let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); // check depth 2 - assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(RpoDigest::from(LEAVES4[0]), tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(RpoDigest::from(LEAVES4[1]), tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(RpoDigest::from(LEAVES4[2]), tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(RpoDigest::from(LEAVES4[3]), tree.get_node(NodeIndex::make(2, 3)).unwrap()); // check depth 1 let (_, node2, node3) = compute_internal_nodes(); @@ -318,10 +319,22 @@ mod tests { let (_, node2, node3) = compute_internal_nodes(); // check depth 2 - assert_eq!(vec![LEAVES4[1], node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(vec![LEAVES4[0], node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(vec![LEAVES4[3], node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(vec![LEAVES4[2], node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap()); + assert_eq!( + vec![RpoDigest::from(LEAVES4[1]), node3], + *tree.get_path(NodeIndex::make(2, 0)).unwrap() + ); + assert_eq!( + vec![RpoDigest::from(LEAVES4[0]), node3], + *tree.get_path(NodeIndex::make(2, 1)).unwrap() + ); + assert_eq!( + vec![RpoDigest::from(LEAVES4[3]), node2], + *tree.get_path(NodeIndex::make(2, 2)).unwrap() + ); + assert_eq!( + vec![RpoDigest::from(LEAVES4[2]), node2], + *tree.get_path(NodeIndex::make(2, 3)).unwrap() + ); // check depth 1 assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap()); @@ -334,7 +347,7 @@ mod tests { // update one leaf let value = 3; - let new_node = int_to_node(9); + let new_node = int_to_leaf(9); let mut expected_leaves = LEAVES8.to_vec(); expected_leaves[value as usize] = new_node; let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap(); @@ -344,7 +357,7 @@ mod tests { // update another leaf let value = 6; - let new_node = int_to_node(10); + let new_node = int_to_leaf(10); expected_leaves[value as usize] = new_node; let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap(); @@ -417,11 +430,13 @@ mod tests { // HELPER FUNCTIONS // -------------------------------------------------------------------------------------------- - fn compute_internal_nodes() -> (Word, Word, Word) { - let node2 = Rpo256::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat()); - let node3 = Rpo256::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat()); + fn compute_internal_nodes() -> (RpoDigest, RpoDigest, RpoDigest) { + let node2 = + Rpo256::hash_elements(&[Word::from(LEAVES4[0]), Word::from(LEAVES4[1])].concat()); + let node3 = + Rpo256::hash_elements(&[Word::from(LEAVES4[2]), Word::from(LEAVES4[3])].concat()); let root = Rpo256::merge(&[node2, node3]); - (root.into(), node2.into(), node3.into()) + (root, node2, node3) } } diff --git a/src/merkle/mmr/accumulator.rs b/src/merkle/mmr/accumulator.rs index ce0ee49..0729c94 100644 --- a/src/merkle/mmr/accumulator.rs +++ b/src/merkle/mmr/accumulator.rs @@ -1,4 +1,7 @@ -use super::{super::Vec, super::ZERO, Felt, MmrProof, Rpo256, Word}; +use super::{ + super::{RpoDigest, Vec, ZERO}, + Felt, MmrProof, Rpo256, Word, +}; #[derive(Debug, Clone, PartialEq)] pub struct MmrPeaks { @@ -25,7 +28,7 @@ pub struct MmrPeaks { /// leaves, starting from the peak with most children, to the one with least. /// /// Invariant: The length of `peaks` must be equal to the number of true bits in `num_leaves`. - pub peaks: Vec, + pub peaks: Vec, } impl MmrPeaks { @@ -38,7 +41,7 @@ impl MmrPeaks { Rpo256::hash_elements(&self.flatten_and_pad_peaks()).into() } - pub fn verify(&self, value: Word, opening: MmrProof) -> bool { + pub fn verify(&self, value: RpoDigest, opening: MmrProof) -> bool { let root = &self.peaks[opening.peak_index()]; opening.merkle_path.verify(opening.relative_pos() as u64, value, root) } @@ -72,7 +75,15 @@ impl MmrPeaks { }; let mut elements = Vec::with_capacity(len); - elements.extend_from_slice(&self.peaks.as_slice().concat()); + elements.extend_from_slice( + &self + .peaks + .as_slice() + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); elements.resize(len, ZERO); elements } diff --git a/src/merkle/mmr/full.rs b/src/merkle/mmr/full.rs index 6737f7b..7ceace3 100644 --- a/src/merkle/mmr/full.rs +++ b/src/merkle/mmr/full.rs @@ -10,10 +10,10 @@ //! depths, i.e. as part of adding adding a new element to the forest the trees with same depth are //! merged, creating a new tree with depth d+1, this process is continued until the property is //! restabilished. -use super::bit::TrueBitPositionIterator; use super::{ - super::{InnerNodeInfo, MerklePath, Vec}, - MmrPeaks, MmrProof, Rpo256, Word, + super::{InnerNodeInfo, MerklePath, RpoDigest, Vec}, + bit::TrueBitPositionIterator, + MmrPeaks, MmrProof, Rpo256, }; use core::fmt::{Display, Formatter}; @@ -38,7 +38,7 @@ pub struct Mmr { /// the elements of every tree in the forest to be stored in the same sequential buffer. It /// also means new elements can be added to the forest, and merging of trees is very cheap with /// no need to copy elements. - pub(super) nodes: Vec, + pub(super) nodes: Vec, } #[derive(Debug, PartialEq, Eq, Copy, Clone)] @@ -129,7 +129,7 @@ impl Mmr { /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element /// has position 0, the second position 1, and so on. - pub fn get(&self, pos: usize) -> Result { + pub fn get(&self, pos: usize) -> Result { // find the target tree responsible for the MMR position let tree_bit = leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?; @@ -153,7 +153,7 @@ impl Mmr { } /// Adds a new element to the MMR. - pub fn add(&mut self, el: Word) { + pub fn add(&mut self, el: RpoDigest) { // Note: every node is also a tree of size 1, adding an element to the forest creates a new // rooted-tree of size 1. This may temporarily break the invariant that every tree in the // forest has different sizes, the loop below will eagerly merge trees of same size and @@ -164,7 +164,7 @@ impl Mmr { let mut right = el; let mut left_tree = 1; while self.forest & left_tree != 0 { - right = *Rpo256::merge(&[self.nodes[left_offset].into(), right.into()]); + right = Rpo256::merge(&[self.nodes[left_offset], right]); self.nodes.push(right); left_offset = left_offset.saturating_sub(nodes_in_forest(left_tree)); @@ -176,7 +176,7 @@ impl Mmr { /// Returns an accumulator representing the current state of the MMR. pub fn accumulator(&self) -> MmrPeaks { - let peaks: Vec = TrueBitPositionIterator::new(self.forest) + let peaks: Vec = TrueBitPositionIterator::new(self.forest) .rev() .map(|bit| nodes_in_forest(1 << bit)) .scan(0, |offset, el| { @@ -212,7 +212,7 @@ impl Mmr { relative_pos: usize, index_offset: usize, mut index: usize, - ) -> (Word, Vec) { + ) -> (RpoDigest, Vec) { // collect the Merkle path let mut tree_depth = tree_bit as usize; let mut path = Vec::with_capacity(tree_depth + 1); @@ -247,7 +247,7 @@ impl Mmr { impl From for Mmr where - T: IntoIterator, + T: IntoIterator, { fn from(values: T) -> Self { let mut mmr = Mmr::new(); diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index aca4ff7..1973aa5 100644 --- a/src/merkle/mmr/tests.rs +++ b/src/merkle/mmr/tests.rs @@ -1,10 +1,14 @@ -use super::bit::TrueBitPositionIterator; -use super::full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest}; use super::{ - super::{InnerNodeInfo, Vec, WORD_SIZE, ZERO}, - Mmr, MmrPeaks, Rpo256, Word, + super::{InnerNodeInfo, Vec, WORD_SIZE}, + bit::TrueBitPositionIterator, + full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest}, + Mmr, MmrPeaks, Rpo256, +}; +use crate::{ + hash::rpo::RpoDigest, + merkle::{int_to_node, MerklePath}, + Felt, Word, }; -use crate::merkle::{int_to_node, MerklePath}; #[test] fn test_position_equal_or_higher_than_leafs_is_never_contained() { @@ -99,7 +103,7 @@ fn test_nodes_in_forest_single_bit() { } } -const LEAVES: [Word; 7] = [ +const LEAVES: [RpoDigest; 7] = [ int_to_node(0), int_to_node(1), int_to_node(2), @@ -114,14 +118,38 @@ fn test_mmr_simple() { let mut postorder = Vec::new(); postorder.push(LEAVES[0]); postorder.push(LEAVES[1]); - postorder.push(*Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat())); + postorder.push(Rpo256::hash_elements( + &[LEAVES[0], LEAVES[1]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + )); postorder.push(LEAVES[2]); postorder.push(LEAVES[3]); - postorder.push(*Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat())); - postorder.push(*Rpo256::hash_elements(&[postorder[2], postorder[5]].concat())); + postorder.push(Rpo256::hash_elements( + &[LEAVES[2], LEAVES[3]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + )); + postorder.push(Rpo256::hash_elements( + &[postorder[2], postorder[5]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + )); postorder.push(LEAVES[4]); postorder.push(LEAVES[5]); - postorder.push(*Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat())); + postorder.push(Rpo256::hash_elements( + &[LEAVES[4], LEAVES[5]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + )); postorder.push(LEAVES[6]); let mut mmr = Mmr::new(); @@ -195,8 +223,20 @@ fn test_mmr_simple() { #[test] fn test_mmr_open() { let mmr: Mmr = LEAVES.into(); - let h01: Word = Rpo256::hash_elements(&LEAVES[0..2].concat()).into(); - let h23: Word = Rpo256::hash_elements(&LEAVES[2..4].concat()).into(); + let h01: RpoDigest = Rpo256::hash_elements( + &LEAVES[0..2] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); + let h23: RpoDigest = Rpo256::hash_elements( + &LEAVES[2..4] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); // node at pos 7 is the root assert!(mmr.open(7).is_err(), "Element 7 is not in the tree, result should be None"); @@ -214,7 +254,7 @@ fn test_mmr_open() { "MmrProof should be valid for the current accumulator." ); - // nodes 4,5 are detph 1 + // nodes 4,5 are depth 1 let root_to_path = MerklePath::new(vec![LEAVES[4]]); let opening = mmr .open(5) @@ -361,10 +401,34 @@ fn test_mmr_inner_nodes() { let mmr: Mmr = LEAVES.into(); let nodes: Vec = mmr.inner_nodes().collect(); - let h01 = *Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat()); - let h23 = *Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat()); - let h0123 = *Rpo256::hash_elements(&[h01, h23].concat()); - let h45 = *Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat()); + let h01 = Rpo256::hash_elements( + &[LEAVES[0], LEAVES[1]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); + let h23 = Rpo256::hash_elements( + &[LEAVES[2], LEAVES[3]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); + let h0123 = Rpo256::hash_elements( + &[h01, h23] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); + let h45 = Rpo256::hash_elements( + &[LEAVES[4], LEAVES[5]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); let postorder = vec![ InnerNodeInfo { value: h01, @@ -396,17 +460,45 @@ fn test_mmr_hash_peaks() { let mmr: Mmr = LEAVES.into(); let peaks = mmr.accumulator(); - let first_peak = *Rpo256::merge(&[ - Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat()), - Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat()), + let first_peak = Rpo256::merge(&[ + Rpo256::hash_elements( + &[LEAVES[0], LEAVES[1]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ), + Rpo256::hash_elements( + &[LEAVES[2], LEAVES[3]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ), ]); - let second_peak = *Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat()); + let second_peak = Rpo256::hash_elements( + &[LEAVES[4], LEAVES[5]] + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat(), + ); let third_peak = LEAVES[6]; // minimum length is 16 let mut expected_peaks = [first_peak, second_peak, third_peak].to_vec(); - expected_peaks.resize(16, [ZERO; WORD_SIZE]); - assert_eq!(peaks.hash_peaks(), *Rpo256::hash_elements(&expected_peaks.as_slice().concat())); + expected_peaks.resize(16, RpoDigest::default()); + assert_eq!( + peaks.hash_peaks(), + *Rpo256::hash_elements( + &expected_peaks + .as_slice() + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat() + ) + ); } #[test] @@ -422,10 +514,17 @@ fn test_mmr_peaks_hash_less_than_16() { // minimum length is 16 let mut expected_peaks = peaks.clone(); - expected_peaks.resize(16, [ZERO; WORD_SIZE]); + expected_peaks.resize(16, RpoDigest::default()); assert_eq!( accumulator.hash_peaks(), - *Rpo256::hash_elements(&expected_peaks.as_slice().concat()) + *Rpo256::hash_elements( + &expected_peaks + .as_slice() + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat() + ) ); } } @@ -441,10 +540,17 @@ fn test_mmr_peaks_hash_odd() { // odd length bigger than 16 is padded to the next even nubmer let mut expected_peaks = peaks.clone(); - expected_peaks.resize(18, [ZERO; WORD_SIZE]); + expected_peaks.resize(18, RpoDigest::default()); assert_eq!( accumulator.hash_peaks(), - *Rpo256::hash_elements(&expected_peaks.as_slice().concat()) + *Rpo256::hash_elements( + &expected_peaks + .as_slice() + .iter() + .map(|digest| digest.into()) + .collect::>() + .concat() + ) ); } diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 631b960..a0d50d6 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -10,7 +10,6 @@ use core::fmt; mod empty_roots; pub use empty_roots::EmptySubtreeRoots; -use empty_roots::EMPTY_WORD; mod index; pub use index::NodeIndex; @@ -44,7 +43,7 @@ pub use node::InnerNodeInfo; #[derive(Clone, Debug, PartialEq, Eq)] pub enum MerkleError { - ConflictingRoots(Vec), + ConflictingRoots(Vec), DepthTooSmall(u8), DepthTooBig(u64), DuplicateValuesForIndex(u64), @@ -54,9 +53,9 @@ pub enum MerkleError { InvalidPath(MerklePath), InvalidNumEntries(usize, usize), NodeNotInSet(NodeIndex), - NodeNotInStore(Word, NodeIndex), + NodeNotInStore(RpoDigest, NodeIndex), NumLeavesNotPowerOfTwo(usize), - RootNotInStore(Word), + RootNotInStore(RpoDigest), } impl fmt::Display for MerkleError { @@ -95,6 +94,11 @@ impl std::error::Error for MerkleError {} // ================================================================================================ #[cfg(test)] -const fn int_to_node(value: u64) -> Word { +const fn int_to_node(value: u64) -> RpoDigest { + RpoDigest::new([Felt::new(value), ZERO, ZERO, ZERO]) +} + +#[cfg(test)] +const fn int_to_leaf(value: u64) -> Word { [Felt::new(value), ZERO, ZERO, ZERO] } diff --git a/src/merkle/node.rs b/src/merkle/node.rs index 681e306..8440af8 100644 --- a/src/merkle/node.rs +++ b/src/merkle/node.rs @@ -1,9 +1,9 @@ -use super::Word; +use crate::hash::rpo::RpoDigest; /// Representation of a node with two children used for iterating over containers. #[derive(Debug, Clone, PartialEq, Eq)] pub struct InnerNodeInfo { - pub value: Word, - pub left: Word, - pub right: Word, + pub value: RpoDigest, + pub left: RpoDigest, + pub right: RpoDigest, } diff --git a/src/merkle/path.rs b/src/merkle/path.rs index e873d84..1842915 100644 --- a/src/merkle/path.rs +++ b/src/merkle/path.rs @@ -1,4 +1,4 @@ -use super::{vec, InnerNodeInfo, MerkleError, NodeIndex, Rpo256, Vec, Word}; +use super::{vec, InnerNodeInfo, MerkleError, NodeIndex, Rpo256, RpoDigest, Vec}; use core::ops::{Deref, DerefMut}; // MERKLE PATH @@ -7,7 +7,7 @@ use core::ops::{Deref, DerefMut}; /// A merkle path container, composed of a sequence of nodes of a Merkle tree. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct MerklePath { - nodes: Vec, + nodes: Vec, } impl MerklePath { @@ -15,7 +15,7 @@ impl MerklePath { // -------------------------------------------------------------------------------------------- /// Creates a new Merkle path from a list of nodes. - pub fn new(nodes: Vec) -> Self { + pub fn new(nodes: Vec) -> Self { Self { nodes } } @@ -28,13 +28,13 @@ impl MerklePath { } /// Computes the merkle root for this opening. - pub fn compute_root(&self, index: u64, node: Word) -> Result { + pub fn compute_root(&self, index: u64, node: RpoDigest) -> Result { let mut index = NodeIndex::new(self.depth(), index)?; let root = self.nodes.iter().copied().fold(node, |node, sibling| { // compute the node and move to the next iteration. - let input = index.build_node(node.into(), sibling.into()); + let input = index.build_node(node, sibling); index.move_up(); - Rpo256::merge(&input).into() + Rpo256::merge(&input) }); Ok(root) } @@ -42,7 +42,7 @@ impl MerklePath { /// Verifies the Merkle opening proof towards the provided root. /// /// Returns `true` if `node` exists at `index` in a Merkle tree with `root`. - pub fn verify(&self, index: u64, node: Word, root: &Word) -> bool { + pub fn verify(&self, index: u64, node: RpoDigest, root: &RpoDigest) -> bool { match self.compute_root(index, node) { Ok(computed_root) => root == &computed_root, Err(_) => false, @@ -55,7 +55,11 @@ impl MerklePath { /// /// # Errors /// Returns an error if the specified index is not valid for this path. - pub fn inner_nodes(&self, index: u64, node: Word) -> Result { + pub fn inner_nodes( + &self, + index: u64, + node: RpoDigest, + ) -> Result { Ok(InnerNodeIterator { nodes: &self.nodes, index: NodeIndex::new(self.depth(), index)?, @@ -64,8 +68,8 @@ impl MerklePath { } } -impl From> for MerklePath { - fn from(path: Vec) -> Self { +impl From> for MerklePath { + fn from(path: Vec) -> Self { Self::new(path) } } @@ -73,7 +77,7 @@ impl From> for MerklePath { impl Deref for MerklePath { // we use `Vec` here instead of slice so we can call vector mutation methods directly from the // merkle path (example: `Vec::remove`). - type Target = Vec; + type Target = Vec; fn deref(&self) -> &Self::Target { &self.nodes @@ -89,15 +93,15 @@ impl DerefMut for MerklePath { // ITERATORS // ================================================================================================ -impl FromIterator for MerklePath { - fn from_iter>(iter: T) -> Self { +impl FromIterator for MerklePath { + fn from_iter>(iter: T) -> Self { Self::new(iter.into_iter().collect()) } } impl IntoIterator for MerklePath { - type Item = Word; - type IntoIter = vec::IntoIter; + type Item = RpoDigest; + type IntoIter = vec::IntoIter; fn into_iter(self) -> Self::IntoIter { self.nodes.into_iter() @@ -106,9 +110,9 @@ impl IntoIterator for MerklePath { /// An iterator over internal nodes of a [MerklePath]. pub struct InnerNodeIterator<'a> { - nodes: &'a Vec, + nodes: &'a Vec, index: NodeIndex, - value: Word, + value: RpoDigest, } impl<'a> Iterator for InnerNodeIterator<'a> { @@ -123,7 +127,7 @@ impl<'a> Iterator for InnerNodeIterator<'a> { (self.value, self.nodes[sibling_pos]) }; - self.value = Rpo256::merge(&[left.into(), right.into()]).into(); + self.value = Rpo256::merge(&[left, right]); self.index.move_up(); Some(InnerNodeInfo { @@ -144,7 +148,7 @@ impl<'a> Iterator for InnerNodeIterator<'a> { #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct ValuePath { /// The node value opening for `path`. - pub value: Word, + pub value: RpoDigest, /// The path from `value` to `root` (exclusive). pub path: MerklePath, } @@ -156,7 +160,7 @@ pub struct ValuePath { #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct RootPath { /// The node value opening for `path`. - pub root: Word, + pub root: RpoDigest, /// The path from `value` to `root` (exclusive). pub path: MerklePath, } diff --git a/src/merkle/path_set.rs b/src/merkle/path_set.rs index ed945fb..322faeb 100644 --- a/src/merkle/path_set.rs +++ b/src/merkle/path_set.rs @@ -1,4 +1,5 @@ -use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, ValuePath, Vec, Word, ZERO}; +use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, ValuePath, Vec}; +use crate::{hash::rpo::RpoDigest, Word}; // MERKLE PATH SET // ================================================================================================ @@ -6,7 +7,7 @@ use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, ValuePath, Vec /// A set of Merkle paths. #[derive(Debug, Clone, PartialEq, Eq)] pub struct MerklePathSet { - root: Word, + root: RpoDigest, total_depth: u8, paths: BTreeMap, } @@ -17,7 +18,7 @@ impl MerklePathSet { /// Returns an empty MerklePathSet. pub fn new(depth: u8) -> Self { - let root = [ZERO; 4]; + let root = RpoDigest::default(); let paths = BTreeMap::new(); Self { @@ -32,7 +33,7 @@ impl MerklePathSet { /// Analogous to `[Self::add_path]`. pub fn with_paths(self, paths: I) -> Result where - I: IntoIterator, + I: IntoIterator, { paths.into_iter().try_fold(self, |mut set, (index, value, path)| { set.add_path(index, value, path)?; @@ -44,7 +45,7 @@ impl MerklePathSet { // -------------------------------------------------------------------------------------------- /// Returns the root to which all paths in this set resolve. - pub const fn root(&self) -> Word { + pub const fn root(&self) -> RpoDigest { self.root } @@ -61,7 +62,7 @@ impl MerklePathSet { /// Returns an error if: /// * The specified index is not valid for the depth of structure. /// * Requested node does not exist in the set. - pub fn get_node(&self, index: NodeIndex) -> Result { + pub fn get_node(&self, index: NodeIndex) -> Result { if index.depth() != self.total_depth { return Err(MerkleError::InvalidDepth { expected: self.total_depth, @@ -84,7 +85,7 @@ impl MerklePathSet { /// * Leaf with the requested path does not exist in the set. pub fn get_leaf(&self, index: u64) -> Result { let index = NodeIndex::new(self.depth(), index)?; - self.get_node(index) + Ok(self.get_node(index)?.into()) } /// Returns a Merkle path to the node at the specified index. The node itself is @@ -150,7 +151,7 @@ impl MerklePathSet { pub fn add_path( &mut self, index_value: u64, - value: Word, + value: RpoDigest, mut path: MerklePath, ) -> Result<(), MerkleError> { let mut index = NodeIndex::new(path.len() as u8, index_value)?; @@ -166,15 +167,15 @@ impl MerklePathSet { path.insert(parity as usize, value); // traverse to the root, updating the nodes - let root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into(); + let root: RpoDigest = Rpo256::merge(&[path[0], path[1]]); let root = path.iter().skip(2).copied().fold(root, |root, hash| { index.move_up(); - Rpo256::merge(&index.build_node(root.into(), hash.into())).into() + Rpo256::merge(&index.build_node(root, hash)) }); // 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 self.root == [ZERO; 4] { + if self.root == RpoDigest::default() { self.root = root; } else if self.root != root { return Err(MerkleError::ConflictingRoots([self.root, root].to_vec())); @@ -191,7 +192,11 @@ impl MerklePathSet { /// # Errors /// Returns an error if: /// * Requested node does not exist in the set. - pub fn update_leaf(&mut self, base_index_value: u64, value: Word) -> Result<(), MerkleError> { + pub fn update_leaf( + &mut self, + base_index_value: u64, + value: RpoDigest, + ) -> Result<(), MerkleError> { let mut index = NodeIndex::new(self.depth(), base_index_value)?; let parity = index.value() & 1; let path_key = index.value() - parity; @@ -203,24 +208,24 @@ impl MerklePathSet { // Fill old_hashes vector ----------------------------------------------------------------- let mut current_index = index; let mut old_hashes = Vec::with_capacity(path.len().saturating_sub(2)); - let mut root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into(); + let mut root: RpoDigest = Rpo256::merge(&[path[0], path[1]]); for hash in path.iter().skip(2).copied() { old_hashes.push(root); current_index.move_up(); - let input = current_index.build_node(hash.into(), root.into()); - root = Rpo256::merge(&input).into(); + let input = current_index.build_node(hash, root); + root = Rpo256::merge(&input); } // Fill new_hashes vector ----------------------------------------------------------------- path[index.is_value_odd() as usize] = value; let mut new_hashes = Vec::with_capacity(path.len().saturating_sub(2)); - let mut new_root: Word = Rpo256::merge(&[path[0].into(), path[1].into()]).into(); + let mut new_root: RpoDigest = Rpo256::merge(&[path[0], path[1]]); for path_hash in path.iter().skip(2).copied() { new_hashes.push(new_root); index.move_up(); - let input = current_index.build_node(path_hash.into(), new_root.into()); - new_root = Rpo256::merge(&input).into(); + let input = current_index.build_node(path_hash, new_root); + new_root = Rpo256::merge(&input); } self.root = new_root; @@ -345,13 +350,13 @@ mod tests { let g = int_to_node(7); let h = int_to_node(8); - let i = Rpo256::merge(&[a.into(), b.into()]); - let j = Rpo256::merge(&[c.into(), d.into()]); - let k = Rpo256::merge(&[e.into(), f.into()]); - let l = Rpo256::merge(&[g.into(), h.into()]); + let i = Rpo256::merge(&[a, b]); + let j = Rpo256::merge(&[c, d]); + let k = Rpo256::merge(&[e, f]); + let l = Rpo256::merge(&[g, h]); - let m = Rpo256::merge(&[i.into(), j.into()]); - let n = Rpo256::merge(&[k.into(), l.into()]); + let m = Rpo256::merge(&[i, j]); + let n = Rpo256::merge(&[k, l]); let root = Rpo256::merge(&[m.into(), n.into()]); @@ -359,31 +364,31 @@ mod tests { let value = b; let index = 1; - let path = MerklePath::new([a.into(), j.into(), n.into()].to_vec()); + let path = MerklePath::new([a, j, n].to_vec()); set.add_path(index, value, path.clone()).unwrap(); - assert_eq!(value, set.get_leaf(index).unwrap()); - assert_eq!(Word::from(root), set.root()); + assert_eq!(*value, set.get_leaf(index).unwrap()); + assert_eq!(RpoDigest::from(root), set.root()); let value = e; let index = 4; let path = MerklePath::new([f.into(), l.into(), m.into()].to_vec()); set.add_path(index, value, path.clone()).unwrap(); - assert_eq!(value, set.get_leaf(index).unwrap()); - assert_eq!(Word::from(root), set.root()); + assert_eq!(*value, set.get_leaf(index).unwrap()); + assert_eq!(RpoDigest::from(root), set.root()); let value = a; let index = 0; let path = MerklePath::new([b.into(), j.into(), n.into()].to_vec()); set.add_path(index, value, path.clone()).unwrap(); - assert_eq!(value, set.get_leaf(index).unwrap()); - assert_eq!(Word::from(root), set.root()); + assert_eq!(*value, set.get_leaf(index).unwrap()); + assert_eq!(RpoDigest::from(root), set.root()); let value = h; let index = 7; let path = MerklePath::new([g.into(), k.into(), m.into()].to_vec()); set.add_path(index, value, path.clone()).unwrap(); - assert_eq!(value, set.get_leaf(index).unwrap()); - assert_eq!(Word::from(root), set.root()); + assert_eq!(*value, set.get_leaf(index).unwrap()); + assert_eq!(RpoDigest::from(root), set.root()); } // HELPER FUNCTIONS @@ -397,11 +402,11 @@ mod tests { /// - 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 { if is_even(node_pos) { - Rpo256::merge(&[node.into(), sibling.into()]).into() + Rpo256::merge(&[node, sibling]) } else { - Rpo256::merge(&[sibling.into(), node.into()]).into() + Rpo256::merge(&[sibling, node]) } } } diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 9d25316..2540ae4 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -1,6 +1,6 @@ use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, - Rpo256, RpoDigest, Vec, Word, EMPTY_WORD, + empty_roots::EMPTY_WORD, BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, + MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word, }; #[cfg(test)] @@ -15,7 +15,7 @@ mod tests; #[derive(Debug, Clone, PartialEq, Eq)] pub struct SimpleSmt { depth: u8, - root: Word, + root: RpoDigest, leaves: BTreeMap, branches: BTreeMap, empty_hashes: Vec, @@ -49,7 +49,7 @@ impl SimpleSmt { } let empty_hashes = EmptySubtreeRoots::empty_hashes(depth).to_vec(); - let root = empty_hashes[0].into(); + let root = empty_hashes[0]; Ok(Self { root, @@ -107,7 +107,7 @@ impl SimpleSmt { // -------------------------------------------------------------------------------------------- /// Returns the root of this Merkle tree. - pub const fn root(&self) -> Word { + pub const fn root(&self) -> RpoDigest { self.root } @@ -121,7 +121,7 @@ impl SimpleSmt { /// # 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. - pub fn get_node(&self, index: NodeIndex) -> Result { + pub fn get_node(&self, index: NodeIndex) -> Result { if index.is_root() { Err(MerkleError::DepthTooSmall(index.depth())) } else if index.depth() > self.depth() { @@ -129,11 +129,12 @@ impl SimpleSmt { } else if index.depth() == self.depth() { // the lookup in empty_hashes could fail only if empty_hashes were not built correctly // by the constructor as we check the depth of the lookup above. - Ok(self - .get_leaf_node(index.value()) - .unwrap_or_else(|| self.empty_hashes[index.depth() as usize].into())) + Ok(RpoDigest::from( + self.get_leaf_node(index.value()) + .unwrap_or_else(|| *self.empty_hashes[index.depth() as usize]), + )) } else { - Ok(self.get_branch_node(&index).parent().into()) + Ok(self.get_branch_node(&index).parent()) } } @@ -143,7 +144,7 @@ impl SimpleSmt { /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}. pub fn get_leaf(&self, index: u64) -> Result { let index = NodeIndex::new(self.depth, index)?; - self.get_node(index) + Ok(self.get_node(index)?.into()) } /// Returns a Merkle path from the node at the specified index to the root. @@ -166,9 +167,9 @@ impl SimpleSmt { index.move_up(); let BranchNode { left, right } = self.get_branch_node(&index); let value = if is_right { left } else { right }; - path.push(*value); + path.push(value); } - Ok(path.into()) + Ok(MerklePath::new(path)) } /// Return a Merkle path from the leaf at the specified index to the root. @@ -193,9 +194,9 @@ impl SimpleSmt { /// Returns an iterator over the inner nodes of this Merkle tree. pub fn inner_nodes(&self) -> impl Iterator + '_ { self.branches.values().map(|e| InnerNodeInfo { - value: e.parent().into(), - left: e.left.into(), - right: e.right.into(), + value: e.parent(), + left: e.left, + right: e.right, }) } @@ -226,7 +227,7 @@ impl SimpleSmt { self.insert_branch_node(index, left, right); value = Rpo256::merge(&[left, right]); } - self.root = value.into(); + self.root = value; Ok(old_value) } diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 0174e15..b12a94d 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -1,6 +1,10 @@ use super::{ - super::{int_to_node, InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt}, - NodeIndex, Rpo256, Vec, Word, EMPTY_WORD, + super::{InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt}, + NodeIndex, Rpo256, Vec, +}; +use crate::{ + merkle::{empty_roots::EMPTY_WORD, int_to_leaf}, + Word, }; // TEST DATA @@ -9,20 +13,20 @@ use super::{ const KEYS4: [u64; 4] = [0, 1, 2, 3]; const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7]; -const VALUES4: [Word; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; +const VALUES4: [Word; 4] = [int_to_leaf(1), int_to_leaf(2), int_to_leaf(3), int_to_leaf(4)]; 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_leaf(1), + int_to_leaf(2), + int_to_leaf(3), + int_to_leaf(4), + int_to_leaf(5), + int_to_leaf(6), + int_to_leaf(7), + int_to_leaf(8), ]; -const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8]; +const ZERO_VALUES8: [Word; 8] = [int_to_leaf(0); 8]; // TESTS // ================================================================================================ @@ -32,7 +36,7 @@ fn build_empty_tree() { // tree of depth 3 let smt = SimpleSmt::new(3).unwrap(); let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap(); - assert_eq!(mt.root(), smt.root()); + assert_eq!(mt.root(), smt.root().into()); } #[test] @@ -42,11 +46,11 @@ fn build_sparse_tree() { // insert single value let key = 6; - let new_node = int_to_node(7); + let new_node = int_to_leaf(7); values[key as usize] = new_node; let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf"); let mt2 = MerkleTree::new(values.clone()).unwrap(); - assert_eq!(mt2.root(), smt.root()); + assert_eq!(mt2.root(), smt.root().into()); assert_eq!( mt2.get_path(NodeIndex::make(3, 6)).unwrap(), smt.get_path(NodeIndex::make(3, 6)).unwrap() @@ -55,11 +59,11 @@ fn build_sparse_tree() { // insert second value at distinct leaf branch let key = 2; - let new_node = int_to_node(3); + let new_node = int_to_leaf(3); values[key as usize] = new_node; let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf"); let mt3 = MerkleTree::new(values).unwrap(); - assert_eq!(mt3.root(), smt.root()); + assert_eq!(mt3.root(), smt.root().into()); assert_eq!( mt3.get_path(NodeIndex::make(3, 2)).unwrap(), smt.get_path(NodeIndex::make(3, 2)).unwrap() @@ -78,16 +82,16 @@ fn test_depth2_tree() { assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap()); // check get_node() - assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(VALUES4[0], *tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(VALUES4[1], *tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(VALUES4[2], *tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(VALUES4[3], *tree.get_node(NodeIndex::make(2, 3)).unwrap()); // check get_path(): depth 2 - assert_eq!(vec![VALUES4[1], node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(vec![VALUES4[0], node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(vec![VALUES4[3], node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(vec![VALUES4[2], node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(vec![VALUES4[1].into(), node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(vec![VALUES4[0].into(), node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(vec![VALUES4[3].into(), node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(vec![VALUES4[2].into(), node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap()); // check get_path(): depth 1 assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap()); @@ -99,10 +103,10 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> { let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap(); // check depth 2 - assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(VALUES4[0], *tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(VALUES4[1], *tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(VALUES4[2], *tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(VALUES4[3], *tree.get_node(NodeIndex::make(2, 3)).unwrap()); // get parent nodes let root = tree.root(); @@ -142,7 +146,7 @@ fn update_leaf() { // update one value let key = 3; - let new_node = int_to_node(9); + let new_node = int_to_leaf(9); let mut expected_values = VALUES8.to_vec(); expected_values[key] = new_node; let expected_tree = MerkleTree::new(expected_values.clone()).unwrap(); @@ -153,7 +157,7 @@ fn update_leaf() { // update another value let key = 6; - let new_node = int_to_node(10); + let new_node = int_to_leaf(10); expected_values[key] = new_node; let expected_tree = MerkleTree::new(expected_values.clone()).unwrap(); @@ -172,34 +176,34 @@ fn small_tree_opening_is_consistent() { // / \ / \ / \ / \ // a b 0 0 c 0 0 d - let z = Word::from(RpoDigest::default()); + let z = EMPTY_WORD; 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 e = RpoDigest::from(Rpo256::merge(&[a.into(), b.into()])); + let f = RpoDigest::from(Rpo256::merge(&[z.into(), z.into()])); + let g = RpoDigest::from(Rpo256::merge(&[c.into(), z.into()])); + let h = RpoDigest::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 i = RpoDigest::from(Rpo256::merge(&[e.into(), f.into()])); + let j = RpoDigest::from(Rpo256::merge(&[g.into(), h.into()])); - let k = Word::from(Rpo256::merge(&[i.into(), j.into()])); + let k = RpoDigest::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(); - assert_eq!(tree.root(), Word::from(k)); + assert_eq!(tree.root(), RpoDigest::from(k)); - let cases: Vec<(u8, u64, Vec)> = vec![ - (3, 0, vec![b, f, j]), - (3, 1, vec![a, f, j]), - (3, 4, vec![z, h, i]), - (3, 7, vec![z, g, i]), + let cases: Vec<(u8, u64, Vec)> = vec![ + (3, 0, vec![b.into(), f, j]), + (3, 1, vec![a.into(), f, j]), + (3, 4, vec![z.into(), h, i]), + (3, 7, vec![z.into(), g, i]), (2, 0, vec![f, j]), (2, 1, vec![e, j]), (2, 2, vec![h, i]), @@ -217,26 +221,26 @@ fn small_tree_opening_is_consistent() { #[test] fn fail_on_duplicates() { - let entries = [(1_u64, int_to_node(1)), (5, int_to_node(2)), (1_u64, int_to_node(3))]; + let entries = [(1_u64, int_to_leaf(1)), (5, int_to_leaf(2)), (1_u64, int_to_leaf(3))]; let smt = SimpleSmt::with_leaves(64, entries); assert!(smt.is_err()); - let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2)), (1_u64, int_to_node(0))]; + let entries = [(1_u64, int_to_leaf(0)), (5, int_to_leaf(2)), (1_u64, int_to_leaf(0))]; let smt = SimpleSmt::with_leaves(64, entries); assert!(smt.is_err()); - let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2)), (1_u64, int_to_node(1))]; + let entries = [(1_u64, int_to_leaf(0)), (5, int_to_leaf(2)), (1_u64, int_to_leaf(1))]; let smt = SimpleSmt::with_leaves(64, entries); assert!(smt.is_err()); - let entries = [(1_u64, int_to_node(1)), (5, int_to_node(2)), (1_u64, int_to_node(0))]; + let entries = [(1_u64, int_to_leaf(1)), (5, int_to_leaf(2)), (1_u64, int_to_leaf(0))]; let smt = SimpleSmt::with_leaves(64, entries); assert!(smt.is_err()); } #[test] fn with_no_duplicates_empty_node() { - let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2))]; + let entries = [(1_u64, int_to_leaf(0)), (5, int_to_leaf(2))]; let smt = SimpleSmt::with_leaves(64, entries); assert!(smt.is_ok()); } @@ -244,9 +248,21 @@ fn with_no_duplicates_empty_node() { // HELPER FUNCTIONS // -------------------------------------------------------------------------------------------- -fn compute_internal_nodes() -> (Word, Word, Word) { - let node2 = Rpo256::hash_elements(&[VALUES4[0], VALUES4[1]].concat()); - let node3 = Rpo256::hash_elements(&[VALUES4[2], VALUES4[3]].concat()); +fn compute_internal_nodes() -> (RpoDigest, RpoDigest, RpoDigest) { + let node2 = Rpo256::hash_elements( + &[VALUES4[0], VALUES4[1]] + .iter() + .map(|digest| *digest) + .collect::>() + .concat(), + ); + let node3 = Rpo256::hash_elements( + &[VALUES4[2], VALUES4[3]] + .iter() + .map(|digest| *digest) + .collect::>() + .concat(), + ); let root = Rpo256::merge(&[node2, node3]); (root.into(), node2.into(), node3.into()) diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index d4c1ffb..a412092 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,6 +1,6 @@ use super::{ mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, - MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, TieredSmt, ValuePath, Vec, Word, + MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, TieredSmt, ValuePath, Vec, }; use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use core::borrow::Borrow; @@ -130,21 +130,20 @@ impl MerkleStore { /// This method can return the following errors: /// - `RootNotInStore` if the `root` is not present in the store. /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. - pub fn get_node(&self, root: Word, index: NodeIndex) -> Result { - let mut hash: RpoDigest = root.into(); + pub fn get_node(&self, root: RpoDigest, index: NodeIndex) -> Result { + let mut hash = root; // corner case: check the root is in the store when called with index `NodeIndex::root()` - self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash.into()))?; + self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?; for i in (0..index.depth()).rev() { - let node = - self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash.into(), index))?; + let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?; let bit = (index.value() >> i) & 1; hash = if bit == 0 { node.left } else { node.right } } - Ok(hash.into()) + Ok(hash) } /// Returns the node at the specified `index` and its opening to the `root`. @@ -155,23 +154,22 @@ impl MerkleStore { /// This method can return the following errors: /// - `RootNotInStore` if the `root` is not present in the store. /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. - pub fn get_path(&self, root: Word, index: NodeIndex) -> Result { - let mut hash: RpoDigest = root.into(); + pub fn get_path(&self, root: RpoDigest, index: NodeIndex) -> Result { + let mut hash = root; let mut path = Vec::with_capacity(index.depth().into()); // corner case: check the root is in the store when called with index `NodeIndex::root()` - self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash.into()))?; + self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?; for i in (0..index.depth()).rev() { - let node = - self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash.into(), index))?; + let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?; let bit = (index.value() >> i) & 1; hash = if bit == 0 { - path.push(node.right.into()); + path.push(node.right); node.left } else { - path.push(node.left.into()); + path.push(node.left); node.right } } @@ -180,7 +178,7 @@ impl MerkleStore { path.reverse(); Ok(ValuePath { - value: hash.into(), + value: hash, path: MerklePath::new(path), }) } @@ -202,7 +200,7 @@ impl MerkleStore { /// information, check [NodeIndex::new]. pub fn get_leaf_depth( &self, - root: Word, + root: RpoDigest, tree_depth: u8, index: u64, ) -> Result { @@ -221,9 +219,9 @@ impl MerkleStore { // check if the root exists, providing the proper error report if it doesn't let empty = EmptySubtreeRoots::empty_hashes(tree_depth); - let mut hash: RpoDigest = root.into(); + let mut hash = root; if !self.nodes.contains_key(&hash) { - return Err(MerkleError::RootNotInStore(hash.into())); + return Err(MerkleError::RootNotInStore(hash)); } // we traverse from root to leaf, so the path is reversed @@ -266,11 +264,11 @@ impl MerkleStore { pub fn subset(&self, roots: I) -> MerkleStore where I: Iterator, - R: Borrow, + R: Borrow, { let mut store = MerkleStore::new(); for root in roots { - let root = RpoDigest::from(*root.borrow()); + let root = *root.borrow(); store.clone_tree_from(root, self); } store @@ -279,9 +277,9 @@ impl MerkleStore { /// Iterator over the inner nodes of the [MerkleStore]. pub fn inner_nodes(&self) -> impl Iterator + '_ { self.nodes.iter().map(|(r, n)| InnerNodeInfo { - value: r.into(), - left: n.left.into(), - right: n.right.into(), + value: *r, + left: n.left, + right: n.right, }) } @@ -294,9 +292,9 @@ impl MerkleStore { I: Iterator, { for node in iter { - let value: RpoDigest = node.value.into(); - let left: RpoDigest = node.left.into(); - let right: RpoDigest = node.right.into(); + let value: RpoDigest = node.value; + let left: RpoDigest = node.left; + let right: RpoDigest = node.right; debug_assert_eq!(Rpo256::merge(&[left, right]), value); self.nodes.insert(value, Node { left, right }); @@ -313,13 +311,13 @@ impl MerkleStore { pub fn add_merkle_path( &mut self, index: u64, - node: Word, + node: RpoDigest, path: MerklePath, - ) -> Result { - let root = path.inner_nodes(index, node)?.fold(Word::default(), |_, node| { - let value: RpoDigest = node.value.into(); - let left: RpoDigest = node.left.into(); - let right: RpoDigest = node.right.into(); + ) -> Result { + let root = path.inner_nodes(index, node)?.fold(RpoDigest::default(), |_, node| { + let value: RpoDigest = node.value; + let left: RpoDigest = node.left; + let right: RpoDigest = node.right; debug_assert_eq!(Rpo256::merge(&[left, right]), value); self.nodes.insert(value, Node { left, right }); @@ -337,7 +335,7 @@ impl MerkleStore { /// For further reference, check [MerkleStore::add_merkle_path]. pub fn add_merkle_paths(&mut self, paths: I) -> Result<(), MerkleError> where - I: IntoIterator, + I: IntoIterator, { for (index_value, node, path) in paths.into_iter() { self.add_merkle_path(index_value, node, path)?; @@ -348,7 +346,10 @@ impl MerkleStore { /// Appends the provided [MerklePathSet] into the store. /// /// For further reference, check [MerkleStore::add_merkle_path]. - pub fn add_merkle_path_set(&mut self, path_set: &MerklePathSet) -> Result { + pub fn add_merkle_path_set( + &mut self, + path_set: &MerklePathSet, + ) -> Result { let root = path_set.root(); for (index, path) in path_set.to_paths() { self.add_merkle_path(index, path.value, path.path)?; @@ -365,9 +366,9 @@ impl MerkleStore { /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. pub fn set_node( &mut self, - mut root: Word, + mut root: RpoDigest, index: NodeIndex, - value: Word, + value: RpoDigest, ) -> Result { let node = value; let ValuePath { value, path } = self.get_path(root, index)?; @@ -383,14 +384,21 @@ impl MerkleStore { /// Merges two elements and adds the resulting node into the store. /// /// Merges arbitrary values. They may be leafs, nodes, or a mixture of both. - pub fn merge_roots(&mut self, root1: Word, root2: Word) -> Result { - let left: RpoDigest = root1.into(); - let right: RpoDigest = root2.into(); - - let parent = Rpo256::merge(&[left, right]); - self.nodes.insert(parent, Node { left, right }); - - Ok(parent.into()) + pub fn merge_roots( + &mut self, + left_root: RpoDigest, + right_root: RpoDigest, + ) -> Result { + let parent = Rpo256::merge(&[left_root, right_root]); + self.nodes.insert( + parent, + Node { + left: left_root, + right: right_root, + }, + ); + + Ok(parent) } // HELPER METHODS @@ -404,7 +412,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); } diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 0bdfa4f..1a7ead6 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -1,10 +1,10 @@ use super::{ - super::EMPTY_WORD, Deserializable, EmptySubtreeRoots, MerkleError, MerklePath, MerkleStore, - NodeIndex, RpoDigest, Serializable, + Deserializable, EmptySubtreeRoots, MerkleError, MerklePath, MerkleStore, NodeIndex, RpoDigest, + Serializable, }; use crate::{ hash::rpo::Rpo256, - merkle::{int_to_node, MerklePathSet, MerkleTree, SimpleSmt}, + merkle::{int_to_leaf, int_to_node, MerklePathSet, MerkleTree, SimpleSmt}, Felt, Word, WORD_SIZE, }; @@ -15,17 +15,17 @@ use std::error::Error; // ================================================================================================ const KEYS4: [u64; 4] = [0, 1, 2, 3]; -const VALUES4: [Word; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; +const VALUES4: [Word; 4] = [int_to_leaf(1), int_to_leaf(2), int_to_leaf(3), int_to_leaf(4)]; 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_leaf(1), + int_to_leaf(2), + int_to_leaf(3), + int_to_leaf(4), + int_to_leaf(5), + int_to_leaf(6), + int_to_leaf(7), + int_to_leaf(8), ]; // TESTS @@ -36,13 +36,13 @@ fn test_root_not_in_store() -> Result<(), MerkleError> { let mtree = MerkleTree::new(VALUES4.to_vec())?; let store = MerkleStore::from(&mtree); assert_eq!( - store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)), - Err(MerkleError::RootNotInStore(VALUES4[0])), + store.get_node(VALUES4[0].into(), NodeIndex::make(mtree.depth(), 0)), + Err(MerkleError::RootNotInStore(VALUES4[0].into())), "Leaf 0 is not a root" ); assert_eq!( - store.get_path(VALUES4[0], NodeIndex::make(mtree.depth(), 0)), - Err(MerkleError::RootNotInStore(VALUES4[0])), + store.get_path(VALUES4[0].into(), NodeIndex::make(mtree.depth(), 0)), + Err(MerkleError::RootNotInStore(VALUES4[0].into())), "Leaf 0 is not a root" ); @@ -58,22 +58,22 @@ fn test_merkle_tree() -> Result<(), MerkleError> { // checks the leaves in the store corresponds to the expected values assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)), - Ok(VALUES4[0]), + Ok(VALUES4[0].into()), "node 0 must be in the tree" ); assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)), - Ok(VALUES4[1]), + Ok(VALUES4[1].into()), "node 1 must be in the tree" ); assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)), - Ok(VALUES4[2]), + Ok(VALUES4[2].into()), "node 2 must be in the tree" ); assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)), - Ok(VALUES4[3]), + Ok(VALUES4[3].into()), "node 3 must be in the tree" ); @@ -104,7 +104,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { // assert the merkle path returned by the store is the same as the one in the tree let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap(); assert_eq!( - VALUES4[0], result.value, + VALUES4[0], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -115,7 +115,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 1)).unwrap(); assert_eq!( - VALUES4[1], result.value, + VALUES4[1], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -126,7 +126,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 2)).unwrap(); assert_eq!( - VALUES4[2], result.value, + VALUES4[2], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -137,7 +137,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 3)).unwrap(); assert_eq!( - VALUES4[3], result.value, + VALUES4[3], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -152,7 +152,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { #[test] fn test_empty_roots() { let store = MerkleStore::default(); - let mut root = RpoDigest::new(EMPTY_WORD); + let mut root = RpoDigest::default(); for depth in 0..255 { root = Rpo256::merge(&[root; 2]); @@ -176,13 +176,17 @@ fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> { let index = NodeIndex::make(depth, 0); let store_path = store.get_path(smt.root(), index)?; let smt_path = smt.get_path(index)?; - assert_eq!(store_path.value, EMPTY_WORD, "the leaf of an empty tree is always ZERO"); + assert_eq!( + store_path.value, + RpoDigest::default(), + "the leaf of an empty tree is always ZERO" + ); assert_eq!( store_path.path, smt_path, "the returned merkle path does not match the computed values" ); assert_eq!( - store_path.path.compute_root(depth.into(), EMPTY_WORD).unwrap(), + store_path.path.compute_root(depth.into(), RpoDigest::default()).unwrap(), smt.root(), "computed root from the path must match the empty tree root" ); @@ -201,16 +205,16 @@ fn test_get_invalid_node() { #[test] fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { let keys2: [u64; 2] = [0, 1]; - let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)]; + let leaves2: [Word; 2] = [int_to_leaf(1), int_to_leaf(2)]; let smt = SimpleSmt::with_leaves(1, keys2.into_iter().zip(leaves2.into_iter())).unwrap(); let store = MerkleStore::from(&smt); let idx = NodeIndex::make(1, 0); - assert_eq!(smt.get_node(idx).unwrap(), leaves2[0]); - assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap()); + assert_eq!(smt.get_node(idx).unwrap(), leaves2[0].into()); + assert_eq!(store.get_node(smt.root().into(), idx).unwrap(), smt.get_node(idx).unwrap()); let idx = NodeIndex::make(1, 1); - assert_eq!(smt.get_node(idx).unwrap(), leaves2[1]); + assert_eq!(smt.get_node(idx).unwrap(), leaves2[1].into()); assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap()); Ok(()) @@ -227,28 +231,28 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { // STORE LEAVES ARE CORRECT ============================================================== // checks the leaves in the store corresponds to the expected values assert_eq!( - store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)), - Ok(VALUES4[0]), + store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 0)), + Ok(VALUES4[0].into()), "node 0 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)), - Ok(VALUES4[1]), + store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 1)), + Ok(VALUES4[1].into()), "node 1 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)), - Ok(VALUES4[2]), + store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 2)), + Ok(VALUES4[2].into()), "node 2 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)), - Ok(VALUES4[3]), + store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 3)), + Ok(VALUES4[3].into()), "node 3 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)), - Ok(EMPTY_WORD), + store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 4)), + Ok(RpoDigest::default()), "unmodified node 4 must be ZERO" ); @@ -284,7 +288,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { // assert the merkle path returned by the store is the same as the one in the tree let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 0)).unwrap(); assert_eq!( - VALUES4[0], result.value, + VALUES4[0], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -295,7 +299,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 1)).unwrap(); assert_eq!( - VALUES4[1], result.value, + VALUES4[1], *result.value, "Value for merkle path at index 1 must match leaf value" ); assert_eq!( @@ -306,7 +310,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 2)).unwrap(); assert_eq!( - VALUES4[2], result.value, + VALUES4[2], *result.value, "Value for merkle path at index 2 must match leaf value" ); assert_eq!( @@ -317,7 +321,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 3)).unwrap(); assert_eq!( - VALUES4[3], result.value, + VALUES4[3], *result.value, "Value for merkle path at index 3 must match leaf value" ); assert_eq!( @@ -328,7 +332,8 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 4)).unwrap(); assert_eq!( - EMPTY_WORD, result.value, + RpoDigest::default(), + result.value, "Value for merkle path at index 4 must match leaf value" ); assert_eq!( @@ -357,10 +362,10 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let p3 = mtree.get_path(NodeIndex::make(2, i3)).unwrap(); let paths = [ - (i0, VALUES4[i0 as usize], p0), - (i1, VALUES4[i1 as usize], p1), - (i2, VALUES4[i2 as usize], p2), - (i3, VALUES4[i3 as usize], p3), + (i0, VALUES4[i0 as usize].into(), p0), + (i1, VALUES4[i1 as usize].into(), p1), + (i2, VALUES4[i2 as usize].into(), p2), + (i3, VALUES4[i3 as usize].into(), p3), ]; let mut store = MerkleStore::default(); @@ -373,22 +378,22 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { // checks the leaves in the store corresponds to the expected values assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 0)), - Ok(VALUES4[0]), + Ok(VALUES4[0].into()), "node 0 must be in the set" ); assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 1)), - Ok(VALUES4[1]), + Ok(VALUES4[1].into()), "node 1 must be in the set" ); assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 2)), - Ok(VALUES4[2]), + Ok(VALUES4[2].into()), "node 2 must be in the set" ); assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 3)), - Ok(VALUES4[3]), + Ok(VALUES4[3].into()), "node 3 must be in the set" ); @@ -419,7 +424,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { // assert the merkle path returned by the store is the same as the one in the set let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 0)).unwrap(); assert_eq!( - VALUES4[0], result.value, + VALUES4[0], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -430,7 +435,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 1)).unwrap(); assert_eq!( - VALUES4[1], result.value, + VALUES4[1], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -441,7 +446,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 2)).unwrap(); assert_eq!( - VALUES4[2], result.value, + VALUES4[2], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -452,7 +457,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 3)).unwrap(); assert_eq!( - VALUES4[3], result.value, + VALUES4[3], *result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -477,7 +482,7 @@ fn wont_open_to_different_depth_root() { for depth in (1..=63).rev() { root = Rpo256::merge(&[root, empty[depth]]); } - let root = Word::from(root); + let root = RpoDigest::from(root); // For this example, the depth of the Merkle tree is 1, as we have only two leaves. Here we // attempt to fetch a node on the maximum depth, and it should fail because the root shouldn't @@ -556,20 +561,20 @@ fn test_constructors() -> Result<(), MerkleError> { let d = 2; let paths = [ - (0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0)).unwrap()), - (1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1)).unwrap()), - (2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2)).unwrap()), - (3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3)).unwrap()), + (0, VALUES4[0].into(), mtree.get_path(NodeIndex::make(d, 0)).unwrap()), + (1, VALUES4[1].into(), mtree.get_path(NodeIndex::make(d, 1)).unwrap()), + (2, VALUES4[2].into(), mtree.get_path(NodeIndex::make(d, 2)).unwrap()), + (3, VALUES4[3].into(), mtree.get_path(NodeIndex::make(d, 3)).unwrap()), ]; let mut store1 = MerkleStore::default(); store1.add_merkle_paths(paths.clone())?; let mut store2 = MerkleStore::default(); - store2.add_merkle_path(0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0))?)?; - store2.add_merkle_path(1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1))?)?; - store2.add_merkle_path(2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2))?)?; - store2.add_merkle_path(3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?; + store2.add_merkle_path(0, VALUES4[0].into(), mtree.get_path(NodeIndex::make(d, 0))?)?; + store2.add_merkle_path(1, VALUES4[1].into(), mtree.get_path(NodeIndex::make(d, 1))?)?; + store2.add_merkle_path(2, VALUES4[2].into(), mtree.get_path(NodeIndex::make(d, 2))?)?; + store2.add_merkle_path(3, VALUES4[3].into(), mtree.get_path(NodeIndex::make(d, 3))?)?; let set = MerklePathSet::new(d).with_paths(paths).unwrap(); for key in [0, 1, 2, 3] { @@ -590,11 +595,11 @@ fn node_path_should_be_truncated_by_midtier_insert() { let key = 0b11010010_11001100_11001100_11001100_11001100_11001100_11001100_11001100_u64; let mut store = MerkleStore::new(); - let root: Word = EmptySubtreeRoots::empty_hashes(64)[0].into(); + let root: RpoDigest = EmptySubtreeRoots::empty_hashes(64)[0]; // insert first node - works as expected let depth = 64; - let node = [Felt::new(key); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(key); WORD_SIZE]); let index = NodeIndex::new(depth, key).unwrap(); let root = store.set_node(root, index, node).unwrap().root; let result = store.get_node(root, index).unwrap(); @@ -607,7 +612,7 @@ fn node_path_should_be_truncated_by_midtier_insert() { let key = key ^ (1 << 63); let key = key >> 8; let depth = 56; - let node = [Felt::new(key); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(key); WORD_SIZE]); let index = NodeIndex::new(depth, key).unwrap(); let root = store.set_node(root, index, node).unwrap().root; let result = store.get_node(root, index).unwrap(); @@ -626,13 +631,13 @@ fn node_path_should_be_truncated_by_midtier_insert() { #[test] fn get_leaf_depth_works_depth_64() { let mut store = MerkleStore::new(); - let mut root: Word = EmptySubtreeRoots::empty_hashes(64)[0].into(); + let mut root: RpoDigest = EmptySubtreeRoots::empty_hashes(64)[0]; let key = u64::MAX; // this will create a rainbow tree and test all opening to depth 64 for d in 0..64 { let k = key & (u64::MAX >> d); - let node = [Felt::new(k); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(k); WORD_SIZE]); let index = NodeIndex::new(64, k).unwrap(); // assert the leaf doesn't exist before the insert. the returned depth should always @@ -649,14 +654,14 @@ fn get_leaf_depth_works_depth_64() { #[test] fn get_leaf_depth_works_with_incremental_depth() { let mut store = MerkleStore::new(); - let mut root: Word = EmptySubtreeRoots::empty_hashes(64)[0].into(); + let mut root: RpoDigest = EmptySubtreeRoots::empty_hashes(64)[0]; // insert some path to the left of the root and assert it let key = 0b01001011_10110110_00001101_01110100_00111011_10101101_00000100_01000001_u64; assert_eq!(0, store.get_leaf_depth(root, 64, key).unwrap()); let depth = 64; let index = NodeIndex::new(depth, key).unwrap(); - let node = [Felt::new(key); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(key); WORD_SIZE]); root = store.set_node(root, index, node).unwrap().root; assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); @@ -665,7 +670,7 @@ fn get_leaf_depth_works_with_incremental_depth() { assert_eq!(1, store.get_leaf_depth(root, 64, key).unwrap()); let depth = 16; let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap(); - let node = [Felt::new(key); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(key); WORD_SIZE]); root = store.set_node(root, index, node).unwrap().root; assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); @@ -673,7 +678,7 @@ fn get_leaf_depth_works_with_incremental_depth() { let key = 0b11001011_10110111_00000000_00000000_00000000_00000000_00000000_00000000_u64; assert_eq!(16, store.get_leaf_depth(root, 64, key).unwrap()); let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap(); - let node = [Felt::new(key); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(key); WORD_SIZE]); root = store.set_node(root, index, node).unwrap().root; assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); @@ -682,7 +687,7 @@ fn get_leaf_depth_works_with_incremental_depth() { assert_eq!(15, store.get_leaf_depth(root, 64, key).unwrap()); let depth = 17; let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap(); - let node = [Felt::new(key); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(key); WORD_SIZE]); root = store.set_node(root, index, node).unwrap().root; assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); } @@ -690,7 +695,7 @@ fn get_leaf_depth_works_with_incremental_depth() { #[test] fn get_leaf_depth_works_with_depth_8() { let mut store = MerkleStore::new(); - let mut root: Word = EmptySubtreeRoots::empty_hashes(8)[0].into(); + let mut root: RpoDigest = EmptySubtreeRoots::empty_hashes(8)[0]; // insert some random, 8 depth keys. `a` diverges from the first bit let a = 0b01101001_u64; @@ -700,7 +705,7 @@ fn get_leaf_depth_works_with_depth_8() { for k in [a, b, c, d] { let index = NodeIndex::new(8, k).unwrap(); - let node = [Felt::new(k); WORD_SIZE]; + let node = RpoDigest::from([Felt::new(k); WORD_SIZE]); root = store.set_node(root, index, node).unwrap().root; } @@ -780,7 +785,7 @@ fn check_mstore_subtree(store: &MerkleStore, subtree: &MerkleTree) { for (i, value) in subtree.leaves() { let index = NodeIndex::new(subtree.depth(), i).unwrap(); let path1 = store.get_path(subtree.root(), index).unwrap(); - assert_eq!(&path1.value, value); + assert_eq!(*path1.value, *value); let path2 = subtree.get_path(index).unwrap(); assert_eq!(path1.path, path2); diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index eaf1cd8..3fb3cdd 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/src/merkle/tiered_smt/mod.rs @@ -1,6 +1,6 @@ use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, - Rpo256, RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO, + empty_roots::EMPTY_WORD, BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, + MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, StarkField, Vec, Word, ZERO, }; use core::cmp; @@ -123,7 +123,7 @@ impl TieredSmt { let mut path = Vec::with_capacity(index.depth() as usize); for _ in 0..index.depth() { let node = self.get_node_unchecked(&index.sibling()); - path.push(node.into()); + path.push(node); index.move_up(); } @@ -200,9 +200,9 @@ impl TieredSmt { self.nodes.iter().filter_map(|(index, node)| { if is_inner_node(index) { Some(InnerNodeInfo { - value: node.into(), - left: self.get_node_unchecked(&index.left_child()).into(), - right: self.get_node_unchecked(&index.right_child()).into(), + value: *node, + left: self.get_node_unchecked(&index.left_child()), + right: self.get_node_unchecked(&index.right_child()), }) } else { None @@ -456,7 +456,7 @@ impl BottomLeaf { let mut elements = Vec::with_capacity(self.values.len() * 2); for (key, val) in self.values.iter() { key.iter().for_each(|&v| elements.push(Felt::new(v))); - elements.extend_from_slice(val); + elements.extend_from_slice(val.as_slice()); } // TODO: hash in domain Rpo256::hash_elements(&elements) diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index f46a4ae..28ac42a 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -66,11 +66,19 @@ fn tsmt_insert_two_16() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(32, raw_a >> 32); let leaf_node_a = build_leaf_node(key_a, val_a, 32); - tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_a, leaf_node_a.into()) + .unwrap() + .root + .into(); let index_b = NodeIndex::make(32, raw_b >> 32); let leaf_node_b = build_leaf_node(key_b, val_b, 32); - tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_b, leaf_node_b.into()) + .unwrap() + .root + .into(); // --- verify that data is consistent between store and tree -------------- @@ -78,12 +86,12 @@ fn tsmt_insert_two_16() { assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -122,11 +130,19 @@ fn tsmt_insert_two_32() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(48, raw_a >> 16); let leaf_node_a = build_leaf_node(key_a, val_a, 48); - tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_a, leaf_node_a.into()) + .unwrap() + .root + .into(); let index_b = NodeIndex::make(48, raw_b >> 16); let leaf_node_b = build_leaf_node(key_b, val_b, 48); - tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_b, leaf_node_b.into()) + .unwrap() + .root + .into(); // --- verify that data is consistent between store and tree -------------- @@ -134,12 +150,12 @@ fn tsmt_insert_two_32() { assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -181,15 +197,27 @@ fn tsmt_insert_three() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(32, raw_a >> 32); let leaf_node_a = build_leaf_node(key_a, val_a, 32); - tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_a, leaf_node_a.into()) + .unwrap() + .root + .into(); let index_b = NodeIndex::make(32, raw_b >> 32); let leaf_node_b = build_leaf_node(key_b, val_b, 32); - tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_b, leaf_node_b.into()) + .unwrap() + .root + .into(); let index_c = NodeIndex::make(32, raw_c >> 32); let leaf_node_c = build_leaf_node(key_c, val_c, 32); - tree_root = store.set_node(tree_root, index_c, leaf_node_c.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_c, leaf_node_c.into()) + .unwrap() + .root + .into(); // --- verify that data is consistent between store and tree -------------- @@ -197,17 +225,17 @@ fn tsmt_insert_three() { assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); assert_eq!(smt.get_value(key_c), val_c); assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c); - let expected_path = store.get_path(tree_root, index_c).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_c).unwrap().path; assert_eq!(smt.get_path(index_c).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -236,13 +264,13 @@ fn tsmt_update() { let mut tree_root = get_init_root(); let index = NodeIndex::make(16, raw >> 48); let leaf_node = build_leaf_node(key, value_b, 16); - tree_root = store.set_node(tree_root, index, leaf_node.into()).unwrap().root; + tree_root = store.set_node(tree_root.into(), index, leaf_node.into()).unwrap().root.into(); assert_eq!(smt.root(), tree_root.into()); assert_eq!(smt.get_value(key), value_b); assert_eq!(smt.get_node(index).unwrap(), leaf_node); - let expected_path = store.get_path(tree_root, index).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index).unwrap().path; assert_eq!(smt.get_path(index).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -281,7 +309,7 @@ fn tsmt_bottom_tier() { // key_b is smaller than key_a. let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]); let mut tree_root = get_init_root(); - tree_root = store.set_node(tree_root, index, leaf_node.into()).unwrap().root; + tree_root = store.set_node(tree_root.into(), index, leaf_node.into()).unwrap().root.into(); // --- verify that data is consistent between store and tree -------------- @@ -291,7 +319,7 @@ fn tsmt_bottom_tier() { assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index).unwrap(), leaf_node); - let expected_path = store.get_path(tree_root, index).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index).unwrap().path; assert_eq!(smt.get_path(index).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -329,11 +357,19 @@ fn tsmt_bottom_tier_two() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(64, raw_a); let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]); - tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_a, leaf_node_a.into()) + .unwrap() + .root + .into(); let index_b = NodeIndex::make(64, raw_b); let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]); - tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root; + tree_root = store + .set_node(tree_root.into(), index_b, leaf_node_b.into()) + .unwrap() + .root + .into(); // --- verify that data is consistent between store and tree -------------- @@ -341,12 +377,12 @@ fn tsmt_bottom_tier_two() { assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; + let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -423,7 +459,7 @@ fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { let mut key = Word::from(key); key[3] = ZERO; elements.extend_from_slice(&key); - elements.extend_from_slice(val); + elements.extend_from_slice(val.as_slice()); } Rpo256::hash_elements(&elements) From 0e0a3fda4f2ad18eda2a79f88262ce8b0d1a43d5 Mon Sep 17 00:00:00 2001 From: tohrnii <100405913+tohrnii@users.noreply.github.com> Date: Tue, 13 Jun 2023 10:28:15 +0100 Subject: [PATCH 11/19] refactor: refactor to clean up and simplify things --- src/merkle/merkle_tree.rs | 76 +++++++--------- src/merkle/mod.rs | 4 + src/merkle/path_set.rs | 48 +++++----- src/merkle/simple_smt/tests.rs | 82 +++++++++-------- src/merkle/store/tests.rs | 157 +++++++++++++++++---------------- src/merkle/tiered_smt/tests.rs | 64 +++----------- 6 files changed, 200 insertions(+), 231 deletions(-) diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 0fc99a7..4e8d154 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -1,6 +1,6 @@ use super::{InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word}; use crate::utils::{string::String, uninit_vector, word_to_hex}; -use core::{fmt, slice}; +use core::{fmt, ops::Deref, slice}; use winter_math::log2; // MERKLE TREE @@ -117,7 +117,11 @@ impl MerkleTree { /// Returns an iterator over the leaves of this [MerkleTree]. pub fn leaves(&self) -> impl Iterator { let leaves_start = self.nodes.len() / 2; - self.nodes.iter().skip(leaves_start).enumerate().map(|(i, v)| (i as u64, &**v)) + self.nodes + .iter() + .skip(leaves_start) + .enumerate() + .map(|(i, v)| (i as u64, v.deref())) } /// Returns n iterator over every inner node of this [MerkleTree]. @@ -256,29 +260,29 @@ pub fn path_to_text(path: &MerklePath) -> Result { mod tests { use super::*; use crate::{ - merkle::{int_to_leaf, InnerNodeInfo}, + merkle::{digests_to_words, int_to_leaf, int_to_node, InnerNodeInfo}, Felt, Word, WORD_SIZE, }; use core::mem::size_of; use proptest::prelude::*; - const LEAVES4: [Word; WORD_SIZE] = - [int_to_leaf(1), int_to_leaf(2), int_to_leaf(3), int_to_leaf(4)]; - - const LEAVES8: [Word; 8] = [ - int_to_leaf(1), - int_to_leaf(2), - int_to_leaf(3), - int_to_leaf(4), - int_to_leaf(5), - int_to_leaf(6), - int_to_leaf(7), - int_to_leaf(8), + const LEAVES4: [RpoDigest; WORD_SIZE] = + [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; + + const LEAVES8: [RpoDigest; 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), ]; #[test] fn build_merkle_tree() { - let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + let tree = super::MerkleTree::new(digests_to_words(&LEAVES4)).unwrap(); assert_eq!(8, tree.nodes.len()); // leaves were copied correctly @@ -297,13 +301,13 @@ mod tests { #[test] fn get_leaf() { - let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + let tree = super::MerkleTree::new(digests_to_words(&LEAVES4)).unwrap(); // check depth 2 - assert_eq!(RpoDigest::from(LEAVES4[0]), tree.get_node(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(RpoDigest::from(LEAVES4[1]), tree.get_node(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(RpoDigest::from(LEAVES4[2]), tree.get_node(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(RpoDigest::from(LEAVES4[3]), tree.get_node(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); // check depth 1 let (_, node2, node3) = compute_internal_nodes(); @@ -314,27 +318,15 @@ mod tests { #[test] fn get_path() { - let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + let tree = super::MerkleTree::new(digests_to_words(&LEAVES4)).unwrap(); let (_, node2, node3) = compute_internal_nodes(); // check depth 2 - assert_eq!( - vec![RpoDigest::from(LEAVES4[1]), node3], - *tree.get_path(NodeIndex::make(2, 0)).unwrap() - ); - assert_eq!( - vec![RpoDigest::from(LEAVES4[0]), node3], - *tree.get_path(NodeIndex::make(2, 1)).unwrap() - ); - assert_eq!( - vec![RpoDigest::from(LEAVES4[3]), node2], - *tree.get_path(NodeIndex::make(2, 2)).unwrap() - ); - assert_eq!( - vec![RpoDigest::from(LEAVES4[2]), node2], - *tree.get_path(NodeIndex::make(2, 3)).unwrap() - ); + assert_eq!(vec![LEAVES4[1], node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(vec![LEAVES4[0], node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(vec![LEAVES4[3], node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(vec![LEAVES4[2], node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap()); // check depth 1 assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap()); @@ -343,13 +335,13 @@ mod tests { #[test] fn update_leaf() { - let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap(); + let mut tree = super::MerkleTree::new(digests_to_words(&LEAVES8)).unwrap(); // update one leaf let value = 3; let new_node = int_to_leaf(9); - let mut expected_leaves = LEAVES8.to_vec(); - expected_leaves[value as usize] = new_node; + let mut expected_leaves = digests_to_words(&LEAVES8); + expected_leaves[value as usize] = new_node.into(); let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap(); tree.update_leaf(value, new_node).unwrap(); @@ -367,7 +359,7 @@ mod tests { #[test] fn nodes() -> Result<(), MerkleError> { - let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + let tree = super::MerkleTree::new(digests_to_words(&LEAVES4)).unwrap(); let root = tree.root(); let l1n0 = tree.get_node(NodeIndex::make(1, 0))?; let l1n1 = tree.get_node(NodeIndex::make(1, 1))?; diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index a0d50d6..61dad43 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -102,3 +102,7 @@ const fn int_to_node(value: u64) -> RpoDigest { const fn int_to_leaf(value: u64) -> Word { [Felt::new(value), ZERO, ZERO, ZERO] } + +pub fn digests_to_words(digests: &[RpoDigest]) -> Vec { + digests.iter().map(|d| d.into()).collect() +} diff --git a/src/merkle/path_set.rs b/src/merkle/path_set.rs index 322faeb..ea77b92 100644 --- a/src/merkle/path_set.rs +++ b/src/merkle/path_set.rs @@ -36,7 +36,7 @@ impl MerklePathSet { I: IntoIterator, { paths.into_iter().try_fold(self, |mut set, (index, value, path)| { - set.add_path(index, value, path)?; + set.add_path(index, value.into(), path)?; Ok(set) }) } @@ -151,7 +151,7 @@ impl MerklePathSet { pub fn add_path( &mut self, index_value: u64, - value: RpoDigest, + value: Word, mut path: MerklePath, ) -> Result<(), MerkleError> { let mut index = NodeIndex::new(path.len() as u8, index_value)?; @@ -164,10 +164,10 @@ impl MerklePathSet { // update the current path let parity = index_value & 1; - path.insert(parity as usize, value); + path.insert(parity as usize, value.into()); // traverse to the root, updating the nodes - let root: RpoDigest = Rpo256::merge(&[path[0], path[1]]); + let root = Rpo256::merge(&[path[0], path[1]]); let root = path.iter().skip(2).copied().fold(root, |root, hash| { index.move_up(); Rpo256::merge(&index.build_node(root, hash)) @@ -192,11 +192,7 @@ impl MerklePathSet { /// # Errors /// Returns an error if: /// * Requested node does not exist in the set. - pub fn update_leaf( - &mut self, - base_index_value: u64, - value: RpoDigest, - ) -> Result<(), MerkleError> { + pub fn update_leaf(&mut self, base_index_value: u64, value: Word) -> Result<(), MerkleError> { let mut index = NodeIndex::new(self.depth(), base_index_value)?; let parity = index.value() & 1; let path_key = index.value() - parity; @@ -208,7 +204,7 @@ impl MerklePathSet { // Fill old_hashes vector ----------------------------------------------------------------- let mut current_index = index; let mut old_hashes = Vec::with_capacity(path.len().saturating_sub(2)); - let mut root: RpoDigest = Rpo256::merge(&[path[0], path[1]]); + let mut root = Rpo256::merge(&[path[0], path[1]]); for hash in path.iter().skip(2).copied() { old_hashes.push(root); current_index.move_up(); @@ -217,10 +213,10 @@ impl MerklePathSet { } // Fill new_hashes vector ----------------------------------------------------------------- - path[index.is_value_odd() as usize] = value; + path[index.is_value_odd() as usize] = value.into(); let mut new_hashes = Vec::with_capacity(path.len().saturating_sub(2)); - let mut new_root: RpoDigest = Rpo256::merge(&[path[0], path[1]]); + let mut new_root = Rpo256::merge(&[path[0], path[1]]); for path_hash in path.iter().skip(2).copied() { new_hashes.push(new_root); index.move_up(); @@ -250,7 +246,7 @@ impl MerklePathSet { #[cfg(test)] mod tests { use super::*; - use crate::merkle::int_to_node; + use crate::merkle::{int_to_leaf, int_to_node}; #[test] fn get_root() { @@ -323,20 +319,20 @@ mod tests { ]) .unwrap(); - let new_hash_6 = int_to_node(100); - let new_hash_5 = int_to_node(55); + let new_hash_6 = int_to_leaf(100); + let new_hash_5 = int_to_leaf(55); set.update_leaf(index_6, new_hash_6).unwrap(); let new_path_4 = set.get_path(NodeIndex::make(depth, index_4)).unwrap(); - let new_hash_67 = calculate_parent_hash(new_hash_6, 14_u64, hash_7); + let new_hash_67 = calculate_parent_hash(new_hash_6.into(), 14_u64, hash_7); assert_eq!(new_hash_67, new_path_4[1]); set.update_leaf(index_5, new_hash_5).unwrap(); let new_path_4 = set.get_path(NodeIndex::make(depth, index_4)).unwrap(); let new_path_6 = set.get_path(NodeIndex::make(depth, index_6)).unwrap(); - let new_hash_45 = calculate_parent_hash(new_hash_5, 13_u64, hash_4); + let new_hash_45 = calculate_parent_hash(new_hash_5.into(), 13_u64, hash_4); assert_eq!(new_hash_45, new_path_6[1]); - assert_eq!(new_hash_5, new_path_4[0]); + assert_eq!(RpoDigest::from(new_hash_5), new_path_4[0]); } #[test] @@ -365,30 +361,30 @@ mod tests { let value = b; let index = 1; let path = MerklePath::new([a, j, n].to_vec()); - set.add_path(index, value, path.clone()).unwrap(); + set.add_path(index, value.into(), path.clone()).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); - assert_eq!(RpoDigest::from(root), set.root()); + assert_eq!(root, set.root()); let value = e; let index = 4; let path = MerklePath::new([f.into(), l.into(), m.into()].to_vec()); - set.add_path(index, value, path.clone()).unwrap(); + set.add_path(index, value.into(), path.clone()).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); - assert_eq!(RpoDigest::from(root), set.root()); + assert_eq!(root, set.root()); let value = a; let index = 0; let path = MerklePath::new([b.into(), j.into(), n.into()].to_vec()); - set.add_path(index, value, path.clone()).unwrap(); + set.add_path(index, value.into(), path.clone()).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); - assert_eq!(RpoDigest::from(root), set.root()); + assert_eq!(root, set.root()); let value = h; let index = 7; let path = MerklePath::new([g.into(), k.into(), m.into()].to_vec()); - set.add_path(index, value, path.clone()).unwrap(); + set.add_path(index, value.into(), path.clone()).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); - assert_eq!(RpoDigest::from(root), set.root()); + assert_eq!(root, set.root()); } // HELPER FUNCTIONS diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index b12a94d..11c9261 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -3,7 +3,7 @@ use super::{ NodeIndex, Rpo256, Vec, }; use crate::{ - merkle::{empty_roots::EMPTY_WORD, int_to_leaf}, + merkle::{digests_to_words, empty_roots::EMPTY_WORD, int_to_leaf, int_to_node}, Word, }; @@ -13,17 +13,17 @@ use crate::{ const KEYS4: [u64; 4] = [0, 1, 2, 3]; const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7]; -const VALUES4: [Word; 4] = [int_to_leaf(1), int_to_leaf(2), int_to_leaf(3), int_to_leaf(4)]; - -const VALUES8: [Word; 8] = [ - int_to_leaf(1), - int_to_leaf(2), - int_to_leaf(3), - int_to_leaf(4), - int_to_leaf(5), - int_to_leaf(6), - int_to_leaf(7), - int_to_leaf(8), +const VALUES4: [RpoDigest; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; + +const VALUES8: [RpoDigest; 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), ]; const ZERO_VALUES8: [Word; 8] = [int_to_leaf(0); 8]; @@ -47,7 +47,7 @@ fn build_sparse_tree() { // insert single value let key = 6; let new_node = int_to_leaf(7); - values[key as usize] = new_node; + values[key as usize] = new_node.into(); let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf"); let mt2 = MerkleTree::new(values.clone()).unwrap(); assert_eq!(mt2.root(), smt.root().into()); @@ -73,7 +73,9 @@ fn build_sparse_tree() { #[test] fn test_depth2_tree() { - let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap(); + let tree = + SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter())) + .unwrap(); // check internal structure let (root, node2, node3) = compute_internal_nodes(); @@ -82,10 +84,10 @@ fn test_depth2_tree() { assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap()); // check get_node() - assert_eq!(VALUES4[0], *tree.get_node(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(VALUES4[1], *tree.get_node(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(VALUES4[2], *tree.get_node(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(VALUES4[3], *tree.get_node(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); // check get_path(): depth 2 assert_eq!(vec![VALUES4[1].into(), node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap()); @@ -100,13 +102,15 @@ fn test_depth2_tree() { #[test] fn test_inner_node_iterator() -> Result<(), MerkleError> { - let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap(); + let tree = + SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter())) + .unwrap(); // check depth 2 - assert_eq!(VALUES4[0], *tree.get_node(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(VALUES4[1], *tree.get_node(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(VALUES4[2], *tree.get_node(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(VALUES4[3], *tree.get_node(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); // get parent nodes let root = tree.root(); @@ -120,19 +124,19 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> { let nodes: Vec = tree.inner_nodes().collect(); let expected = vec![ InnerNodeInfo { - value: root.into(), - left: l1n0.into(), - right: l1n1.into(), + value: root, + left: l1n0, + right: l1n1, }, InnerNodeInfo { - value: l1n0.into(), - left: l2n0.into(), - right: l2n1.into(), + value: l1n0, + left: l2n0, + right: l2n1, }, InnerNodeInfo { - value: l1n1.into(), - left: l2n2.into(), - right: l2n3.into(), + value: l1n1, + left: l2n2, + right: l2n3, }, ]; assert_eq!(nodes, expected); @@ -142,18 +146,20 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> { #[test] fn update_leaf() { - let mut tree = SimpleSmt::with_leaves(3, KEYS8.into_iter().zip(VALUES8.into_iter())).unwrap(); + let mut tree = + SimpleSmt::with_leaves(3, KEYS8.into_iter().zip(digests_to_words(&VALUES8).into_iter())) + .unwrap(); // update one value let key = 3; let new_node = int_to_leaf(9); - let mut expected_values = VALUES8.to_vec(); + let mut expected_values = digests_to_words(&VALUES8); expected_values[key] = new_node; let expected_tree = MerkleTree::new(expected_values.clone()).unwrap(); let old_leaf = tree.update_leaf(key as u64, new_node).unwrap(); assert_eq!(expected_tree.root(), tree.root); - assert_eq!(old_leaf, VALUES8[key]); + assert_eq!(old_leaf, *VALUES8[key]); // update another value let key = 6; @@ -163,7 +169,7 @@ fn update_leaf() { let old_leaf = tree.update_leaf(key as u64, new_node).unwrap(); assert_eq!(expected_tree.root(), tree.root); - assert_eq!(old_leaf, VALUES8[key]); + assert_eq!(old_leaf, *VALUES8[key]); } #[test] @@ -252,14 +258,14 @@ fn compute_internal_nodes() -> (RpoDigest, RpoDigest, RpoDigest) { let node2 = Rpo256::hash_elements( &[VALUES4[0], VALUES4[1]] .iter() - .map(|digest| *digest) + .map(|digest| digest.into()) .collect::>() .concat(), ); let node3 = Rpo256::hash_elements( &[VALUES4[2], VALUES4[3]] .iter() - .map(|digest| *digest) + .map(|digest| digest.into()) .collect::>() .concat(), ); diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 1a7ead6..9101f1b 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -4,7 +4,7 @@ use super::{ }; use crate::{ hash::rpo::Rpo256, - merkle::{int_to_leaf, int_to_node, MerklePathSet, MerkleTree, SimpleSmt}, + merkle::{digests_to_words, int_to_leaf, int_to_node, MerklePathSet, MerkleTree, SimpleSmt}, Felt, Word, WORD_SIZE, }; @@ -15,17 +15,17 @@ use std::error::Error; // ================================================================================================ const KEYS4: [u64; 4] = [0, 1, 2, 3]; -const VALUES4: [Word; 4] = [int_to_leaf(1), int_to_leaf(2), int_to_leaf(3), int_to_leaf(4)]; - -const VALUES8: [Word; 8] = [ - int_to_leaf(1), - int_to_leaf(2), - int_to_leaf(3), - int_to_leaf(4), - int_to_leaf(5), - int_to_leaf(6), - int_to_leaf(7), - int_to_leaf(8), +const VALUES4: [RpoDigest; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; + +const VALUES8: [RpoDigest; 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 @@ -33,16 +33,16 @@ const VALUES8: [Word; 8] = [ #[test] fn test_root_not_in_store() -> Result<(), MerkleError> { - let mtree = MerkleTree::new(VALUES4.to_vec())?; + let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; let store = MerkleStore::from(&mtree); assert_eq!( - store.get_node(VALUES4[0].into(), NodeIndex::make(mtree.depth(), 0)), - Err(MerkleError::RootNotInStore(VALUES4[0].into())), + store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)), + Err(MerkleError::RootNotInStore(VALUES4[0])), "Leaf 0 is not a root" ); assert_eq!( - store.get_path(VALUES4[0].into(), NodeIndex::make(mtree.depth(), 0)), - Err(MerkleError::RootNotInStore(VALUES4[0].into())), + store.get_path(VALUES4[0], NodeIndex::make(mtree.depth(), 0)), + Err(MerkleError::RootNotInStore(VALUES4[0])), "Leaf 0 is not a root" ); @@ -51,29 +51,29 @@ fn test_root_not_in_store() -> Result<(), MerkleError> { #[test] fn test_merkle_tree() -> Result<(), MerkleError> { - let mtree = MerkleTree::new(VALUES4.to_vec())?; + let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; let store = MerkleStore::from(&mtree); // STORE LEAVES ARE CORRECT ------------------------------------------------------------------- // checks the leaves in the store corresponds to the expected values assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)), - Ok(VALUES4[0].into()), + Ok(VALUES4[0]), "node 0 must be in the tree" ); assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)), - Ok(VALUES4[1].into()), + Ok(VALUES4[1]), "node 1 must be in the tree" ); assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)), - Ok(VALUES4[2].into()), + Ok(VALUES4[2]), "node 2 must be in the tree" ); assert_eq!( store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)), - Ok(VALUES4[3].into()), + Ok(VALUES4[3]), "node 3 must be in the tree" ); @@ -104,7 +104,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { // assert the merkle path returned by the store is the same as the one in the tree let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap(); assert_eq!( - VALUES4[0], *result.value, + VALUES4[0], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -115,7 +115,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 1)).unwrap(); assert_eq!( - VALUES4[1], *result.value, + VALUES4[1], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -126,7 +126,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 2)).unwrap(); assert_eq!( - VALUES4[2], *result.value, + VALUES4[2], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -137,7 +137,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 3)).unwrap(); assert_eq!( - VALUES4[3], *result.value, + VALUES4[3], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -157,7 +157,7 @@ fn test_empty_roots() { for depth in 0..255 { root = Rpo256::merge(&[root; 2]); assert!( - store.get_node(root.into(), NodeIndex::make(0, 0)).is_ok(), + store.get_node(root, NodeIndex::make(0, 0)).is_ok(), "The root of the empty tree of depth {depth} must be registered" ); } @@ -197,7 +197,8 @@ fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> { #[test] fn test_get_invalid_node() { - let mtree = MerkleTree::new(VALUES4.to_vec()).expect("creating a merkle tree must work"); + let mtree = + MerkleTree::new(digests_to_words(&VALUES4)).expect("creating a merkle tree must work"); let store = MerkleStore::from(&mtree); let _ = store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)); } @@ -211,7 +212,7 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { let idx = NodeIndex::make(1, 0); assert_eq!(smt.get_node(idx).unwrap(), leaves2[0].into()); - assert_eq!(store.get_node(smt.root().into(), idx).unwrap(), smt.get_node(idx).unwrap()); + assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap()); let idx = NodeIndex::make(1, 1); assert_eq!(smt.get_node(idx).unwrap(), leaves2[1].into()); @@ -222,36 +223,38 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { #[test] fn test_sparse_merkle_tree() -> Result<(), MerkleError> { - let smt = - SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, KEYS4.into_iter().zip(VALUES4.into_iter())) - .unwrap(); + let smt = SimpleSmt::with_leaves( + SimpleSmt::MAX_DEPTH, + KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()), + ) + .unwrap(); let store = MerkleStore::from(&smt); // STORE LEAVES ARE CORRECT ============================================================== // checks the leaves in the store corresponds to the expected values assert_eq!( - store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 0)), - Ok(VALUES4[0].into()), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)), + Ok(VALUES4[0]), "node 0 must be in the tree" ); assert_eq!( - store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 1)), - Ok(VALUES4[1].into()), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)), + Ok(VALUES4[1]), "node 1 must be in the tree" ); assert_eq!( - store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 2)), - Ok(VALUES4[2].into()), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)), + Ok(VALUES4[2]), "node 2 must be in the tree" ); assert_eq!( - store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 3)), - Ok(VALUES4[3].into()), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)), + Ok(VALUES4[3]), "node 3 must be in the tree" ); assert_eq!( - store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 4)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)), Ok(RpoDigest::default()), "unmodified node 4 must be ZERO" ); @@ -288,7 +291,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { // assert the merkle path returned by the store is the same as the one in the tree let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 0)).unwrap(); assert_eq!( - VALUES4[0], *result.value, + VALUES4[0], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -299,7 +302,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 1)).unwrap(); assert_eq!( - VALUES4[1], *result.value, + VALUES4[1], result.value, "Value for merkle path at index 1 must match leaf value" ); assert_eq!( @@ -310,7 +313,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 2)).unwrap(); assert_eq!( - VALUES4[2], *result.value, + VALUES4[2], result.value, "Value for merkle path at index 2 must match leaf value" ); assert_eq!( @@ -321,7 +324,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 3)).unwrap(); assert_eq!( - VALUES4[3], *result.value, + VALUES4[3], result.value, "Value for merkle path at index 3 must match leaf value" ); assert_eq!( @@ -347,7 +350,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { #[test] fn test_add_merkle_paths() -> Result<(), MerkleError> { - let mtree = MerkleTree::new(VALUES4.to_vec())?; + let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; let i0 = 0; let p0 = mtree.get_path(NodeIndex::make(2, i0)).unwrap(); @@ -362,10 +365,10 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let p3 = mtree.get_path(NodeIndex::make(2, i3)).unwrap(); let paths = [ - (i0, VALUES4[i0 as usize].into(), p0), - (i1, VALUES4[i1 as usize].into(), p1), - (i2, VALUES4[i2 as usize].into(), p2), - (i3, VALUES4[i3 as usize].into(), p3), + (i0, VALUES4[i0 as usize], p0), + (i1, VALUES4[i1 as usize], p1), + (i2, VALUES4[i2 as usize], p2), + (i3, VALUES4[i3 as usize], p3), ]; let mut store = MerkleStore::default(); @@ -378,22 +381,22 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { // checks the leaves in the store corresponds to the expected values assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 0)), - Ok(VALUES4[0].into()), + Ok(VALUES4[0]), "node 0 must be in the set" ); assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 1)), - Ok(VALUES4[1].into()), + Ok(VALUES4[1]), "node 1 must be in the set" ); assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 2)), - Ok(VALUES4[2].into()), + Ok(VALUES4[2]), "node 2 must be in the set" ); assert_eq!( store.get_node(set.root(), NodeIndex::make(set.depth(), 3)), - Ok(VALUES4[3].into()), + Ok(VALUES4[3]), "node 3 must be in the set" ); @@ -424,7 +427,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { // assert the merkle path returned by the store is the same as the one in the set let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 0)).unwrap(); assert_eq!( - VALUES4[0], *result.value, + VALUES4[0], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -435,7 +438,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 1)).unwrap(); assert_eq!( - VALUES4[1], *result.value, + VALUES4[1], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -446,7 +449,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 2)).unwrap(); assert_eq!( - VALUES4[2], *result.value, + VALUES4[2], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -457,7 +460,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 3)).unwrap(); assert_eq!( - VALUES4[3], *result.value, + VALUES4[3], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( @@ -517,7 +520,7 @@ fn store_path_opens_from_leaf() { let mtree = MerkleTree::new(vec![a, b, c, d, e, f, g, h]).unwrap(); let store = MerkleStore::from(&mtree); - let path = store.get_path(root.into(), NodeIndex::make(3, 1)).unwrap().path; + let path = store.get_path(root, NodeIndex::make(3, 1)).unwrap().path; let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec()); assert_eq!(path, expected); @@ -525,7 +528,7 @@ fn store_path_opens_from_leaf() { #[test] fn test_set_node() -> Result<(), MerkleError> { - let mtree = MerkleTree::new(VALUES4.to_vec())?; + let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; let mut store = MerkleStore::from(&mtree); let value = int_to_node(42); let index = NodeIndex::make(mtree.depth(), 0); @@ -537,7 +540,7 @@ fn test_set_node() -> Result<(), MerkleError> { #[test] fn test_constructors() -> Result<(), MerkleError> { - let mtree = MerkleTree::new(VALUES4.to_vec())?; + let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; let store = MerkleStore::from(&mtree); let depth = mtree.depth(); @@ -549,7 +552,11 @@ fn test_constructors() -> Result<(), MerkleError> { } let depth = 32; - let smt = SimpleSmt::with_leaves(depth, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap(); + let smt = SimpleSmt::with_leaves( + depth, + KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()), + ) + .unwrap(); let store = MerkleStore::from(&smt); let depth = smt.depth(); @@ -561,20 +568,20 @@ fn test_constructors() -> Result<(), MerkleError> { let d = 2; let paths = [ - (0, VALUES4[0].into(), mtree.get_path(NodeIndex::make(d, 0)).unwrap()), - (1, VALUES4[1].into(), mtree.get_path(NodeIndex::make(d, 1)).unwrap()), - (2, VALUES4[2].into(), mtree.get_path(NodeIndex::make(d, 2)).unwrap()), - (3, VALUES4[3].into(), mtree.get_path(NodeIndex::make(d, 3)).unwrap()), + (0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0)).unwrap()), + (1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1)).unwrap()), + (2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2)).unwrap()), + (3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3)).unwrap()), ]; let mut store1 = MerkleStore::default(); store1.add_merkle_paths(paths.clone())?; let mut store2 = MerkleStore::default(); - store2.add_merkle_path(0, VALUES4[0].into(), mtree.get_path(NodeIndex::make(d, 0))?)?; - store2.add_merkle_path(1, VALUES4[1].into(), mtree.get_path(NodeIndex::make(d, 1))?)?; - store2.add_merkle_path(2, VALUES4[2].into(), mtree.get_path(NodeIndex::make(d, 2))?)?; - store2.add_merkle_path(3, VALUES4[3].into(), mtree.get_path(NodeIndex::make(d, 3))?)?; + store2.add_merkle_path(0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0))?)?; + store2.add_merkle_path(1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1))?)?; + store2.add_merkle_path(2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2))?)?; + store2.add_merkle_path(3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?; let set = MerklePathSet::new(d).with_paths(paths).unwrap(); for key in [0, 1, 2, 3] { @@ -744,16 +751,16 @@ fn get_leaf_depth_works_with_depth_8() { #[test] fn mstore_subset() { // add a Merkle tree of depth 3 to the store - let mtree = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let mtree = MerkleTree::new(digests_to_words(&VALUES8)).unwrap(); let mut store = MerkleStore::default(); let empty_store_num_nodes = store.nodes.len(); store.extend(mtree.inner_nodes()); // build 3 subtrees contained within the above Merkle tree; note that subtree2 is a subset // of subtree1 - let subtree1 = MerkleTree::new(VALUES8[..4].to_vec()).unwrap(); - let subtree2 = MerkleTree::new(VALUES8[2..4].to_vec()).unwrap(); - let subtree3 = MerkleTree::new(VALUES8[6..].to_vec()).unwrap(); + let subtree1 = MerkleTree::new(digests_to_words(&VALUES8[..4])).unwrap(); + let subtree2 = MerkleTree::new(digests_to_words(&VALUES8[2..4])).unwrap(); + let subtree3 = MerkleTree::new(digests_to_words(&VALUES8[6..])).unwrap(); // --- extract all 3 subtrees --------------------------------------------- @@ -798,7 +805,7 @@ fn check_mstore_subtree(store: &MerkleStore, subtree: &MerkleTree) { #[cfg(feature = "std")] #[test] fn test_serialization() -> Result<(), Box> { - let mtree = MerkleTree::new(VALUES4.to_vec())?; + let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; let store = MerkleStore::from(&mtree); let decoded = MerkleStore::read_from_bytes(&store.to_bytes()).expect("deserialization failed"); assert_eq!(store, decoded); diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index 28ac42a..b4fbd77 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -17,7 +17,7 @@ fn tsmt_insert_one() { // 16 most significant bits of the key let index = NodeIndex::make(16, raw >> 48); let leaf_node = build_leaf_node(key, value, 16); - let tree_root = store.set_node(smt.root().into(), index, leaf_node.into()).unwrap().root; + let tree_root = store.set_node(smt.root(), index, leaf_node).unwrap().root; smt.insert(key, value); @@ -66,19 +66,11 @@ fn tsmt_insert_two_16() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(32, raw_a >> 32); let leaf_node_a = build_leaf_node(key_a, val_a, 32); - tree_root = store - .set_node(tree_root.into(), index_a, leaf_node_a.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; let index_b = NodeIndex::make(32, raw_b >> 32); let leaf_node_b = build_leaf_node(key_b, val_b, 32); - tree_root = store - .set_node(tree_root.into(), index_b, leaf_node_b.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; // --- verify that data is consistent between store and tree -------------- @@ -130,19 +122,11 @@ fn tsmt_insert_two_32() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(48, raw_a >> 16); let leaf_node_a = build_leaf_node(key_a, val_a, 48); - tree_root = store - .set_node(tree_root.into(), index_a, leaf_node_a.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; let index_b = NodeIndex::make(48, raw_b >> 16); let leaf_node_b = build_leaf_node(key_b, val_b, 48); - tree_root = store - .set_node(tree_root.into(), index_b, leaf_node_b.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; // --- verify that data is consistent between store and tree -------------- @@ -197,27 +181,15 @@ fn tsmt_insert_three() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(32, raw_a >> 32); let leaf_node_a = build_leaf_node(key_a, val_a, 32); - tree_root = store - .set_node(tree_root.into(), index_a, leaf_node_a.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; let index_b = NodeIndex::make(32, raw_b >> 32); let leaf_node_b = build_leaf_node(key_b, val_b, 32); - tree_root = store - .set_node(tree_root.into(), index_b, leaf_node_b.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; let index_c = NodeIndex::make(32, raw_c >> 32); let leaf_node_c = build_leaf_node(key_c, val_c, 32); - tree_root = store - .set_node(tree_root.into(), index_c, leaf_node_c.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_c, leaf_node_c).unwrap().root; // --- verify that data is consistent between store and tree -------------- @@ -264,7 +236,7 @@ fn tsmt_update() { let mut tree_root = get_init_root(); let index = NodeIndex::make(16, raw >> 48); let leaf_node = build_leaf_node(key, value_b, 16); - tree_root = store.set_node(tree_root.into(), index, leaf_node.into()).unwrap().root.into(); + tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root; assert_eq!(smt.root(), tree_root.into()); @@ -309,7 +281,7 @@ fn tsmt_bottom_tier() { // key_b is smaller than key_a. let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]); let mut tree_root = get_init_root(); - tree_root = store.set_node(tree_root.into(), index, leaf_node.into()).unwrap().root.into(); + tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root; // --- verify that data is consistent between store and tree -------------- @@ -357,19 +329,11 @@ fn tsmt_bottom_tier_two() { let mut tree_root = get_init_root(); let index_a = NodeIndex::make(64, raw_a); let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]); - tree_root = store - .set_node(tree_root.into(), index_a, leaf_node_a.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; let index_b = NodeIndex::make(64, raw_b); let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]); - tree_root = store - .set_node(tree_root.into(), index_b, leaf_node_b.into()) - .unwrap() - .root - .into(); + tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; // --- verify that data is consistent between store and tree -------------- @@ -442,8 +406,8 @@ fn tsmt_node_not_available() { // HELPER FUNCTIONS // ================================================================================================ -fn get_init_root() -> Word { - EmptySubtreeRoots::empty_hashes(64)[0].into() +fn get_init_root() -> RpoDigest { + EmptySubtreeRoots::empty_hashes(64)[0] } fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest { From b9def61e28c52ec7d956d5a276d2fa4590431dc5 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Tue, 13 Jun 2023 16:14:07 +0300 Subject: [PATCH 12/19] 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()); +} From 049ae32cbff25c16905ed7a43150b46598789af0 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Tue, 13 Jun 2023 14:40:31 -0700 Subject: [PATCH 13/19] chore: clean up test code --- src/merkle/mmr/tests.rs | 140 ++++++--------------------------- src/merkle/mod.rs | 3 +- src/merkle/simple_smt/tests.rs | 48 +++++------ src/merkle/store/tests.rs | 9 +-- src/merkle/tiered_smt/tests.rs | 36 ++++----- 5 files changed, 66 insertions(+), 170 deletions(-) diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index 1973aa5..cb367fc 100644 --- a/src/merkle/mmr/tests.rs +++ b/src/merkle/mmr/tests.rs @@ -1,5 +1,5 @@ use super::{ - super::{InnerNodeInfo, Vec, WORD_SIZE}, + super::{InnerNodeInfo, Vec}, bit::TrueBitPositionIterator, full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest}, Mmr, MmrPeaks, Rpo256, @@ -118,38 +118,14 @@ fn test_mmr_simple() { let mut postorder = Vec::new(); postorder.push(LEAVES[0]); postorder.push(LEAVES[1]); - postorder.push(Rpo256::hash_elements( - &[LEAVES[0], LEAVES[1]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - )); + postorder.push(Rpo256::merge(&[LEAVES[0], LEAVES[1]])); postorder.push(LEAVES[2]); postorder.push(LEAVES[3]); - postorder.push(Rpo256::hash_elements( - &[LEAVES[2], LEAVES[3]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - )); - postorder.push(Rpo256::hash_elements( - &[postorder[2], postorder[5]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - )); + postorder.push(Rpo256::merge(&[LEAVES[2], LEAVES[3]])); + postorder.push(Rpo256::merge(&[postorder[2], postorder[5]])); postorder.push(LEAVES[4]); postorder.push(LEAVES[5]); - postorder.push(Rpo256::hash_elements( - &[LEAVES[4], LEAVES[5]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - )); + postorder.push(Rpo256::merge(&[LEAVES[4], LEAVES[5]])); postorder.push(LEAVES[6]); let mut mmr = Mmr::new(); @@ -223,20 +199,8 @@ fn test_mmr_simple() { #[test] fn test_mmr_open() { let mmr: Mmr = LEAVES.into(); - let h01: RpoDigest = Rpo256::hash_elements( - &LEAVES[0..2] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); - let h23: RpoDigest = Rpo256::hash_elements( - &LEAVES[2..4] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); + let h01 = Rpo256::merge(&[LEAVES[0], LEAVES[1]]); + let h23 = Rpo256::merge(&[LEAVES[2], LEAVES[3]]); // node at pos 7 is the root assert!(mmr.open(7).is_err(), "Element 7 is not in the tree, result should be None"); @@ -401,34 +365,10 @@ fn test_mmr_inner_nodes() { let mmr: Mmr = LEAVES.into(); let nodes: Vec = mmr.inner_nodes().collect(); - let h01 = Rpo256::hash_elements( - &[LEAVES[0], LEAVES[1]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); - let h23 = Rpo256::hash_elements( - &[LEAVES[2], LEAVES[3]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); - let h0123 = Rpo256::hash_elements( - &[h01, h23] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); - let h45 = Rpo256::hash_elements( - &[LEAVES[4], LEAVES[5]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); + let h01 = Rpo256::merge(&[LEAVES[0], LEAVES[1]]); + let h23 = Rpo256::merge(&[LEAVES[2], LEAVES[3]]); + let h0123 = Rpo256::merge(&[h01, h23]); + let h45 = Rpo256::merge(&[LEAVES[4], LEAVES[5]]); let postorder = vec![ InnerNodeInfo { value: h01, @@ -461,28 +401,10 @@ fn test_mmr_hash_peaks() { let peaks = mmr.accumulator(); let first_peak = Rpo256::merge(&[ - Rpo256::hash_elements( - &[LEAVES[0], LEAVES[1]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ), - Rpo256::hash_elements( - &[LEAVES[2], LEAVES[3]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ), + Rpo256::merge(&[LEAVES[0], LEAVES[1]]), + Rpo256::merge(&[LEAVES[2], LEAVES[3]]), ]); - let second_peak = Rpo256::hash_elements( - &[LEAVES[4], LEAVES[5]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); + let second_peak = Rpo256::merge(&[LEAVES[4], LEAVES[5]]); let third_peak = LEAVES[6]; // minimum length is 16 @@ -490,14 +412,7 @@ fn test_mmr_hash_peaks() { expected_peaks.resize(16, RpoDigest::default()); assert_eq!( peaks.hash_peaks(), - *Rpo256::hash_elements( - &expected_peaks - .as_slice() - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat() - ) + *Rpo256::hash_elements(&digests_to_elements(&expected_peaks)) ); } @@ -517,14 +432,7 @@ fn test_mmr_peaks_hash_less_than_16() { expected_peaks.resize(16, RpoDigest::default()); assert_eq!( accumulator.hash_peaks(), - *Rpo256::hash_elements( - &expected_peaks - .as_slice() - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat() - ) + *Rpo256::hash_elements(&digests_to_elements(&expected_peaks)) ); } } @@ -543,14 +451,7 @@ fn test_mmr_peaks_hash_odd() { expected_peaks.resize(18, RpoDigest::default()); assert_eq!( accumulator.hash_peaks(), - *Rpo256::hash_elements( - &expected_peaks - .as_slice() - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat() - ) + *Rpo256::hash_elements(&digests_to_elements(&expected_peaks)) ); } @@ -582,3 +483,10 @@ mod property_tests { } } } + +// HELPER FUNCTIONS +// ================================================================================================ + +fn digests_to_elements(digests: &[RpoDigest]) -> Vec { + digests.iter().flat_map(|v| Word::from(v)).collect() +} diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 61dad43..8601b60 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -103,6 +103,7 @@ const fn int_to_leaf(value: u64) -> Word { [Felt::new(value), ZERO, ZERO, ZERO] } -pub fn digests_to_words(digests: &[RpoDigest]) -> Vec { +#[cfg(test)] +fn digests_to_words(digests: &[RpoDigest]) -> Vec { digests.iter().map(|d| d.into()).collect() } diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 11c9261..ce6b8f3 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -36,7 +36,7 @@ fn build_empty_tree() { // tree of depth 3 let smt = SimpleSmt::new(3).unwrap(); let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap(); - assert_eq!(mt.root(), smt.root().into()); + assert_eq!(mt.root(), smt.root()); } #[test] @@ -47,10 +47,10 @@ fn build_sparse_tree() { // insert single value let key = 6; let new_node = int_to_leaf(7); - values[key as usize] = new_node.into(); + values[key as usize] = new_node; let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf"); let mt2 = MerkleTree::new(values.clone()).unwrap(); - assert_eq!(mt2.root(), smt.root().into()); + assert_eq!(mt2.root(), smt.root()); assert_eq!( mt2.get_path(NodeIndex::make(3, 6)).unwrap(), smt.get_path(NodeIndex::make(3, 6)).unwrap() @@ -63,7 +63,7 @@ fn build_sparse_tree() { values[key as usize] = new_node; let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf"); let mt3 = MerkleTree::new(values).unwrap(); - assert_eq!(mt3.root(), smt.root().into()); + assert_eq!(mt3.root(), smt.root()); assert_eq!( mt3.get_path(NodeIndex::make(3, 2)).unwrap(), smt.get_path(NodeIndex::make(3, 2)).unwrap() @@ -90,10 +90,10 @@ fn test_depth2_tree() { assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); // check get_path(): depth 2 - assert_eq!(vec![VALUES4[1].into(), node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap()); - assert_eq!(vec![VALUES4[0].into(), node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap()); - assert_eq!(vec![VALUES4[3].into(), node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap()); - assert_eq!(vec![VALUES4[2].into(), node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap()); + assert_eq!(vec![VALUES4[1], node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(vec![VALUES4[0], node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(vec![VALUES4[3], node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(vec![VALUES4[2], node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap()); // check get_path(): depth 1 assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap()); @@ -189,15 +189,15 @@ fn small_tree_opening_is_consistent() { let c = Word::from(Rpo256::merge(&[b.into(); 2])); let d = Word::from(Rpo256::merge(&[c.into(); 2])); - let e = RpoDigest::from(Rpo256::merge(&[a.into(), b.into()])); - let f = RpoDigest::from(Rpo256::merge(&[z.into(), z.into()])); - let g = RpoDigest::from(Rpo256::merge(&[c.into(), z.into()])); - let h = RpoDigest::from(Rpo256::merge(&[z.into(), d.into()])); + let e = Rpo256::merge(&[a.into(), b.into()]); + let f = Rpo256::merge(&[z.into(), z.into()]); + let g = Rpo256::merge(&[c.into(), z.into()]); + let h = Rpo256::merge(&[z.into(), d.into()]); - let i = RpoDigest::from(Rpo256::merge(&[e.into(), f.into()])); - let j = RpoDigest::from(Rpo256::merge(&[g.into(), h.into()])); + let i = Rpo256::merge(&[e, f]); + let j = Rpo256::merge(&[g, h]); - let k = RpoDigest::from(Rpo256::merge(&[i.into(), j.into()])); + let k = Rpo256::merge(&[i, j]); let depth = 3; let entries = vec![(0, a), (1, b), (4, c), (7, d)]; @@ -255,21 +255,9 @@ fn with_no_duplicates_empty_node() { // -------------------------------------------------------------------------------------------- fn compute_internal_nodes() -> (RpoDigest, RpoDigest, RpoDigest) { - let node2 = Rpo256::hash_elements( - &[VALUES4[0], VALUES4[1]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); - let node3 = Rpo256::hash_elements( - &[VALUES4[2], VALUES4[3]] - .iter() - .map(|digest| digest.into()) - .collect::>() - .concat(), - ); + let node2 = Rpo256::merge(&[VALUES4[0], VALUES4[1]]); + let node3 = Rpo256::merge(&[VALUES4[2], VALUES4[3]]); let root = Rpo256::merge(&[node2, node3]); - (root.into(), node2.into(), node3.into()) + (root, node2, node3) } diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 9101f1b..1bc5db9 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -485,7 +485,6 @@ fn wont_open_to_different_depth_root() { for depth in (1..=63).rev() { root = Rpo256::merge(&[root, empty[depth]]); } - let root = RpoDigest::from(root); // For this example, the depth of the Merkle tree is 1, as we have only two leaves. Here we // attempt to fetch a node on the maximum depth, and it should fail because the root shouldn't @@ -513,16 +512,16 @@ fn store_path_opens_from_leaf() { let k = Rpo256::merge(&[e.into(), f.into()]); let l = Rpo256::merge(&[g.into(), h.into()]); - let m = Rpo256::merge(&[i.into(), j.into()]); - let n = Rpo256::merge(&[k.into(), l.into()]); + let m = Rpo256::merge(&[i, j]); + let n = Rpo256::merge(&[k, l]); - let root = Rpo256::merge(&[m.into(), n.into()]); + let root = Rpo256::merge(&[m, n]); let mtree = MerkleTree::new(vec![a, b, c, d, e, f, g, h]).unwrap(); let store = MerkleStore::from(&mtree); let path = store.get_path(root, NodeIndex::make(3, 1)).unwrap().path; - let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec()); + let expected = MerklePath::new([a.into(), j, n].to_vec()); assert_eq!(path, expected); } diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index b4fbd77..894c628 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -21,7 +21,7 @@ fn tsmt_insert_one() { smt.insert(key, value); - assert_eq!(smt.root(), tree_root.into()); + assert_eq!(smt.root(), tree_root); // make sure the value was inserted, and the node is at the expected index assert_eq!(smt.get_value(key), value); @@ -74,16 +74,16 @@ fn tsmt_insert_two_16() { // --- verify that data is consistent between store and tree -------------- - assert_eq!(smt.root(), tree_root.into()); + assert_eq!(smt.root(), tree_root); assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; + let expected_path = store.get_path(tree_root, index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; + let expected_path = store.get_path(tree_root, index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -130,16 +130,16 @@ fn tsmt_insert_two_32() { // --- verify that data is consistent between store and tree -------------- - assert_eq!(smt.root(), tree_root.into()); + assert_eq!(smt.root(), tree_root); assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; + let expected_path = store.get_path(tree_root, index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; + let expected_path = store.get_path(tree_root, index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -193,21 +193,21 @@ fn tsmt_insert_three() { // --- verify that data is consistent between store and tree -------------- - assert_eq!(smt.root(), tree_root.into()); + assert_eq!(smt.root(), tree_root); assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; + let expected_path = store.get_path(tree_root, index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; + let expected_path = store.get_path(tree_root, index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); assert_eq!(smt.get_value(key_c), val_c); assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c); - let expected_path = store.get_path(tree_root.into(), index_c).unwrap().path; + let expected_path = store.get_path(tree_root, index_c).unwrap().path; assert_eq!(smt.get_path(index_c).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -238,11 +238,11 @@ fn tsmt_update() { let leaf_node = build_leaf_node(key, value_b, 16); tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root; - assert_eq!(smt.root(), tree_root.into()); + assert_eq!(smt.root(), tree_root); assert_eq!(smt.get_value(key), value_b); assert_eq!(smt.get_node(index).unwrap(), leaf_node); - let expected_path = store.get_path(tree_root.into(), index).unwrap().path; + let expected_path = store.get_path(tree_root, index).unwrap().path; assert_eq!(smt.get_path(index).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -285,13 +285,13 @@ fn tsmt_bottom_tier() { // --- verify that data is consistent between store and tree -------------- - assert_eq!(smt.root(), tree_root.into()); + assert_eq!(smt.root(), tree_root); assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index).unwrap(), leaf_node); - let expected_path = store.get_path(tree_root.into(), index).unwrap().path; + let expected_path = store.get_path(tree_root, index).unwrap().path; assert_eq!(smt.get_path(index).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of @@ -337,16 +337,16 @@ fn tsmt_bottom_tier_two() { // --- verify that data is consistent between store and tree -------------- - assert_eq!(smt.root(), tree_root.into()); + assert_eq!(smt.root(), tree_root); assert_eq!(smt.get_value(key_a), val_a); assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path; + let expected_path = store.get_path(tree_root, index_a).unwrap().path; assert_eq!(smt.get_path(index_a).unwrap(), expected_path); assert_eq!(smt.get_value(key_b), val_b); assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path; + let expected_path = store.get_path(tree_root, index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); // make sure inner nodes match - the store contains more entries because it keeps track of From fe5cac9edc06afe9f8872ebf551b80c0958ab4d3 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Tue, 13 Jun 2023 22:43:08 -0700 Subject: [PATCH 14/19] fix: compilation errors --- src/hash/rpo/tests.rs | 6 +- src/merkle/index.rs | 2 +- src/merkle/merkle_tree.rs | 8 +-- src/merkle/mmr/tests.rs | 8 +-- src/merkle/partial_mt/mod.rs | 21 ++++--- src/merkle/partial_mt/tests.rs | 105 ++++++++++++++++----------------- src/merkle/path_set.rs | 16 ++--- src/merkle/simple_smt/tests.rs | 2 +- src/merkle/tiered_smt/tests.rs | 4 +- 9 files changed, 84 insertions(+), 88 deletions(-) diff --git a/src/hash/rpo/tests.rs b/src/hash/rpo/tests.rs index 8b2a258..d0f6889 100644 --- a/src/hash/rpo/tests.rs +++ b/src/hash/rpo/tests.rs @@ -206,7 +206,7 @@ fn sponge_bytes_with_remainder_length_wont_panic() { // size. // // this is a preliminary test to the fuzzy-stress of proptest. - Rpo256::hash(&vec![0; 113]); + Rpo256::hash(&[0; 113]); } #[test] @@ -230,8 +230,8 @@ fn sponge_zeroes_collision() { proptest! { #[test] - fn rpo256_wont_panic_with_arbitrary_input(ref vec in any::>()) { - Rpo256::hash(&vec); + fn rpo256_wont_panic_with_arbitrary_input(ref bytes in any::>()) { + Rpo256::hash(bytes); } } diff --git a/src/merkle/index.rs b/src/merkle/index.rs index 564d981..f17216f 100644 --- a/src/merkle/index.rs +++ b/src/merkle/index.rs @@ -190,7 +190,7 @@ mod tests { if value > (1 << depth) { // round up depth += 1; } - NodeIndex::new(depth, value.into()).unwrap() + NodeIndex::new(depth, value).unwrap() } } diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 4e8d154..cfb61bc 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -287,7 +287,7 @@ mod tests { // leaves were copied correctly for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) { - assert_eq!(*a, RpoDigest::from(*b)); + assert_eq!(a, b); } let (root, node2, node3) = compute_internal_nodes(); @@ -341,7 +341,7 @@ mod tests { let value = 3; let new_node = int_to_leaf(9); let mut expected_leaves = digests_to_words(&LEAVES8); - expected_leaves[value as usize] = new_node.into(); + expected_leaves[value as usize] = new_node; let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap(); tree.update_leaf(value, new_node).unwrap(); @@ -408,8 +408,8 @@ mod tests { let digest = RpoDigest::from(word); // assert the addresses are different - let word_ptr = (&word).as_ptr() as *const u8; - let digest_ptr = (&digest).as_ptr() as *const u8; + let word_ptr = word.as_ptr() as *const u8; + let digest_ptr = digest.as_ptr() as *const u8; assert_ne!(word_ptr, digest_ptr); // compare the bytes representation diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index cb367fc..13722e2 100644 --- a/src/merkle/mmr/tests.rs +++ b/src/merkle/mmr/tests.rs @@ -439,15 +439,15 @@ fn test_mmr_peaks_hash_less_than_16() { #[test] fn test_mmr_peaks_hash_odd() { - let peaks: Vec<_> = (0..=17).map(|i| int_to_node(i)).collect(); + let peaks: Vec<_> = (0..=17).map(int_to_node).collect(); let accumulator = MmrPeaks { num_leaves: (1 << peaks.len()) - 1, peaks: peaks.clone(), }; - // odd length bigger than 16 is padded to the next even nubmer - let mut expected_peaks = peaks.clone(); + // odd length bigger than 16 is padded to the next even number + let mut expected_peaks = peaks; expected_peaks.resize(18, RpoDigest::default()); assert_eq!( accumulator.hash_peaks(), @@ -488,5 +488,5 @@ mod property_tests { // ================================================================================================ fn digests_to_elements(digests: &[RpoDigest]) -> Vec { - digests.iter().flat_map(|v| Word::from(v)).collect() + digests.iter().flat_map(Word::from).collect() } diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index 3807de3..8800662 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -1,6 +1,5 @@ use super::{ - BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec, - Word, EMPTY_WORD, + BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec, ZERO, }; use crate::utils::{format, string::String, word_to_hex}; use core::fmt; @@ -15,7 +14,7 @@ mod tests; const ROOT_INDEX: NodeIndex = NodeIndex::root(); /// An RpoDigest consisting of 4 ZERO elements. -const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD); +const EMPTY_DIGEST: RpoDigest = RpoDigest::new([ZERO; 4]); // PARTIAL MERKLE TREE // ================================================================================================ @@ -50,7 +49,7 @@ impl PartialMerkleTree { // CONSTRUCTORS // -------------------------------------------------------------------------------------------- - /// Returns a new emply [PartialMerkleTree]. + /// Returns a new empty [PartialMerkleTree]. pub fn new() -> Self { PartialMerkleTree { max_depth: 0, @@ -108,7 +107,7 @@ impl PartialMerkleTree { paths.push(( leaf, ValuePath { - value: *self.get_node(leaf).expect("Failed to get leaf node"), + value: self.get_node(leaf).expect("Failed to get leaf node"), path: self.get_path(leaf).expect("Failed to get path"), }, )); @@ -142,7 +141,7 @@ impl PartialMerkleTree { index.move_up(); let sibling = self.nodes.get(&sibling_index).cloned().expect("Sibling node not in the map"); - path.push(Word::from(sibling)); + path.push(sibling); } Ok(MerklePath::new(path)) } @@ -189,11 +188,11 @@ impl PartialMerkleTree { // 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()); + self.nodes.insert(sibling_node_index, path[0]); // traverse to the root, updating the nodes let mut index_value = index_value; - let node = Rpo256::merge(&index_value.build_node(value, path[0].into())); + let node = Rpo256::merge(&index_value.build_node(value, path[0])); let root = path.iter().skip(1).copied().fold(node, |node, hash| { index_value.move_up(); // insert calculated node to the nodes map @@ -215,11 +214,11 @@ impl PartialMerkleTree { // - 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() { + if self.nodes.insert(sibling_node, hash).is_none() { self.leaves.insert(sibling_node); } - Rpo256::merge(&index_value.build_node(node, hash.into())) + Rpo256::merge(&index_value.build_node(node, hash)) }); // if the path set is empty (the root is all ZEROs), set the root to the root of the added @@ -227,7 +226,7 @@ impl PartialMerkleTree { 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())); + return Err(MerkleError::ConflictingRoots([self.root(), root].to_vec())); } Ok(()) diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index ba62896..91638da 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}, - ValuePath, Vec, Word, + super::{digests_to_words, int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree}, + RpoDigest, ValuePath, Vec, }; // TEST DATA @@ -18,7 +18,7 @@ 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] = [ +const VALUES8: [RpoDigest; 8] = [ int_to_node(30), int_to_node(31), int_to_node(32), @@ -44,22 +44,21 @@ const VALUES8: [Word; 8] = [ // (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 +// NodeIndex(3, 5) will be labeled 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(); + let mt = MerkleTree::new(digests_to_words(&VALUES8)).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(); + let pmt = PartialMerkleTree::with_paths([(3, path33.value, path33.path)]).unwrap(); - assert_eq!(pmt.root(), expected_root.into()); + assert_eq!(pmt.root(), expected_root); } /// This test checks correctness of the `add_path()` and `get_path()` functions. First it creates a @@ -67,7 +66,7 @@ fn get_root() { /// 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(); + let mt = MerkleTree::new(digests_to_words(&VALUES8)).unwrap(); let expected_root = mt.root(); let ms = MerkleStore::from(&mt); @@ -76,10 +75,8 @@ fn add_and_get_paths() { let expected_path22 = ms.get_path(expected_root, NODE22).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(); + pmt.add_path(3, expected_path33.value, expected_path33.path.clone()).unwrap(); + pmt.add_path(2, expected_path22.value, expected_path22.path.clone()).unwrap(); let path33 = pmt.get_path(NODE33).unwrap(); let path22 = pmt.get_path(NODE22).unwrap(); @@ -87,58 +84,58 @@ fn add_and_get_paths() { assert_eq!(expected_path33.path, path33); assert_eq!(expected_path22.path, path22); - assert_eq!(expected_root, *actual_root); + 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(); + let mt = MerkleTree::new(digests_to_words(&VALUES8)).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(); + let pmt = PartialMerkleTree::with_paths([(3, path33.value, 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()); + 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()); } /// 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 mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let mt = MerkleTree::new(digests_to_words(&VALUES8)).unwrap(); let root = mt.root(); 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 mut pmt = PartialMerkleTree::with_paths([(3, path33.value, path33.path)]).unwrap(); let new_value32 = int_to_node(132); let expected_root = ms.set_node(root, NODE32, new_value32).unwrap().root; - pmt.update_leaf(NODE32, new_value32.into()).unwrap(); + pmt.update_leaf(NODE32, new_value32).unwrap(); let actual_root = pmt.root(); - assert_eq!(expected_root, *actual_root); + assert_eq!(expected_root, actual_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.into()).unwrap(); + pmt.update_leaf(NODE20, new_value20).unwrap(); let actual_root = pmt.root(); - assert_eq!(expected_root, *actual_root); + assert_eq!(expected_root, actual_root); } /// 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(); + let mt = MerkleTree::new(digests_to_words(&VALUES8)).unwrap(); let expected_root = mt.root(); let ms = MerkleStore::from(&mt); @@ -147,8 +144,8 @@ fn get_paths() { 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(); + pmt.add_path(3, path33.value, path33.path).unwrap(); + pmt.add_path(2, path22.value, path22.path).unwrap(); // After PMT creation with path33 (33; 32, 20, 11) and path22 (22; 23, 10) we will have this // tree: // @@ -170,7 +167,7 @@ fn get_paths() { ( leaf, ValuePath { - value: mt.get_node(leaf).unwrap().into(), + value: mt.get_node(leaf).unwrap(), path: mt.get_path(leaf).unwrap(), }, ) @@ -185,7 +182,7 @@ fn get_paths() { // Checks correctness of leaves determination when using the `leaves()` function. #[test] fn leaves() { - let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let mt = MerkleTree::new(digests_to_words(&VALUES8)).unwrap(); let expected_root = mt.root(); let ms = MerkleStore::from(&mt); @@ -193,7 +190,7 @@ fn leaves() { 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 mut pmt = PartialMerkleTree::with_paths([(3, path33.value, path33.path)]).unwrap(); // After PMT creation with path33 (33; 32, 20, 11) we will have this tree: // // ______root______ @@ -206,17 +203,17 @@ fn leaves() { // // 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(); - let value32 = mt.get_node(NODE32).unwrap().into(); - let value33 = mt.get_node(NODE33).unwrap().into(); + let value11 = mt.get_node(NODE11).unwrap(); + let value20 = mt.get_node(NODE20).unwrap(); + let value32 = mt.get_node(NODE32).unwrap(); + let value33 = mt.get_node(NODE33).unwrap(); let leaves = vec![(NODE11, value11), (NODE20, value20), (NODE32, value32), (NODE33, value33)]; - let expected_leaves = leaves.iter().map(|&tuple| tuple); + let expected_leaves = leaves.iter().copied(); assert!(expected_leaves.eq(pmt.leaves())); - pmt.add_path(2, path22.value.into(), path22.path).unwrap(); + pmt.add_path(2, path22.value, path22.path).unwrap(); // After adding the path22 (22; 23, 10) to the existing PMT we will have this tree: // // ______root______ @@ -229,11 +226,11 @@ fn leaves() { // // 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(); - 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 value20 = mt.get_node(NODE20).unwrap(); + let value22 = mt.get_node(NODE22).unwrap(); + let value23 = mt.get_node(NODE23).unwrap(); + let value32 = mt.get_node(NODE32).unwrap(); + let value33 = mt.get_node(NODE33).unwrap(); let leaves = vec![ (NODE20, value20), @@ -243,7 +240,7 @@ fn leaves() { (NODE33, value33), ]; - let expected_leaves = leaves.iter().map(|&tuple| tuple); + let expected_leaves = leaves.iter().copied(); assert!(expected_leaves.eq(pmt.leaves())); } @@ -254,22 +251,22 @@ fn err_add_path() { 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(); + pmt.add_path(3, int_to_node(6), path33).unwrap(); - assert!(pmt.add_path(2, int_to_node(7).into(), path22).is_err()); + assert!(pmt.add_path(2, int_to_node(7), 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 mt = MerkleTree::new(digests_to_words(&VALUES8)).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(); + let pmt = PartialMerkleTree::with_paths([(3, path33.value, path33.path)]).unwrap(); assert!(pmt.get_node(NODE22).is_err()); assert!(pmt.get_node(NODE23).is_err()); @@ -280,14 +277,14 @@ fn err_get_node() { /// 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 mt = MerkleTree::new(digests_to_words(&VALUES8)).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(); + let pmt = PartialMerkleTree::with_paths([(3, path33.value, path33.path)]).unwrap(); assert!(pmt.get_path(NODE22).is_err()); assert!(pmt.get_path(NODE23).is_err()); @@ -297,17 +294,17 @@ fn err_get_path() { #[test] fn err_update_leaf() { - let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let mt = MerkleTree::new(digests_to_words(&VALUES8)).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(); + let mut pmt = PartialMerkleTree::with_paths([(3, path33.value, 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()); + assert!(pmt.update_leaf(NODE22, int_to_node(22)).is_err()); + assert!(pmt.update_leaf(NODE23, int_to_node(23)).is_err()); + assert!(pmt.update_leaf(NODE30, int_to_node(30)).is_err()); + assert!(pmt.update_leaf(NODE31, int_to_node(31)).is_err()); } diff --git a/src/merkle/path_set.rs b/src/merkle/path_set.rs index ea77b92..169e073 100644 --- a/src/merkle/path_set.rs +++ b/src/merkle/path_set.rs @@ -354,35 +354,35 @@ mod tests { let m = Rpo256::merge(&[i, j]); let n = Rpo256::merge(&[k, l]); - let root = Rpo256::merge(&[m.into(), n.into()]); + let root = Rpo256::merge(&[m, n]); let mut set = MerklePathSet::new(3); let value = b; let index = 1; let path = MerklePath::new([a, j, n].to_vec()); - set.add_path(index, value.into(), path.clone()).unwrap(); + set.add_path(index, value.into(), path).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); assert_eq!(root, set.root()); let value = e; let index = 4; - let path = MerklePath::new([f.into(), l.into(), m.into()].to_vec()); - set.add_path(index, value.into(), path.clone()).unwrap(); + let path = MerklePath::new([f, l, m].to_vec()); + set.add_path(index, value.into(), path).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); assert_eq!(root, set.root()); let value = a; let index = 0; - let path = MerklePath::new([b.into(), j.into(), n.into()].to_vec()); - set.add_path(index, value.into(), path.clone()).unwrap(); + let path = MerklePath::new([b, j, n].to_vec()); + set.add_path(index, value.into(), path).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); assert_eq!(root, set.root()); let value = h; let index = 7; - let path = MerklePath::new([g.into(), k.into(), m.into()].to_vec()); - set.add_path(index, value.into(), path.clone()).unwrap(); + let path = MerklePath::new([g, k, m].to_vec()); + set.add_path(index, value.into(), path).unwrap(); assert_eq!(*value, set.get_leaf(index).unwrap()); assert_eq!(root, set.root()); } diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index ce6b8f3..d86f98b 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -203,7 +203,7 @@ fn small_tree_opening_is_consistent() { let entries = vec![(0, a), (1, b), (4, c), (7, d)]; let tree = SimpleSmt::with_leaves(depth, entries).unwrap(); - assert_eq!(tree.root(), RpoDigest::from(k)); + assert_eq!(tree.root(), k); let cases: Vec<(u8, u64, Vec)> = vec![ (3, 0, vec![b.into(), f, j]), diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index 894c628..fbf26ea 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -432,10 +432,10 @@ fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { fn get_non_empty_nodes(store: &MerkleStore) -> Vec { store .inner_nodes() - .filter(|node| !is_empty_subtree(&RpoDigest::from(node.value))) + .filter(|node| !is_empty_subtree(&node.value)) .collect::>() } fn is_empty_subtree(node: &RpoDigest) -> bool { - EmptySubtreeRoots::empty_hashes(255).contains(&node) + EmptySubtreeRoots::empty_hashes(255).contains(node) } From 4215e83ae5f995011261cbb1fd140ea861f05956 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Tue, 13 Jun 2023 22:53:14 -0700 Subject: [PATCH 15/19] feat: add EMPTY_VALUE const to SMTs --- src/merkle/simple_smt/mod.rs | 13 ++++++++----- src/merkle/tiered_smt/mod.rs | 15 +++++++++------ 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 2540ae4..c8da302 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -1,6 +1,6 @@ use super::{ - empty_roots::EMPTY_WORD, BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, - MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word, + BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, + Rpo256, RpoDigest, Vec, Word, }; #[cfg(test)] @@ -31,6 +31,9 @@ impl SimpleSmt { /// Maximum supported depth. pub const MAX_DEPTH: u8 = 64; + /// Value of an empty leaf. + pub const EMPTY_VALUE: Word = super::empty_roots::EMPTY_WORD; + // CONSTRUCTORS // -------------------------------------------------------------------------------------------- @@ -91,12 +94,12 @@ impl SimpleSmt { let mut empty_entries = BTreeSet::new(); for (key, value) in entries { let old_value = tree.update_leaf(key, value)?; - if old_value != EMPTY_WORD || empty_entries.contains(&key) { + if old_value != Self::EMPTY_VALUE || empty_entries.contains(&key) { return Err(MerkleError::DuplicateValuesForIndex(key)); } // if we've processed an empty entry, add the key to the set of empty entry keys, and // if this key was already in the set, return an error - if value == EMPTY_WORD && !empty_entries.insert(key) { + if value == Self::EMPTY_VALUE && !empty_entries.insert(key) { return Err(MerkleError::DuplicateValuesForIndex(key)); } } @@ -210,7 +213,7 @@ impl SimpleSmt { /// # Errors /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}. pub fn update_leaf(&mut self, index: u64, value: Word) -> Result { - let old_value = self.insert_leaf_node(index, value).unwrap_or(EMPTY_WORD); + let old_value = self.insert_leaf_node(index, value).unwrap_or(Self::EMPTY_VALUE); // if the old value and new value are the same, there is nothing to update if value == old_value { diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index 3fb3cdd..3269517 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/src/merkle/tiered_smt/mod.rs @@ -1,6 +1,6 @@ use super::{ - empty_roots::EMPTY_WORD, BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, - MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, StarkField, Vec, Word, ZERO, + BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, + Rpo256, RpoDigest, StarkField, Vec, Word, ZERO, }; use core::cmp; @@ -54,6 +54,9 @@ impl TieredSmt { /// Maximum node depth. This is also the bottom tier of the tree. const MAX_DEPTH: u8 = 64; + /// Value of an empty leaf. + pub const EMPTY_VALUE: Word = super::empty_roots::EMPTY_WORD; + // CONSTRUCTORS // -------------------------------------------------------------------------------------------- @@ -74,12 +77,12 @@ impl TieredSmt { let mut empty_entries = BTreeSet::new(); for (key, value) in entries { let old_value = tree.insert(key, value); - if old_value != EMPTY_WORD || empty_entries.contains(&key) { + if old_value != Self::EMPTY_VALUE || empty_entries.contains(&key) { return Err(MerkleError::DuplicateValuesForKey(key)); } // if we've processed an empty entry, add the key to the set of empty entry keys, and // if this key was already in the set, return an error - if value == EMPTY_WORD && !empty_entries.insert(key) { + if value == Self::EMPTY_VALUE && !empty_entries.insert(key) { return Err(MerkleError::DuplicateValuesForKey(key)); } } @@ -136,7 +139,7 @@ impl TieredSmt { pub fn get_value(&self, key: RpoDigest) -> Word { match self.values.get(&key) { Some(value) => *value, - None => EMPTY_WORD, + None => Self::EMPTY_VALUE, } } @@ -149,7 +152,7 @@ impl TieredSmt { /// If the value for the specified key was not previously set, [ZERO; 4] is returned. pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word { // insert the value into the key-value map, and if nothing has changed, return - let old_value = self.values.insert(key, value).unwrap_or(EMPTY_WORD); + let old_value = self.values.insert(key, value).unwrap_or(Self::EMPTY_VALUE); if old_value == value { return old_value; } From 679a30e02ee864b3af7d964d8135fbb9e99dd52a Mon Sep 17 00:00:00 2001 From: frisitano Date: Fri, 16 Jun 2023 11:42:49 +0100 Subject: [PATCH 16/19] feat: introduce recorder objects --- src/data.rs | 307 +++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/merkle/mmr/full.rs | 1 + src/merkle/mod.rs | 6 +- src/merkle/partial_mt/mod.rs | 3 +- src/merkle/path.rs | 6 + src/merkle/store/mod.rs | 240 ++++++++++++++++++--------- src/merkle/store/tests.rs | 61 ++++++- src/utils.rs | 2 +- 9 files changed, 541 insertions(+), 86 deletions(-) create mode 100644 src/data.rs diff --git a/src/data.rs b/src/data.rs new file mode 100644 index 0000000..cd82502 --- /dev/null +++ b/src/data.rs @@ -0,0 +1,307 @@ +use super::utils::{ + collections::{btree_map::IntoIter, BTreeMap, BTreeSet}, + Box, +}; +use core::{ + cell::RefCell, + iter::{Chain, Filter}, +}; + +// KEY-VALUE MAP TRAIT +// ================================================================================================ +/// A trait that defines the interface for a key-value map. +pub trait KvMap { + fn get(&self, key: &K) -> Option<&V>; + fn contains_key(&self, key: &K) -> bool; + fn len(&self) -> usize; + fn is_empty(&self) -> bool { + self.len() == 0 + } + fn iter(&self) -> Box + '_>; + fn insert(&mut self, key: K, value: V) -> Option; +} + +// RECORDING MAP +// ================================================================================================ + +/// A [RecordingMap] that records read requests to the underlying key-value map. +/// The data recorder is used to generate a proof for read requests. +/// +/// The [RecordingMap] is composed of three parts: +/// - `data`: which contains the initial key-value pairs from the underlying data set. +/// - `delta`: which contains key-value pairs which have been created after instantiation. +/// - `updated_keys`: which tracks keys from `data` which have been updated in `delta`. +/// - `trace`: which contains the keys from the initial data set (`data`) that are read. +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct RecordingMap { + data: BTreeMap, + delta: BTreeMap, + updated_keys: BTreeSet, + trace: RefCell>, +} + +impl RecordingMap { + // CONSTRUCTOR + // -------------------------------------------------------------------------------------------- + /// Returns a new [RecordingMap] instance initialized with the provided key-value pairs. + /// ([BTreeMap]). + pub fn new(init: impl IntoIterator) -> Self { + RecordingMap { + data: init.into_iter().collect(), + delta: BTreeMap::new(), + updated_keys: BTreeSet::new(), + trace: RefCell::new(BTreeSet::new()), + } + } + + // FINALIZER + // -------------------------------------------------------------------------------------------- + /// Consumes the [DataRecorder] and returns a [BTreeMap] containing the key-value pairs from + /// the initial data set that were read during recording. + pub fn into_proof(self) -> BTreeMap { + self.data + .into_iter() + .filter(|(k, _)| self.trace.borrow().contains(k)) + .collect::>() + } +} + +impl KvMap for RecordingMap { + // ACCESSORS + // -------------------------------------------------------------------------------------------- + /// Returns a reference to the value associated with the given key if the value exists. If the + /// key is part of the initial data set, the key access is recorded. + fn get(&self, key: &K) -> Option<&V> { + if let Some(value) = self.delta.get(key) { + return Some(value); + } + + match self.data.get(key) { + None => None, + Some(value) => { + self.trace.borrow_mut().insert(key.clone()); + Some(value) + } + } + } + + /// Returns a boolean to indicate whether the given key exists in the data set. If the key is + /// part of the initial data set, the key access is recorded. + fn contains_key(&self, key: &K) -> bool { + if self.delta.contains_key(key) { + return true; + } + + match self.data.contains_key(key) { + true => { + self.trace.borrow_mut().insert(key.clone()); + true + } + false => false, + } + } + + /// Returns the number of key-value pairs in the data set. + fn len(&self) -> usize { + self.data.len() + self.delta.len() - self.updated_keys.len() + } + + /// Returns an iterator over the key-value pairs in the data set. + fn iter(&self) -> Box + '_> { + Box::new( + self.data + .iter() + .filter(|(k, _)| !self.updated_keys.contains(k)) + .chain(self.delta.iter()), + ) + } + + // MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Inserts a key-value pair into the data set. If the key already exists in the data set, the + /// value is updated and the old value is returned. + fn insert(&mut self, key: K, value: V) -> Option { + if let Some(value) = self.delta.insert(key.clone(), value) { + return Some(value); + } + + match self.data.get(&key) { + None => None, + Some(value) => { + self.trace.borrow_mut().insert(key.clone()); + self.updated_keys.insert(key); + Some(value.clone()) + } + } + } +} + +// RECORDING MAP TRAIT IMPLS +// ================================================================================================ + +impl Extend<(K, V)> for RecordingMap { + fn extend>(&mut self, iter: T) { + iter.into_iter().for_each(move |(k, v)| { + self.insert(k, v); + }); + } +} + +impl Default for RecordingMap { + fn default() -> Self { + RecordingMap::new(BTreeMap::new()) + } +} + +impl IntoIterator for RecordingMap { + type Item = (K, V); + type IntoIter = + Chain, Box bool>>, IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + #[allow(clippy::type_complexity)] + let filter_updated: Box bool> = + Box::new(move |(k, _)| !self.updated_keys.contains(k)); + let data_iter = self.data.into_iter().filter(filter_updated); + let updates_iter = self.delta.into_iter(); + + data_iter.chain(updates_iter) + } +} + +// BTREE MAP `KvMap` IMPLEMENTATION +// ================================================================================================ +impl KvMap for BTreeMap { + fn get(&self, key: &K) -> Option<&V> { + self.get(key) + } + + fn contains_key(&self, key: &K) -> bool { + self.contains_key(key) + } + + fn len(&self) -> usize { + self.len() + } + + fn iter(&self) -> Box + '_> { + Box::new(self.iter()) + } + + fn insert(&mut self, key: K, value: V) -> Option { + self.insert(key, value) + } +} + +// TESTS +// ================================================================================================ +#[cfg(test)] +mod test_recorder { + use super::*; + + const ITEMS: [(u64, u64); 5] = [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]; + + #[test] + fn test_get_item() { + // instantiate a recording map + let map = RecordingMap::new(ITEMS.to_vec()); + + // get a few items + let get_items = [0, 1, 2]; + for key in get_items.iter() { + map.get(key); + } + + // convert the map into a proof + let proof = map.into_proof(); + + // check that the proof contains the expected values + for (key, value) in ITEMS.iter() { + match get_items.contains(key) { + true => assert_eq!(proof.get(key), Some(value)), + false => assert_eq!(proof.get(key), None), + } + } + } + + #[test] + fn test_contains_key() { + // instantiate a recording map + let map = RecordingMap::new(ITEMS.to_vec()); + + // check if the map contains a few items + let get_items = [0, 1, 2]; + for key in get_items.iter() { + map.contains_key(key); + } + + // convert the map into a proof + let proof = map.into_proof(); + + // check that the proof contains the expected values + for (key, _) in ITEMS.iter() { + match get_items.contains(key) { + true => assert_eq!(proof.contains_key(key), true), + false => assert_eq!(proof.contains_key(key), false), + } + } + } + + #[test] + fn test_len() { + // instantiate a recording map + let mut map = RecordingMap::new(ITEMS.to_vec()); + // length of the map should be equal to the number of items + assert_eq!(map.len(), ITEMS.len()); + + // inserting entry with key that already exists should not change the length + map.insert(4, 5); + assert_eq!(map.len(), ITEMS.len()); + + // inserting entry with new key should increase the length + map.insert(5, 5); + assert_eq!(map.len(), ITEMS.len() + 1); + + // get some items so that they are saved in the trace + let get_items = [0, 1, 2]; + for key in get_items.iter() { + map.contains_key(key); + } + + // Note: The length reported by the proof will be different to the length originally + // reported by the map. + let proof = map.into_proof(); + + // length of the proof should be equal to get_items + 1. The extra item is the original + // value at key = 4u64 + assert_eq!(proof.len(), get_items.len() + 1); + } + + #[test] + fn test_iter() { + let mut map = RecordingMap::new(ITEMS.to_vec()); + assert!(map.iter().all(|(x, y)| ITEMS.contains(&(*x, *y)))); + + // when inserting entry with key that already exists the iterator should return the new value + let new_value = 5; + map.insert(4, new_value); + assert_eq!(map.iter().count(), ITEMS.len()); + assert!(map.iter().all(|(x, y)| if x == &4 { + y == &new_value + } else { + ITEMS.contains(&(*x, *y)) + })); + } + + #[test] + fn test_is_empty() { + // instantiate an empty recording map + let empty_map: RecordingMap = RecordingMap::default(); + assert!(empty_map.is_empty()); + + // instantiate a non-empty recording map + let map = RecordingMap::new(ITEMS.to_vec()); + assert!(!map.is_empty()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 7c7d753..f0dc7ee 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ #[cfg_attr(test, macro_use)] extern crate alloc; +pub mod data; pub mod hash; pub mod merkle; pub mod utils; diff --git a/src/merkle/mmr/full.rs b/src/merkle/mmr/full.rs index 7ceace3..d2fbbeb 100644 --- a/src/merkle/mmr/full.rs +++ b/src/merkle/mmr/full.rs @@ -28,6 +28,7 @@ use std::error::Error; /// /// Since this is a full representation of the MMR, elements are never removed and the MMR will /// grow roughly `O(2n)` in number of leaf elements. +#[derive(Debug, Clone)] pub struct Mmr { /// Refer to the `forest` method documentation for details of the semantics of this value. pub(super) forest: usize, diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 2617787..00e8ffc 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -1,4 +1,5 @@ use super::{ + data::{KvMap, RecordingMap}, hash::rpo::{Rpo256, RpoDigest}, utils::collections::{vec, BTreeMap, BTreeSet, Vec}, Felt, StarkField, Word, WORD_SIZE, ZERO, @@ -33,7 +34,10 @@ mod mmr; pub use mmr::{Mmr, MmrPeaks, MmrProof}; mod store; -pub use store::MerkleStore; +pub use store::{ + GenericMerkleStore, MerkleMap, MerkleMapT, MerkleStore, RecordingMerkleMap, + RecordingMerkleStore, +}; mod node; pub use node::InnerNodeInfo; diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index 8800662..3558c9f 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -154,7 +154,8 @@ impl PartialMerkleTree { self.leaves.iter().map(|&leaf| { ( leaf, - self.get_node(leaf).expect(&format!("Leaf with {leaf} is not in the nodes map")), + self.get_node(leaf) + .unwrap_or_else(|_| panic!("Leaf with {leaf} is not in the nodes map")), ) }) } diff --git a/src/merkle/path.rs b/src/merkle/path.rs index 1842915..975bc68 100644 --- a/src/merkle/path.rs +++ b/src/merkle/path.rs @@ -68,6 +68,12 @@ impl MerklePath { } } +impl From for Vec { + fn from(path: MerklePath) -> Self { + path.nodes + } +} + impl From> for MerklePath { fn from(path: Vec) -> Self { Self::new(path) diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index a412092..c85d4fd 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,6 +1,7 @@ use super::{ - mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, - MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, TieredSmt, ValuePath, Vec, + mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, KvMap, MerkleError, MerklePath, + MerklePathSet, MerkleTree, NodeIndex, RecordingMap, RootPath, Rpo256, RpoDigest, SimpleSmt, + TieredSmt, ValuePath, Vec, }; use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use core::borrow::Borrow; @@ -8,12 +9,56 @@ use core::borrow::Borrow; #[cfg(test)] mod tests; +// TRAIT / TYPE DECLARATIONS +// ================================================================================================ +/// A supertrait that defines the required traits for a type to be used as a data map backend for +/// the [GenericMerkleStore] +pub trait MerkleMapT: + KvMap + + Extend<(RpoDigest, Node)> + + FromIterator<(RpoDigest, Node)> + + IntoIterator +{ +} + +// MERKLE STORE +// ------------------------------------------------------------------------------------------------ + +/// Type that represents a standard MerkleStore. +pub type MerkleStore = GenericMerkleStore; + +/// Declaration of a BTreeMap that uses a [RpoDigest] as a key and a [Node] as the value. This type +/// is used as a data backend for the standard [GenericMerkleStore]. +pub type MerkleMap = BTreeMap; + +/// Implementation of [MerkleMapT] trait on [MerkleMap]. +impl MerkleMapT for MerkleMap {} + +// RECORDING MERKLE STORE +// ------------------------------------------------------------------------------------------------ + +/// Type that represents a MerkleStore with recording capabilities. +pub type RecordingMerkleStore = GenericMerkleStore; + +/// Declaration of a [RecordingMap] that uses a [RpoDigest] as a key and a [Node] as the value. +/// This type is used as a data backend for the recording [GenericMerkleStore]. +pub type RecordingMerkleMap = RecordingMap; + +/// Implementation of [MerkleMapT] on [RecordingMerkleMap]. +impl MerkleMapT for RecordingMerkleMap {} + +// NODE DEFINITION +// ================================================================================================ + #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)] pub struct Node { left: RpoDigest, right: RpoDigest, } +// MERKLE STORE IMPLEMENTATION +// ================================================================================================ + /// An in-memory data store for Merkelized data. /// /// This is a in memory data store for Merkle trees, this store allows all the nodes of multiple @@ -51,9 +96,8 @@ pub struct Node { /// let tree2 = MerkleTree::new(vec![A, B, C, D, E, F, G, H1]).unwrap(); /// /// // populates the store with two merkle trees, common nodes are shared -/// store -/// .extend(tree1.inner_nodes()) -/// .extend(tree2.inner_nodes()); +/// store.extend(tree1.inner_nodes()); +/// store.extend(tree2.inner_nodes()); /// /// // every leaf except the last are the same /// for i in 0..7 { @@ -78,41 +122,25 @@ pub struct Node { /// assert_eq!(store.num_internal_nodes() - 255, 10); /// ``` #[derive(Debug, Clone, Eq, PartialEq)] -pub struct MerkleStore { - nodes: BTreeMap, +pub struct GenericMerkleStore { + nodes: T, } -impl Default for MerkleStore { +impl Default for GenericMerkleStore { fn default() -> Self { Self::new() } } -impl MerkleStore { +impl GenericMerkleStore { // CONSTRUCTORS // -------------------------------------------------------------------------------------------- - /// Creates an empty `MerkleStore` instance. - pub fn new() -> MerkleStore { + /// Creates an empty `GenericMerkleStore` instance. + pub fn new() -> GenericMerkleStore { // pre-populate the store with the empty hashes - let subtrees = EmptySubtreeRoots::empty_hashes(255); - let nodes = subtrees - .iter() - .rev() - .copied() - .zip(subtrees.iter().rev().skip(1).copied()) - .map(|(child, parent)| { - ( - parent, - Node { - left: child, - right: child, - }, - ) - }) - .collect(); - - MerkleStore { nodes } + let nodes = empty_hashes().into_iter().collect(); + GenericMerkleStore { nodes } } // PUBLIC ACCESSORS @@ -261,12 +289,12 @@ impl MerkleStore { /// nodes which are descendants of the specified roots. /// /// The roots for which no descendants exist in this Merkle store are ignored. - pub fn subset(&self, roots: I) -> MerkleStore + pub fn subset(&self, roots: I) -> GenericMerkleStore where I: Iterator, R: Borrow, { - let mut store = MerkleStore::new(); + let mut store = GenericMerkleStore::new(); for root in roots { let root = *root.borrow(); store.clone_tree_from(root, self); @@ -274,7 +302,7 @@ impl MerkleStore { store } - /// Iterator over the inner nodes of the [MerkleStore]. + /// Iterator over the inner nodes of the [GenericMerkleStore]. pub fn inner_nodes(&self) -> impl Iterator + '_ { self.nodes.iter().map(|(r, n)| InnerNodeInfo { value: *r, @@ -286,23 +314,6 @@ impl MerkleStore { // STATE MUTATORS // -------------------------------------------------------------------------------------------- - /// Adds a sequence of nodes yielded by the provided iterator into the store. - pub fn extend(&mut self, iter: I) -> &mut MerkleStore - where - I: Iterator, - { - for node in iter { - let value: RpoDigest = node.value; - let left: RpoDigest = node.left; - let right: RpoDigest = node.right; - - debug_assert_eq!(Rpo256::merge(&[left, right]), value); - self.nodes.insert(value, Node { left, right }); - } - - self - } - /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the /// new root. /// @@ -332,7 +343,7 @@ impl MerkleStore { /// This will compute the sibling elements for each Merkle `path` and include all the nodes /// into the store. /// - /// For further reference, check [MerkleStore::add_merkle_path]. + /// For further reference, check [GenericMerkleStore::add_merkle_path]. pub fn add_merkle_paths(&mut self, paths: I) -> Result<(), MerkleError> where I: IntoIterator, @@ -345,7 +356,7 @@ impl MerkleStore { /// Appends the provided [MerklePathSet] into the store. /// - /// For further reference, check [MerkleStore::add_merkle_path]. + /// For further reference, check [GenericMerkleStore::add_merkle_path]. pub fn add_merkle_path_set( &mut self, path_set: &MerklePathSet, @@ -420,55 +431,126 @@ impl MerkleStore { } } +// RECORDING MERKLE STORE FINALIZER +// =============================================================================================== + +impl RecordingMerkleStore { + /// Consumes the [DataRecorder] and returns a [BTreeMap] containing the key-value pairs from + /// the initial data set that were read during recording. + pub fn into_proof(self) -> MerkleMap { + self.nodes.into_proof() + } +} + +// EMPTY HASHES +// ================================================================================================ +/// Creates empty hashes for all the subtrees of a tree with a max depth of 255. +fn empty_hashes() -> impl IntoIterator { + let subtrees = EmptySubtreeRoots::empty_hashes(255); + subtrees.iter().rev().copied().zip(subtrees.iter().rev().skip(1).copied()).map( + |(child, parent)| { + ( + parent, + Node { + left: child, + right: child, + }, + ) + }, + ) +} + +/// Consumes an iterator of [InnerNodeInfo] and returns an iterator of `(value, node)` tuples +/// which includes the nodes associate with roots of empty subtrees up to a depth of 255. +fn combine_nodes_with_empty_hashes( + nodes: impl IntoIterator, +) -> impl Iterator { + nodes + .into_iter() + .map(|info| { + ( + info.value, + Node { + left: info.left, + right: info.right, + }, + ) + }) + .chain(empty_hashes().into_iter()) +} + // CONVERSIONS // ================================================================================================ -impl From<&MerkleTree> for MerkleStore { +impl From<&MerkleTree> for GenericMerkleStore { fn from(value: &MerkleTree) -> Self { - let mut store = MerkleStore::new(); - store.extend(value.inner_nodes()); - store + let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); + GenericMerkleStore { nodes } } } -impl From<&SimpleSmt> for MerkleStore { +impl From<&SimpleSmt> for GenericMerkleStore { fn from(value: &SimpleSmt) -> Self { - let mut store = MerkleStore::new(); - store.extend(value.inner_nodes()); - store + let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); + GenericMerkleStore { nodes } } } -impl From<&Mmr> for MerkleStore { +impl From<&Mmr> for GenericMerkleStore { fn from(value: &Mmr) -> Self { - let mut store = MerkleStore::new(); - store.extend(value.inner_nodes()); - store + let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); + GenericMerkleStore { nodes } } } -impl From<&TieredSmt> for MerkleStore { +impl From<&TieredSmt> for GenericMerkleStore { fn from(value: &TieredSmt) -> Self { - let mut store = MerkleStore::new(); - store.extend(value.inner_nodes()); - store + let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); + GenericMerkleStore { nodes } } } -impl FromIterator for MerkleStore { - fn from_iter>(iter: T) -> Self { - let mut store = MerkleStore::new(); - store.extend(iter.into_iter()); - store +impl FromIterator for GenericMerkleStore { + fn from_iter>(iter: I) -> Self { + let nodes = combine_nodes_with_empty_hashes(iter).collect(); + GenericMerkleStore { nodes } + } +} + +impl From for RecordingMerkleStore { + fn from(value: MerkleStore) -> Self { + GenericMerkleStore { + nodes: RecordingMerkleMap::new(value.nodes.into_iter()), + } + } +} + +impl FromIterator<(RpoDigest, Node)> for RecordingMerkleMap { + fn from_iter>(iter: T) -> Self { + RecordingMerkleMap::new(iter) + } +} + +impl From for MerkleStore { + fn from(value: MerkleMap) -> Self { + GenericMerkleStore { nodes: value } } } // ITERATORS // ================================================================================================ -impl Extend for MerkleStore { - fn extend>(&mut self, iter: T) { - self.extend(iter.into_iter()); +impl Extend for GenericMerkleStore { + fn extend>(&mut self, iter: I) { + self.nodes.extend(iter.into_iter().map(|info| { + ( + info.value, + Node { + left: info.left, + right: info.right, + }, + ) + })); } } @@ -490,7 +572,7 @@ impl Deserializable for Node { } } -impl Serializable for MerkleStore { +impl Serializable for GenericMerkleStore { fn write_into(&self, target: &mut W) { target.write_u64(self.nodes.len() as u64); @@ -501,10 +583,10 @@ impl Serializable for MerkleStore { } } -impl Deserializable for MerkleStore { +impl Deserializable for GenericMerkleStore { fn read_from(source: &mut R) -> Result { let len = source.read_u64()?; - let mut nodes: BTreeMap = BTreeMap::new(); + let mut nodes: MerkleMap = BTreeMap::new(); for _ in 0..len { let key = RpoDigest::read_from(source)?; @@ -512,6 +594,6 @@ impl Deserializable for MerkleStore { nodes.insert(key, value); } - Ok(MerkleStore { nodes }) + Ok(GenericMerkleStore { nodes }) } } diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 1bc5db9..56b124f 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -1,13 +1,16 @@ use super::{ - Deserializable, EmptySubtreeRoots, MerkleError, MerklePath, MerkleStore, NodeIndex, RpoDigest, - Serializable, + EmptySubtreeRoots, MerkleError, MerklePath, MerkleStore, NodeIndex, RecordingMerkleStore, + RpoDigest, }; use crate::{ hash::rpo::Rpo256, merkle::{digests_to_words, int_to_leaf, int_to_node, MerklePathSet, MerkleTree, SimpleSmt}, - Felt, Word, WORD_SIZE, + Felt, Word, ONE, WORD_SIZE, ZERO, }; +#[cfg(feature = "std")] +use super::{Deserializable, Serializable}; + #[cfg(feature = "std")] use std::error::Error; @@ -17,6 +20,7 @@ use std::error::Error; const KEYS4: [u64; 4] = [0, 1, 2, 3]; const VALUES4: [RpoDigest; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; +const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7]; const VALUES8: [RpoDigest; 8] = [ int_to_node(1), int_to_node(2), @@ -34,7 +38,7 @@ const VALUES8: [RpoDigest; 8] = [ #[test] fn test_root_not_in_store() -> Result<(), MerkleError> { let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; - let store = MerkleStore::from(&mtree); + let store = MerkleStore::default(); assert_eq!( store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)), Err(MerkleError::RootNotInStore(VALUES4[0])), @@ -810,3 +814,52 @@ fn test_serialization() -> Result<(), Box> { assert_eq!(store, decoded); Ok(()) } + +// MERKLE RECORDER +// ================================================================================================ +#[test] +fn test_recorder() { + // instantiate recorder from MerkleTree and SimpleSmt + let mtree = MerkleTree::new(digests_to_words(&VALUES4)).unwrap(); + let smtree = SimpleSmt::with_leaves( + 64, + KEYS8.into_iter().zip(VALUES8.into_iter().map(|x| x.into()).rev()), + ) + .unwrap(); + let mut recorder: RecordingMerkleStore = + mtree.inner_nodes().chain(smtree.inner_nodes()).collect(); + + // get nodes from both trees and make sure they are correct + let index_0 = NodeIndex::new(mtree.depth(), 0).unwrap(); + let node = recorder.get_node(mtree.root(), index_0).unwrap(); + assert_eq!(node, mtree.get_node(index_0).unwrap()); + + let index_1 = NodeIndex::new(smtree.depth(), 1).unwrap(); + let node = recorder.get_node(smtree.root(), index_1).unwrap(); + assert_eq!(node, smtree.get_node(index_1).unwrap()); + + // insert a value and assert that when we request it next time it is accurate + let new_value = [ZERO, ZERO, ONE, ONE].into(); + let index_2 = NodeIndex::new(smtree.depth(), 2).unwrap(); + let root = recorder.set_node(smtree.root(), index_2, new_value).unwrap().root; + assert_eq!(recorder.get_node(root, index_2).unwrap(), new_value); + + // construct the proof + let proof = recorder.into_proof(); + let merkle_store: MerkleStore = proof.into(); + + // make sure the proof contains all nodes from both trees + let node = merkle_store.get_node(mtree.root(), index_0).unwrap(); + assert_eq!(node, mtree.get_node(index_0).unwrap()); + + let node = merkle_store.get_node(smtree.root(), index_1).unwrap(); + assert_eq!(node, smtree.get_node(index_1).unwrap()); + + let node = merkle_store.get_node(smtree.root(), index_2).unwrap(); + assert_eq!(node, smtree.get_leaf(index_2.value()).unwrap().into()); + + // assert that is doesnt contain nodes that were not recorded + let not_recorded_index = NodeIndex::new(smtree.depth(), 4).unwrap(); + assert!(merkle_store.get_node(smtree.root(), not_recorded_index).is_err()); + assert!(smtree.get_node(not_recorded_index).is_ok()); +} diff --git a/src/utils.rs b/src/utils.rs index b9b7849..e350b69 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -11,7 +11,7 @@ pub use std::format; // RE-EXPORTS // ================================================================================================ pub use winter_utils::{ - collections, string, uninit_vector, ByteReader, ByteWriter, Deserializable, + collections, string, uninit_vector, Box, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable, SliceReader, }; From f08644e4dfc220febecb6c4413cdaf13d0fc1e53 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Fri, 23 Jun 2023 23:07:24 -0700 Subject: [PATCH 17/19] refactor: simplify recording MerkleStore structure --- benches/store.rs | 2 +- src/lib.rs | 1 - src/merkle/mod.rs | 8 +- src/merkle/partial_mt/tests.rs | 5 +- src/merkle/store/mod.rs | 234 ++++++++++++---------------- src/merkle/store/tests.rs | 10 +- src/{data.rs => utils/kv_map.rs} | 251 +++++++++++++++++-------------- src/{utils.rs => utils/mod.rs} | 17 ++- 8 files changed, 256 insertions(+), 272 deletions(-) rename src/{data.rs => utils/kv_map.rs} (62%) rename src/{utils.rs => utils/mod.rs} (55%) diff --git a/benches/store.rs b/benches/store.rs index f3d9cfa..d6da04b 100644 --- a/benches/store.rs +++ b/benches/store.rs @@ -1,5 +1,5 @@ use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion}; -use miden_crypto::merkle::{MerkleStore, MerkleTree, NodeIndex, SimpleSmt}; +use miden_crypto::merkle::{DefaultMerkleStore as MerkleStore, MerkleTree, NodeIndex, SimpleSmt}; use miden_crypto::Word; use miden_crypto::{hash::rpo::RpoDigest, Felt}; use rand_utils::{rand_array, rand_value}; diff --git a/src/lib.rs b/src/lib.rs index f0dc7ee..7c7d753 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,7 +4,6 @@ #[cfg_attr(test, macro_use)] extern crate alloc; -pub mod data; pub mod hash; pub mod merkle; pub mod utils; diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 00e8ffc..ded7313 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -1,7 +1,6 @@ use super::{ - data::{KvMap, RecordingMap}, hash::rpo::{Rpo256, RpoDigest}, - utils::collections::{vec, BTreeMap, BTreeSet, Vec}, + utils::collections::{vec, BTreeMap, BTreeSet, KvMap, RecordingMap, Vec}, Felt, StarkField, Word, WORD_SIZE, ZERO, }; use core::fmt; @@ -34,10 +33,7 @@ mod mmr; pub use mmr::{Mmr, MmrPeaks, MmrProof}; mod store; -pub use store::{ - GenericMerkleStore, MerkleMap, MerkleMapT, MerkleStore, RecordingMerkleMap, - RecordingMerkleStore, -}; +pub use store::{DefaultMerkleStore, MerkleStore, RecordingMerkleStore}; mod node; pub use node::InnerNodeInfo; diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index 91638da..ed5281f 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -1,5 +1,8 @@ use super::{ - super::{digests_to_words, int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree}, + super::{ + digests_to_words, int_to_node, DefaultMerkleStore as MerkleStore, MerkleTree, NodeIndex, + PartialMerkleTree, + }, RpoDigest, ValuePath, Vec, }; diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index c85d4fd..ab69ecc 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -9,46 +9,14 @@ use core::borrow::Borrow; #[cfg(test)] mod tests; -// TRAIT / TYPE DECLARATIONS -// ================================================================================================ -/// A supertrait that defines the required traits for a type to be used as a data map backend for -/// the [GenericMerkleStore] -pub trait MerkleMapT: - KvMap - + Extend<(RpoDigest, Node)> - + FromIterator<(RpoDigest, Node)> - + IntoIterator -{ -} - // MERKLE STORE -// ------------------------------------------------------------------------------------------------ - -/// Type that represents a standard MerkleStore. -pub type MerkleStore = GenericMerkleStore; - -/// Declaration of a BTreeMap that uses a [RpoDigest] as a key and a [Node] as the value. This type -/// is used as a data backend for the standard [GenericMerkleStore]. -pub type MerkleMap = BTreeMap; - -/// Implementation of [MerkleMapT] trait on [MerkleMap]. -impl MerkleMapT for MerkleMap {} - -// RECORDING MERKLE STORE -// ------------------------------------------------------------------------------------------------ - -/// Type that represents a MerkleStore with recording capabilities. -pub type RecordingMerkleStore = GenericMerkleStore; - -/// Declaration of a [RecordingMap] that uses a [RpoDigest] as a key and a [Node] as the value. -/// This type is used as a data backend for the recording [GenericMerkleStore]. -pub type RecordingMerkleMap = RecordingMap; +// ================================================================================================ -/// Implementation of [MerkleMapT] on [RecordingMerkleMap]. -impl MerkleMapT for RecordingMerkleMap {} +/// A default [MerkleStore] which uses a simple [BTreeMap] as the backing storage. +pub type DefaultMerkleStore = MerkleStore>; -// NODE DEFINITION -// ================================================================================================ +/// A [MerkleStore] with recording capabilities which uses [RecordingMap] as the backing storage. +pub type RecordingMerkleStore = MerkleStore>; #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)] pub struct Node { @@ -56,9 +24,6 @@ pub struct Node { right: RpoDigest, } -// MERKLE STORE IMPLEMENTATION -// ================================================================================================ - /// An in-memory data store for Merkelized data. /// /// This is a in memory data store for Merkle trees, this store allows all the nodes of multiple @@ -87,7 +52,7 @@ pub struct Node { /// # let T1 = MerkleTree::new([A, B, C, D, E, F, G, H1].to_vec()).expect("even number of leaves provided"); /// # let ROOT0 = T0.root(); /// # let ROOT1 = T1.root(); -/// let mut store = MerkleStore::new(); +/// let mut store: MerkleStore = MerkleStore::new(); /// /// // the store is initialized with the SMT empty nodes /// assert_eq!(store.num_internal_nodes(), 255); @@ -122,25 +87,25 @@ pub struct Node { /// assert_eq!(store.num_internal_nodes() - 255, 10); /// ``` #[derive(Debug, Clone, Eq, PartialEq)] -pub struct GenericMerkleStore { +pub struct MerkleStore = BTreeMap> { nodes: T, } -impl Default for GenericMerkleStore { +impl> Default for MerkleStore { fn default() -> Self { Self::new() } } -impl GenericMerkleStore { +impl> MerkleStore { // CONSTRUCTORS // -------------------------------------------------------------------------------------------- - /// Creates an empty `GenericMerkleStore` instance. - pub fn new() -> GenericMerkleStore { + /// Creates an empty `MerkleStore` instance. + pub fn new() -> MerkleStore { // pre-populate the store with the empty hashes let nodes = empty_hashes().into_iter().collect(); - GenericMerkleStore { nodes } + MerkleStore { nodes } } // PUBLIC ACCESSORS @@ -154,10 +119,10 @@ impl GenericMerkleStore { /// Returns the node at `index` rooted on the tree `root`. /// /// # Errors - /// /// This method can return the following errors: /// - `RootNotInStore` if the `root` is not present in the store. - /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. + /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in + /// the store. pub fn get_node(&self, root: RpoDigest, index: NodeIndex) -> Result { let mut hash = root; @@ -181,7 +146,8 @@ impl GenericMerkleStore { /// # Errors /// This method can return the following errors: /// - `RootNotInStore` if the `root` is not present in the store. - /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. + /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in + /// the store. pub fn get_path(&self, root: RpoDigest, index: NodeIndex) -> Result { let mut hash = root; let mut path = Vec::with_capacity(index.depth().into()); @@ -225,7 +191,7 @@ impl GenericMerkleStore { /// - The path from the root continues to a depth greater than `tree_depth`. /// - The provided `tree_depth` is greater than `64. /// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more - /// information, check [NodeIndex::new]. + /// information, check [NodeIndex::new]. pub fn get_leaf_depth( &self, root: RpoDigest, @@ -289,12 +255,12 @@ impl GenericMerkleStore { /// nodes which are descendants of the specified roots. /// /// The roots for which no descendants exist in this Merkle store are ignored. - pub fn subset(&self, roots: I) -> GenericMerkleStore + pub fn subset(&self, roots: I) -> MerkleStore where I: Iterator, R: Borrow, { - let mut store = GenericMerkleStore::new(); + let mut store = MerkleStore::new(); for root in roots { let root = *root.borrow(); store.clone_tree_from(root, self); @@ -302,7 +268,7 @@ impl GenericMerkleStore { store } - /// Iterator over the inner nodes of the [GenericMerkleStore]. + /// Iterator over the inner nodes of the [MerkleStore]. pub fn inner_nodes(&self) -> impl Iterator + '_ { self.nodes.iter().map(|(r, n)| InnerNodeInfo { value: *r, @@ -343,7 +309,7 @@ impl GenericMerkleStore { /// This will compute the sibling elements for each Merkle `path` and include all the nodes /// into the store. /// - /// For further reference, check [GenericMerkleStore::add_merkle_path]. + /// For further reference, check [MerkleStore::add_merkle_path]. pub fn add_merkle_paths(&mut self, paths: I) -> Result<(), MerkleError> where I: IntoIterator, @@ -356,7 +322,7 @@ impl GenericMerkleStore { /// Appends the provided [MerklePathSet] into the store. /// - /// For further reference, check [GenericMerkleStore::add_merkle_path]. + /// For further reference, check [MerkleStore::add_merkle_path]. pub fn add_merkle_path_set( &mut self, path_set: &MerklePathSet, @@ -371,10 +337,10 @@ impl GenericMerkleStore { /// Sets a node to `value`. /// /// # Errors - /// /// This method can return the following errors: /// - `RootNotInStore` if the `root` is not present in the store. - /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. + /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in + /// the store. pub fn set_node( &mut self, mut root: RpoDigest, @@ -412,6 +378,14 @@ impl GenericMerkleStore { Ok(parent) } + // DESTRUCTURING + // -------------------------------------------------------------------------------------------- + + /// Returns the inner storage of this MerkleStore while consuming `self`. + pub fn into_inner(self) -> T { + self.nodes + } + // HELPER METHODS // -------------------------------------------------------------------------------------------- @@ -431,116 +405,62 @@ impl GenericMerkleStore { } } -// RECORDING MERKLE STORE FINALIZER -// =============================================================================================== - -impl RecordingMerkleStore { - /// Consumes the [DataRecorder] and returns a [BTreeMap] containing the key-value pairs from - /// the initial data set that were read during recording. - pub fn into_proof(self) -> MerkleMap { - self.nodes.into_proof() - } -} - -// EMPTY HASHES -// ================================================================================================ -/// Creates empty hashes for all the subtrees of a tree with a max depth of 255. -fn empty_hashes() -> impl IntoIterator { - let subtrees = EmptySubtreeRoots::empty_hashes(255); - subtrees.iter().rev().copied().zip(subtrees.iter().rev().skip(1).copied()).map( - |(child, parent)| { - ( - parent, - Node { - left: child, - right: child, - }, - ) - }, - ) -} - -/// Consumes an iterator of [InnerNodeInfo] and returns an iterator of `(value, node)` tuples -/// which includes the nodes associate with roots of empty subtrees up to a depth of 255. -fn combine_nodes_with_empty_hashes( - nodes: impl IntoIterator, -) -> impl Iterator { - nodes - .into_iter() - .map(|info| { - ( - info.value, - Node { - left: info.left, - right: info.right, - }, - ) - }) - .chain(empty_hashes().into_iter()) -} - // CONVERSIONS // ================================================================================================ -impl From<&MerkleTree> for GenericMerkleStore { +impl> From<&MerkleTree> for MerkleStore { fn from(value: &MerkleTree) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); - GenericMerkleStore { nodes } + Self { nodes } } } -impl From<&SimpleSmt> for GenericMerkleStore { +impl> From<&SimpleSmt> for MerkleStore { fn from(value: &SimpleSmt) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); - GenericMerkleStore { nodes } + Self { nodes } } } -impl From<&Mmr> for GenericMerkleStore { +impl> From<&Mmr> for MerkleStore { fn from(value: &Mmr) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); - GenericMerkleStore { nodes } + Self { nodes } } } -impl From<&TieredSmt> for GenericMerkleStore { +impl> From<&TieredSmt> for MerkleStore { fn from(value: &TieredSmt) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); - GenericMerkleStore { nodes } - } -} - -impl FromIterator for GenericMerkleStore { - fn from_iter>(iter: I) -> Self { - let nodes = combine_nodes_with_empty_hashes(iter).collect(); - GenericMerkleStore { nodes } + Self { nodes } } } -impl From for RecordingMerkleStore { - fn from(value: MerkleStore) -> Self { - GenericMerkleStore { - nodes: RecordingMerkleMap::new(value.nodes.into_iter()), - } +impl> From for MerkleStore { + fn from(values: T) -> Self { + let nodes = values.into_iter().chain(empty_hashes().into_iter()).collect(); + Self { nodes } } } -impl FromIterator<(RpoDigest, Node)> for RecordingMerkleMap { - fn from_iter>(iter: T) -> Self { - RecordingMerkleMap::new(iter) +impl> FromIterator for MerkleStore { + fn from_iter>(iter: I) -> Self { + let nodes = combine_nodes_with_empty_hashes(iter.into_iter()).collect(); + Self { nodes } } } -impl From for MerkleStore { - fn from(value: MerkleMap) -> Self { - GenericMerkleStore { nodes: value } +impl> FromIterator<(RpoDigest, Node)> for MerkleStore { + fn from_iter>(iter: I) -> Self { + let nodes = iter.into_iter().chain(empty_hashes().into_iter()).collect(); + Self { nodes } } } // ITERATORS // ================================================================================================ -impl Extend for GenericMerkleStore { +impl> Extend for MerkleStore { fn extend>(&mut self, iter: I) { self.nodes.extend(iter.into_iter().map(|info| { ( @@ -572,7 +492,7 @@ impl Deserializable for Node { } } -impl Serializable for GenericMerkleStore { +impl> Serializable for MerkleStore { fn write_into(&self, target: &mut W) { target.write_u64(self.nodes.len() as u64); @@ -583,17 +503,55 @@ impl Serializable for GenericMerkleStore { } } -impl Deserializable for GenericMerkleStore { +impl> Deserializable for MerkleStore { fn read_from(source: &mut R) -> Result { let len = source.read_u64()?; - let mut nodes: MerkleMap = BTreeMap::new(); + let mut nodes: Vec<(RpoDigest, Node)> = Vec::with_capacity(len as usize); for _ in 0..len { let key = RpoDigest::read_from(source)?; let value = Node::read_from(source)?; - nodes.insert(key, value); + nodes.push((key, value)); } - Ok(GenericMerkleStore { nodes }) + Ok(nodes.into_iter().collect()) } } + +// HELPER FUNCTIONS +// ================================================================================================ + +/// Creates empty hashes for all the subtrees of a tree with a max depth of 255. +fn empty_hashes() -> impl IntoIterator { + let subtrees = EmptySubtreeRoots::empty_hashes(255); + subtrees.iter().rev().copied().zip(subtrees.iter().rev().skip(1).copied()).map( + |(child, parent)| { + ( + parent, + Node { + left: child, + right: child, + }, + ) + }, + ) +} + +/// Consumes an iterator of [InnerNodeInfo] and returns an iterator of `(value, node)` tuples +/// which includes the nodes associate with roots of empty subtrees up to a depth of 255. +fn combine_nodes_with_empty_hashes( + nodes: impl IntoIterator, +) -> impl Iterator { + nodes + .into_iter() + .map(|info| { + ( + info.value, + Node { + left: info.left, + right: info.right, + }, + ) + }) + .chain(empty_hashes().into_iter()) +} diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 56b124f..c6f346f 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -1,6 +1,6 @@ use super::{ - EmptySubtreeRoots, MerkleError, MerklePath, MerkleStore, NodeIndex, RecordingMerkleStore, - RpoDigest, + DefaultMerkleStore as MerkleStore, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, + RecordingMerkleStore, RpoDigest, }; use crate::{ hash::rpo::Rpo256, @@ -38,7 +38,7 @@ const VALUES8: [RpoDigest; 8] = [ #[test] fn test_root_not_in_store() -> Result<(), MerkleError> { let mtree = MerkleTree::new(digests_to_words(&VALUES4))?; - let store = MerkleStore::default(); + let store = MerkleStore::from(&mtree); assert_eq!( store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)), Err(MerkleError::RootNotInStore(VALUES4[0])), @@ -826,6 +826,7 @@ fn test_recorder() { KEYS8.into_iter().zip(VALUES8.into_iter().map(|x| x.into()).rev()), ) .unwrap(); + let mut recorder: RecordingMerkleStore = mtree.inner_nodes().chain(smtree.inner_nodes()).collect(); @@ -845,7 +846,8 @@ fn test_recorder() { assert_eq!(recorder.get_node(root, index_2).unwrap(), new_value); // construct the proof - let proof = recorder.into_proof(); + let rec_map = recorder.into_inner(); + let proof = rec_map.into_proof(); let merkle_store: MerkleStore = proof.into(); // make sure the proof contains all nodes from both trees diff --git a/src/data.rs b/src/utils/kv_map.rs similarity index 62% rename from src/data.rs rename to src/utils/kv_map.rs index cd82502..d9b453d 100644 --- a/src/data.rs +++ b/src/utils/kv_map.rs @@ -1,43 +1,70 @@ -use super::utils::{ +use core::cell::RefCell; +use winter_utils::{ collections::{btree_map::IntoIter, BTreeMap, BTreeSet}, Box, }; -use core::{ - cell::RefCell, - iter::{Chain, Filter}, -}; // KEY-VALUE MAP TRAIT // ================================================================================================ + /// A trait that defines the interface for a key-value map. -pub trait KvMap { +pub trait KvMap: + Extend<(K, V)> + FromIterator<(K, V)> + IntoIterator +{ fn get(&self, key: &K) -> Option<&V>; fn contains_key(&self, key: &K) -> bool; fn len(&self) -> usize; fn is_empty(&self) -> bool { self.len() == 0 } - fn iter(&self) -> Box + '_>; fn insert(&mut self, key: K, value: V) -> Option; + + fn iter(&self) -> Box + '_>; +} + +// BTREE MAP `KvMap` IMPLEMENTATION +// ================================================================================================ + +impl KvMap for BTreeMap { + fn get(&self, key: &K) -> Option<&V> { + self.get(key) + } + + fn contains_key(&self, key: &K) -> bool { + self.contains_key(key) + } + + fn len(&self) -> usize { + self.len() + } + + fn insert(&mut self, key: K, value: V) -> Option { + self.insert(key, value) + } + + fn iter(&self) -> Box + '_> { + Box::new(self.iter()) + } } // RECORDING MAP // ================================================================================================ /// A [RecordingMap] that records read requests to the underlying key-value map. +/// /// The data recorder is used to generate a proof for read requests. /// /// The [RecordingMap] is composed of three parts: -/// - `data`: which contains the initial key-value pairs from the underlying data set. -/// - `delta`: which contains key-value pairs which have been created after instantiation. -/// - `updated_keys`: which tracks keys from `data` which have been updated in `delta`. -/// - `trace`: which contains the keys from the initial data set (`data`) that are read. -#[derive(Debug, Clone, Eq, PartialEq)] +/// - `data`: which contains the current set of key-value pairs in the map. +/// - `updates`: which tracks keys for which values have been since the map was instantiated. +/// updates include both insertions and updates of values under existing keys. +/// - `trace`: which contains the key-value pairs from the original data which have been accesses +/// since the map was instantiated. +#[derive(Debug, Default, Clone, Eq, PartialEq)] pub struct RecordingMap { data: BTreeMap, - delta: BTreeMap, - updated_keys: BTreeSet, - trace: RefCell>, + updates: BTreeSet, + trace: RefCell>, } impl RecordingMap { @@ -48,97 +75,87 @@ impl RecordingMap { pub fn new(init: impl IntoIterator) -> Self { RecordingMap { data: init.into_iter().collect(), - delta: BTreeMap::new(), - updated_keys: BTreeSet::new(), - trace: RefCell::new(BTreeSet::new()), + updates: BTreeSet::new(), + trace: RefCell::new(BTreeMap::new()), } } // FINALIZER // -------------------------------------------------------------------------------------------- - /// Consumes the [DataRecorder] and returns a [BTreeMap] containing the key-value pairs from + + /// Consumes the [RecordingMap] and returns a [BTreeMap] containing the key-value pairs from /// the initial data set that were read during recording. pub fn into_proof(self) -> BTreeMap { - self.data - .into_iter() - .filter(|(k, _)| self.trace.borrow().contains(k)) - .collect::>() + self.trace.take() + } + + // TEST HELPERS + // -------------------------------------------------------------------------------------------- + + #[cfg(test)] + pub fn trace_len(&self) -> usize { + self.trace.borrow().len() + } + + #[cfg(test)] + pub fn updates_len(&self) -> usize { + self.updates.len() } } impl KvMap for RecordingMap { - // ACCESSORS + // PUBLIC ACCESSORS // -------------------------------------------------------------------------------------------- - /// Returns a reference to the value associated with the given key if the value exists. If the - /// key is part of the initial data set, the key access is recorded. - fn get(&self, key: &K) -> Option<&V> { - if let Some(value) = self.delta.get(key) { - return Some(value); - } - match self.data.get(key) { - None => None, - Some(value) => { - self.trace.borrow_mut().insert(key.clone()); - Some(value) + /// Returns a reference to the value associated with the given key if the value exists. + /// + /// If the key is part of the initial data set, the key access is recorded. + fn get(&self, key: &K) -> Option<&V> { + self.data.get(key).map(|value| { + if !self.updates.contains(key) { + self.trace.borrow_mut().insert(key.clone(), value.clone()); } - } + value + }) } - /// Returns a boolean to indicate whether the given key exists in the data set. If the key is - /// part of the initial data set, the key access is recorded. + /// Returns a boolean to indicate whether the given key exists in the data set. + /// + /// If the key is part of the initial data set, the key access is recorded. fn contains_key(&self, key: &K) -> bool { - if self.delta.contains_key(key) { - return true; - } - - match self.data.contains_key(key) { - true => { - self.trace.borrow_mut().insert(key.clone()); - true - } - false => false, - } + self.get(key).is_some() } /// Returns the number of key-value pairs in the data set. fn len(&self) -> usize { - self.data.len() + self.delta.len() - self.updated_keys.len() - } - - /// Returns an iterator over the key-value pairs in the data set. - fn iter(&self) -> Box + '_> { - Box::new( - self.data - .iter() - .filter(|(k, _)| !self.updated_keys.contains(k)) - .chain(self.delta.iter()), - ) + self.data.len() } // MUTATORS // -------------------------------------------------------------------------------------------- - /// Inserts a key-value pair into the data set. If the key already exists in the data set, the - /// value is updated and the old value is returned. + /// Inserts a key-value pair into the data set. + /// + /// If the key already exists in the data set, the value is updated and the old value is + /// returned. fn insert(&mut self, key: K, value: V) -> Option { - if let Some(value) = self.delta.insert(key.clone(), value) { - return Some(value); - } - - match self.data.get(&key) { - None => None, - Some(value) => { - self.trace.borrow_mut().insert(key.clone()); - self.updated_keys.insert(key); - Some(value.clone()) + let new_update = self.updates.insert(key.clone()); + self.data.insert(key.clone(), value).map(|old_value| { + if new_update { + self.trace.borrow_mut().insert(key, old_value.clone()); } - } + old_value + }) } -} -// RECORDING MAP TRAIT IMPLS -// ================================================================================================ + // ITERATION + // -------------------------------------------------------------------------------------------- + + /// Returns an iterator over the key-value pairs in the data set. + fn iter(&self) -> Box + '_> { + Box::new(self.data.iter()) + } +} impl Extend<(K, V)> for RecordingMap { fn extend>(&mut self, iter: T) { @@ -148,56 +165,26 @@ impl Extend<(K, V)> for RecordingMap { } } -impl Default for RecordingMap { - fn default() -> Self { - RecordingMap::new(BTreeMap::new()) +impl FromIterator<(K, V)> for RecordingMap { + fn from_iter>(iter: T) -> Self { + Self::new(iter) } } -impl IntoIterator for RecordingMap { +impl IntoIterator for RecordingMap { type Item = (K, V); - type IntoIter = - Chain, Box bool>>, IntoIter>; + type IntoIter = IntoIter; fn into_iter(self) -> Self::IntoIter { - #[allow(clippy::type_complexity)] - let filter_updated: Box bool> = - Box::new(move |(k, _)| !self.updated_keys.contains(k)); - let data_iter = self.data.into_iter().filter(filter_updated); - let updates_iter = self.delta.into_iter(); - - data_iter.chain(updates_iter) - } -} - -// BTREE MAP `KvMap` IMPLEMENTATION -// ================================================================================================ -impl KvMap for BTreeMap { - fn get(&self, key: &K) -> Option<&V> { - self.get(key) - } - - fn contains_key(&self, key: &K) -> bool { - self.contains_key(key) - } - - fn len(&self) -> usize { - self.len() - } - - fn iter(&self) -> Box + '_> { - Box::new(self.iter()) - } - - fn insert(&mut self, key: K, value: V) -> Option { - self.insert(key, value) + self.data.into_iter() } } // TESTS // ================================================================================================ + #[cfg(test)] -mod test_recorder { +mod tests { use super::*; const ITEMS: [(u64, u64); 5] = [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]; @@ -255,19 +242,49 @@ mod test_recorder { // length of the map should be equal to the number of items assert_eq!(map.len(), ITEMS.len()); - // inserting entry with key that already exists should not change the length + // inserting entry with key that already exists should not change the length, but it does + // add entries to the trace and update sets map.insert(4, 5); assert_eq!(map.len(), ITEMS.len()); + assert_eq!(map.trace_len(), 1); + assert_eq!(map.updates_len(), 1); - // inserting entry with new key should increase the length + // inserting entry with new key should increase the length; it should also record the key + // as an updated key, but the trace length does not change since old values were not touched map.insert(5, 5); assert_eq!(map.len(), ITEMS.len() + 1); + assert_eq!(map.trace_len(), 1); + assert_eq!(map.updates_len(), 2); + + // get some items so that they are saved in the trace; this should record original items + // in the trace, but should not affect the set of updates + let get_items = [0, 1, 2]; + for key in get_items.iter() { + map.contains_key(key); + } + assert_eq!(map.trace_len(), 4); + assert_eq!(map.updates_len(), 2); - // get some items so that they are saved in the trace + // read the same items again, this should not have any effect on either length, trace, or + // the set of updates let get_items = [0, 1, 2]; for key in get_items.iter() { map.contains_key(key); } + assert_eq!(map.trace_len(), 4); + assert_eq!(map.updates_len(), 2); + + // read a newly inserted item; this should not affect either length, trace, or the set of + // updates + let _val = map.get(&5).unwrap(); + assert_eq!(map.trace_len(), 4); + assert_eq!(map.updates_len(), 2); + + // update a newly inserted item; this should not affect either length, trace, or the set + // of updates + map.insert(5, 11); + assert_eq!(map.trace_len(), 4); + assert_eq!(map.updates_len(), 2); // Note: The length reported by the proof will be different to the length originally // reported by the map. diff --git a/src/utils.rs b/src/utils/mod.rs similarity index 55% rename from src/utils.rs rename to src/utils/mod.rs index e350b69..7804420 100644 --- a/src/utils.rs +++ b/src/utils/mod.rs @@ -1,5 +1,4 @@ -use super::Word; -use crate::utils::string::String; +use super::{utils::string::String, Word}; use core::fmt::{self, Write}; #[cfg(not(feature = "std"))] @@ -8,13 +7,23 @@ pub use alloc::format; #[cfg(feature = "std")] pub use std::format; +mod kv_map; + // RE-EXPORTS // ================================================================================================ pub use winter_utils::{ - collections, string, uninit_vector, Box, ByteReader, ByteWriter, Deserializable, - DeserializationError, Serializable, SliceReader, + string, uninit_vector, Box, ByteReader, ByteWriter, Deserializable, DeserializationError, + Serializable, SliceReader, }; +pub mod collections { + pub use super::kv_map::*; + pub use winter_utils::collections::*; +} + +// UTILITY FUNCTIONS +// ================================================================================================ + /// Converts a [Word] into hex. pub fn word_to_hex(w: &Word) -> Result { let mut s = String::new(); From b2d6866d4159d5051930d1f82bf3c1911f4079d3 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Sun, 25 Jun 2023 01:42:21 -0700 Subject: [PATCH 18/19] refactor: rename Merkle store Node into StoreNode --- src/merkle/mod.rs | 2 +- src/merkle/store/mod.rs | 58 ++++++++++++++++++++--------------------- 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index ded7313..3e1c9d9 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -33,7 +33,7 @@ mod mmr; pub use mmr::{Mmr, MmrPeaks, MmrProof}; mod store; -pub use store::{DefaultMerkleStore, MerkleStore, RecordingMerkleStore}; +pub use store::{DefaultMerkleStore, MerkleStore, RecordingMerkleStore, StoreNode}; mod node; pub use node::InnerNodeInfo; diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index ab69ecc..fdba5ed 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -13,13 +13,13 @@ mod tests; // ================================================================================================ /// A default [MerkleStore] which uses a simple [BTreeMap] as the backing storage. -pub type DefaultMerkleStore = MerkleStore>; +pub type DefaultMerkleStore = MerkleStore>; /// A [MerkleStore] with recording capabilities which uses [RecordingMap] as the backing storage. -pub type RecordingMerkleStore = MerkleStore>; +pub type RecordingMerkleStore = MerkleStore>; #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)] -pub struct Node { +pub struct StoreNode { left: RpoDigest, right: RpoDigest, } @@ -87,17 +87,17 @@ pub struct Node { /// assert_eq!(store.num_internal_nodes() - 255, 10); /// ``` #[derive(Debug, Clone, Eq, PartialEq)] -pub struct MerkleStore = BTreeMap> { +pub struct MerkleStore = BTreeMap> { nodes: T, } -impl> Default for MerkleStore { +impl> Default for MerkleStore { fn default() -> Self { Self::new() } } -impl> MerkleStore { +impl> MerkleStore { // CONSTRUCTORS // -------------------------------------------------------------------------------------------- @@ -297,7 +297,7 @@ impl> MerkleStore { let right: RpoDigest = node.right; debug_assert_eq!(Rpo256::merge(&[left, right]), value); - self.nodes.insert(value, Node { left, right }); + self.nodes.insert(value, StoreNode { left, right }); node.value }); @@ -369,7 +369,7 @@ impl> MerkleStore { let parent = Rpo256::merge(&[left_root, right_root]); self.nodes.insert( parent, - Node { + StoreNode { left: left_root, right: right_root, }, @@ -408,50 +408,50 @@ impl> MerkleStore { // CONVERSIONS // ================================================================================================ -impl> From<&MerkleTree> for MerkleStore { +impl> From<&MerkleTree> for MerkleStore { fn from(value: &MerkleTree) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); Self { nodes } } } -impl> From<&SimpleSmt> for MerkleStore { +impl> From<&SimpleSmt> for MerkleStore { fn from(value: &SimpleSmt) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); Self { nodes } } } -impl> From<&Mmr> for MerkleStore { +impl> From<&Mmr> for MerkleStore { fn from(value: &Mmr) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); Self { nodes } } } -impl> From<&TieredSmt> for MerkleStore { +impl> From<&TieredSmt> for MerkleStore { fn from(value: &TieredSmt) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); Self { nodes } } } -impl> From for MerkleStore { +impl> From for MerkleStore { fn from(values: T) -> Self { let nodes = values.into_iter().chain(empty_hashes().into_iter()).collect(); Self { nodes } } } -impl> FromIterator for MerkleStore { +impl> FromIterator for MerkleStore { fn from_iter>(iter: I) -> Self { let nodes = combine_nodes_with_empty_hashes(iter.into_iter()).collect(); Self { nodes } } } -impl> FromIterator<(RpoDigest, Node)> for MerkleStore { - fn from_iter>(iter: I) -> Self { +impl> FromIterator<(RpoDigest, StoreNode)> for MerkleStore { + fn from_iter>(iter: I) -> Self { let nodes = iter.into_iter().chain(empty_hashes().into_iter()).collect(); Self { nodes } } @@ -460,12 +460,12 @@ impl> FromIterator<(RpoDigest, Node)> for MerkleStore< // ITERATORS // ================================================================================================ -impl> Extend for MerkleStore { +impl> Extend for MerkleStore { fn extend>(&mut self, iter: I) { self.nodes.extend(iter.into_iter().map(|info| { ( info.value, - Node { + StoreNode { left: info.left, right: info.right, }, @@ -477,22 +477,22 @@ impl> Extend for MerkleStore { // SERIALIZATION // ================================================================================================ -impl Serializable for Node { +impl Serializable for StoreNode { fn write_into(&self, target: &mut W) { self.left.write_into(target); self.right.write_into(target); } } -impl Deserializable for Node { +impl Deserializable for StoreNode { fn read_from(source: &mut R) -> Result { let left = RpoDigest::read_from(source)?; let right = RpoDigest::read_from(source)?; - Ok(Node { left, right }) + Ok(StoreNode { left, right }) } } -impl> Serializable for MerkleStore { +impl> Serializable for MerkleStore { fn write_into(&self, target: &mut W) { target.write_u64(self.nodes.len() as u64); @@ -503,14 +503,14 @@ impl> Serializable for MerkleStore { } } -impl> Deserializable for MerkleStore { +impl> Deserializable for MerkleStore { fn read_from(source: &mut R) -> Result { let len = source.read_u64()?; - let mut nodes: Vec<(RpoDigest, Node)> = Vec::with_capacity(len as usize); + let mut nodes: Vec<(RpoDigest, StoreNode)> = Vec::with_capacity(len as usize); for _ in 0..len { let key = RpoDigest::read_from(source)?; - let value = Node::read_from(source)?; + let value = StoreNode::read_from(source)?; nodes.push((key, value)); } @@ -522,13 +522,13 @@ impl> Deserializable for MerkleStore { // ================================================================================================ /// Creates empty hashes for all the subtrees of a tree with a max depth of 255. -fn empty_hashes() -> impl IntoIterator { +fn empty_hashes() -> impl IntoIterator { let subtrees = EmptySubtreeRoots::empty_hashes(255); subtrees.iter().rev().copied().zip(subtrees.iter().rev().skip(1).copied()).map( |(child, parent)| { ( parent, - Node { + StoreNode { left: child, right: child, }, @@ -541,13 +541,13 @@ fn empty_hashes() -> impl IntoIterator { /// which includes the nodes associate with roots of empty subtrees up to a depth of 255. fn combine_nodes_with_empty_hashes( nodes: impl IntoIterator, -) -> impl Iterator { +) -> impl Iterator { nodes .into_iter() .map(|info| { ( info.value, - Node { + StoreNode { left: info.left, right: info.right, }, From 858f95d4a12f4fc922a76f5edcae07bba6a7f149 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Sun, 25 Jun 2023 01:54:34 -0700 Subject: [PATCH 19/19] chore: update changelog --- CHANGELOG.md | 6 +++++- README.md | 5 +++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f61b43..6f83519 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,8 @@ -## 0.6.0 (TBD) +## 0.6.0 (2023-06-25) + +* [BREAKING] Added support for recording capabilities for `MerkleStore` (#162). +* [BREAKING] Refactored Merkle struct APIs to use `RpoDigest` instead of `Word` (#157). +* Added initial implementation of `PartialMerkleTree` (#156). ## 0.5.0 (2023-05-26) diff --git a/README.md b/README.md index 19ffab0..b0bbdfe 100644 --- a/README.md +++ b/README.md @@ -12,10 +12,11 @@ For performance benchmarks of these hash functions and their comparison to other ## Merkle [Merkle module](./src/merkle/) provides a set of data structures related to Merkle trees. All these data structures are implemented using the RPO hash function described above. The data structures are: +* `Mmr`: a Merkle mountain range structure designed to function as an append-only log. * `MerkleTree`: a regular fully-balanced binary Merkle tree. The depth of this tree can be at most 64. * `MerklePathSet`: a collection of Merkle authentication paths all resolving to the same root. The length of the paths can be at most 64. -* `MerkleStore`: a collection of Merkle trees of different heights designed to efficiently store trees with common subtrees. -* `Mmr`: a Merkle mountain range structure designed to function as an append-only log. +* `MerkleStore`: a collection of Merkle trees of different heights designed to efficiently store trees with common subtrees. When instantiated with `RecordingMap`, a Merkle store records all accesses to the original data. +* `PartialMerkleTree`: a partial view of a Merkle tree where some sub-trees may not be known. This is similar to a collection of Merkle paths all resolving to the same root. The length of the paths can be at most 64. * `SimpleSmt`: a Sparse Merkle Tree (with no compaction), mapping 64-bit keys to 4-element values. * `TieredSmt`: a Sparse Merkle tree (with compaction), mapping 4-element keys to 4-element values.