Introduce SparseMerkleTree trait (#245)

src/merkle/smt/mod.rs

@@ -0,0 +1,238 @@
+use winter_math::StarkField;
+
+use crate::{
+    hash::rpo::{Rpo256, RpoDigest},
+    Word,
+};
+
+use super::{MerkleError, MerklePath, NodeIndex, Vec};
+
+mod simple;
+pub use simple::SimpleSmt;
+
+// CONSTANTS
+// ================================================================================================
+
+/// Minimum supported depth.
+pub const SMT_MIN_DEPTH: u8 = 1;
+
+/// Maximum supported depth.
+pub const SMT_MAX_DEPTH: u8 = 64;
+
+// SPARSE MERKLE TREE
+// ================================================================================================
+
+/// An abstract description of a sparse Merkle tree.
+///
+/// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed
+/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's
+/// value was not explicitly set, then its value is the default value. Typically, the vast majority
+/// of leaves will store the default value (hence it is "sparse"), and therefore the internal
+/// representation of the tree will only keep track of the leaves that have a different value from
+/// the default.
+///
+/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the
+/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a
+/// tree is just a single value, and is probably a programming mistake.
+///
+/// Every key maps to one leaf. If there are as many keys as there are leaves, then
+/// [Self::Leaf] should be the same type as [Self::Value], as is the case with
+/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`]
+/// must accomodate all keys that map to the same leaf.
+///
+/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
+pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
+    /// The type for a key
+    type Key: Clone;
+    /// The type for a value
+    type Value: Clone + PartialEq;
+    /// The type for a leaf
+    type Leaf;
+    /// The type for an opening (i.e. a "proof") of a leaf
+    type Opening: From<(MerklePath, Self::Leaf)>;
+
+    /// The default value used to compute the hash of empty leaves
+    const EMPTY_VALUE: Self::Value;
+
+    // PROVIDED METHODS
+    // ---------------------------------------------------------------------------------------------
+
+    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
+    /// path to the leaf, as well as the leaf itself.
+    fn open(&self, key: &Self::Key) -> Self::Opening {
+        let leaf = self.get_leaf(key);
+
+        let mut index: NodeIndex = {
+            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(key);
+            leaf_index.into()
+        };
+
+        let merkle_path = {
+            let mut path = Vec::with_capacity(index.depth() as usize);
+            for _ in 0..index.depth() {
+                let is_right = index.is_value_odd();
+                index.move_up();
+                let InnerNode { left, right } = self.get_inner_node(index);
+                let value = if is_right { left } else { right };
+                path.push(value);
+            }
+
+            MerklePath::new(path)
+        };
+
+        (merkle_path, leaf).into()
+    }
+
+    /// Inserts a value at the specified key, returning the previous value associated with that key.
+    /// Recall that by definition, any key that hasn't been updated is associated with
+    /// [`Self::EMPTY_VALUE`].
+    ///
+    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
+    /// updating the root itself.
+    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value {
+        let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE);
+
+        // if the old value and new value are the same, there is nothing to update
+        if value == old_value {
+            return value;
+        }
+
+        let leaf = self.get_leaf(&key);
+        let node_index = {
+            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
+            leaf_index.into()
+        };
+
+        self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
+
+        old_value
+    }
+
+    /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
+    /// `node_hash_at_index` is the hash of the node stored at index.
+    fn recompute_nodes_from_index_to_root(
+        &mut self,
+        mut index: NodeIndex,
+        node_hash_at_index: RpoDigest,
+    ) {
+        let mut value = node_hash_at_index;
+        for _ in 0..index.depth() {
+            let is_right = index.is_value_odd();
+            index.move_up();
+            let InnerNode { left, right } = self.get_inner_node(index);
+            let (left, right) = if is_right { (left, value) } else { (value, right) };
+            self.insert_inner_node(index, InnerNode { left, right });
+            value = Rpo256::merge(&[left, right]);
+        }
+        self.set_root(value);
+    }
+
+    // REQUIRED METHODS
+    // ---------------------------------------------------------------------------------------------
+
+    /// The root of the tree
+    fn root(&self) -> RpoDigest;
+
+    /// Sets the root of the tree
+    fn set_root(&mut self, root: RpoDigest);
+
+    /// Retrieves an inner node at the given index
+    fn get_inner_node(&self, index: NodeIndex) -> InnerNode;
+
+    /// Inserts an inner node at the given index
+    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode);
+
+    /// Inserts a leaf node, and returns the value at the key if already exists
+    fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
+
+    /// Returns the leaf at the specified index.
+    fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
+
+    /// Returns the hash of a leaf
+    fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
+
+    /// Maps a key to a leaf index
+    fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
+}
+
+// INNER NODE
+// ================================================================================================
+
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub(crate) struct InnerNode {
+    pub left: RpoDigest,
+    pub right: RpoDigest,
+}
+
+impl InnerNode {
+    pub fn hash(&self) -> RpoDigest {
+        Rpo256::merge(&[self.left, self.right])
+    }
+}
+
+// LEAF INDEX
+// ================================================================================================
+
+/// The index of a leaf, at a depth known at compile-time.
+#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub struct LeafIndex<const DEPTH: u8> {
+    index: NodeIndex,
+}
+
+impl<const DEPTH: u8> LeafIndex<DEPTH> {
+    pub fn new(value: u64) -> Result<Self, MerkleError> {
+        if DEPTH < SMT_MIN_DEPTH {
+            return Err(MerkleError::DepthTooSmall(DEPTH));
+        }
+
+        Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
+    }
+
+    pub fn value(&self) -> u64 {
+        self.index.value()
+    }
+}
+
+impl LeafIndex<SMT_MAX_DEPTH> {
+    pub const fn new_max_depth(value: u64) -> Self {
+        LeafIndex {
+            index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
+        }
+    }
+}
+
+impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
+    fn from(value: LeafIndex<DEPTH>) -> Self {
+        value.index
+    }
+}
+
+impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
+    type Error = MerkleError;
+
+    fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
+        if node_index.depth() != DEPTH {
+            return Err(MerkleError::InvalidDepth {
+                expected: DEPTH,
+                provided: node_index.depth(),
+            });
+        }
+
+        Self::new(node_index.value())
+    }
+}
+
+impl From<Word> for LeafIndex<SMT_MAX_DEPTH> {
+    fn from(value: Word) -> Self {
+        // We use the most significant `Felt` of a `Word` as the leaf index.
+        Self::new_max_depth(value[3].as_int())
+    }
+}
+
+impl From<RpoDigest> for LeafIndex<SMT_MAX_DEPTH> {
+    fn from(value: RpoDigest) -> Self {
+        Word::from(value).into()
+    }
+}

src/merkle/smt/simple/mod.rs

@@ -0,0 +1,325 @@
+use crate::{
+    merkle::{EmptySubtreeRoots, InnerNodeInfo, MerkleTreeDelta, StoreNode, ValuePath},
+    utils::collections::TryApplyDiff,
+    EMPTY_WORD,
+};
+
+use super::{
+    InnerNode, LeafIndex, MerkleError, NodeIndex, RpoDigest, SparseMerkleTree, Word, SMT_MAX_DEPTH,
+    SMT_MIN_DEPTH,
+};
+use crate::utils::collections::{BTreeMap, BTreeSet};
+
+#[cfg(test)]
+mod tests;
+
+// SPARSE MERKLE TREE
+// ================================================================================================
+
+/// 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)]
+#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+pub struct SimpleSmt<const DEPTH: u8> {
+    root: RpoDigest,
+    leaves: BTreeMap<u64, Word>,
+    inner_nodes: BTreeMap<NodeIndex, InnerNode>,
+}
+
+impl<const DEPTH: u8> SimpleSmt<DEPTH> {
+    // CONSTANTS
+    // --------------------------------------------------------------------------------------------
+
+    /// The default value used to compute the hash of empty leaves
+    pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<DEPTH>>::EMPTY_VALUE;
+
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a new [SimpleSmt].
+    ///
+    /// All leaves in the returned tree are set to [ZERO; 4].
+    ///
+    /// # Errors
+    /// Returns an error if DEPTH is 0 or is greater than 64.
+    pub fn new() -> Result<Self, MerkleError> {
+        // validate the range of the depth.
+        if DEPTH < SMT_MIN_DEPTH {
+            return Err(MerkleError::DepthTooSmall(DEPTH));
+        } else if SMT_MAX_DEPTH < DEPTH {
+            return Err(MerkleError::DepthTooBig(DEPTH as u64));
+        }
+
+        let root = *EmptySubtreeRoots::entry(DEPTH, 0);
+
+        Ok(Self {
+            root,
+            leaves: BTreeMap::new(),
+            inner_nodes: BTreeMap::new(),
+        })
+    }
+
+    /// Returns a new [SimpleSmt] instantiated 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:
+    /// - If the depth is 0 or is greater than 64.
+    /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
+    /// - The provided entries contain multiple values for the same key.
+    pub fn with_leaves(
+        entries: impl IntoIterator<Item = (u64, Word)>,
+    ) -> Result<Self, MerkleError> {
+        // create an empty tree
+        let mut tree = Self::new()?;
+
+        // compute the max number of entries. We use an upper bound of depth 63 because we consider
+        // passing in a vector of size 2^64 infeasible.
+        let max_num_entries = 2_usize.pow(DEPTH.min(63).into());
+
+        // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
+        // entries with the empty value need additional tracking.
+        let mut key_set_to_zero = BTreeSet::new();
+
+        for (idx, (key, value)) in entries.into_iter().enumerate() {
+            if idx >= max_num_entries {
+                return Err(MerkleError::InvalidNumEntries(max_num_entries));
+            }
+
+            let old_value = tree.insert(LeafIndex::<DEPTH>::new(key)?, value);
+
+            if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
+                return Err(MerkleError::DuplicateValuesForIndex(key));
+            }
+
+            if value == Self::EMPTY_VALUE {
+                key_set_to_zero.insert(key);
+            };
+        }
+        Ok(tree)
+    }
+
+    /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
+    /// starting at index 0.
+    pub fn with_contiguous_leaves(
+        entries: impl IntoIterator<Item = Word>,
+    ) -> Result<Self, MerkleError> {
+        Self::with_leaves(
+            entries
+                .into_iter()
+                .enumerate()
+                .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
+        )
+    }
+
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the depth of the tree
+    pub const fn depth(&self) -> u8 {
+        DEPTH
+    }
+
+    /// Returns the root of the tree
+    pub fn root(&self) -> RpoDigest {
+        <Self as SparseMerkleTree<DEPTH>>::root(self)
+    }
+
+    /// Returns the leaf at the specified index.
+    pub fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
+        <Self as SparseMerkleTree<DEPTH>>::get_leaf(self, key)
+    }
+
+    /// Returns a node at the specified index.
+    ///
+    /// # Errors
+    /// Returns an error if the specified index has depth set to 0 or the depth is greater than
+    /// the depth of this Merkle tree.
+    pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
+        if index.is_root() {
+            Err(MerkleError::DepthTooSmall(index.depth()))
+        } else if index.depth() > DEPTH {
+            Err(MerkleError::DepthTooBig(index.depth() as u64))
+        } else if index.depth() == DEPTH {
+            let leaf = self.get_leaf(&LeafIndex::<DEPTH>::try_from(index)?);
+
+            Ok(leaf.into())
+        } else {
+            Ok(self.get_inner_node(index).hash())
+        }
+    }
+
+    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
+    /// path to the leaf, as well as the leaf itself.
+    pub fn open(&self, key: &LeafIndex<DEPTH>) -> ValuePath {
+        <Self as SparseMerkleTree<DEPTH>>::open(self, key)
+    }
+
+    // ITERATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns an iterator over the leaves of this [SimpleSmt].
+    pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
+        self.leaves.iter().map(|(i, w)| (*i, w))
+    }
+
+    /// Returns an iterator over the inner nodes of this [SimpleSmt].
+    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
+        self.inner_nodes.values().map(|e| InnerNodeInfo {
+            value: e.hash(),
+            left: e.left,
+            right: e.right,
+        })
+    }
+
+    // STATE MUTATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Inserts a value at the specified key, returning the previous value associated with that key.
+    /// Recall that by definition, any key that hasn't been updated is associated with
+    /// [`EMPTY_WORD`].
+    ///
+    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
+    /// updating the root itself.
+    pub fn insert(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Word {
+        <Self as SparseMerkleTree<DEPTH>>::insert(self, key, value)
+    }
+
+    /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
+    /// computed as `DEPTH - SUBTREE_DEPTH`.
+    ///
+    /// Returns the new root.
+    pub fn set_subtree<const SUBTREE_DEPTH: u8>(
+        &mut self,
+        subtree_insertion_index: u64,
+        subtree: SimpleSmt<SUBTREE_DEPTH>,
+    ) -> Result<RpoDigest, MerkleError> {
+        if SUBTREE_DEPTH > DEPTH {
+            return Err(MerkleError::InvalidSubtreeDepth {
+                subtree_depth: SUBTREE_DEPTH,
+                tree_depth: DEPTH,
+            });
+        }
+
+        // Verify that `subtree_insertion_index` is valid.
+        let subtree_root_insertion_depth = DEPTH - SUBTREE_DEPTH;
+        let subtree_root_index =
+            NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
+
+        // add leaves
+        // --------------
+
+        // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
+        // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
+        // valid as they are. However, consider what happens when we insert at
+        // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
+        // you can see it as there's a full subtree sitting on its left. In general, for
+        // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
+        // to insert, so we need to adjust all its leaves by `i * 2^d`.
+        let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(SUBTREE_DEPTH.into());
+        for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
+            let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
+            debug_assert!(new_leaf_idx < 2_u64.pow(DEPTH.into()));
+
+            self.leaves.insert(new_leaf_idx, *leaf_value);
+        }
+
+        // add subtree's branch nodes (which includes the root)
+        // --------------
+        for (branch_idx, branch_node) in subtree.inner_nodes {
+            let new_branch_idx = {
+                let new_depth = subtree_root_insertion_depth + branch_idx.depth();
+                let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
+                    + branch_idx.value();
+
+                NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
+            };
+
+            self.inner_nodes.insert(new_branch_idx, branch_node);
+        }
+
+        // recompute nodes starting from subtree root
+        // --------------
+        self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
+
+        Ok(self.root)
+    }
+}
+
+impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
+    type Key = LeafIndex<DEPTH>;
+    type Value = Word;
+    type Leaf = Word;
+    type Opening = ValuePath;
+
+    const EMPTY_VALUE: Self::Value = EMPTY_WORD;
+
+    fn root(&self) -> RpoDigest {
+        self.root
+    }
+
+    fn set_root(&mut self, root: RpoDigest) {
+        self.root = root;
+    }
+
+    fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
+        self.inner_nodes.get(&index).cloned().unwrap_or_else(|| {
+            let node = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
+
+            InnerNode { left: *node, right: *node }
+        })
+    }
+
+    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
+        self.inner_nodes.insert(index, inner_node);
+    }
+
+    fn insert_value(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Option<Word> {
+        self.leaves.insert(key.value(), value)
+    }
+
+    fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
+        // 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.
+        let leaf_pos = key.value();
+
+        match self.leaves.get(&leaf_pos) {
+            Some(word) => *word,
+            None => Word::from(*EmptySubtreeRoots::entry(DEPTH, DEPTH)),
+        }
+    }
+
+    fn hash_leaf(leaf: &Word) -> RpoDigest {
+        // `SimpleSmt` takes the leaf value itself as the hash
+        leaf.into()
+    }
+
+    fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
+        *key
+    }
+}
+
+// TRY APPLY DIFF
+// ================================================================================================
+impl<const DEPTH: u8> TryApplyDiff<RpoDigest, StoreNode> for SimpleSmt<DEPTH> {
+    type Error = MerkleError;
+    type DiffType = MerkleTreeDelta;
+
+    fn try_apply(&mut self, diff: MerkleTreeDelta) -> Result<(), MerkleError> {
+        if diff.depth() != DEPTH {
+            return Err(MerkleError::InvalidDepth { expected: DEPTH, provided: diff.depth() });
+        }
+
+        for slot in diff.cleared_slots() {
+            self.insert(LeafIndex::<DEPTH>::new(*slot)?, Self::EMPTY_VALUE);
+        }
+
+        for (slot, value) in diff.updated_slots() {
+            self.insert(LeafIndex::<DEPTH>::new(*slot)?, *value);
+        }
+
+        Ok(())
+    }
+}

src/merkle/smt/simple/tests.rs

@@ -0,0 +1,437 @@
+use super::{
+    super::{MerkleError, RpoDigest, SimpleSmt},
+    NodeIndex,
+};
+use crate::{
+    hash::rpo::Rpo256,
+    merkle::{
+        digests_to_words, int_to_leaf, int_to_node, smt::SparseMerkleTree, EmptySubtreeRoots,
+        InnerNodeInfo, LeafIndex, MerkleTree,
+    },
+    utils::collections::Vec,
+    Word, EMPTY_WORD,
+};
+
+// TEST DATA
+// ================================================================================================
+
+const KEYS4: [u64; 4] = [0, 1, 2, 3];
+const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
+
+const VALUES4: [RpoDigest; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)];
+
+const VALUES8: [RpoDigest; 8] = [
+    int_to_node(1),
+    int_to_node(2),
+    int_to_node(3),
+    int_to_node(4),
+    int_to_node(5),
+    int_to_node(6),
+    int_to_node(7),
+    int_to_node(8),
+];
+
+const ZERO_VALUES8: [Word; 8] = [int_to_leaf(0); 8];
+
+// TESTS
+// ================================================================================================
+
+#[test]
+fn build_empty_tree() {
+    // tree of depth 3
+    let smt = SimpleSmt::<3>::new().unwrap();
+    let mt = MerkleTree::new(ZERO_VALUES8).unwrap();
+    assert_eq!(mt.root(), smt.root());
+}
+
+#[test]
+fn build_sparse_tree() {
+    const DEPTH: u8 = 3;
+    let mut smt = SimpleSmt::<DEPTH>::new().unwrap();
+    let mut values = ZERO_VALUES8.to_vec();
+
+    // insert single value
+    let key = 6;
+    let new_node = int_to_leaf(7);
+    values[key as usize] = new_node;
+    let old_value = smt.insert(LeafIndex::<DEPTH>::new(key).unwrap(), new_node);
+    let mt2 = MerkleTree::new(values.clone()).unwrap();
+    assert_eq!(mt2.root(), smt.root());
+    assert_eq!(
+        mt2.get_path(NodeIndex::make(3, 6)).unwrap(),
+        smt.open(&LeafIndex::<3>::new(6).unwrap()).path
+    );
+    assert_eq!(old_value, EMPTY_WORD);
+
+    // insert second value at distinct leaf branch
+    let key = 2;
+    let new_node = int_to_leaf(3);
+    values[key as usize] = new_node;
+    let old_value = smt.insert(LeafIndex::<DEPTH>::new(key).unwrap(), new_node);
+    let mt3 = MerkleTree::new(values).unwrap();
+    assert_eq!(mt3.root(), smt.root());
+    assert_eq!(
+        mt3.get_path(NodeIndex::make(3, 2)).unwrap(),
+        smt.open(&LeafIndex::<3>::new(2).unwrap()).path
+    );
+    assert_eq!(old_value, EMPTY_WORD);
+}
+
+/// Tests that [`SimpleSmt::with_contiguous_leaves`] works as expected
+#[test]
+fn build_contiguous_tree() {
+    let tree_with_leaves =
+        SimpleSmt::<2>::with_leaves([0, 1, 2, 3].into_iter().zip(digests_to_words(&VALUES4)))
+            .unwrap();
+
+    let tree_with_contiguous_leaves =
+        SimpleSmt::<2>::with_contiguous_leaves(digests_to_words(&VALUES4)).unwrap();
+
+    assert_eq!(tree_with_leaves, tree_with_contiguous_leaves);
+}
+
+#[test]
+fn test_depth2_tree() {
+    let tree =
+        SimpleSmt::<2>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4))).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());
+
+    // 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());
+
+    // check get_path(): depth 2
+    assert_eq!(vec![VALUES4[1], node3], *tree.open(&LeafIndex::<2>::new(0).unwrap()).path);
+    assert_eq!(vec![VALUES4[0], node3], *tree.open(&LeafIndex::<2>::new(1).unwrap()).path);
+    assert_eq!(vec![VALUES4[3], node2], *tree.open(&LeafIndex::<2>::new(2).unwrap()).path);
+    assert_eq!(vec![VALUES4[2], node2], *tree.open(&LeafIndex::<2>::new(3).unwrap()).path);
+}
+
+#[test]
+fn test_inner_node_iterator() -> Result<(), MerkleError> {
+    let tree =
+        SimpleSmt::<2>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4))).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());
+
+    // get parent nodes
+    let root = tree.root();
+    let l1n0 = tree.get_node(NodeIndex::make(1, 0))?;
+    let l1n1 = tree.get_node(NodeIndex::make(1, 1))?;
+    let l2n0 = tree.get_node(NodeIndex::make(2, 0))?;
+    let l2n1 = tree.get_node(NodeIndex::make(2, 1))?;
+    let l2n2 = tree.get_node(NodeIndex::make(2, 2))?;
+    let l2n3 = tree.get_node(NodeIndex::make(2, 3))?;
+
+    let nodes: Vec<InnerNodeInfo> = tree.inner_nodes().collect();
+    let expected = vec![
+        InnerNodeInfo { value: root, left: l1n0, right: l1n1 },
+        InnerNodeInfo { value: l1n0, left: l2n0, right: l2n1 },
+        InnerNodeInfo { value: l1n1, left: l2n2, right: l2n3 },
+    ];
+    assert_eq!(nodes, expected);
+
+    Ok(())
+}
+
+#[test]
+fn update_leaf() {
+    const DEPTH: u8 = 3;
+    let mut tree =
+        SimpleSmt::<DEPTH>::with_leaves(KEYS8.into_iter().zip(digests_to_words(&VALUES8))).unwrap();
+
+    // update one value
+    let key = 3;
+    let new_node = int_to_leaf(9);
+    let mut expected_values = digests_to_words(&VALUES8);
+    expected_values[key] = new_node;
+    let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
+
+    let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
+    assert_eq!(expected_tree.root(), tree.root);
+    assert_eq!(old_leaf, *VALUES8[key]);
+
+    // update another value
+    let key = 6;
+    let new_node = int_to_leaf(10);
+    expected_values[key] = new_node;
+    let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
+
+    let old_leaf = tree.insert(LeafIndex::<DEPTH>::new(key as u64).unwrap(), new_node);
+    assert_eq!(expected_tree.root(), tree.root);
+    assert_eq!(old_leaf, *VALUES8[key]);
+}
+
+#[test]
+fn small_tree_opening_is_consistent() {
+    //        ____k____
+    //       /         \
+    //     _i_         _j_
+    //    /   \       /   \
+    //   e     f     g     h
+    //  / \   / \   / \   / \
+    // a   b 0   0 c   0 0   d
+
+    let z = EMPTY_WORD;
+
+    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
+    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
+    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
+    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
+
+    let e = Rpo256::merge(&[a.into(), b.into()]);
+    let f = Rpo256::merge(&[z.into(), z.into()]);
+    let g = Rpo256::merge(&[c.into(), z.into()]);
+    let h = Rpo256::merge(&[z.into(), d.into()]);
+
+    let i = Rpo256::merge(&[e, f]);
+    let j = Rpo256::merge(&[g, h]);
+
+    let k = Rpo256::merge(&[i, j]);
+
+    let entries = vec![(0, a), (1, b), (4, c), (7, d)];
+    let tree = SimpleSmt::<3>::with_leaves(entries).unwrap();
+
+    assert_eq!(tree.root(), k);
+
+    let cases: Vec<(u64, Vec<RpoDigest>)> = vec![
+        (0, vec![b.into(), f, j]),
+        (1, vec![a.into(), f, j]),
+        (4, vec![z.into(), h, i]),
+        (7, vec![z.into(), g, i]),
+    ];
+
+    for (key, path) in cases {
+        let opening = tree.open(&LeafIndex::<3>::new(key).unwrap());
+
+        assert_eq!(path, *opening.path);
+    }
+}
+
+#[test]
+fn test_simplesmt_fail_on_duplicates() {
+    let values = [
+        // same key, same value
+        (int_to_leaf(1), int_to_leaf(1)),
+        // same key, different values
+        (int_to_leaf(1), int_to_leaf(2)),
+        // same key, set to zero
+        (EMPTY_WORD, int_to_leaf(1)),
+        // same key, re-set to zero
+        (int_to_leaf(1), EMPTY_WORD),
+        // same key, set to zero twice
+        (EMPTY_WORD, EMPTY_WORD),
+    ];
+
+    for (first, second) in values.iter() {
+        // consecutive
+        let entries = [(1, *first), (1, *second)];
+        let smt = SimpleSmt::<64>::with_leaves(entries);
+        assert_eq!(smt.unwrap_err(), MerkleError::DuplicateValuesForIndex(1));
+
+        // not consecutive
+        let entries = [(1, *first), (5, int_to_leaf(5)), (1, *second)];
+        let smt = SimpleSmt::<64>::with_leaves(entries);
+        assert_eq!(smt.unwrap_err(), MerkleError::DuplicateValuesForIndex(1));
+    }
+}
+
+#[test]
+fn with_no_duplicates_empty_node() {
+    let entries = [(1_u64, int_to_leaf(0)), (5, int_to_leaf(2))];
+    let smt = SimpleSmt::<64>::with_leaves(entries);
+    assert!(smt.is_ok());
+}
+
+#[test]
+fn test_simplesmt_with_leaves_nonexisting_leaf() {
+    // TESTING WITH EMPTY WORD
+    // --------------------------------------------------------------------------------------------
+
+    // Depth 1 has 2 leaf. Position is 0-indexed, position 2 doesn't exist.
+    let leaves = [(2, EMPTY_WORD)];
+    let result = SimpleSmt::<1>::with_leaves(leaves);
+    assert!(result.is_err());
+
+    // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist.
+    let leaves = [(4, EMPTY_WORD)];
+    let result = SimpleSmt::<2>::with_leaves(leaves);
+    assert!(result.is_err());
+
+    // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist.
+    let leaves = [(8, EMPTY_WORD)];
+    let result = SimpleSmt::<3>::with_leaves(leaves);
+    assert!(result.is_err());
+
+    // TESTING WITH A VALUE
+    // --------------------------------------------------------------------------------------------
+    let value = int_to_node(1);
+
+    // Depth 1 has 2 leaves. Position is 0-indexed, position 2 doesn't exist.
+    let leaves = [(2, *value)];
+    let result = SimpleSmt::<1>::with_leaves(leaves);
+    assert!(result.is_err());
+
+    // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist.
+    let leaves = [(4, *value)];
+    let result = SimpleSmt::<2>::with_leaves(leaves);
+    assert!(result.is_err());
+
+    // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist.
+    let leaves = [(8, *value)];
+    let result = SimpleSmt::<3>::with_leaves(leaves);
+    assert!(result.is_err());
+}
+
+#[test]
+fn test_simplesmt_set_subtree() {
+    // Final Tree:
+    //
+    //        ____k____
+    //       /         \
+    //     _i_         _j_
+    //    /   \       /   \
+    //   e     f     g     h
+    //  / \   / \   / \   / \
+    // a   b 0   0 c   0 0   d
+
+    let z = EMPTY_WORD;
+
+    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
+    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
+    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
+    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
+
+    let e = Rpo256::merge(&[a.into(), b.into()]);
+    let f = Rpo256::merge(&[z.into(), z.into()]);
+    let g = Rpo256::merge(&[c.into(), z.into()]);
+    let h = Rpo256::merge(&[z.into(), d.into()]);
+
+    let i = Rpo256::merge(&[e, f]);
+    let j = Rpo256::merge(&[g, h]);
+
+    let k = Rpo256::merge(&[i, j]);
+
+    // subtree:
+    //   g
+    //  / \
+    // c   0
+    let subtree = {
+        let entries = vec![(0, c)];
+        SimpleSmt::<1>::with_leaves(entries).unwrap()
+    };
+
+    // insert subtree
+    const TREE_DEPTH: u8 = 3;
+    let tree = {
+        let entries = vec![(0, a), (1, b), (7, d)];
+        let mut tree = SimpleSmt::<TREE_DEPTH>::with_leaves(entries).unwrap();
+
+        tree.set_subtree(2, subtree).unwrap();
+
+        tree
+    };
+
+    assert_eq!(tree.root(), k);
+    assert_eq!(tree.get_leaf(&LeafIndex::<TREE_DEPTH>::new(4).unwrap()), c);
+    assert_eq!(tree.get_inner_node(NodeIndex::new_unchecked(2, 2)).hash(), g);
+}
+
+/// Ensures that an invalid input node index into `set_subtree()` incurs no mutation of the tree
+#[test]
+fn test_simplesmt_set_subtree_unchanged_for_wrong_index() {
+    // Final Tree:
+    //
+    //        ____k____
+    //       /         \
+    //     _i_         _j_
+    //    /   \       /   \
+    //   e     f     g     h
+    //  / \   / \   / \   / \
+    // a   b 0   0 c   0 0   d
+
+    let z = EMPTY_WORD;
+
+    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
+    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
+    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
+    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
+
+    // subtree:
+    //   g
+    //  / \
+    // c   0
+    let subtree = {
+        let entries = vec![(0, c)];
+        SimpleSmt::<1>::with_leaves(entries).unwrap()
+    };
+
+    let mut tree = {
+        let entries = vec![(0, a), (1, b), (7, d)];
+        SimpleSmt::<3>::with_leaves(entries).unwrap()
+    };
+    let tree_root_before_insertion = tree.root();
+
+    // insert subtree
+    assert!(tree.set_subtree(500, subtree).is_err());
+
+    assert_eq!(tree.root(), tree_root_before_insertion);
+}
+
+/// We insert an empty subtree that has the same depth as the original tree
+#[test]
+fn test_simplesmt_set_subtree_entire_tree() {
+    // Initial Tree:
+    //
+    //        ____k____
+    //       /         \
+    //     _i_         _j_
+    //    /   \       /   \
+    //   e     f     g     h
+    //  / \   / \   / \   / \
+    // a   b 0   0 c   0 0   d
+
+    let z = EMPTY_WORD;
+
+    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
+    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
+    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
+    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
+
+    // subtree: E3
+    const DEPTH: u8 = 3;
+    let subtree = { SimpleSmt::<DEPTH>::with_leaves(Vec::new()).unwrap() };
+    assert_eq!(subtree.root(), *EmptySubtreeRoots::entry(DEPTH, 0));
+
+    // insert subtree
+    let mut tree = {
+        let entries = vec![(0, a), (1, b), (4, c), (7, d)];
+        SimpleSmt::<3>::with_leaves(entries).unwrap()
+    };
+
+    tree.set_subtree(0, subtree).unwrap();
+
+    assert_eq!(tree.root(), *EmptySubtreeRoots::entry(DEPTH, 0));
+}
+
+// HELPER FUNCTIONS
+// --------------------------------------------------------------------------------------------
+
+fn compute_internal_nodes() -> (RpoDigest, RpoDigest, RpoDigest) {
+    let node2 = Rpo256::merge(&[VALUES4[0], VALUES4[1]]);
+    let node3 = Rpo256::merge(&[VALUES4[2], VALUES4[3]]);
+    let root = Rpo256::merge(&[node2, node3]);
+
+    (root, node2, node3)
+}