From 867b772d9a3f894ece86c629a2adcb78f5d41b05 Mon Sep 17 00:00:00 2001 From: Victor Lopez Date: Thu, 16 Mar 2023 19:50:59 +0100 Subject: [PATCH 1/3] fix: merkle store panics on bounds Prior to this commit, the MerkleStore panicked under certain bounds. It will prevent such panics by using checked operations. ilog2, for instance, will panic when the operand is zero. However, there is a documentation rule enforcing the merkle tree to be size at least 2. If this rule is checked, then the panic is impossible. --- src/bit.rs | 4 +- src/merkle/mod.rs | 2 +- src/merkle/path_set.rs | 8 ++ src/merkle/store.rs | 181 ++++++++++++++++++++++++++--------------- 4 files changed, 127 insertions(+), 68 deletions(-) diff --git a/src/bit.rs b/src/bit.rs index a58be2b..5eb2577 100644 --- a/src/bit.rs +++ b/src/bit.rs @@ -19,7 +19,7 @@ impl BitIterator { let mask = bitmask(n); let ones = self.mask.trailing_ones(); let mask_position = ones; - self.mask ^= mask << mask_position; + self.mask ^= mask.checked_shl(mask_position).unwrap_or(0); self } @@ -31,7 +31,7 @@ impl BitIterator { let mask = bitmask(n); let ones = self.mask.leading_ones(); let mask_position = u64::BITS - ones - n; - self.mask ^= mask << mask_position; + self.mask ^= mask.checked_shl(mask_position).unwrap_or(0); self } } diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 3db75dd..550eb0c 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -35,7 +35,7 @@ pub use store::MerkleStore; // ERRORS // ================================================================================================ -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub enum MerkleError { ConflictingRoots(Vec), DepthTooSmall(u8), diff --git a/src/merkle/path_set.rs b/src/merkle/path_set.rs index 0b9d85c..b483949 100644 --- a/src/merkle/path_set.rs +++ b/src/merkle/path_set.rs @@ -57,6 +57,14 @@ impl MerklePathSet { self.total_depth } + /// Returns all the leaf indexes of this path set. + pub fn indexes(&self) -> impl Iterator + '_ { + self.paths + .keys() + .copied() + .map(|index| NodeIndex::new(self.total_depth, index)) + } + /// Returns a node at the specified index. /// /// # Errors diff --git a/src/merkle/store.rs b/src/merkle/store.rs index 5b5c88e..12c38c9 100644 --- a/src/merkle/store.rs +++ b/src/merkle/store.rs @@ -4,8 +4,8 @@ //! (leaves or internal) to live as long as necessary and without duplication, this allows the //! implementation of efficient persistent data structures use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerkleTree, NodeIndex, Rpo256, - RpoDigest, SimpleSmt, Vec, Word, + BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree, + NodeIndex, Rpo256, RpoDigest, SimpleSmt, Vec, Word, }; #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)] @@ -81,6 +81,51 @@ impl MerkleStore { Ok(self) } + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns the node at `index` rooted on the tree `root`. + /// + /// # Errors + /// + /// This will return `NodeNotInStorage` if the element is not present in the store. + pub fn get_node(&self, root: Word, index: NodeIndex) -> Result { + let mut hash: RpoDigest = root.into(); + + // Check the root is in the storage when called with `NodeIndex::root()` + self.nodes + .get(&hash) + .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?; + + for bit in index.bit_iterator().rev() { + let node = self + .nodes + .get(&hash) + .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?; + hash = if bit { node.right } else { node.left } + } + + Ok(hash.into()) + } + + /// Returns the path for the node at `index` rooted on the tree `root`. + /// + /// The path starts at the sibling of the target leaf. + /// + /// # Errors + /// + /// This will return `NodeNotInStorage` if the element is not present in the store. + pub fn get_path(&self, root: Word, mut index: NodeIndex) -> Result { + let mut path = Vec::with_capacity(index.depth().saturating_sub(1) as usize); + while index.depth() > 0 { + let sibling = index.sibling(); + index.move_up(); + let node = self.get_node(root, sibling)?; + path.push(node); + } + Ok(MerklePath::new(path)) + } + // STATE MUTATORS // -------------------------------------------------------------------------------------------- @@ -99,6 +144,10 @@ impl MerkleStore { I: IntoIterator, { let leaves: Vec<_> = leaves.into_iter().collect(); + if leaves.len() < 2 { + return Err(MerkleError::DepthTooSmall(leaves.len() as u8)); + } + let layers = leaves.len().ilog2(); let tree = MerkleTree::new(leaves)?; @@ -225,74 +274,24 @@ impl MerkleStore { Ok(roots.iter().next().unwrap().into()) } - // PUBLIC ACCESSORS - // -------------------------------------------------------------------------------------------- - - /// Returns the node at `index` rooted on the tree `root`. - /// - /// # Errors - /// - /// This will return `NodeNotInStorage` if the element is not present in the store. - pub fn get_node(&self, root: Word, index: NodeIndex) -> Result { - let mut hash: RpoDigest = root.into(); - - // Check the root is in the storage when called with `NodeIndex::root()` - self.nodes - .get(&hash) - .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?; - - for bit in index.bit_iterator().rev() { - let node = self - .nodes - .get(&hash) - .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?; - hash = if bit { node.right } else { node.left } - } - - Ok(hash.into()) - } - - /// Returns the path for the node at `index` rooted on the tree `root`. - /// - /// # Errors - /// - /// This will return `NodeNotInStorage` if the element is not present in the store. - pub fn get_path( - &self, - root: Word, - index: NodeIndex, - ) -> Result<(Word, MerklePath), MerkleError> { - let mut hash: RpoDigest = root.into(); - let mut path = Vec::new(); - let node = RpoDigest::default(); - for bit in index.bit_iterator() { - let node = self - .nodes - .get(&hash) - .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?; - - hash = if bit { - path.push(node.left.into()); - node.right - } else { - path.push(node.right.into()); - node.left - } - } - - Ok((node.into(), MerklePath::new(path))) + /// Appends the provided [MerklePathSet] into the store. + pub fn add_merkle_path_set(&mut self, path_set: &MerklePathSet) -> Result { + let root = path_set.root(); + path_set.indexes().try_fold(root, |_, index| { + let node = path_set.get_node(index)?; + let path = path_set.get_path(index)?; + self.add_merkle_path(index.value(), node, path) + }) } - // DATA MUTATORS - // -------------------------------------------------------------------------------------------- - pub fn set_node( &mut self, root: Word, index: NodeIndex, value: Word, ) -> Result { - let (current_node, path) = self.get_path(root, index)?; + let current_node = self.get_node(root, index)?; + let path = self.get_path(root, index)?; if current_node != value { self.add_merkle_path(index.value(), value, path) } else { @@ -331,9 +330,12 @@ impl MerkleStore { #[cfg(test)] mod test { - use super::{MerkleError, MerkleStore, MerkleTree, NodeIndex, SimpleSmt, Word}; - use crate::merkle::int_to_node; - use crate::merkle::MerklePathSet; + use super::*; + use crate::{ + hash::rpo::Rpo256, + merkle::{int_to_node, MerklePathSet}, + Felt, Word, + }; const KEYS4: [u64; 4] = [0, 1, 2, 3]; const LEAVES4: [Word; 4] = [ @@ -494,4 +496,53 @@ mod test { Ok(()) } + + #[test] + fn wont_open_to_different_depth_root() { + let empty = EmptySubtreeRoots::empty_hashes(64); + let a = [Felt::new(1); 4]; + let b = [Felt::new(2); 4]; + + // compute the root for a different depth + let mut root = Rpo256::merge(&[a.into(), b.into()]); + for depth in (1..=63).rev() { + root = Rpo256::merge(&[root, empty[depth]]); + } + let root = Word::from(root); + + let store = MerkleStore::default().with_merkle_tree([a, b]).unwrap(); + let index = NodeIndex::root(); + let err = store.get_node(root, index).err().unwrap(); + assert_eq!(err, MerkleError::NodeNotInStorage(root, index)); + } + + #[test] + fn store_path_opens_from_leaf() { + let a = [Felt::new(1); 4]; + let b = [Felt::new(2); 4]; + let c = [Felt::new(3); 4]; + let d = [Felt::new(4); 4]; + let e = [Felt::new(5); 4]; + let f = [Felt::new(6); 4]; + let g = [Felt::new(7); 4]; + let h = [Felt::new(8); 4]; + + 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 m = Rpo256::merge(&[i.into(), j.into()]); + let n = Rpo256::merge(&[k.into(), l.into()]); + + let root = Rpo256::merge(&[m.into(), n.into()]); + + let store = MerkleStore::default() + .with_merkle_tree([a, b, c, d, e, f, g, h]) + .unwrap(); + let path = store.get_path(root.into(), NodeIndex::new(3, 1)).unwrap(); + + let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec()); + assert_eq!(path, expected); + } } From 8cb245dc1f36ef1b1d6ce884bc130d191247f139 Mon Sep 17 00:00:00 2001 From: Victor Lopez Date: Thu, 16 Mar 2023 19:50:59 +0100 Subject: [PATCH 2/3] bugfix: reverse merkle path to match other structures The store builds the path from root to leaf, this updates the code to return a path from leaf to root, as it is done by the other structures. This also added custom error for missing root. --- src/merkle/mod.rs | 6 +- src/merkle/{store.rs => store/mod.rs} | 295 +++------------ src/merkle/store/tests.rs | 497 ++++++++++++++++++++++++++ 3 files changed, 555 insertions(+), 243 deletions(-) rename src/merkle/{store.rs => store/mod.rs} (53%) create mode 100644 src/merkle/store/tests.rs diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 550eb0c..411dd7e 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -40,13 +40,14 @@ pub enum MerkleError { ConflictingRoots(Vec), DepthTooSmall(u8), DepthTooBig(u64), - NodeNotInStorage(Word, NodeIndex), + NodeNotInStore(Word, NodeIndex), NumLeavesNotPowerOfTwo(usize), InvalidIndex(NodeIndex), InvalidDepth { expected: u8, provided: u8 }, InvalidPath(MerklePath), InvalidEntriesCount(usize, usize), NodeNotInSet(u64), + RootNotInStore(Word), } impl fmt::Display for MerkleError { @@ -70,7 +71,8 @@ impl fmt::Display for MerkleError { InvalidPath(_path) => write!(f, "the provided path is not valid"), InvalidEntriesCount(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"), NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"), - NodeNotInStorage(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the storage", hash, index.value(), index.depth(),), + NodeNotInStore(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the storage", hash, index.value(), index.depth(),), + RootNotInStore(root) => write!(f, "the root {:?} is not in the storage", root), } } } diff --git a/src/merkle/store.rs b/src/merkle/store/mod.rs similarity index 53% rename from src/merkle/store.rs rename to src/merkle/store/mod.rs index 12c38c9..961d4b8 100644 --- a/src/merkle/store.rs +++ b/src/merkle/store/mod.rs @@ -8,6 +8,9 @@ use super::{ NodeIndex, Rpo256, RpoDigest, SimpleSmt, Vec, Word, }; +#[cfg(test)] +mod tests; + #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)] pub struct Node { left: RpoDigest, @@ -88,42 +91,67 @@ impl MerkleStore { /// /// # Errors /// - /// This will return `NodeNotInStorage` if the element is not present in the store. + /// This method can return the following errors: + /// - `RootNotInStore` if the `root` is not present in the storage. + /// - `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(); - // Check the root is in the storage when called with `NodeIndex::root()` + // corner case: check the root is in the storage when called with index `NodeIndex::root()` self.nodes .get(&hash) - .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?; + .ok_or(MerkleError::RootNotInStore(hash.into()))?; for bit in index.bit_iterator().rev() { let node = self .nodes .get(&hash) - .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?; + .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?; hash = if bit { node.right } else { node.left } } Ok(hash.into()) } - /// Returns the path for the node at `index` rooted on the tree `root`. + /// Returns the node at the specified `index` and its opening to the `root`. /// /// The path starts at the sibling of the target leaf. /// /// # Errors /// - /// This will return `NodeNotInStorage` if the element is not present in the store. - pub fn get_path(&self, root: Word, mut index: NodeIndex) -> Result { - let mut path = Vec::with_capacity(index.depth().saturating_sub(1) as usize); - while index.depth() > 0 { - let sibling = index.sibling(); - index.move_up(); - let node = self.get_node(root, sibling)?; - path.push(node); + /// This method can return the following errors: + /// - `RootNotInStore` if the `root` is not present in the storage. + /// - `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<(Word, MerklePath), MerkleError> { + let mut hash: RpoDigest = root.into(); + let mut path = Vec::with_capacity(index.depth().into()); + + // corner case: check the root is in the storage when called with index `NodeIndex::root()` + self.nodes + .get(&hash) + .ok_or(MerkleError::RootNotInStore(hash.into()))?; + + for bit in index.bit_iterator().rev() { + let node = self + .nodes + .get(&hash) + .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?; + + hash = if bit { + path.push(node.left.into()); + node.right + } else { + path.push(node.right.into()); + node.left + } } - Ok(MerklePath::new(path)) + + path.reverse(); + Ok((hash.into(), MerklePath::new(path))) } // STATE MUTATORS @@ -269,8 +297,6 @@ impl MerkleStore { self.add_merkle_path(index_value, node, path)?; } - // Returns the parent of the last paths (assumes all paths have the same parent) or empty - // The length of unique_roots is checked above, so this wont panic Ok(roots.iter().next().unwrap().into()) } @@ -284,16 +310,22 @@ impl MerkleStore { }) } + /// Sets a node to `value`. + /// + /// # Errors + /// + /// This method can return the following errors: + /// - `RootNotInStore` if the `root` is not present in the storage. + /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. pub fn set_node( &mut self, root: Word, index: NodeIndex, value: Word, ) -> Result { - let current_node = self.get_node(root, index)?; - let path = self.get_path(root, index)?; - if current_node != value { - self.add_merkle_path(index.value(), value, path) + let result = self.get_path(root, index)?; + if result.0 != value { + self.add_merkle_path(index.value(), value, result.1) } else { Ok(root) } @@ -304,12 +336,12 @@ impl MerkleStore { let root2: RpoDigest = root2.into(); if !self.nodes.contains_key(&root1) { - Err(MerkleError::NodeNotInStorage( + Err(MerkleError::NodeNotInStore( root1.into(), NodeIndex::new(0, 0), )) } else if !self.nodes.contains_key(&root1) { - Err(MerkleError::NodeNotInStorage( + Err(MerkleError::NodeNotInStore( root2.into(), NodeIndex::new(0, 0), )) @@ -327,222 +359,3 @@ impl MerkleStore { } } } - -#[cfg(test)] -mod test { - use super::*; - use crate::{ - hash::rpo::Rpo256, - merkle::{int_to_node, MerklePathSet}, - Felt, Word, - }; - - const KEYS4: [u64; 4] = [0, 1, 2, 3]; - const LEAVES4: [Word; 4] = [ - int_to_node(1), - int_to_node(2), - int_to_node(3), - int_to_node(4), - ]; - - #[test] - fn test_add_merkle_tree() -> Result<(), MerkleError> { - let mut store = MerkleStore::default(); - - let mtree = MerkleTree::new(LEAVES4.to_vec())?; - store.add_merkle_tree(LEAVES4.to_vec())?; - - assert!( - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)) - .is_ok(), - "node 0 must be in the tree" - ); - assert!( - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)) - .is_ok(), - "node 1 must be in the tree" - ); - assert!( - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)) - .is_ok(), - "node 2 must be in the tree" - ); - assert!( - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)) - .is_ok(), - "node 3 must be in the tree" - ); - - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)) - .expect("node 0 must be in tree"); - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)) - .expect("node 1 must be in tree"); - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)) - .expect("node 2 must be in tree"); - store - .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)) - .expect("node 3 must be in tree"); - - Ok(()) - } - - #[test] - fn test_get_invalid_node() { - let mut store = MerkleStore::default(); - let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work"); - store - .add_merkle_tree(LEAVES4.to_vec()) - .expect("adding a merkle tree to the store must work"); - let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)); - } - - #[test] - fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { - let mut store = MerkleStore::default(); - let keys2: [u64; 2] = [0, 1]; - let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)]; - store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?; - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) - .unwrap() - .with_leaves(keys2.into_iter().zip(leaves2.into_iter())) - .unwrap(); - - let idx = NodeIndex::new(1, 0); - assert_eq!( - store.get_node(smt.root(), idx).unwrap(), - smt.get_node(&idx).unwrap() - ); - - Ok(()) - } - - #[test] - fn test_add_sparse_merkle_tree() -> Result<(), MerkleError> { - let mut store = MerkleStore::default(); - store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?; - - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) - .unwrap() - .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter())) - .unwrap(); - - let idx = NodeIndex::new(1, 0); - assert_eq!( - store.get_node(smt.root(), idx).unwrap(), - smt.get_node(&idx).unwrap() - ); - let idx = NodeIndex::new(1, 1); - assert_eq!( - store.get_node(smt.root(), idx).unwrap(), - smt.get_node(&idx).unwrap() - ); - - Ok(()) - } - - #[test] - fn test_add_merkle_paths() -> Result<(), MerkleError> { - let mut store = MerkleStore::default(); - let mtree = MerkleTree::new(LEAVES4.to_vec())?; - - let i0 = 0; - let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap(); - - let i1 = 1; - let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap(); - - let i2 = 2; - let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap(); - - let i3 = 3; - let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap(); - - let paths = [ - (i0, LEAVES4[i0 as usize], p0), - (i1, LEAVES4[i1 as usize], p1), - (i2, LEAVES4[i2 as usize], p2), - (i3, LEAVES4[i3 as usize], p3), - ]; - - store - .add_merkle_paths(paths.clone()) - .expect("the valid paths must work"); - - let set = MerklePathSet::new(3).with_paths(paths).unwrap(); - - assert_eq!( - set.get_node(NodeIndex::new(3, 0)).unwrap(), - store.get_node(set.root(), NodeIndex::new(2, 0b00)).unwrap(), - ); - assert_eq!( - set.get_node(NodeIndex::new(3, 1)).unwrap(), - store.get_node(set.root(), NodeIndex::new(2, 0b01)).unwrap(), - ); - assert_eq!( - set.get_node(NodeIndex::new(3, 2)).unwrap(), - store.get_node(set.root(), NodeIndex::new(2, 0b10)).unwrap(), - ); - assert_eq!( - set.get_node(NodeIndex::new(3, 3)).unwrap(), - store.get_node(set.root(), NodeIndex::new(2, 0b11)).unwrap(), - ); - - Ok(()) - } - - #[test] - fn wont_open_to_different_depth_root() { - let empty = EmptySubtreeRoots::empty_hashes(64); - let a = [Felt::new(1); 4]; - let b = [Felt::new(2); 4]; - - // compute the root for a different depth - let mut root = Rpo256::merge(&[a.into(), b.into()]); - for depth in (1..=63).rev() { - root = Rpo256::merge(&[root, empty[depth]]); - } - let root = Word::from(root); - - let store = MerkleStore::default().with_merkle_tree([a, b]).unwrap(); - let index = NodeIndex::root(); - let err = store.get_node(root, index).err().unwrap(); - assert_eq!(err, MerkleError::NodeNotInStorage(root, index)); - } - - #[test] - fn store_path_opens_from_leaf() { - let a = [Felt::new(1); 4]; - let b = [Felt::new(2); 4]; - let c = [Felt::new(3); 4]; - let d = [Felt::new(4); 4]; - let e = [Felt::new(5); 4]; - let f = [Felt::new(6); 4]; - let g = [Felt::new(7); 4]; - let h = [Felt::new(8); 4]; - - 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 m = Rpo256::merge(&[i.into(), j.into()]); - let n = Rpo256::merge(&[k.into(), l.into()]); - - let root = Rpo256::merge(&[m.into(), n.into()]); - - let store = MerkleStore::default() - .with_merkle_tree([a, b, c, d, e, f, g, h]) - .unwrap(); - let path = store.get_path(root.into(), NodeIndex::new(3, 1)).unwrap(); - - let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec()); - assert_eq!(path, expected); - } -} diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs new file mode 100644 index 0000000..a60adb4 --- /dev/null +++ b/src/merkle/store/tests.rs @@ -0,0 +1,497 @@ +use super::*; +use crate::{ + hash::rpo::Rpo256, + merkle::{int_to_node, MerklePathSet}, + Felt, Word, +}; + +const KEYS4: [u64; 4] = [0, 1, 2, 3]; +const LEAVES4: [Word; 4] = [ + int_to_node(1), + int_to_node(2), + int_to_node(3), + int_to_node(4), +]; + +#[test] +fn test_root_not_in_storage() -> Result<(), MerkleError> { + let mtree = MerkleTree::new(LEAVES4.to_vec())?; + let store = MerkleStore::default().with_merkle_tree(LEAVES4)?; + assert_eq!( + store.get_node(LEAVES4[0], NodeIndex::new(mtree.depth(), 0)), + Err(MerkleError::RootNotInStore(LEAVES4[0])), + "Leaf 0 is not a root" + ); + assert_eq!( + store.get_path(LEAVES4[0], NodeIndex::new(mtree.depth(), 0)), + Err(MerkleError::RootNotInStore(LEAVES4[0])), + "Leaf 0 is not a root" + ); + + Ok(()) +} + +#[test] +fn test_merkle_tree() -> Result<(), MerkleError> { + let mut store = MerkleStore::default(); + + let mtree = MerkleTree::new(LEAVES4.to_vec())?; + store.add_merkle_tree(LEAVES4.to_vec())?; + + // STORE LEAVES ARE CORRECT ============================================================== + // checks the leaves in the store corresponds to the expected values + assert_eq!( + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)), + Ok(LEAVES4[0]), + "node 0 must be in the tree" + ); + assert_eq!( + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)), + Ok(LEAVES4[1]), + "node 1 must be in the tree" + ); + assert_eq!( + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)), + Ok(LEAVES4[2]), + "node 2 must be in the tree" + ); + assert_eq!( + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)), + Ok(LEAVES4[3]), + "node 3 must be in the tree" + ); + + // STORE LEAVES MATCH TREE =============================================================== + // sanity check the values returned by the store and the tree + assert_eq!( + mtree.get_node(NodeIndex::new(mtree.depth(), 0)), + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)), + "node 0 must be the same for both MerkleTree and MerkleStore" + ); + assert_eq!( + mtree.get_node(NodeIndex::new(mtree.depth(), 1)), + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)), + "node 1 must be the same for both MerkleTree and MerkleStore" + ); + assert_eq!( + mtree.get_node(NodeIndex::new(mtree.depth(), 2)), + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)), + "node 2 must be the same for both MerkleTree and MerkleStore" + ); + assert_eq!( + mtree.get_node(NodeIndex::new(mtree.depth(), 3)), + store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)), + "node 3 must be the same for both MerkleTree and MerkleStore" + ); + + // STORE MERKLE PATH MATCHS ============================================================== + // 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::new(mtree.depth(), 0)) + .unwrap(); + assert_eq!( + LEAVES4[0], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + mtree.get_path(NodeIndex::new(mtree.depth(), 0)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 1)) + .unwrap(); + assert_eq!( + LEAVES4[1], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + mtree.get_path(NodeIndex::new(mtree.depth(), 1)), + Ok(result.1), + "merkle path for index 1 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 2)) + .unwrap(); + assert_eq!( + LEAVES4[2], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + mtree.get_path(NodeIndex::new(mtree.depth(), 2)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 3)) + .unwrap(); + assert_eq!( + LEAVES4[3], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + mtree.get_path(NodeIndex::new(mtree.depth(), 3)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + Ok(()) +} + +#[test] +fn test_get_invalid_node() { + let mut store = MerkleStore::default(); + let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work"); + store + .add_merkle_tree(LEAVES4.to_vec()) + .expect("adding a merkle tree to the store must work"); + let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)); +} + +#[test] +fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { + let mut store = MerkleStore::default(); + let keys2: [u64; 2] = [0, 1]; + let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)]; + store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?; + let smt = SimpleSmt::new(1) + .unwrap() + .with_leaves(keys2.into_iter().zip(leaves2.into_iter())) + .unwrap(); + + let idx = NodeIndex::new(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() + ); + + let idx = NodeIndex::new(1, 1); + assert_eq!(smt.get_node(&idx).unwrap(), leaves2[1]); + assert_eq!( + store.get_node(smt.root(), idx).unwrap(), + smt.get_node(&idx).unwrap() + ); + + Ok(()) +} + +#[test] +fn test_sparse_merkle_tree() -> Result<(), MerkleError> { + let mut store = MerkleStore::default(); + store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?; + + let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) + .unwrap() + .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter())) + .unwrap(); + + // STORE LEAVES ARE CORRECT ============================================================== + // checks the leaves in the store corresponds to the expected values + assert_eq!( + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 0)), + Ok(LEAVES4[0]), + "node 0 must be in the tree" + ); + assert_eq!( + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 1)), + Ok(LEAVES4[1]), + "node 1 must be in the tree" + ); + assert_eq!( + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 2)), + Ok(LEAVES4[2]), + "node 2 must be in the tree" + ); + assert_eq!( + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 3)), + Ok(LEAVES4[3]), + "node 3 must be in the tree" + ); + + // STORE LEAVES MATCH TREE =============================================================== + // sanity check the values returned by the store and the tree + assert_eq!( + smt.get_node(&NodeIndex::new(smt.depth(), 0)), + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 0)), + "node 0 must be the same for both SparseMerkleTree and MerkleStore" + ); + assert_eq!( + smt.get_node(&NodeIndex::new(smt.depth(), 1)), + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 1)), + "node 1 must be the same for both SparseMerkleTree and MerkleStore" + ); + assert_eq!( + smt.get_node(&NodeIndex::new(smt.depth(), 2)), + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 2)), + "node 2 must be the same for both SparseMerkleTree and MerkleStore" + ); + assert_eq!( + smt.get_node(&NodeIndex::new(smt.depth(), 3)), + store.get_node(smt.root(), NodeIndex::new(smt.depth(), 3)), + "node 3 must be the same for both SparseMerkleTree and MerkleStore" + ); + + // STORE MERKLE PATH MATCHS ============================================================== + // 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::new(smt.depth(), 0)) + .unwrap(); + assert_eq!( + LEAVES4[0], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + smt.get_path(NodeIndex::new(smt.depth(), 0)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(smt.root(), NodeIndex::new(smt.depth(), 1)) + .unwrap(); + assert_eq!( + LEAVES4[1], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + smt.get_path(NodeIndex::new(smt.depth(), 1)), + Ok(result.1), + "merkle path for index 1 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(smt.root(), NodeIndex::new(smt.depth(), 2)) + .unwrap(); + assert_eq!( + LEAVES4[2], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + smt.get_path(NodeIndex::new(smt.depth(), 2)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(smt.root(), NodeIndex::new(smt.depth(), 3)) + .unwrap(); + assert_eq!( + LEAVES4[3], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + smt.get_path(NodeIndex::new(smt.depth(), 3)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + Ok(()) +} + +#[test] +fn test_add_merkle_paths() -> Result<(), MerkleError> { + let mtree = MerkleTree::new(LEAVES4.to_vec())?; + + let i0 = 0; + let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap(); + + let i1 = 1; + let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap(); + + let i2 = 2; + let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap(); + + let i3 = 3; + let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap(); + + let paths = [ + (i0, LEAVES4[i0 as usize], p0), + (i1, LEAVES4[i1 as usize], p1), + (i2, LEAVES4[i2 as usize], p2), + (i3, LEAVES4[i3 as usize], p3), + ]; + + let mut store = MerkleStore::default(); + store + .add_merkle_paths(paths.clone()) + .expect("the valid paths must work"); + + let depth = 3; + let set = MerklePathSet::new(depth).with_paths(paths).unwrap(); + + // STORE LEAVES ARE CORRECT ============================================================== + // checks the leaves in the store corresponds to the expected values + assert_eq!( + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 0)), + Ok(LEAVES4[0]), + "node 0 must be in the set" + ); + assert_eq!( + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 1)), + Ok(LEAVES4[1]), + "node 1 must be in the set" + ); + assert_eq!( + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 2)), + Ok(LEAVES4[2]), + "node 2 must be in the set" + ); + assert_eq!( + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 3)), + Ok(LEAVES4[3]), + "node 3 must be in the set" + ); + + // STORE LEAVES MATCH SET ================================================================ + // sanity check the values returned by the store and the set + assert_eq!( + set.get_node(NodeIndex::new(set.depth(), 0)), + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 0)), + "node 0 must be the same for both SparseMerkleTree and MerkleStore" + ); + assert_eq!( + set.get_node(NodeIndex::new(set.depth(), 1)), + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 1)), + "node 1 must be the same for both SparseMerkleTree and MerkleStore" + ); + assert_eq!( + set.get_node(NodeIndex::new(set.depth(), 2)), + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 2)), + "node 2 must be the same for both SparseMerkleTree and MerkleStore" + ); + assert_eq!( + set.get_node(NodeIndex::new(set.depth(), 3)), + store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 3)), + "node 3 must be the same for both SparseMerkleTree and MerkleStore" + ); + + // STORE MERKLE PATH MATCHS ============================================================== + // 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::new(set.depth() - 1, 0)) + .unwrap(); + assert_eq!( + LEAVES4[0], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + set.get_path(NodeIndex::new(set.depth(), 0)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(set.root(), NodeIndex::new(set.depth() - 1, 1)) + .unwrap(); + assert_eq!( + LEAVES4[1], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + set.get_path(NodeIndex::new(set.depth(), 1)), + Ok(result.1), + "merkle path for index 1 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(set.root(), NodeIndex::new(set.depth() - 1, 2)) + .unwrap(); + assert_eq!( + LEAVES4[2], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + set.get_path(NodeIndex::new(set.depth(), 2)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + let result = store + .get_path(set.root(), NodeIndex::new(set.depth() - 1, 3)) + .unwrap(); + assert_eq!( + LEAVES4[3], result.0, + "Value for merkle path at index 0 must match leaf value" + ); + assert_eq!( + set.get_path(NodeIndex::new(set.depth(), 3)), + Ok(result.1), + "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" + ); + + Ok(()) +} + +#[test] +fn wont_open_to_different_depth_root() { + let empty = EmptySubtreeRoots::empty_hashes(64); + let a = [Felt::new(1); 4]; + let b = [Felt::new(2); 4]; + + // Compute the root for a different depth. We cherry-pick this specific depth to prevent a + // regression to a bug in the past that allowed the user to fetch a node at a depth lower than + // the inserted path of a Merkle tree. + let mut root = Rpo256::merge(&[a.into(), b.into()]); + for depth in (1..=63).rev() { + root = Rpo256::merge(&[root, empty[depth]]); + } + let root = Word::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 + // exist for the set. + let store = MerkleStore::default().with_merkle_tree([a, b]).unwrap(); + let index = NodeIndex::root(); + let err = store.get_node(root, index).err().unwrap(); + assert_eq!(err, MerkleError::RootNotInStore(root)); +} + +#[test] +fn store_path_opens_from_leaf() { + let a = [Felt::new(1); 4]; + let b = [Felt::new(2); 4]; + let c = [Felt::new(3); 4]; + let d = [Felt::new(4); 4]; + let e = [Felt::new(5); 4]; + let f = [Felt::new(6); 4]; + let g = [Felt::new(7); 4]; + let h = [Felt::new(8); 4]; + + 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 m = Rpo256::merge(&[i.into(), j.into()]); + let n = Rpo256::merge(&[k.into(), l.into()]); + + let root = Rpo256::merge(&[m.into(), n.into()]); + + let store = MerkleStore::default() + .with_merkle_tree([a, b, c, d, e, f, g, h]) + .unwrap(); + let path = store.get_path(root.into(), NodeIndex::new(3, 1)).unwrap(); + + let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec()); + assert_eq!(path.1, expected); +} + +#[test] +fn test_set_node() -> Result<(), MerkleError> { + let mtree = MerkleTree::new(LEAVES4.to_vec())?; + let mut store = MerkleStore::default().with_merkle_tree(LEAVES4)?; + let value = int_to_node(42); + let index = NodeIndex::new(mtree.depth(), 0); + let new_root = store.set_node(mtree.root(), index, value)?; + assert_eq!( + store.get_node(new_root, index), + Ok(value), + "Value must have changed" + ); + + Ok(()) +} From 17eb8d78d3ec9d4fa5e501f2c104cdcff77fffe9 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Mon, 20 Mar 2023 21:55:57 +0100 Subject: [PATCH 3/3] chore: storage -> store --- src/merkle/mod.rs | 4 ++-- src/merkle/store/mod.rs | 18 +++++++++--------- src/merkle/store/tests.rs | 2 +- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 411dd7e..0afbce9 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -71,8 +71,8 @@ impl fmt::Display for MerkleError { InvalidPath(_path) => write!(f, "the provided path is not valid"), InvalidEntriesCount(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"), NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"), - NodeNotInStore(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the storage", hash, index.value(), index.depth(),), - RootNotInStore(root) => write!(f, "the root {:?} is not in the storage", root), + NodeNotInStore(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the store", hash, index.value(), index.depth(),), + RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root), } } } diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 961d4b8..5610429 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -92,12 +92,12 @@ impl MerkleStore { /// # Errors /// /// This method can return the following errors: - /// - `RootNotInStore` if the `root` is not present in the storage. + /// - `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(); - // corner case: check the root is in the storage when called with index `NodeIndex::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()))?; @@ -120,7 +120,7 @@ impl MerkleStore { /// # Errors /// /// This method can return the following errors: - /// - `RootNotInStore` if the `root` is not present in the storage. + /// - `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, @@ -130,7 +130,7 @@ impl MerkleStore { let mut hash: RpoDigest = root.into(); let mut path = Vec::with_capacity(index.depth().into()); - // corner case: check the root is in the storage when called with index `NodeIndex::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()))?; @@ -160,7 +160,7 @@ impl MerkleStore { /// Adds all the nodes of a Merkle tree represented by `leaves`. /// /// This will instantiate a Merkle tree using `leaves` and include all the nodes into the - /// storage. + /// store. /// /// # Errors /// @@ -208,7 +208,7 @@ impl MerkleStore { /// Adds all the nodes of a Sparse Merkle tree represented by `entries`. /// /// This will instantiate a Sparse Merkle tree using `entries` and include all the nodes into - /// the storage. + /// the store. /// /// # Errors /// @@ -236,7 +236,7 @@ impl MerkleStore { /// Adds all the nodes of a Merkle path represented by `path`. /// /// This will compute the sibling elements determined by the Merkle `path` and `node`, and - /// include all the nodes into the storage. + /// include all the nodes into the store. pub fn add_merkle_path( &mut self, index_value: u64, @@ -270,7 +270,7 @@ impl MerkleStore { /// Adds all the nodes of multiple Merkle paths into the store. /// /// This will compute the sibling elements for each Merkle `path` and include all the nodes - /// into the storage. + /// into the store. /// /// # Errors /// @@ -315,7 +315,7 @@ impl MerkleStore { /// # Errors /// /// This method can return the following errors: - /// - `RootNotInStore` if the `root` is not present in the storage. + /// - `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 set_node( &mut self, diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index a60adb4..c469553 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -14,7 +14,7 @@ const LEAVES4: [Word; 4] = [ ]; #[test] -fn test_root_not_in_storage() -> Result<(), MerkleError> { +fn test_root_not_in_store() -> Result<(), MerkleError> { let mtree = MerkleTree::new(LEAVES4.to_vec())?; let store = MerkleStore::default().with_merkle_tree(LEAVES4)?; assert_eq!(