From fcf03478ba1cb17ea15c35b6e833002486bb1442 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Sun, 7 May 2023 00:21:18 -0700 Subject: [PATCH] refactor: update SimpleSmt interfaces --- benches/smt.rs | 14 ++- benches/store.rs | 24 ++--- src/merkle/empty_roots.rs | 10 +- src/merkle/mod.rs | 9 +- src/merkle/simple_smt/mod.rs | 189 +++++++++++++++++---------------- src/merkle/simple_smt/tests.rs | 167 +++++++++-------------------- src/merkle/store/tests.rs | 17 +-- 7 files changed, 183 insertions(+), 247 deletions(-) diff --git a/benches/smt.rs b/benches/smt.rs index 3c63c04..44e3ea5 100644 --- a/benches/smt.rs +++ b/benches/smt.rs @@ -18,8 +18,8 @@ fn smt_rpo(c: &mut Criterion) { (i, word) }) .collect(); - let tree = SimpleSmt::new(depth).unwrap().with_leaves(entries).unwrap(); - trees.push(tree); + let tree = SimpleSmt::with_leaves(depth, entries).unwrap(); + trees.push((tree, count)); } } @@ -29,10 +29,9 @@ fn smt_rpo(c: &mut Criterion) { let mut insert = c.benchmark_group(format!("smt update_leaf")); - for tree in trees.iter_mut() { + for (tree, count) in trees.iter_mut() { let depth = tree.depth(); - let count = tree.leaves_count() as u64; - let key = count >> 2; + let key = *count >> 2; insert.bench_with_input( format!("simple smt(depth:{depth},count:{count})"), &(key, leaf), @@ -48,10 +47,9 @@ fn smt_rpo(c: &mut Criterion) { let mut path = c.benchmark_group(format!("smt get_leaf_path")); - for tree in trees.iter_mut() { + for (tree, count) in trees.iter_mut() { let depth = tree.depth(); - let count = tree.leaves_count() as u64; - let key = count >> 2; + let key = *count >> 2; path.bench_with_input( format!("simple smt(depth:{depth},count:{count})"), &key, diff --git a/benches/store.rs b/benches/store.rs index 3ad79c4..aaa9806 100644 --- a/benches/store.rs +++ b/benches/store.rs @@ -104,10 +104,7 @@ fn get_leaf_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) - .unwrap() - .with_leaves(smt_leaves.clone()) - .unwrap(); + let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap(); let store = MerkleStore::from(&smt); let depth = smt.depth(); let root = smt.root(); @@ -215,10 +212,7 @@ fn get_node_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) - .unwrap() - .with_leaves(smt_leaves.clone()) - .unwrap(); + let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap(); let store = MerkleStore::from(&smt); let root = smt.root(); let half_depth = smt.depth() / 2; @@ -292,10 +286,7 @@ fn get_leaf_path_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) - .unwrap() - .with_leaves(smt_leaves.clone()) - .unwrap(); + let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap(); let store = MerkleStore::from(&smt); let depth = smt.depth(); let root = smt.root(); @@ -361,7 +352,7 @@ fn new(c: &mut Criterion) { .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>() }, - |l| black_box(SimpleSmt::new(SimpleSmt::MAX_DEPTH).unwrap().with_leaves(l)), + |l| black_box(SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, l)), BatchSize::SmallInput, ) }); @@ -376,7 +367,7 @@ fn new(c: &mut Criterion) { .collect::>() }, |l| { - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH).unwrap().with_leaves(l).unwrap(); + let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, l).unwrap(); black_box(MerkleStore::from(&smt)); }, BatchSize::SmallInput, @@ -442,10 +433,7 @@ fn update_leaf_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let mut smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) - .unwrap() - .with_leaves(smt_leaves.clone()) - .unwrap(); + let mut smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap(); let mut store = MerkleStore::from(&smt); let depth = smt.depth(); let root = smt.root(); diff --git a/src/merkle/empty_roots.rs b/src/merkle/empty_roots.rs index 7f8c4a2..b1b0b30 100644 --- a/src/merkle/empty_roots.rs +++ b/src/merkle/empty_roots.rs @@ -1,6 +1,12 @@ -use super::{Felt, RpoDigest, WORD_SIZE, ZERO}; +use super::{Felt, RpoDigest, Word, WORD_SIZE, ZERO}; use core::slice; +// CONSTANTS +// ================================================================================================ + +/// A word consisting of 4 ZERO elements. +pub const EMPTY_WORD: Word = [ZERO; WORD_SIZE]; + // EMPTY NODES SUBTREES // ================================================================================================ @@ -1570,7 +1576,7 @@ fn all_depths_opens_to_zero() { assert_eq!(depth as usize + 1, subtree.len()); // assert the opening is zero - let initial = RpoDigest::new([ZERO; WORD_SIZE]); + let initial = RpoDigest::new(EMPTY_WORD); assert_eq!(initial, subtree.remove(0)); // compute every node of the path manually and compare with the output diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index aca3dde..d894e9f 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -1,6 +1,6 @@ use super::{ hash::rpo::{Rpo256, RpoDigest}, - utils::collections::{vec, BTreeMap, Vec}, + utils::collections::{vec, BTreeMap, BTreeSet, Vec}, Felt, StarkField, Word, WORD_SIZE, ZERO, }; use core::fmt; @@ -10,6 +10,7 @@ use core::fmt; mod empty_roots; pub use empty_roots::EmptySubtreeRoots; +use empty_roots::EMPTY_WORD; mod index; pub use index::NodeIndex; @@ -43,12 +44,13 @@ pub enum MerkleError { ConflictingRoots(Vec), DepthTooSmall(u8), DepthTooBig(u64), + DuplicateValuesForKey(u64), NodeNotInStore(Word, NodeIndex), NumLeavesNotPowerOfTwo(usize), InvalidIndex { depth: u8, value: u64 }, InvalidDepth { expected: u8, provided: u8 }, InvalidPath(MerklePath), - InvalidEntriesCount(usize, usize), + InvalidNumEntries(usize, usize), NodeNotInSet(u64), RootNotInStore(Word), } @@ -60,6 +62,7 @@ impl fmt::Display for MerkleError { ConflictingRoots(roots) => write!(f, "the merkle paths roots do not match {roots:?}"), DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"), DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"), + DuplicateValuesForKey(key) => write!(f, "multiple values provided for key {key}"), NumLeavesNotPowerOfTwo(leaves) => { write!(f, "the leaves count {leaves} is not a power of 2") } @@ -72,7 +75,7 @@ impl fmt::Display for MerkleError { "the provided depth {provided} is not valid for {expected}" ), 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}"), + InvalidNumEntries(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 store", hash, index.value(), index.depth(),), RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root), diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 1f29fe0..6b0a6c2 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -1,6 +1,6 @@ use super::{ - BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, - RpoDigest, Vec, Word, + BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, + Rpo256, RpoDigest, Vec, Word, EMPTY_WORD, }; #[cfg(test)] @@ -10,6 +10,7 @@ mod tests; // ================================================================================================ /// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction. +/// /// The root of the tree is recomputed on each new leaf update. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SimpleSmt { @@ -20,18 +21,6 @@ pub struct SimpleSmt { empty_hashes: Vec, } -#[derive(Debug, Default, Clone, PartialEq, Eq)] -struct BranchNode { - left: RpoDigest, - right: RpoDigest, -} - -impl BranchNode { - fn parent(&self) -> RpoDigest { - Rpo256::merge(&[self.left, self.right]) - } -} - impl SimpleSmt { // CONSTANTS // -------------------------------------------------------------------------------------------- @@ -45,7 +34,12 @@ impl SimpleSmt { // CONSTRUCTORS // -------------------------------------------------------------------------------------------- - /// Creates a new simple SMT with the provided depth. + /// Returns a new [SimpleSmt] instantiated with the specified depth. + /// + /// All leaves in the returned tree are set to [ZERO; 4]. + /// + /// # Errors + /// Returns an error if the depth is 0 or is greater than 64. pub fn new(depth: u8) -> Result { // validate the range of the depth. if depth < Self::MIN_DEPTH { @@ -66,36 +60,47 @@ impl SimpleSmt { }) } - /// Appends the provided entries as leaves of the tree. + /// Returns a new [SimpleSmt] instantiated with the specified depth and with leaves + /// set as specified by the provided entries. /// - /// # Errors + /// All leaves omitted from the entries list are set to [ZERO; 4]. /// - /// The function will fail if the provided entries count exceed the maximum tree capacity, that - /// is `2^{depth}`. - pub fn with_leaves(mut self, entries: R) -> Result + /// # 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(depth: u8, entries: R) -> Result where R: IntoIterator, I: Iterator + ExactSizeIterator, { - // check if the leaves count will fit the depth setup - let mut entries = entries.into_iter(); - let max = 1 << self.depth.min(63); + // create an empty tree + let mut tree = Self::new(depth)?; + + // 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 << tree.depth.min(63); if entries.len() > max { - return Err(MerkleError::InvalidEntriesCount(max, entries.len())); + return Err(MerkleError::InvalidNumEntries(max, entries.len())); } - // append leaves and return - entries.try_for_each(|(key, leaf)| self.insert_leaf(key, leaf))?; - Ok(self) - } - - /// Replaces the internal empty digests used when a given depth doesn't contain a node. - pub fn with_empty_subtrees(mut self, hashes: I) -> Self - where - I: IntoIterator, - { - self.replace_empty_subtrees(hashes.into_iter().collect()); - self + // append leaves to the tree returning an error if a duplicate entry for the same key + // is found + 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) { + 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) { + return Err(MerkleError::DuplicateValuesForKey(key)); + } + } + Ok(tree) } // PUBLIC ACCESSORS @@ -111,40 +116,43 @@ impl SimpleSmt { self.depth } - // PROVIDERS - // -------------------------------------------------------------------------------------------- - - /// Returns the set count of the keys of the leaves. - pub fn leaves_count(&self) -> usize { - self.leaves.len() - } - /// Returns a node at the specified index. /// /// # Errors - /// Returns an error if: - /// * The specified depth is greater than the depth of the tree. + /// 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 { if index.is_root() { Err(MerkleError::DepthTooSmall(index.depth())) } else if index.depth() > self.depth() { Err(MerkleError::DepthTooBig(index.depth() as u64)) } else if index.depth() == self.depth() { - self.get_leaf_node(index.value()) - .or_else(|| self.empty_hashes.get(index.depth() as usize).copied().map(Word::from)) - .ok_or(MerkleError::NodeNotInSet(index.value())) + // 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())) } else { - let branch_node = self.get_branch_node(&index); - Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into()) + Ok(self.get_branch_node(&index).parent().into()) } } - /// Returns a Merkle path from the node at the specified key to the root. The node itself is - /// not included in the path. + /// Returns a value of the leaf at the specified index. /// /// # Errors - /// Returns an error if: - /// * The specified depth is greater than the depth of the tree. + /// 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) + } + + /// 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. pub fn get_path(&self, mut index: NodeIndex) -> Result { if index.is_root() { return Err(MerkleError::DepthTooSmall(index.depth())); @@ -163,18 +171,18 @@ impl SimpleSmt { Ok(path.into()) } - /// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not - /// included in the path. + /// Return a Merkle path from the leaf at the specified index to the root. + /// + /// The leaf itself is not included in the path. /// /// # Errors - /// Returns an error if: - /// * The specified key does not exist as a leaf node. - pub fn get_leaf_path(&self, key: u64) -> Result { - let index = NodeIndex::new(self.depth(), key)?; + /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}. + pub fn get_leaf_path(&self, index: u64) -> Result { + let index = NodeIndex::new(self.depth(), index)?; self.get_path(index) } - /// Iterator over the inner nodes of the [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(), @@ -186,27 +194,21 @@ impl SimpleSmt { // STATE MUTATORS // -------------------------------------------------------------------------------------------- - /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the - /// tree. + /// 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 specified key is not a valid leaf index for this tree. - pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> { - let index = NodeIndex::new(self.depth(), key)?; - if !self.check_leaf_node_exists(key) { - return Err(MerkleError::NodeNotInSet(index.value())); - } - self.insert_leaf(key, value)?; - - Ok(()) - } + /// 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); - /// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree - pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> { - self.insert_leaf_node(key, value); + // if the old value and new value are the same, there is nothing to update + if value == old_value { + return Ok(value); + } - // TODO consider using a map `index |-> word` instead of `index |-> (word, word)` - let mut index = NodeIndex::new(self.depth(), key)?; + let mut index = NodeIndex::new(self.depth(), index)?; let mut value = RpoDigest::from(value); for _ in 0..index.depth() { let is_right = index.is_value_odd(); @@ -217,26 +219,18 @@ impl SimpleSmt { value = Rpo256::merge(&[left, right]); } self.root = value.into(); - Ok(()) + Ok(old_value) } // HELPER METHODS // -------------------------------------------------------------------------------------------- - fn replace_empty_subtrees(&mut self, hashes: Vec) { - self.empty_hashes = hashes; - } - - fn check_leaf_node_exists(&self, key: u64) -> bool { - self.leaves.contains_key(&key) - } - fn get_leaf_node(&self, key: u64) -> Option { self.leaves.get(&key).copied() } - fn insert_leaf_node(&mut self, key: u64, node: Word) { - self.leaves.insert(key, node); + fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option { + self.leaves.insert(key, node) } fn get_branch_node(&self, index: &NodeIndex) -> BranchNode { @@ -254,3 +248,18 @@ impl SimpleSmt { self.branches.insert(index, branch); } } + +// BRANCH NODE +// ================================================================================================ + +#[derive(Debug, Default, Clone, PartialEq, Eq)] +struct BranchNode { + left: RpoDigest, + right: RpoDigest, +} + +impl BranchNode { + fn parent(&self) -> RpoDigest { + Rpo256::merge(&[self.left, self.right]) + } +} diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 582e003..0174e15 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -1,9 +1,10 @@ use super::{ super::{int_to_node, InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt}, - NodeIndex, Rpo256, Vec, Word, + NodeIndex, Rpo256, Vec, Word, EMPTY_WORD, }; -use proptest::prelude::*; -use rand_utils::prng_array; + +// TEST DATA +// ================================================================================================ const KEYS4: [u64; 4] = [0, 1, 2, 3]; const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7]; @@ -23,25 +24,17 @@ const VALUES8: [Word; 8] = [ const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8]; +// TESTS +// ================================================================================================ + #[test] 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()); } -#[test] -fn empty_digests_are_consistent() { - let depth = 5; - let root = SimpleSmt::new(depth).unwrap().root(); - let computed: [RpoDigest; 2] = (0..depth).fold([Default::default(); 2], |state, _| { - let digest = Rpo256::merge(&state); - [digest; 2] - }); - - assert_eq!(Word::from(computed[0]), root); -} - #[test] fn build_sparse_tree() { let mut smt = SimpleSmt::new(3).unwrap(); @@ -51,80 +44,59 @@ fn build_sparse_tree() { let key = 6; let new_node = int_to_node(7); values[key as usize] = new_node; - smt.insert_leaf(key, new_node).expect("Failed to insert leaf"); + 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.get_path(NodeIndex::make(3, 6)).unwrap(), smt.get_path(NodeIndex::make(3, 6)).unwrap() ); + assert_eq!(old_value, EMPTY_WORD); // insert second value at distinct leaf branch let key = 2; let new_node = int_to_node(3); values[key as usize] = new_node; - smt.insert_leaf(key, new_node).expect("Failed to insert leaf"); + 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.get_path(NodeIndex::make(3, 2)).unwrap(), smt.get_path(NodeIndex::make(3, 2)).unwrap() ); + assert_eq!(old_value, EMPTY_WORD); } #[test] -fn build_full_tree() { - let tree = SimpleSmt::new(2) - .unwrap() - .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter())) - .unwrap(); +fn test_depth2_tree() { + let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap(); + // check internal structure let (root, node2, node3) = compute_internal_nodes(); assert_eq!(root, tree.root()); assert_eq!(node2, tree.get_node(NodeIndex::make(1, 0)).unwrap()); assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap()); -} - -#[test] -fn get_values() { - let tree = SimpleSmt::new(2) - .unwrap() - .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter())) - .unwrap(); - // check depth 2 + // 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()); -} -#[test] -fn get_path() { - let tree = SimpleSmt::new(2) - .unwrap() - .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter())) - .unwrap(); - - let (_, node2, node3) = compute_internal_nodes(); - - // check depth 2 + // 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()); - // check depth 1 + // check get_path(): depth 1 assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap()); assert_eq!(vec![node2], *tree.get_path(NodeIndex::make(1, 1)).unwrap()); } #[test] -fn test_parent_node_iterator() -> Result<(), MerkleError> { - let tree = SimpleSmt::new(2) - .unwrap() - .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter())) - .unwrap(); +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()); @@ -166,35 +138,28 @@ fn test_parent_node_iterator() -> Result<(), MerkleError> { #[test] fn update_leaf() { - let mut tree = SimpleSmt::new(3) - .unwrap() - .with_leaves(KEYS8.into_iter().zip(VALUES8.into_iter())) - .unwrap(); + let mut tree = SimpleSmt::with_leaves(3, KEYS8.into_iter().zip(VALUES8.into_iter())).unwrap(); // update one value let key = 3; let new_node = int_to_node(9); let mut expected_values = VALUES8.to_vec(); expected_values[key] = new_node; - let expected_tree = SimpleSmt::new(3) - .unwrap() - .with_leaves(KEYS8.into_iter().zip(expected_values.clone().into_iter())) - .unwrap(); + let expected_tree = MerkleTree::new(expected_values.clone()).unwrap(); - tree.update_leaf(key as u64, new_node).unwrap(); - assert_eq!(expected_tree.root, tree.root); + 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]); // update another value let key = 6; let new_node = int_to_node(10); expected_values[key] = new_node; - let expected_tree = SimpleSmt::new(3) - .unwrap() - .with_leaves(KEYS8.into_iter().zip(expected_values.into_iter())) - .unwrap(); + let expected_tree = MerkleTree::new(expected_values.clone()).unwrap(); - tree.update_leaf(key as u64, new_node).unwrap(); - assert_eq!(expected_tree.root, tree.root); + 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]); } #[test] @@ -226,7 +191,7 @@ fn small_tree_opening_is_consistent() { let depth = 3; let entries = vec![(0, a), (1, b), (4, c), (7, d)]; - let tree = SimpleSmt::new(depth).unwrap().with_leaves(entries).unwrap(); + let tree = SimpleSmt::with_leaves(depth, entries).unwrap(); assert_eq!(tree.root(), Word::from(k)); @@ -250,56 +215,30 @@ fn small_tree_opening_is_consistent() { } } -proptest! { - #[test] - fn arbitrary_openings_single_leaf( - depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH, - key in prop::num::u64::ANY, - leaf in prop::num::u64::ANY, - ) { - let mut tree = SimpleSmt::new(depth).unwrap(); - - let key = key % (1 << depth as u64); - let leaf = int_to_node(leaf); - - tree.insert_leaf(key, leaf.into()).unwrap(); - tree.get_leaf_path(key).unwrap(); - - // traverse to root, fetching all paths - for d in 1..depth { - let k = key >> (depth - d); - tree.get_path(NodeIndex::make(d, k)).unwrap(); - } - } +#[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 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 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 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 smt = SimpleSmt::with_leaves(64, entries); + assert!(smt.is_err()); +} - #[test] - fn arbitrary_openings_multiple_leaves( - depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH, - count in 2u8..10u8, - ref seed in any::<[u8; 32]>() - ) { - let mut tree = SimpleSmt::new(depth).unwrap(); - let mut seed = *seed; - let leaves = (1 << depth) - 1; - - for _ in 0..count { - seed = prng_array(seed); - - let mut key = [0u8; 8]; - let mut leaf = [0u8; 8]; - - key.copy_from_slice(&seed[..8]); - leaf.copy_from_slice(&seed[8..16]); - - let key = u64::from_le_bytes(key); - let key = key % leaves; - let leaf = u64::from_le_bytes(leaf); - let leaf = int_to_node(leaf); - - tree.insert_leaf(key, leaf).unwrap(); - tree.get_leaf_path(key).unwrap(); - } - } +#[test] +fn with_no_duplicates_empty_node() { + let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2))]; + let smt = SimpleSmt::with_leaves(64, entries); + assert!(smt.is_ok()); } // HELPER FUNCTIONS diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 02e12b2..3c2c3e0 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -183,10 +183,7 @@ fn test_get_invalid_node() { 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 smt = SimpleSmt::new(1) - .unwrap() - .with_leaves(keys2.into_iter().zip(leaves2.into_iter())) - .unwrap(); + 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); @@ -202,10 +199,9 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { #[test] fn test_sparse_merkle_tree() -> Result<(), MerkleError> { - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) - .unwrap() - .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter())) - .unwrap(); + let smt = + SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, KEYS4.into_iter().zip(LEAVES4.into_iter())) + .unwrap(); let store = MerkleStore::from(&smt); @@ -526,10 +522,7 @@ fn test_constructors() -> Result<(), MerkleError> { } let depth = 32; - let smt = SimpleSmt::new(depth) - .unwrap() - .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter())) - .unwrap(); + let smt = SimpleSmt::with_leaves(depth, KEYS4.into_iter().zip(LEAVES4.into_iter())).unwrap(); let store = MerkleStore::from(&smt); let depth = smt.depth();