Merge pull request #93 from 0xPolygonMiden/hacka-add-merkle-store

Add merkle store

src/bit.rs

@@ -0,0 +1,169 @@
+/// Yields the bits of a `u64`.
+pub struct BitIterator {
+    /// The value that is being iterated bit-wise
+    value: u64,
+    /// True bits in the `mask` are the bits that have been visited.
+    mask: u64,
+}
+
+impl BitIterator {
+    pub fn new(value: u64) -> BitIterator {
+        BitIterator { value, mask: 0 }
+    }
+
+    /// An efficient skip implementation.
+    ///
+    /// Note: The compiler is smart enough to translate a `skip(n)` into a single shift instruction
+    /// if the code is inlined, however inlining does not always happen.
+    pub fn skip_front(mut self, n: u32) -> Self {
+        let mask = bitmask(n);
+        let ones = self.mask.trailing_ones();
+        let mask_position = ones;
+        self.mask ^= mask << mask_position;
+        self
+    }
+
+    /// An efficient skip from the back.
+    ///
+    /// Note: The compiler is smart enough to translate a `skip(n)` into a single shift instruction
+    /// if the code is inlined, however inlining does not always happen.
+    pub fn skip_back(mut self, n: u32) -> Self {
+        let mask = bitmask(n);
+        let ones = self.mask.leading_ones();
+        let mask_position = u64::BITS - ones - n;
+        self.mask ^= mask << mask_position;
+        self
+    }
+}
+
+impl Iterator for BitIterator {
+    type Item = bool;
+
+    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
+        // trailing_ones is implemented with trailing_zeros, and the zeros are computed with the
+        // intrinsic cttz. [Rust 1.67.0] x86 uses the `bsf` instruction. AArch64 uses the `rbit
+        // clz` instructions.
+        let ones = self.mask.trailing_ones();
+
+        if ones == u64::BITS {
+            None
+        } else {
+            let bit_position = ones;
+            let mask = 1 << bit_position;
+            self.mask ^= mask;
+            let bit = self.value & mask;
+            Some(bit != 0)
+        }
+    }
+}
+
+impl DoubleEndedIterator for BitIterator {
+    fn next_back(&mut self) -> Option<<Self as Iterator>::Item> {
+        // leading_ones is implemented with leading_zeros, and the zeros are computed with the
+        // intrinsic ctlz. [Rust 1.67.0] x86 uses the `bsr` instruction. AArch64 uses the `clz`
+        // instruction.
+        let ones = self.mask.leading_ones();
+
+        if ones == u64::BITS {
+            None
+        } else {
+            let bit_position = u64::BITS - ones - 1;
+            let mask = 1 << bit_position;
+            self.mask ^= mask;
+            let bit = self.value & mask;
+            Some(bit != 0)
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::BitIterator;
+
+    #[test]
+    fn test_bit_iterator() {
+        let v = 0b1;
+        let mut it = BitIterator::new(v);
+        assert!(it.next().unwrap(), "first bit is true");
+        assert!(it.all(|v| v == false), "every other value is false");
+
+        let v = 0b10;
+        let mut it = BitIterator::new(v);
+        assert!(!it.next().unwrap(), "first bit is false");
+        assert!(it.next().unwrap(), "first bit is true");
+        assert!(it.all(|v| v == false), "every other value is false");
+
+        let v = 0b10;
+        let mut it = BitIterator::new(v);
+        assert!(!it.next_back().unwrap(), "last bit is false");
+        assert!(!it.next().unwrap(), "first bit is false");
+        assert!(it.next().unwrap(), "first bit is true");
+        assert!(it.all(|v| v == false), "every other value is false");
+    }
+
+    #[test]
+    fn test_bit_iterator_skip() {
+        let v = 0b1;
+        let mut it = BitIterator::new(v).skip_front(1);
+        assert!(it.all(|v| v == false), "every other value is false");
+
+        let v = 0b10;
+        let mut it = BitIterator::new(v).skip_front(1);
+        assert!(it.next().unwrap(), "first bit is true");
+        assert!(it.all(|v| v == false), "every other value is false");
+
+        let high_bit = 0b1 << (u64::BITS - 1);
+        let mut it = BitIterator::new(high_bit).skip_back(1);
+        assert!(it.all(|v| v == false), "every other value is false");
+
+        let v = 0b10;
+        let mut it = BitIterator::new(v).skip_back(1);
+        assert!(!it.next_back().unwrap(), "last bit is false");
+        assert!(!it.next().unwrap(), "first bit is false");
+        assert!(it.next().unwrap(), "first bit is true");
+        assert!(it.all(|v| v == false), "every other value is false");
+    }
+
+    #[test]
+    fn test_skip_all() {
+        let v = 0b1;
+        let mut it = BitIterator::new(v).skip_front(u64::BITS);
+        assert!(it.next().is_none(), "iterator must be exhausted");
+
+        let v = 0b1;
+        let mut it = BitIterator::new(v).skip_back(u64::BITS);
+        assert!(it.next().is_none(), "iterator must be exhausted");
+    }
+
+    #[test]
+    fn test_bit_iterator_count_bits_after_skip() {
+        let any_value = 0b1;
+        for s in 0..u64::BITS {
+            let it = BitIterator::new(any_value).skip_front(s);
+            assert_eq!(it.count() as u32, u64::BITS - s)
+        }
+
+        let any_value = 0b1;
+        for s in 1..u64::BITS {
+            let it = BitIterator::new(any_value).skip_back(s);
+            assert_eq!(it.count() as u32, u64::BITS - s)
+        }
+    }
+
+    #[test]
+    fn test_bit_iterator_rev() {
+        let v = 0b1;
+        let mut it = BitIterator::new(v).rev();
+        assert!(it.nth(63).unwrap(), "the last value is true");
+    }
+}
+
+// UTILITIES
+// ===============================================================================================
+
+fn bitmask(s: u32) -> u64 {
+    match 1u64.checked_shl(s) {
+        Some(r) => r - 1,
+        None => u64::MAX,
+    }
+}

src/hash/rpo/digest.rs

@@ -73,12 +73,24 @@ impl From<[Felt; DIGEST_SIZE]> for RpoDigest {
     }
 }
 
+impl From<&RpoDigest> for [Felt; DIGEST_SIZE] {
+    fn from(value: &RpoDigest) -> Self {
+        value.0
+    }
+}
+
 impl From<RpoDigest> for [Felt; DIGEST_SIZE] {
     fn from(value: RpoDigest) -> Self {
         value.0
     }
 }
 
+impl From<&RpoDigest> for [u8; 32] {
+    fn from(value: &RpoDigest) -> Self {
+        value.as_bytes()
+    }
+}
+
 impl From<RpoDigest> for [u8; 32] {
     fn from(value: RpoDigest) -> Self {
         value.as_bytes()

src/lib.rs

@@ -4,6 +4,7 @@
 #[cfg_attr(test, macro_use)]
 extern crate alloc;
 
+mod bit;
 pub mod hash;
 pub mod merkle;
 

src/merkle/index.rs

@@ -1,4 +1,5 @@
 use super::{Felt, MerkleError, RpoDigest, StarkField};
+use crate::bit::BitIterator;
 
 // NODE INDEX
 // ================================================================================================
@@ -97,6 +98,19 @@ impl NodeIndex {
         self.depth == 0
     }
 
+    /// Returns a bit iterator for the `value`.
+    ///
+    /// Bits read from left-to-right represent which internal node's child should be visited to
+    /// arrive at the leaf. From the right-to-left the bit represent the position the hash of the
+    /// current element should go.
+    ///
+    /// Additionally, the value that is not visisted are the sibling values necessary for a Merkle
+    /// opening.
+    pub fn bit_iterator(&self) -> BitIterator {
+        let depth: u32 = self.depth.into();
+        BitIterator::new(self.value).skip_back(u64::BITS - depth)
+    }
+
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 

src/merkle/merkle_tree.rs

@@ -9,7 +9,7 @@ use winter_math::log2;
 /// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two).
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct MerkleTree {
-    nodes: Vec<Word>,
+    pub(crate) nodes: Vec<Word>,
 }
 
 impl MerkleTree {
@@ -108,6 +108,8 @@ impl MerkleTree {
             index.move_up();
         }
 
+        debug_assert!(index.is_root(), "the path must include the root");
+
         Ok(path.into())
     }
 

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;
@@ -29,13 +29,18 @@ pub use simple_smt::SimpleSmt;
 mod mmr;
 pub use mmr::{Mmr, MmrPeaks};
 
+mod store;
+pub use store::MerkleStore;
+
 // ERRORS
 // ================================================================================================
 
 #[derive(Clone, Debug)]
 pub enum MerkleError {
+    ConflictingRoots(Vec<Word>),
     DepthTooSmall(u8),
     DepthTooBig(u64),
+    NodeNotInStorage(Word, NodeIndex),
     NumLeavesNotPowerOfTwo(usize),
     InvalidIndex(NodeIndex),
     InvalidDepth { expected: u8, provided: u8 },
@@ -48,6 +53,7 @@ impl fmt::Display for MerkleError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         use MerkleError::*;
         match self {
+            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"),
             NumLeavesNotPowerOfTwo(leaves) => {
@@ -64,6 +70,7 @@ 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(),),
         }
     }
 }

src/merkle/simple_smt/mod.rs

@@ -15,7 +15,7 @@ mod tests;
 pub struct SimpleSmt {
     root: Word,
     depth: u8,
-    store: Store,
+    pub(crate) store: Store,
 }
 
 impl SimpleSmt {
@@ -207,17 +207,17 @@ impl SimpleSmt {
 /// respectively. Hashes for blank subtrees at each layer are stored in `empty_hashes`, beginning
 /// with the root hash of an empty tree, and ending with the zero value of a leaf node.
 #[derive(Debug, Clone, PartialEq, Eq)]
-struct Store {
-    branches: BTreeMap<NodeIndex, BranchNode>,
+pub(crate) struct Store {
+    pub(crate) branches: BTreeMap<NodeIndex, BranchNode>,
     leaves: BTreeMap<u64, Word>,
-    empty_hashes: Vec<RpoDigest>,
+    pub(crate) empty_hashes: Vec<RpoDigest>,
     depth: u8,
 }
 
 #[derive(Debug, Default, Clone, PartialEq, Eq)]
-struct BranchNode {
-    left: RpoDigest,
-    right: RpoDigest,
+pub(crate) struct BranchNode {
+    pub(crate) left: RpoDigest,
+    pub(crate) right: RpoDigest,
 }
 
 impl Store {

src/merkle/store.rs

@@ -0,0 +1,461 @@
+//! An in-memory data store for Merkle-lized data
+//!
+//! This is a in memory data store for Merkle trees, this store allows all the nodes of a tree
+//! (leaves or internal) to live as long as necessary and without duplication, this allows the
+//! implementation of efficient persistent data structures
+use super::{
+    BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerkleTree, NodeIndex, Rpo256,
+    RpoDigest, SimpleSmt, Vec, Word,
+};
+
+#[derive(Debug)]
+pub struct Node {
+    left: RpoDigest,
+    right: RpoDigest,
+}
+
+pub struct MerkleStore {
+    nodes: BTreeMap<RpoDigest, Node>,
+}
+
+impl Default for MerkleStore {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl MerkleStore {
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Creates an empty `MerkleStore` instance.
+    pub fn new() -> MerkleStore {
+        let mut nodes = BTreeMap::new();
+
+        // pre-populate the store with the empty hashes
+        let subtrees = EmptySubtreeRoots::empty_hashes(64);
+        for (child, parent) in subtrees.iter().zip(subtrees.iter().skip(1)) {
+            nodes.insert(
+                *parent,
+                Node {
+                    left: *child,
+                    right: *child,
+                },
+            );
+        }
+        MerkleStore { nodes }
+    }
+
+    /// 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.
+    ///
+    /// # Errors
+    ///
+    /// This method may return the following errors:
+    /// - `DepthTooSmall` if leaves is empty or contains only 1 element
+    /// - `NumLeavesNotPowerOfTwo` if the number of leaves is not a power-of-two
+    pub fn add_merkle_tree(&mut self, leaves: Vec<Word>) -> Result<Word, MerkleError> {
+        let layers = leaves.len().ilog2();
+        let tree = MerkleTree::new(leaves)?;
+
+        let mut depth = 0;
+        let mut parent_offset = 1;
+        let mut child_offset = 2;
+        while depth < layers {
+            let layer_size = 1usize << depth;
+            for _ in 0..layer_size {
+                // merkle tree is using level form representation, so left and right siblings are
+                // next to each other
+                let left = tree.nodes[child_offset];
+                let right = tree.nodes[child_offset + 1];
+                self.nodes.insert(
+                    tree.nodes[parent_offset].into(),
+                    Node {
+                        left: left.into(),
+                        right: right.into(),
+                    },
+                );
+                parent_offset += 1;
+                child_offset += 2;
+            }
+            depth += 1;
+        }
+
+        Ok(tree.nodes[1])
+    }
+
+    /// 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.
+    ///
+    /// # Errors
+    ///
+    /// This will return `InvalidEntriesCount` if the length of `entries` is not `63`.
+    pub fn add_sparse_merkle_tree<R, I>(&mut self, entries: R) -> Result<Word, MerkleError>
+    where
+        R: IntoIterator<IntoIter = I>,
+        I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
+    {
+        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)?.with_leaves(entries)?;
+        for branch in smt.store.branches.values() {
+            let parent = Rpo256::merge(&[branch.left, branch.right]);
+            self.nodes.insert(
+                parent,
+                Node {
+                    left: branch.left,
+                    right: branch.right,
+                },
+            );
+        }
+
+        Ok(smt.root())
+    }
+
+    /// 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.
+    pub fn add_merkle_path(
+        &mut self,
+        index_value: u64,
+        node: Word,
+        path: MerklePath,
+    ) -> Result<Word, MerkleError> {
+        let mut node = node;
+        let mut index = NodeIndex::new(self.nodes.len() as u8, index_value);
+
+        for sibling in path {
+            let (left, right) = match index.is_value_odd() {
+                true => (sibling, node),
+                false => (node, sibling),
+            };
+            let parent = Rpo256::merge(&[left.into(), right.into()]);
+            self.nodes.insert(
+                parent,
+                Node {
+                    left: left.into(),
+                    right: right.into(),
+                },
+            );
+
+            index.move_up();
+            node = parent.into();
+        }
+
+        Ok(node)
+    }
+
+    /// 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.
+    ///
+    /// # Errors
+    ///
+    /// Every path must resolve to the same root, otherwise this will return an `ConflictingRoots`
+    /// error.
+    pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<Word, MerkleError>
+    where
+        I: IntoIterator<Item = (u64, Word, MerklePath)>,
+    {
+        let paths: Vec<(u64, Word, MerklePath)> = paths.into_iter().collect();
+
+        let roots: BTreeSet<RpoDigest> = paths
+            .iter()
+            .map(|(index, node, path)| path.compute_root(*index, *node).into())
+            .collect();
+
+        if roots.len() != 1 {
+            return Err(MerkleError::ConflictingRoots(
+                roots.iter().map(|v| Word::from(*v)).collect(),
+            ));
+        }
+
+        for (index_value, node, path) in paths {
+            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
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the node at `index` rooted on the tree `root`.
+    ///
+    /// # Errors
+    ///
+    /// This will return `NodeNotInStorage` if the element is not present in the store.
+    pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
+        let mut hash: RpoDigest = root.into();
+        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)?;
+        if current_node != value {
+            self.add_merkle_path(index.value(), value, path)
+        } else {
+            Ok(root)
+        }
+    }
+
+    pub fn merge_roots(&mut self, root1: Word, root2: Word) -> Result<Word, MerkleError> {
+        let root1: RpoDigest = root1.into();
+        let root2: RpoDigest = root2.into();
+
+        if !self.nodes.contains_key(&root1) {
+            Err(MerkleError::NodeNotInStorage(
+                root1.into(),
+                NodeIndex::new(0, 0),
+            ))
+        } else if !self.nodes.contains_key(&root1) {
+            Err(MerkleError::NodeNotInStorage(
+                root2.into(),
+                NodeIndex::new(0, 0),
+            ))
+        } else {
+            let parent: Word = Rpo256::merge(&[root1, root2]).into();
+            self.nodes.insert(
+                parent.into(),
+                Node {
+                    left: root1,
+                    right: root2,
+                },
+            );
+
+            Ok(parent)
+        }
+    }
+}
+
+#[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_node_returns_self_for_root() {
+        let store = MerkleStore::default();
+        let root_idx = NodeIndex::new(0, 0);
+
+        // the root does not need any lookups in the storage itself, so the value is just returned
+        assert_eq!(store.get_node(LEAVES4[0], root_idx).unwrap(), LEAVES4[0]);
+        assert_eq!(store.get_node(LEAVES4[1], root_idx).unwrap(), LEAVES4[1]);
+        assert_eq!(store.get_node(LEAVES4[2], root_idx).unwrap(), LEAVES4[2]);
+        assert_eq!(store.get_node(LEAVES4[3], root_idx).unwrap(), LEAVES4[3]);
+    }
+
+    #[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(())
+    }
+}