refactor: update SimpleSmt interfaces

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();
+            let tree = SimpleSmt::with_leaves(depth, entries).unwrap();
-            trees.push(tree);
+            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,

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::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
+        let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
-            .unwrap()
-            .with_leaves(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::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
+        let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
-            .unwrap()
-            .with_leaves(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::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
+        let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
-            .unwrap()
-            .with_leaves(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::<Vec<(u64, Word)>>()
                 },
-                |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::<Vec<(u64, Word)>>()
                 },
                 |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::<Vec<(u64, Word)>>();
-        let mut smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
+        let mut smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
-            .unwrap()
-            .with_leaves(smt_leaves.clone())
-            .unwrap();
         let mut store = MerkleStore::from(&smt);
         let depth = smt.depth();
         let root = smt.root();

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

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<Word>),
     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),

src/merkle/simple_smt/mod.rs

@@ -1,6 +1,6 @@
 use super::{
-    BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256,
+    BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex,
-    RpoDigest, Vec, Word,
+    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<RpoDigest>,
 }
 
-#[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<Self, MerkleError> {
         // 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.
+    ///
+    /// All leaves omitted from the entries list are set to [ZERO; 4].
     ///
     /// # Errors
-    ///
+    /// Returns an error if:
-    /// The function will fail if the provided entries count exceed the maximum tree capacity, that
+    /// - If the depth is 0 or is greater than 64.
-    /// is `2^{depth}`.
+    /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
-    pub fn with_leaves<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
+    /// - The provided entries contain multiple values for the same key.
+    pub fn with_leaves<R, I>(depth: u8, entries: R) -> Result<Self, MerkleError>
     where
         R: IntoIterator<IntoIter = I>,
         I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
     {
-        // check if the leaves count will fit the depth setup
+        // create an empty tree
-        let mut entries = entries.into_iter();
+        let mut tree = Self::new(depth)?;
-        let max = 1 << self.depth.min(63);
+
+        // 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
+        // append leaves to the tree returning an error if a duplicate entry for the same key
-        entries.try_for_each(|(key, leaf)| self.insert_leaf(key, leaf))?;
+        // is found
-        Ok(self)
+        let mut empty_entries = BTreeSet::new();
-    }
+        for (key, value) in entries {
-
+            let old_value = tree.update_leaf(key, value)?;
-    /// Replaces the internal empty digests used when a given depth doesn't contain a node.
+            if old_value != EMPTY_WORD || empty_entries.contains(&key) {
-    pub fn with_empty_subtrees<I>(mut self, hashes: I) -> Self
+                return Err(MerkleError::DuplicateValuesForKey(key));
-    where
+            }
-        I: IntoIterator<Item = RpoDigest>,
+            // 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
-        self.replace_empty_subtrees(hashes.into_iter().collect());
+            if value == EMPTY_WORD && !empty_entries.insert(key) {
-        self
+                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:
+    /// Returns an error if the specified index has depth set to 0 or the depth is greater than
-    /// * The specified depth is greater than the depth of the tree.
+    /// the depth of this Merkle tree.
     pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
         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())
+            // the lookup in empty_hashes could fail only if empty_hashes were not built correctly
-                .or_else(|| self.empty_hashes.get(index.depth() as usize).copied().map(Word::from))
+            // by the constructor as we check the depth of the lookup above.
-                .ok_or(MerkleError::NodeNotInSet(index.value()))
+            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(self.get_branch_node(&index).parent().into())
-            Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
         }
     }
 
-    /// Returns a Merkle path from the node at the specified key to the root. The node itself is
+    /// Returns a value of the leaf at the specified index.
-    /// not included in the path.
     ///
     /// # Errors
-    /// Returns an error if:
+    /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
-    /// * The specified depth is greater than the depth of the tree.
+    pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
+        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<MerklePath, MerkleError> {
         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
+    /// Return a Merkle path from the leaf at the specified index to the root.
-    /// included in the path.
+    ///
+    /// The leaf itself is not included in the path.
     ///
     /// # Errors
-    /// Returns an error if:
+    /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
-    /// * The specified key does not exist as a leaf node.
+    pub fn get_leaf_path(&self, index: u64) -> Result<MerklePath, MerkleError> {
-    pub fn get_leaf_path(&self, key: u64) -> Result<MerklePath, MerkleError> {
+        let index = NodeIndex::new(self.depth(), index)?;
-        let index = NodeIndex::new(self.depth(), key)?;
         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<Item = InnerNodeInfo> + '_ {
         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
+    /// Updates value of the leaf at the specified index returning the old leaf value.
-    /// tree.
+    ///
+    /// 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.
+    /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
-    pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
+    pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
-        let index = NodeIndex::new(self.depth(), key)?;
+        let old_value = self.insert_leaf_node(index, value).unwrap_or(EMPTY_WORD);
-        if !self.check_leaf_node_exists(key) {
+
-            return Err(MerkleError::NodeNotInSet(index.value()));
+        // if the old value and new value are the same, there is nothing to update
+        if value == old_value {
+            return Ok(value);
         }
-        self.insert_leaf(key, value)?;
 
-        Ok(())
+        let mut index = NodeIndex::new(self.depth(), index)?;
-    }
-
-    /// 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);
-
-        // TODO consider using a map `index |-> word` instead of `index |-> (word, word)`
-        let mut index = NodeIndex::new(self.depth(), key)?;
         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<RpoDigest>) {
-        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<Word> {
         self.leaves.get(&key).copied()
     }
 
-    fn insert_leaf_node(&mut self, key: u64, node: Word) {
+    fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
-        self.leaves.insert(key, node);
+        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])
+    }
+}

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() {
+fn test_depth2_tree() {
-    let tree = SimpleSmt::new(2)
+    let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap();
-        .unwrap()
-        .with_leaves(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]
+    // check get_node()
-fn get_values() {
-    let tree = SimpleSmt::new(2)
-        .unwrap()
-        .with_leaves(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());
-}
 
-#[test]
+    // check get_path(): depth 2
-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
     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> {
+fn test_inner_node_iterator() -> Result<(), MerkleError> {
-    let tree = SimpleSmt::new(2)
+    let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap();
-        .unwrap()
-        .with_leaves(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)
+    let mut tree = SimpleSmt::with_leaves(3, KEYS8.into_iter().zip(VALUES8.into_iter())).unwrap();
-        .unwrap()
-        .with_leaves(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)
+    let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
-        .unwrap()
-        .with_leaves(KEYS8.into_iter().zip(expected_values.clone().into_iter()))
-        .unwrap();
 
-    tree.update_leaf(key as u64, new_node).unwrap();
+    let old_leaf = tree.update_leaf(key as u64, new_node).unwrap();
-    assert_eq!(expected_tree.root, tree.root);
+    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)
+    let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
-        .unwrap()
-        .with_leaves(KEYS8.into_iter().zip(expected_values.into_iter()))
-        .unwrap();
 
-    tree.update_leaf(key as u64, new_node).unwrap();
+    let old_leaf = tree.update_leaf(key as u64, new_node).unwrap();
-    assert_eq!(expected_tree.root, tree.root);
+    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]
-    #[test]
+fn fail_on_duplicates() {
-    fn arbitrary_openings_single_leaf(
+    let entries = [(1_u64, int_to_node(1)), (5, int_to_node(2)), (1_u64, int_to_node(3))];
-        depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
+    let smt = SimpleSmt::with_leaves(64, entries);
-        key in prop::num::u64::ANY,
+    assert!(smt.is_err());
-        leaf in prop::num::u64::ANY,
-    ) {
-        let mut tree = SimpleSmt::new(depth).unwrap();
 
-        let key = key % (1 << depth as u64);
+    let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2)), (1_u64, int_to_node(0))];
-        let leaf = int_to_node(leaf);
+    let smt = SimpleSmt::with_leaves(64, entries);
+    assert!(smt.is_err());
 
-        tree.insert_leaf(key, leaf.into()).unwrap();
+    let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2)), (1_u64, int_to_node(1))];
-        tree.get_leaf_path(key).unwrap();
+    let smt = SimpleSmt::with_leaves(64, entries);
+    assert!(smt.is_err());
 
-        // traverse to root, fetching all paths
+    let entries = [(1_u64, int_to_node(1)), (5, int_to_node(2)), (1_u64, int_to_node(0))];
-        for d in 1..depth {
+    let smt = SimpleSmt::with_leaves(64, entries);
-            let k = key >> (depth - d);
+    assert!(smt.is_err());
-            tree.get_path(NodeIndex::make(d, k)).unwrap();
+}
-        }
-    }
 
-    #[test]
+#[test]
-    fn arbitrary_openings_multiple_leaves(
+fn with_no_duplicates_empty_node() {
-        depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
+    let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2))];
-        count in 2u8..10u8,
+    let smt = SimpleSmt::with_leaves(64, entries);
-        ref seed in any::<[u8; 32]>()
+    assert!(smt.is_ok());
-    ) {
-        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();
-        }
-    }
 }
 
 // HELPER FUNCTIONS

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)
+    let smt = SimpleSmt::with_leaves(1, keys2.into_iter().zip(leaves2.into_iter())).unwrap();
-        .unwrap()
-        .with_leaves(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)
+    let smt =
-        .unwrap()
+        SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, KEYS4.into_iter().zip(LEAVES4.into_iter()))
-        .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
+            .unwrap();
-        .unwrap();
 
     let store = MerkleStore::from(&smt);
 
@@ -526,10 +522,7 @@ fn test_constructors() -> Result<(), MerkleError> {
     }
 
     let depth = 32;
-    let smt = SimpleSmt::new(depth)
+    let smt = SimpleSmt::with_leaves(depth, KEYS4.into_iter().zip(LEAVES4.into_iter())).unwrap();
-        .unwrap()
-        .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
-        .unwrap();
     let store = MerkleStore::from(&smt);
     let depth = smt.depth();