Merge pull request #95 from 0xPolygonMiden/vlopes11-fix-merkle-store-bounds

fix: merkle store panics on bounds

src/bit.rs

@@ -19,7 +19,7 @@ impl BitIterator {
         let mask = bitmask(n);
         let ones = self.mask.trailing_ones();
         let mask_position = ones;
-        self.mask ^= mask << mask_position;
+        self.mask ^= mask.checked_shl(mask_position).unwrap_or(0);
         self
     }
 
@@ -31,7 +31,7 @@ impl BitIterator {
         let mask = bitmask(n);
         let ones = self.mask.leading_ones();
         let mask_position = u64::BITS - ones - n;
-        self.mask ^= mask << mask_position;
+        self.mask ^= mask.checked_shl(mask_position).unwrap_or(0);
         self
     }
 }

src/merkle/mod.rs

@@ -35,18 +35,19 @@ pub use store::MerkleStore;
 // ERRORS
 // ================================================================================================
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum MerkleError {
     ConflictingRoots(Vec<Word>),
     DepthTooSmall(u8),
     DepthTooBig(u64),
-    NodeNotInStorage(Word, NodeIndex),
+    NodeNotInStore(Word, NodeIndex),
     NumLeavesNotPowerOfTwo(usize),
     InvalidIndex(NodeIndex),
     InvalidDepth { expected: u8, provided: u8 },
     InvalidPath(MerklePath),
     InvalidEntriesCount(usize, usize),
     NodeNotInSet(u64),
+    RootNotInStore(Word),
 }
 
 impl fmt::Display for MerkleError {
@@ -70,7 +71,8 @@ impl fmt::Display for MerkleError {
             InvalidPath(_path) => write!(f, "the provided path is not valid"),
             InvalidEntriesCount(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"),
             NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"),
-            NodeNotInStorage(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the storage", hash, index.value(), index.depth(),),
+            NodeNotInStore(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the store", hash, index.value(), index.depth(),),
+            RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root),
         }
     }
 }

src/merkle/path_set.rs

@@ -57,6 +57,14 @@ impl MerklePathSet {
         self.total_depth
     }
 
+    /// Returns all the leaf indexes of this path set.
+    pub fn indexes(&self) -> impl Iterator<Item = NodeIndex> + '_ {
+        self.paths
+            .keys()
+            .copied()
+            .map(|index| NodeIndex::new(self.total_depth, index))
+    }
+
     /// Returns a node at the specified index.
     ///
     /// # Errors

src/merkle/store/mod.rs

@@ -4,10 +4,13 @@
 //! (leaves or internal) to live as long as necessary and without duplication, this allows the
 //! implementation of efficient persistent data structures
 use super::{
-    BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerkleTree, NodeIndex, Rpo256,
+    BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree,
-    RpoDigest, SimpleSmt, Vec, Word,
+    NodeIndex, Rpo256, RpoDigest, SimpleSmt, Vec, Word,
 };
 
+#[cfg(test)]
+mod tests;
+
 #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
 pub struct Node {
     left: RpoDigest,
@@ -81,13 +84,83 @@ impl MerkleStore {
         Ok(self)
     }
 
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the node at `index` rooted on the tree `root`.
+    ///
+    /// # Errors
+    ///
+    /// This method can return the following errors:
+    /// - `RootNotInStore` if the `root` is not present in the store.
+    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
+    pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
+        let mut hash: RpoDigest = root.into();
+
+        // corner case: check the root is in the store when called with index `NodeIndex::root()`
+        self.nodes
+            .get(&hash)
+            .ok_or(MerkleError::RootNotInStore(hash.into()))?;
+
+        for bit in index.bit_iterator().rev() {
+            let node = self
+                .nodes
+                .get(&hash)
+                .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?;
+            hash = if bit { node.right } else { node.left }
+        }
+
+        Ok(hash.into())
+    }
+
+    /// Returns the node at the specified `index` and its opening to the `root`.
+    ///
+    /// The path starts at the sibling of the target leaf.
+    ///
+    /// # Errors
+    ///
+    /// This method can return the following errors:
+    /// - `RootNotInStore` if the `root` is not present in the store.
+    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
+    pub fn get_path(
+        &self,
+        root: Word,
+        index: NodeIndex,
+    ) -> Result<(Word, MerklePath), MerkleError> {
+        let mut hash: RpoDigest = root.into();
+        let mut path = Vec::with_capacity(index.depth().into());
+
+        // corner case: check the root is in the store when called with index `NodeIndex::root()`
+        self.nodes
+            .get(&hash)
+            .ok_or(MerkleError::RootNotInStore(hash.into()))?;
+
+        for bit in index.bit_iterator().rev() {
+            let node = self
+                .nodes
+                .get(&hash)
+                .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?;
+
+            hash = if bit {
+                path.push(node.left.into());
+                node.right
+            } else {
+                path.push(node.right.into());
+                node.left
+            }
+        }
+
+        path.reverse();
+        Ok((hash.into(), MerklePath::new(path)))
+    }
+
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
     /// Adds all the nodes of a Merkle tree represented by `leaves`.
     ///
     /// This will instantiate a Merkle tree using `leaves` and include all the nodes into the
-    /// storage.
+    /// store.
     ///
     /// # Errors
     ///
@@ -99,6 +172,10 @@ impl MerkleStore {
         I: IntoIterator<Item = Word>,
     {
         let leaves: Vec<_> = leaves.into_iter().collect();
+        if leaves.len() < 2 {
+            return Err(MerkleError::DepthTooSmall(leaves.len() as u8));
+        }
+
         let layers = leaves.len().ilog2();
         let tree = MerkleTree::new(leaves)?;
 
@@ -131,7 +208,7 @@ impl MerkleStore {
     /// Adds all the nodes of a Sparse Merkle tree represented by `entries`.
     ///
     /// This will instantiate a Sparse Merkle tree using `entries` and include all the nodes into
-    /// the storage.
+    /// the store.
     ///
     /// # Errors
     ///
@@ -159,7 +236,7 @@ impl MerkleStore {
     /// Adds all the nodes of a Merkle path represented by `path`.
     ///
     /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
-    /// include all the nodes into the storage.
+    /// include all the nodes into the store.
     pub fn add_merkle_path(
         &mut self,
         index_value: u64,
@@ -193,7 +270,7 @@ impl MerkleStore {
     /// Adds all the nodes of multiple Merkle paths into the store.
     ///
     /// This will compute the sibling elements for each Merkle `path` and include all the nodes
-    /// into the storage.
+    /// into the store.
     ///
     /// # Errors
     ///
@@ -220,81 +297,35 @@ impl MerkleStore {
             self.add_merkle_path(index_value, node, path)?;
         }
 
-        // Returns the parent of the last paths (assumes all paths have the same parent) or empty
-        // The length of unique_roots is checked above, so this wont panic
         Ok(roots.iter().next().unwrap().into())
     }
 
-    // PUBLIC ACCESSORS
+    /// Appends the provided [MerklePathSet] into the store.
-    // --------------------------------------------------------------------------------------------
+    pub fn add_merkle_path_set(&mut self, path_set: &MerklePathSet) -> Result<Word, MerkleError> {
+        let root = path_set.root();
+        path_set.indexes().try_fold(root, |_, index| {
+            let node = path_set.get_node(index)?;
+            let path = path_set.get_path(index)?;
+            self.add_merkle_path(index.value(), node, path)
+        })
+    }
 
-    /// Returns the node at `index` rooted on the tree `root`.
+    /// Sets a node to `value`.
     ///
     /// # Errors
     ///
-    /// This will return `NodeNotInStorage` if the element is not present in the store.
+    /// This method can return the following errors:
-    pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
+    /// - `RootNotInStore` if the `root` is not present in the store.
-        let mut hash: RpoDigest = root.into();
+    /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
-
-        // Check the root is in the storage when called with `NodeIndex::root()`
-        self.nodes
-            .get(&hash)
-            .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
-
-        for bit in index.bit_iterator().rev() {
-            let node = self
-                .nodes
-                .get(&hash)
-                .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
-            hash = if bit { node.right } else { node.left }
-        }
-
-        Ok(hash.into())
-    }
-
-    /// Returns the path for the node at `index` rooted on the tree `root`.
-    ///
-    /// # Errors
-    ///
-    /// This will return `NodeNotInStorage` if the element is not present in the store.
-    pub fn get_path(
-        &self,
-        root: Word,
-        index: NodeIndex,
-    ) -> Result<(Word, MerklePath), MerkleError> {
-        let mut hash: RpoDigest = root.into();
-        let mut path = Vec::new();
-        let node = RpoDigest::default();
-        for bit in index.bit_iterator() {
-            let node = self
-                .nodes
-                .get(&hash)
-                .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
-
-            hash = if bit {
-                path.push(node.left.into());
-                node.right
-            } else {
-                path.push(node.right.into());
-                node.left
-            }
-        }
-
-        Ok((node.into(), MerklePath::new(path)))
-    }
-
-    // DATA MUTATORS
-    // --------------------------------------------------------------------------------------------
-
     pub fn set_node(
         &mut self,
         root: Word,
         index: NodeIndex,
         value: Word,
     ) -> Result<Word, MerkleError> {
-        let (current_node, path) = self.get_path(root, index)?;
+        let result = self.get_path(root, index)?;
-        if current_node != value {
+        if result.0 != value {
-            self.add_merkle_path(index.value(), value, path)
+            self.add_merkle_path(index.value(), value, result.1)
         } else {
             Ok(root)
         }
@@ -305,12 +336,12 @@ impl MerkleStore {
         let root2: RpoDigest = root2.into();
 
         if !self.nodes.contains_key(&root1) {
-            Err(MerkleError::NodeNotInStorage(
+            Err(MerkleError::NodeNotInStore(
                 root1.into(),
                 NodeIndex::new(0, 0),
             ))
         } else if !self.nodes.contains_key(&root1) {
-            Err(MerkleError::NodeNotInStorage(
+            Err(MerkleError::NodeNotInStore(
                 root2.into(),
                 NodeIndex::new(0, 0),
             ))
@@ -328,170 +359,3 @@ impl MerkleStore {
         }
     }
 }
-
-#[cfg(test)]
-mod test {
-    use super::{MerkleError, MerkleStore, MerkleTree, NodeIndex, SimpleSmt, Word};
-    use crate::merkle::int_to_node;
-    use crate::merkle::MerklePathSet;
-
-    const KEYS4: [u64; 4] = [0, 1, 2, 3];
-    const LEAVES4: [Word; 4] = [
-        int_to_node(1),
-        int_to_node(2),
-        int_to_node(3),
-        int_to_node(4),
-    ];
-
-    #[test]
-    fn test_add_merkle_tree() -> Result<(), MerkleError> {
-        let mut store = MerkleStore::default();
-
-        let mtree = MerkleTree::new(LEAVES4.to_vec())?;
-        store.add_merkle_tree(LEAVES4.to_vec())?;
-
-        assert!(
-            store
-                .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
-                .is_ok(),
-            "node 0 must be in the tree"
-        );
-        assert!(
-            store
-                .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
-                .is_ok(),
-            "node 1 must be in the tree"
-        );
-        assert!(
-            store
-                .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
-                .is_ok(),
-            "node 2 must be in the tree"
-        );
-        assert!(
-            store
-                .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
-                .is_ok(),
-            "node 3 must be in the tree"
-        );
-
-        store
-            .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
-            .expect("node 0 must be in tree");
-        store
-            .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
-            .expect("node 1 must be in tree");
-        store
-            .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
-            .expect("node 2 must be in tree");
-        store
-            .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
-            .expect("node 3 must be in tree");
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_get_invalid_node() {
-        let mut store = MerkleStore::default();
-        let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work");
-        store
-            .add_merkle_tree(LEAVES4.to_vec())
-            .expect("adding a merkle tree to the store must work");
-        let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3));
-    }
-
-    #[test]
-    fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
-        let mut store = MerkleStore::default();
-        let keys2: [u64; 2] = [0, 1];
-        let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)];
-        store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?;
-        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
-            .unwrap()
-            .with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
-            .unwrap();
-
-        let idx = NodeIndex::new(1, 0);
-        assert_eq!(
-            store.get_node(smt.root(), idx).unwrap(),
-            smt.get_node(&idx).unwrap()
-        );
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_add_sparse_merkle_tree() -> Result<(), MerkleError> {
-        let mut store = MerkleStore::default();
-        store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
-
-        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
-            .unwrap()
-            .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
-            .unwrap();
-
-        let idx = NodeIndex::new(1, 0);
-        assert_eq!(
-            store.get_node(smt.root(), idx).unwrap(),
-            smt.get_node(&idx).unwrap()
-        );
-        let idx = NodeIndex::new(1, 1);
-        assert_eq!(
-            store.get_node(smt.root(), idx).unwrap(),
-            smt.get_node(&idx).unwrap()
-        );
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_add_merkle_paths() -> Result<(), MerkleError> {
-        let mut store = MerkleStore::default();
-        let mtree = MerkleTree::new(LEAVES4.to_vec())?;
-
-        let i0 = 0;
-        let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap();
-
-        let i1 = 1;
-        let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap();
-
-        let i2 = 2;
-        let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap();
-
-        let i3 = 3;
-        let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap();
-
-        let paths = [
-            (i0, LEAVES4[i0 as usize], p0),
-            (i1, LEAVES4[i1 as usize], p1),
-            (i2, LEAVES4[i2 as usize], p2),
-            (i3, LEAVES4[i3 as usize], p3),
-        ];
-
-        store
-            .add_merkle_paths(paths.clone())
-            .expect("the valid paths must work");
-
-        let set = MerklePathSet::new(3).with_paths(paths).unwrap();
-
-        assert_eq!(
-            set.get_node(NodeIndex::new(3, 0)).unwrap(),
-            store.get_node(set.root(), NodeIndex::new(2, 0b00)).unwrap(),
-        );
-        assert_eq!(
-            set.get_node(NodeIndex::new(3, 1)).unwrap(),
-            store.get_node(set.root(), NodeIndex::new(2, 0b01)).unwrap(),
-        );
-        assert_eq!(
-            set.get_node(NodeIndex::new(3, 2)).unwrap(),
-            store.get_node(set.root(), NodeIndex::new(2, 0b10)).unwrap(),
-        );
-        assert_eq!(
-            set.get_node(NodeIndex::new(3, 3)).unwrap(),
-            store.get_node(set.root(), NodeIndex::new(2, 0b11)).unwrap(),
-        );
-
-        Ok(())
-    }
-}

src/merkle/store/tests.rs

@@ -0,0 +1,497 @@
+use super::*;
+use crate::{
+    hash::rpo::Rpo256,
+    merkle::{int_to_node, MerklePathSet},
+    Felt, Word,
+};
+
+const KEYS4: [u64; 4] = [0, 1, 2, 3];
+const LEAVES4: [Word; 4] = [
+    int_to_node(1),
+    int_to_node(2),
+    int_to_node(3),
+    int_to_node(4),
+];
+
+#[test]
+fn test_root_not_in_store() -> Result<(), MerkleError> {
+    let mtree = MerkleTree::new(LEAVES4.to_vec())?;
+    let store = MerkleStore::default().with_merkle_tree(LEAVES4)?;
+    assert_eq!(
+        store.get_node(LEAVES4[0], NodeIndex::new(mtree.depth(), 0)),
+        Err(MerkleError::RootNotInStore(LEAVES4[0])),
+        "Leaf 0 is not a root"
+    );
+    assert_eq!(
+        store.get_path(LEAVES4[0], NodeIndex::new(mtree.depth(), 0)),
+        Err(MerkleError::RootNotInStore(LEAVES4[0])),
+        "Leaf 0 is not a root"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_merkle_tree() -> Result<(), MerkleError> {
+    let mut store = MerkleStore::default();
+
+    let mtree = MerkleTree::new(LEAVES4.to_vec())?;
+    store.add_merkle_tree(LEAVES4.to_vec())?;
+
+    // STORE LEAVES ARE CORRECT ==============================================================
+    // checks the leaves in the store corresponds to the expected values
+    assert_eq!(
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)),
+        Ok(LEAVES4[0]),
+        "node 0 must be in the tree"
+    );
+    assert_eq!(
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)),
+        Ok(LEAVES4[1]),
+        "node 1 must be in the tree"
+    );
+    assert_eq!(
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)),
+        Ok(LEAVES4[2]),
+        "node 2 must be in the tree"
+    );
+    assert_eq!(
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)),
+        Ok(LEAVES4[3]),
+        "node 3 must be in the tree"
+    );
+
+    // STORE LEAVES MATCH TREE ===============================================================
+    // sanity check the values returned by the store and the tree
+    assert_eq!(
+        mtree.get_node(NodeIndex::new(mtree.depth(), 0)),
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)),
+        "node 0 must be the same for both MerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        mtree.get_node(NodeIndex::new(mtree.depth(), 1)),
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)),
+        "node 1 must be the same for both MerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        mtree.get_node(NodeIndex::new(mtree.depth(), 2)),
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)),
+        "node 2 must be the same for both MerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        mtree.get_node(NodeIndex::new(mtree.depth(), 3)),
+        store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)),
+        "node 3 must be the same for both MerkleTree and MerkleStore"
+    );
+
+    // STORE MERKLE PATH MATCHS ==============================================================
+    // assert the merkle path returned by the store is the same as the one in the tree
+    let result = store
+        .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 0))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[0], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        mtree.get_path(NodeIndex::new(mtree.depth(), 0)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 1))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[1], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        mtree.get_path(NodeIndex::new(mtree.depth(), 1)),
+        Ok(result.1),
+        "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 2))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[2], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        mtree.get_path(NodeIndex::new(mtree.depth(), 2)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 3))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[3], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        mtree.get_path(NodeIndex::new(mtree.depth(), 3)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_get_invalid_node() {
+    let mut store = MerkleStore::default();
+    let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work");
+    store
+        .add_merkle_tree(LEAVES4.to_vec())
+        .expect("adding a merkle tree to the store must work");
+    let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3));
+}
+
+#[test]
+fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
+    let mut store = MerkleStore::default();
+    let keys2: [u64; 2] = [0, 1];
+    let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)];
+    store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?;
+    let smt = SimpleSmt::new(1)
+        .unwrap()
+        .with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
+        .unwrap();
+
+    let idx = NodeIndex::new(1, 0);
+    assert_eq!(smt.get_node(&idx).unwrap(), leaves2[0]);
+    assert_eq!(
+        store.get_node(smt.root(), idx).unwrap(),
+        smt.get_node(&idx).unwrap()
+    );
+
+    let idx = NodeIndex::new(1, 1);
+    assert_eq!(smt.get_node(&idx).unwrap(), leaves2[1]);
+    assert_eq!(
+        store.get_node(smt.root(), idx).unwrap(),
+        smt.get_node(&idx).unwrap()
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
+    let mut store = MerkleStore::default();
+    store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
+
+    let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
+        .unwrap()
+        .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
+        .unwrap();
+
+    // STORE LEAVES ARE CORRECT ==============================================================
+    // checks the leaves in the store corresponds to the expected values
+    assert_eq!(
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 0)),
+        Ok(LEAVES4[0]),
+        "node 0 must be in the tree"
+    );
+    assert_eq!(
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 1)),
+        Ok(LEAVES4[1]),
+        "node 1 must be in the tree"
+    );
+    assert_eq!(
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 2)),
+        Ok(LEAVES4[2]),
+        "node 2 must be in the tree"
+    );
+    assert_eq!(
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 3)),
+        Ok(LEAVES4[3]),
+        "node 3 must be in the tree"
+    );
+
+    // STORE LEAVES MATCH TREE ===============================================================
+    // sanity check the values returned by the store and the tree
+    assert_eq!(
+        smt.get_node(&NodeIndex::new(smt.depth(), 0)),
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 0)),
+        "node 0 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        smt.get_node(&NodeIndex::new(smt.depth(), 1)),
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 1)),
+        "node 1 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        smt.get_node(&NodeIndex::new(smt.depth(), 2)),
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 2)),
+        "node 2 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        smt.get_node(&NodeIndex::new(smt.depth(), 3)),
+        store.get_node(smt.root(), NodeIndex::new(smt.depth(), 3)),
+        "node 3 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+
+    // STORE MERKLE PATH MATCHS ==============================================================
+    // assert the merkle path returned by the store is the same as the one in the tree
+    let result = store
+        .get_path(smt.root(), NodeIndex::new(smt.depth(), 0))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[0], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        smt.get_path(NodeIndex::new(smt.depth(), 0)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(smt.root(), NodeIndex::new(smt.depth(), 1))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[1], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        smt.get_path(NodeIndex::new(smt.depth(), 1)),
+        Ok(result.1),
+        "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(smt.root(), NodeIndex::new(smt.depth(), 2))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[2], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        smt.get_path(NodeIndex::new(smt.depth(), 2)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(smt.root(), NodeIndex::new(smt.depth(), 3))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[3], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        smt.get_path(NodeIndex::new(smt.depth(), 3)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_add_merkle_paths() -> Result<(), MerkleError> {
+    let mtree = MerkleTree::new(LEAVES4.to_vec())?;
+
+    let i0 = 0;
+    let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap();
+
+    let i1 = 1;
+    let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap();
+
+    let i2 = 2;
+    let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap();
+
+    let i3 = 3;
+    let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap();
+
+    let paths = [
+        (i0, LEAVES4[i0 as usize], p0),
+        (i1, LEAVES4[i1 as usize], p1),
+        (i2, LEAVES4[i2 as usize], p2),
+        (i3, LEAVES4[i3 as usize], p3),
+    ];
+
+    let mut store = MerkleStore::default();
+    store
+        .add_merkle_paths(paths.clone())
+        .expect("the valid paths must work");
+
+    let depth = 3;
+    let set = MerklePathSet::new(depth).with_paths(paths).unwrap();
+
+    // STORE LEAVES ARE CORRECT ==============================================================
+    // checks the leaves in the store corresponds to the expected values
+    assert_eq!(
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 0)),
+        Ok(LEAVES4[0]),
+        "node 0 must be in the set"
+    );
+    assert_eq!(
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 1)),
+        Ok(LEAVES4[1]),
+        "node 1 must be in the set"
+    );
+    assert_eq!(
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 2)),
+        Ok(LEAVES4[2]),
+        "node 2 must be in the set"
+    );
+    assert_eq!(
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 3)),
+        Ok(LEAVES4[3]),
+        "node 3 must be in the set"
+    );
+
+    // STORE LEAVES MATCH SET ================================================================
+    // sanity check the values returned by the store and the set
+    assert_eq!(
+        set.get_node(NodeIndex::new(set.depth(), 0)),
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 0)),
+        "node 0 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        set.get_node(NodeIndex::new(set.depth(), 1)),
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 1)),
+        "node 1 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        set.get_node(NodeIndex::new(set.depth(), 2)),
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 2)),
+        "node 2 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+    assert_eq!(
+        set.get_node(NodeIndex::new(set.depth(), 3)),
+        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 3)),
+        "node 3 must be the same for both SparseMerkleTree and MerkleStore"
+    );
+
+    // STORE MERKLE PATH MATCHS ==============================================================
+    // assert the merkle path returned by the store is the same as the one in the set
+    let result = store
+        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 0))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[0], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        set.get_path(NodeIndex::new(set.depth(), 0)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 1))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[1], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        set.get_path(NodeIndex::new(set.depth(), 1)),
+        Ok(result.1),
+        "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 2))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[2], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        set.get_path(NodeIndex::new(set.depth(), 2)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    let result = store
+        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 3))
+        .unwrap();
+    assert_eq!(
+        LEAVES4[3], result.0,
+        "Value for merkle path at index 0 must match leaf value"
+    );
+    assert_eq!(
+        set.get_path(NodeIndex::new(set.depth(), 3)),
+        Ok(result.1),
+        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn wont_open_to_different_depth_root() {
+    let empty = EmptySubtreeRoots::empty_hashes(64);
+    let a = [Felt::new(1); 4];
+    let b = [Felt::new(2); 4];
+
+    // Compute the root for a different depth. We cherry-pick this specific depth to prevent a
+    // regression to a bug in the past that allowed the user to fetch a node at a depth lower than
+    // the inserted path of a Merkle tree.
+    let mut root = Rpo256::merge(&[a.into(), b.into()]);
+    for depth in (1..=63).rev() {
+        root = Rpo256::merge(&[root, empty[depth]]);
+    }
+    let root = Word::from(root);
+
+    // For this example, the depth of the Merkle tree is 1, as we have only two leaves. Here we
+    // attempt to fetch a node on the maximum depth, and it should fail because the root shouldn't
+    // exist for the set.
+    let store = MerkleStore::default().with_merkle_tree([a, b]).unwrap();
+    let index = NodeIndex::root();
+    let err = store.get_node(root, index).err().unwrap();
+    assert_eq!(err, MerkleError::RootNotInStore(root));
+}
+
+#[test]
+fn store_path_opens_from_leaf() {
+    let a = [Felt::new(1); 4];
+    let b = [Felt::new(2); 4];
+    let c = [Felt::new(3); 4];
+    let d = [Felt::new(4); 4];
+    let e = [Felt::new(5); 4];
+    let f = [Felt::new(6); 4];
+    let g = [Felt::new(7); 4];
+    let h = [Felt::new(8); 4];
+
+    let i = Rpo256::merge(&[a.into(), b.into()]);
+    let j = Rpo256::merge(&[c.into(), d.into()]);
+    let k = Rpo256::merge(&[e.into(), f.into()]);
+    let l = Rpo256::merge(&[g.into(), h.into()]);
+
+    let m = Rpo256::merge(&[i.into(), j.into()]);
+    let n = Rpo256::merge(&[k.into(), l.into()]);
+
+    let root = Rpo256::merge(&[m.into(), n.into()]);
+
+    let store = MerkleStore::default()
+        .with_merkle_tree([a, b, c, d, e, f, g, h])
+        .unwrap();
+    let path = store.get_path(root.into(), NodeIndex::new(3, 1)).unwrap();
+
+    let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec());
+    assert_eq!(path.1, expected);
+}
+
+#[test]
+fn test_set_node() -> Result<(), MerkleError> {
+    let mtree = MerkleTree::new(LEAVES4.to_vec())?;
+    let mut store = MerkleStore::default().with_merkle_tree(LEAVES4)?;
+    let value = int_to_node(42);
+    let index = NodeIndex::new(mtree.depth(), 0);
+    let new_root = store.set_node(mtree.root(), index, value)?;
+    assert_eq!(
+        store.get_node(new_root, index),
+        Ok(value),
+        "Value must have changed"
+    );
+
+    Ok(())
+}