Merge pull request #152 from 0xPolygonMiden/bobbin-tsmt

Basic Tiered MST

src/hash/rpo/digest.rs

@@ -2,7 +2,7 @@ use super::{Digest, Felt, StarkField, DIGEST_SIZE, ZERO};
 use crate::utils::{
     string::String, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
 };
-use core::{cmp::Ordering, ops::Deref};
+use core::{cmp::Ordering, fmt::Display, ops::Deref};
 
 // DIGEST TRAIT IMPLEMENTATIONS
 // ================================================================================================
@@ -85,6 +85,28 @@ impl From for [Felt; DIGEST_SIZE] {
     }
 }
 
+impl From<&RpoDigest> for [u64; DIGEST_SIZE] {
+    fn from(value: &RpoDigest) -> Self {
+        [
+            value.0[0].as_int(),
+            value.0[1].as_int(),
+            value.0[2].as_int(),
+            value.0[3].as_int(),
+        ]
+    }
+}
+
+impl From<RpoDigest> for [u64; DIGEST_SIZE] {
+    fn from(value: RpoDigest) -> Self {
+        [
+            value.0[0].as_int(),
+            value.0[1].as_int(),
+            value.0[2].as_int(),
+            value.0[3].as_int(),
+        ]
+    }
+}
+
 impl From<&RpoDigest> for [u8; 32] {
     fn from(value: &RpoDigest) -> Self {
         value.as_bytes()
@@ -134,6 +156,15 @@ impl PartialOrd for RpoDigest {
     }
 }
 
+impl Display for RpoDigest {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        for byte in self.as_bytes() {
+            write!(f, "{byte:02x}")?;
+        }
+        Ok(())
+    }
+}
+
 // TESTS
 // ================================================================================================
 

src/merkle/index.rs

@@ -1,4 +1,5 @@
 use super::{Felt, MerkleError, RpoDigest, StarkField};
+use core::fmt::Display;
 
 // NODE INDEX
 // ================================================================================================
@@ -40,6 +41,12 @@ impl NodeIndex {
         }
     }
 
+    /// Creates a new node index without checking its validity.
+    pub const fn new_unchecked(depth: u8, value: u64) -> Self {
+        debug_assert!((64 - value.leading_zeros()) <= depth as u32);
+        Self { depth, value }
+    }
+
     /// Creates a new node index for testing purposes.
     ///
     /// # Panics
@@ -117,11 +124,26 @@ impl NodeIndex {
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
-    /// Traverse one level towards the root, decrementing the depth by `1`.
-    pub fn move_up(&mut self) -> &mut Self {
+    /// Traverses one level towards the root, decrementing the depth by `1`.
+    pub fn move_up(&mut self) {
         self.depth = self.depth.saturating_sub(1);
         self.value >>= 1;
-        self
+    }
+
+    /// Traverses towards the root until the specified depth is reached.
+    ///
+    /// Assumes that the specified depth is smaller than the current depth.
+    pub fn move_up_to(&mut self, depth: u8) {
+        debug_assert!(depth < self.depth);
+        let delta = self.depth.saturating_sub(depth);
+        self.depth = self.depth.saturating_sub(delta);
+        self.value >>= delta as u32;
+    }
+}
+
+impl Display for NodeIndex {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "depth={}, value={}", self.depth, self.value)
     }
 }
 

src/merkle/mod.rs

@@ -27,6 +27,9 @@ pub use path_set::MerklePathSet;
 mod simple_smt;
 pub use simple_smt::SimpleSmt;
 
+mod tiered_smt;
+pub use tiered_smt::TieredSmt;
+
 mod mmr;
 pub use mmr::{Mmr, MmrPeaks, MmrProof};
 
@@ -44,14 +47,15 @@ pub enum MerkleError {
     ConflictingRoots(Vec<Word>),
     DepthTooSmall(u8),
     DepthTooBig(u64),
-    DuplicateValuesForKey(u64),
-    NodeNotInStore(Word, NodeIndex),
-    NumLeavesNotPowerOfTwo(usize),
+    DuplicateValuesForIndex(u64),
+    DuplicateValuesForKey(RpoDigest),
     InvalidIndex { depth: u8, value: u64 },
     InvalidDepth { expected: u8, provided: u8 },
     InvalidPath(MerklePath),
     InvalidNumEntries(usize, usize),
-    NodeNotInSet(u64),
+    NodeNotInSet(NodeIndex),
+    NodeNotInStore(Word, NodeIndex),
+    NumLeavesNotPowerOfTwo(usize),
     RootNotInStore(Word),
 }
 
@@ -62,10 +66,8 @@ impl fmt::Display for MerkleError {
             ConflictingRoots(roots) => write!(f, "the merkle paths roots do not match {roots:?}"),
             DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"),
             DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"),
+            DuplicateValuesForIndex(key) => write!(f, "multiple values provided for key {key}"),
             DuplicateValuesForKey(key) => write!(f, "multiple values provided for key {key}"),
-            NumLeavesNotPowerOfTwo(leaves) => {
-                write!(f, "the leaves count {leaves} is not a power of 2")
-            }
             InvalidIndex{ depth, value} => write!(
                 f,
                 "the index value {value} is not valid for the depth {depth}"
@@ -76,8 +78,11 @@ impl fmt::Display for MerkleError {
             ),
             InvalidPath(_path) => write!(f, "the provided path is not valid"),
             InvalidNumEntries(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"),
-            NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"),
-            NodeNotInStore(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the store", hash, index.value(), index.depth(),),
+            NodeNotInSet(index) => write!(f, "the node with index ({index}) is not in the set"),
+            NodeNotInStore(hash, index) => write!(f, "the node {hash:?} with index ({index}) is not in the store"),
+            NumLeavesNotPowerOfTwo(leaves) => {
+                write!(f, "the leaves count {leaves} is not a power of 2")
+            }
             RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root),
         }
     }

src/merkle/path_set.rs

@@ -73,7 +73,7 @@ impl MerklePathSet {
         let path_key = index.value() - parity;
         self.paths
             .get(&path_key)
-            .ok_or(MerkleError::NodeNotInSet(path_key))
+            .ok_or(MerkleError::NodeNotInSet(index))
             .map(|path| path[parity as usize])
     }
 
@@ -104,11 +104,8 @@ impl MerklePathSet {
 
         let parity = index.value() & 1;
         let path_key = index.value() - parity;
-        let mut path = self
-            .paths
-            .get(&path_key)
-            .cloned()
-            .ok_or(MerkleError::NodeNotInSet(index.value()))?;
+        let mut path =
+            self.paths.get(&path_key).cloned().ok_or(MerkleError::NodeNotInSet(index))?;
         path.remove(parity as usize);
         Ok(path)
     }
@@ -200,7 +197,7 @@ impl MerklePathSet {
         let path_key = index.value() - parity;
         let path = match self.paths.get_mut(&path_key) {
             Some(path) => path,
-            None => return Err(MerkleError::NodeNotInSet(base_index_value)),
+            None => return Err(MerkleError::NodeNotInSet(index)),
         };
 
         // Fill old_hashes vector -----------------------------------------------------------------

src/merkle/simple_smt/mod.rs

@@ -92,12 +92,12 @@ impl SimpleSmt {
         for (key, value) in entries {
             let old_value = tree.update_leaf(key, value)?;
             if old_value != EMPTY_WORD || empty_entries.contains(&key) {
-                return Err(MerkleError::DuplicateValuesForKey(key));
+                return Err(MerkleError::DuplicateValuesForIndex(key));
             }
             // if we've processed an empty entry, add the key to the set of empty entry keys, and
             // if this key was already in the set, return an error
             if value == EMPTY_WORD && !empty_entries.insert(key) {
-                return Err(MerkleError::DuplicateValuesForKey(key));
+                return Err(MerkleError::DuplicateValuesForIndex(key));
             }
         }
         Ok(tree)

src/merkle/tiered_smt/mod.rs

@@ -0,0 +1,409 @@
+use super::{
+    BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, MerkleError, MerklePath, NodeIndex, Rpo256,
+    RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO,
+};
+use core::cmp;
+
+#[cfg(test)]
+mod tests;
+
+// TIERED SPARSE MERKLE TREE
+// ================================================================================================
+
+/// Tiered (compacted) Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and
+/// values are represented by 4 field elements.
+///
+/// Leaves in the tree can exist only on specific depths called "tiers". These depths are: 16, 32,
+/// 48, and 64. Initially, when a tree is empty, it is equivalent to an empty Sparse Merkle tree
+/// of depth 64 (i.e., leaves at depth 64 are set to [ZERO; 4]). As non-empty values are inserted
+/// into the tree they are added to the first available tier.
+///
+/// For example, when the first key-value is inserted, it will be stored in a node at depth 16
+/// such that the first 16 bits of the key determine the position of the node at depth 16. If
+/// another value with a key sharing the same 16-bit prefix is inserted, both values move into
+/// the next tier (depth 32). This process is repeated until values end up at tier 64. If multiple
+/// values have keys with a common 64-bit prefix, such key-value pairs are stored in a sorted list
+/// at the last tier (depth = 64).
+///
+/// To differentiate between internal and leaf nodes, node values are computed as follows:
+/// - Internal nodes: hash(left_child, right_child).
+/// - Leaf node at depths 16, 32, or 64: hash(rem_key, value, domain=depth).
+/// - Leaf node at depth 64: hash([rem_key_0, value_0, ..., rem_key_n, value_n, domain=64]).
+///
+/// Where rem_key is computed by replacing d most significant bits of the key with zeros where d
+/// is depth (i.e., for a leaf at depth 16, we replace 16 most significant bits of the key with 0).
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct TieredSmt {
+    root: RpoDigest,
+    nodes: BTreeMap<NodeIndex, RpoDigest>,
+    upper_leaves: BTreeMap<NodeIndex, RpoDigest>, // node_index |-> key map
+    bottom_leaves: BTreeMap<u64, BottomLeaf>,     // leaves of depth 64
+    values: BTreeMap<RpoDigest, Word>,
+}
+
+impl TieredSmt {
+    // CONSTANTS
+    // --------------------------------------------------------------------------------------------
+
+    /// The number of levels between tiers.
+    const TIER_SIZE: u8 = 16;
+
+    /// Depths at which leaves can exist in a tiered SMT.
+    const TIER_DEPTHS: [u8; 4] = [16, 32, 48, 64];
+
+    /// Maximum node depth. This is also the bottom tier of the tree.
+    const MAX_DEPTH: u8 = 64;
+
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a new [TieredSmt] instantiated with the specified key-value pairs.
+    ///
+    /// # Errors
+    /// Returns an error if the provided entries contain multiple values for the same key.
+    pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
+    where
+        R: IntoIterator<IntoIter = I>,
+        I: Iterator<Item = (RpoDigest, Word)> + ExactSizeIterator,
+    {
+        // create an empty tree
+        let mut tree = Self::default();
+
+        // append leaves to the tree returning an error if a duplicate entry for the same key
+        // is found
+        let mut empty_entries = BTreeSet::new();
+        for (key, value) in entries {
+            let old_value = tree.insert(key, value);
+            if old_value != EMPTY_WORD || empty_entries.contains(&key) {
+                return Err(MerkleError::DuplicateValuesForKey(key));
+            }
+            // if we've processed an empty entry, add the key to the set of empty entry keys, and
+            // if this key was already in the set, return an error
+            if value == EMPTY_WORD && !empty_entries.insert(key) {
+                return Err(MerkleError::DuplicateValuesForKey(key));
+            }
+        }
+        Ok(tree)
+    }
+
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the root of this Merkle tree.
+    pub const fn root(&self) -> RpoDigest {
+        self.root
+    }
+
+    /// Returns a node at the specified index.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The specified index depth is 0 or greater than 64.
+    /// - The node with the specified index does not exists in the Merkle tree. This is possible
+    ///   when a leaf node with the same index prefix exists at a tier higher than the requested
+    ///   node.
+    pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
+        self.validate_node_access(index)?;
+        Ok(self.get_node_unchecked(&index))
+    }
+
+    /// Returns a Merkle path from the node at the specified index to the root.
+    ///
+    /// The node itself is not included in the path.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The specified index depth is 0 or greater than 64.
+    /// - The node with the specified index does not exists in the Merkle tree. This is possible
+    ///   when a leaf node with the same index prefix exists at a tier higher than the node to
+    ///   which the path is requested.
+    pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
+        self.validate_node_access(index)?;
+
+        let mut path = Vec::with_capacity(index.depth() as usize);
+        for _ in 0..index.depth() {
+            let node = self.get_node_unchecked(&index.sibling());
+            path.push(node.into());
+            index.move_up();
+        }
+
+        Ok(path.into())
+    }
+
+    /// Returns the value associated with the specified key.
+    ///
+    /// If nothing was inserted into this tree for the specified key, [ZERO; 4] is returned.
+    pub fn get_value(&self, key: RpoDigest) -> Word {
+        match self.values.get(&key) {
+            Some(value) => *value,
+            None => EMPTY_WORD,
+        }
+    }
+
+    // STATE MUTATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Inserts the provided value into the tree under the specified key and returns the value
+    /// previously stored under this key.
+    ///
+    /// If the value for the specified key was not previously set, [ZERO; 4] is returned.
+    pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
+        // insert the value into the key-value map, and if nothing has changed, return
+        let old_value = self.values.insert(key, value).unwrap_or(EMPTY_WORD);
+        if old_value == value {
+            return old_value;
+        }
+
+        // determine the index for the value node; this index could have 3 different meanings:
+        // - it points to a root of an empty subtree (excluding depth = 64); in this case, we can
+        //   replace the node with the value node immediately.
+        // - it points to a node at the bottom tier (i.e., depth = 64); in this case, we need to
+        //   process bottom-tier insertion which will be handled by insert_node().
+        // - it points to a leaf node; this node could be a node with the same key or a different
+        //   key with a common prefix; in the latter case, we'll need to move the leaf to a lower
+        //   tier; for this scenario the `leaf_key` will contain the key of the leaf node
+        let (mut index, leaf_key) = self.get_insert_location(&key);
+
+        // if the returned index points to a leaf, and this leaf is for a different key, we need
+        // to move the leaf to a lower tier
+        if let Some(other_key) = leaf_key {
+            if other_key != key {
+                // determine how far down the tree should we move the existing leaf
+                let common_prefix_len = get_common_prefix_tier(&key, &other_key);
+                let depth = cmp::min(common_prefix_len + Self::TIER_SIZE, Self::MAX_DEPTH);
+
+                // move the leaf to the new location; this requires first removing the existing
+                // index, re-computing node value, and inserting the node at a new location
+                let other_index = key_to_index(&other_key, depth);
+                let other_value = *self.values.get(&other_key).expect("no value for other key");
+                self.upper_leaves.remove(&index).expect("other node key not in map");
+                self.insert_node(other_index, other_key, other_value);
+
+                // the new leaf also needs to move down to the same tier
+                index = key_to_index(&key, depth);
+            }
+        }
+
+        // insert the node and return the old value
+        self.insert_node(index, key, value);
+        old_value
+    }
+
+    // HELPER METHODS
+    // --------------------------------------------------------------------------------------------
+
+    /// Checks if the specified index is valid in the context of this Merkle tree.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The specified index depth is 0 or greater than 64.
+    /// - The node for the specified index does not exists in the Merkle tree. This is possible
+    ///   when an ancestors of the specified index is a leaf node.
+    fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> {
+        if index.is_root() {
+            return Err(MerkleError::DepthTooSmall(index.depth()));
+        } else if index.depth() > Self::MAX_DEPTH {
+            return Err(MerkleError::DepthTooBig(index.depth() as u64));
+        } else {
+            // make sure that there are no leaf nodes in the ancestors of the index; since leaf
+            // nodes can live at specific depth, we just need to check these depths.
+            let tier = get_index_tier(&index);
+            let mut tier_index = index;
+            for &depth in Self::TIER_DEPTHS[..tier].iter().rev() {
+                tier_index.move_up_to(depth);
+                if self.upper_leaves.contains_key(&tier_index) {
+                    return Err(MerkleError::NodeNotInSet(index));
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Returns a node at the specified index. If the node does not exist at this index, a root
+    /// for an empty subtree at the index's depth is returned.
+    ///
+    /// Unlike [TieredSmt::get_node()] this does not perform any checks to verify that the returned
+    /// node is valid in the context of this tree.
+    fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest {
+        match self.nodes.get(index) {
+            Some(node) => *node,
+            None => EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[index.depth() as usize],
+        }
+    }
+
+    /// Returns an index at which a node for the specified key should be inserted. If a leaf node
+    /// already exists at that index, returns the key associated with that leaf node.
+    ///
+    /// In case the index falls into the bottom tier (depth = 64), leaf node key is not returned
+    /// as the bottom tier may contain multiple key-value pairs in the same leaf.
+    fn get_insert_location(&self, key: &RpoDigest) -> (NodeIndex, Option<RpoDigest>) {
+        // traverse the tree from the root down checking nodes at tiers 16, 32, and 48. Return if
+        // a node at any of the tiers is either a leaf or a root of an empty subtree.
+        let mse = Word::from(key)[3].as_int();
+        for depth in (Self::TIER_DEPTHS[0]..Self::MAX_DEPTH).step_by(Self::TIER_SIZE as usize) {
+            let index = NodeIndex::new_unchecked(depth, mse >> (Self::MAX_DEPTH - depth));
+            if let Some(leaf_key) = self.upper_leaves.get(&index) {
+                return (index, Some(*leaf_key));
+            } else if !self.nodes.contains_key(&index) {
+                return (index, None);
+            }
+        }
+
+        // if we got here, that means all of the nodes checked so far are internal nodes, and
+        // the new node would need to be inserted in the bottom tier.
+        let index = NodeIndex::new_unchecked(Self::MAX_DEPTH, mse);
+        (index, None)
+    }
+
+    /// Inserts the provided key-value pair at the specified index and updates the root of this
+    /// Merkle tree by recomputing the path to the root.
+    fn insert_node(&mut self, mut index: NodeIndex, key: RpoDigest, value: Word) {
+        let depth = index.depth();
+
+        // insert the key into index-key map and compute the new value of the node
+        let mut node = if index.depth() == Self::MAX_DEPTH {
+            // for the bottom tier, we add the key-value pair to the existing leaf, or create a
+            // new leaf with this key-value pair
+            self.bottom_leaves
+                .entry(index.value())
+                .and_modify(|leaves| leaves.add_value(key, value))
+                .or_insert(BottomLeaf::new(key, value))
+                .hash()
+        } else {
+            // for the upper tiers, we just update the index-key map and compute the value of the
+            // node
+            self.upper_leaves.insert(index, key);
+            // the node value is computed as: hash(remaining_key || value, domain = depth)
+            let remaining_path = get_remaining_path(key, depth.into());
+            Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into())
+        };
+
+        // insert the node and update the path from the node to the root
+        for _ in 0..index.depth() {
+            self.nodes.insert(index, node);
+            let sibling = self.get_node_unchecked(&index.sibling());
+            node = Rpo256::merge(&index.build_node(node, sibling));
+            index.move_up();
+        }
+
+        // update the root
+        self.root = node;
+    }
+}
+
+impl Default for TieredSmt {
+    fn default() -> Self {
+        Self {
+            root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0],
+            nodes: BTreeMap::new(),
+            upper_leaves: BTreeMap::new(),
+            bottom_leaves: BTreeMap::new(),
+            values: BTreeMap::new(),
+        }
+    }
+}
+
+// HELPER FUNCTIONS
+// ================================================================================================
+
+/// Returns the remaining path for the specified key at the specified depth.
+///
+/// Remaining path is computed by setting n most significant bits of the key to zeros, where n is
+/// the specified depth.
+fn get_remaining_path(key: RpoDigest, depth: u32) -> RpoDigest {
+    let mut key = Word::from(key);
+    key[3] = if depth == 64 {
+        ZERO
+    } else {
+        // remove `depth` bits from the most significant key element
+        ((key[3].as_int() << depth) >> depth).into()
+    };
+    key.into()
+}
+
+/// Returns index for the specified key inserted at the specified depth.
+///
+/// The value for the key is computed by taking n most significant bits from the most significant
+/// element of the key, where n is the specified depth.
+fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex {
+    let mse = Word::from(key)[3].as_int();
+    let value = match depth {
+        16 | 32 | 48 | 64 => mse >> ((TieredSmt::MAX_DEPTH - depth) as u32),
+        _ => unreachable!("invalid depth: {depth}"),
+    };
+    NodeIndex::new_unchecked(depth, value)
+}
+
+/// Returns tiered common prefix length between the most significant elements of the provided keys.
+///
+/// Specifically:
+/// - returns 64 if the most significant elements are equal.
+/// - returns 48 if the common prefix is between 48 and 63 bits.
+/// - returns 32 if the common prefix is between 32 and 47 bits.
+/// - returns 16 if the common prefix is between 16 and 31 bits.
+/// - returns 0 if the common prefix is fewer than 16 bits.
+fn get_common_prefix_tier(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
+    let e1 = Word::from(key1)[3].as_int();
+    let e2 = Word::from(key2)[3].as_int();
+    let ex = (e1 ^ e2).leading_zeros() as u8;
+    (ex / 16) * 16
+}
+
+/// Returns a tier for the specified index.
+///
+/// The tiers are defined as follows:
+/// - Tier 0: depth 0 through 16 (inclusive).
+/// - Tier 1: depth 17 through 32 (inclusive).
+/// - Tier 2: depth 33 through 48 (inclusive).
+/// - Tier 3: depth 49 through 64 (inclusive).
+const fn get_index_tier(index: &NodeIndex) -> usize {
+    debug_assert!(index.depth() <= TieredSmt::MAX_DEPTH, "invalid depth");
+    match index.depth() {
+        0..=16 => 0,
+        17..=32 => 1,
+        33..=48 => 2,
+        _ => 3,
+    }
+}
+
+// BOTTOM LEAF
+// ================================================================================================
+
+/// Stores contents of the bottom leaf (i.e., leaf at depth = 64) in a [TieredSmt].
+///
+/// Bottom leaf can contain one or more key-value pairs all sharing the same 64-bit key prefix.
+/// The values are sorted by key to make sure the structure of the leaf is independent of the
+/// insertion order. This guarantees that a leaf with the same set of key-value pairs always has
+/// the same hash value.
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct BottomLeaf {
+    values: BTreeMap<[u64; 4], Word>,
+}
+
+impl BottomLeaf {
+    /// Returns a new [BottomLeaf] with a single key-value pair added.
+    pub fn new(key: RpoDigest, value: Word) -> Self {
+        let mut values = BTreeMap::new();
+        let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32);
+        values.insert(key.into(), value);
+        Self { values }
+    }
+
+    /// Adds a new key-value pair to this leaf.
+    pub fn add_value(&mut self, key: RpoDigest, value: Word) {
+        let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32);
+        self.values.insert(key.into(), value);
+    }
+
+    /// Computes a hash of this leaf.
+    pub fn hash(&self) -> RpoDigest {
+        let mut elements = Vec::with_capacity(self.values.len() * 2);
+        for (key, val) in self.values.iter() {
+            key.iter().for_each(|&v| elements.push(Felt::new(v)));
+            elements.extend_from_slice(val);
+        }
+        // TODO: hash in domain
+        Rpo256::hash_elements(&elements)
+    }
+}

src/merkle/tiered_smt/tests.rs

@@ -0,0 +1,368 @@
+use super::{
+    super::{
+        super::{ONE, ZERO},
+        Felt, MerkleStore, WORD_SIZE,
+    },
+    get_remaining_path, EmptySubtreeRoots, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word,
+};
+
+#[test]
+fn tsmt_insert_one() {
+    let mut smt = TieredSmt::default();
+    let mut store = MerkleStore::default();
+
+    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
+    let value = [ONE; WORD_SIZE];
+
+    // since the tree is empty, the first node will be inserted at depth 16 and the index will be
+    // 16 most significant bits of the key
+    let index = NodeIndex::make(16, raw >> 48);
+    let leaf_node = build_leaf_node(key, value, 16);
+    let tree_root = store.set_node(smt.root().into(), index, leaf_node.into()).unwrap().root;
+
+    smt.insert(key, value);
+
+    assert_eq!(smt.root(), tree_root.into());
+
+    // make sure the value was inserted, and the node is at the expected index
+    assert_eq!(smt.get_value(key), value);
+    assert_eq!(smt.get_node(index).unwrap(), leaf_node);
+
+    // make sure the paths we get from the store and the tree match
+    let expected_path = store.get_path(tree_root, index).unwrap();
+    assert_eq!(smt.get_path(index).unwrap(), expected_path.path);
+}
+
+#[test]
+fn tsmt_insert_two_16() {
+    let mut smt = TieredSmt::default();
+    let mut store = MerkleStore::default();
+
+    // --- insert the first value ---------------------------------------------
+    let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
+    let val_a = [ONE; WORD_SIZE];
+    smt.insert(key_a, val_a);
+
+    // --- insert the second value --------------------------------------------
+    // the key for this value has the same 16-bit prefix as the key for the first value,
+    // thus, on insertions, both values should be pushed to depth 32 tier
+    let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
+    let val_b = [Felt::new(2); WORD_SIZE];
+    smt.insert(key_b, val_b);
+
+    // --- build Merkle store with equivalent data ----------------------------
+    let mut tree_root = get_init_root();
+    let index_a = NodeIndex::make(32, raw_a >> 32);
+    let leaf_node_a = build_leaf_node(key_a, val_a, 32);
+    tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
+
+    let index_b = NodeIndex::make(32, raw_b >> 32);
+    let leaf_node_b = build_leaf_node(key_b, val_b, 32);
+    tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
+
+    // --- verify that data is consistent between store and tree --------------
+
+    assert_eq!(smt.root(), tree_root.into());
+
+    assert_eq!(smt.get_value(key_a), val_a);
+    assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
+    let expected_path = store.get_path(tree_root, index_a).unwrap().path;
+    assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
+
+    assert_eq!(smt.get_value(key_b), val_b);
+    assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
+    let expected_path = store.get_path(tree_root, index_b).unwrap().path;
+    assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
+}
+
+#[test]
+fn tsmt_insert_two_32() {
+    let mut smt = TieredSmt::default();
+    let mut store = MerkleStore::default();
+
+    // --- insert the first value ---------------------------------------------
+    let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
+    let val_a = [ONE; WORD_SIZE];
+    smt.insert(key_a, val_a);
+
+    // --- insert the second value --------------------------------------------
+    // the key for this value has the same 32-bit prefix as the key for the first value,
+    // thus, on insertions, both values should be pushed to depth 48 tier
+    let raw_b = 0b_10101010_10101010_00011111_11111111_00010110_10010011_11100000_00000000_u64;
+    let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
+    let val_b = [Felt::new(2); WORD_SIZE];
+    smt.insert(key_b, val_b);
+
+    // --- build Merkle store with equivalent data ----------------------------
+    let mut tree_root = get_init_root();
+    let index_a = NodeIndex::make(48, raw_a >> 16);
+    let leaf_node_a = build_leaf_node(key_a, val_a, 48);
+    tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
+
+    let index_b = NodeIndex::make(48, raw_b >> 16);
+    let leaf_node_b = build_leaf_node(key_b, val_b, 48);
+    tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
+
+    // --- verify that data is consistent between store and tree --------------
+
+    assert_eq!(smt.root(), tree_root.into());
+
+    assert_eq!(smt.get_value(key_a), val_a);
+    assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
+    let expected_path = store.get_path(tree_root, index_a).unwrap().path;
+    assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
+
+    assert_eq!(smt.get_value(key_b), val_b);
+    assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
+    let expected_path = store.get_path(tree_root, index_b).unwrap().path;
+    assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
+}
+
+#[test]
+fn tsmt_insert_three() {
+    let mut smt = TieredSmt::default();
+    let mut store = MerkleStore::default();
+
+    // --- insert the first value ---------------------------------------------
+    let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
+    let val_a = [ONE; WORD_SIZE];
+    smt.insert(key_a, val_a);
+
+    // --- insert the second value --------------------------------------------
+    // the key for this value has the same 16-bit prefix as the key for the first value,
+    // thus, on insertions, both values should be pushed to depth 32 tier
+    let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
+    let val_b = [Felt::new(2); WORD_SIZE];
+    smt.insert(key_b, val_b);
+
+    // --- insert the third value ---------------------------------------------
+    // the key for this value has the same 16-bit prefix as the keys for the first two,
+    // values; thus, on insertions, it will be inserted into depth 32 tier, but will not
+    // affect locations of the other two values
+    let raw_c = 0b_10101010_10101010_11011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
+    let val_c = [Felt::new(3); WORD_SIZE];
+    smt.insert(key_c, val_c);
+
+    // --- build Merkle store with equivalent data ----------------------------
+    let mut tree_root = get_init_root();
+    let index_a = NodeIndex::make(32, raw_a >> 32);
+    let leaf_node_a = build_leaf_node(key_a, val_a, 32);
+    tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
+
+    let index_b = NodeIndex::make(32, raw_b >> 32);
+    let leaf_node_b = build_leaf_node(key_b, val_b, 32);
+    tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
+
+    let index_c = NodeIndex::make(32, raw_c >> 32);
+    let leaf_node_c = build_leaf_node(key_c, val_c, 32);
+    tree_root = store.set_node(tree_root, index_c, leaf_node_c.into()).unwrap().root;
+
+    // --- verify that data is consistent between store and tree --------------
+
+    assert_eq!(smt.root(), tree_root.into());
+
+    assert_eq!(smt.get_value(key_a), val_a);
+    assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
+    let expected_path = store.get_path(tree_root, index_a).unwrap().path;
+    assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
+
+    assert_eq!(smt.get_value(key_b), val_b);
+    assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
+    let expected_path = store.get_path(tree_root, index_b).unwrap().path;
+    assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
+
+    assert_eq!(smt.get_value(key_c), val_c);
+    assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c);
+    let expected_path = store.get_path(tree_root, index_c).unwrap().path;
+    assert_eq!(smt.get_path(index_c).unwrap(), expected_path);
+}
+
+#[test]
+fn tsmt_update() {
+    let mut smt = TieredSmt::default();
+    let mut store = MerkleStore::default();
+
+    // --- insert a value into the tree ---------------------------------------
+    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
+    let value_a = [ONE; WORD_SIZE];
+    smt.insert(key, value_a);
+
+    // --- update the value ---------------------------------------------------
+    let value_b = [Felt::new(2); WORD_SIZE];
+    smt.insert(key, value_b);
+
+    // --- verify consistency -------------------------------------------------
+    let mut tree_root = get_init_root();
+    let index = NodeIndex::make(16, raw >> 48);
+    let leaf_node = build_leaf_node(key, value_b, 16);
+    tree_root = store.set_node(tree_root, index, leaf_node.into()).unwrap().root;
+
+    assert_eq!(smt.root(), tree_root.into());
+
+    assert_eq!(smt.get_value(key), value_b);
+    assert_eq!(smt.get_node(index).unwrap(), leaf_node);
+    let expected_path = store.get_path(tree_root, index).unwrap().path;
+    assert_eq!(smt.get_path(index).unwrap(), expected_path);
+}
+
+// BOTTOM TIER TESTS
+// ================================================================================================
+
+#[test]
+fn tsmt_bottom_tier() {
+    let mut smt = TieredSmt::default();
+    let mut store = MerkleStore::default();
+
+    // common prefix for the keys
+    let prefix = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+
+    // --- insert the first value ---------------------------------------------
+    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(prefix)]);
+    let val_a = [ONE; WORD_SIZE];
+    smt.insert(key_a, val_a);
+
+    // --- insert the second value --------------------------------------------
+    // this key has the same 64-bit prefix and thus both values should end up in the same
+    // node at depth 64
+    let key_b = RpoDigest::from([ZERO, ONE, ONE, Felt::new(prefix)]);
+    let val_b = [Felt::new(2); WORD_SIZE];
+    smt.insert(key_b, val_b);
+
+    // --- build Merkle store with equivalent data ----------------------------
+    let index = NodeIndex::make(64, prefix);
+    // to build bottom leaf we sort by key starting with the least significant element, thus
+    // key_b is smaller than key_a.
+    let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]);
+    let mut tree_root = get_init_root();
+    tree_root = store.set_node(tree_root, index, leaf_node.into()).unwrap().root;
+
+    // --- verify that data is consistent between store and tree --------------
+
+    assert_eq!(smt.root(), tree_root.into());
+
+    assert_eq!(smt.get_value(key_a), val_a);
+    assert_eq!(smt.get_value(key_b), val_b);
+
+    assert_eq!(smt.get_node(index).unwrap(), leaf_node);
+    let expected_path = store.get_path(tree_root, index).unwrap().path;
+    assert_eq!(smt.get_path(index).unwrap(), expected_path);
+}
+
+#[test]
+fn tsmt_bottom_tier_two() {
+    let mut smt = TieredSmt::default();
+    let mut store = MerkleStore::default();
+
+    // --- insert the first value ---------------------------------------------
+    let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
+    let val_a = [ONE; WORD_SIZE];
+    smt.insert(key_a, val_a);
+
+    // --- insert the second value --------------------------------------------
+    // the key for this value has the same 48-bit prefix as the key for the first value,
+    // thus, on insertions, both should end up in different nodes at depth 64
+    let raw_b = 0b_10101010_10101010_00011111_11111111_10010110_10010011_01100000_00000000_u64;
+    let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
+    let val_b = [Felt::new(2); WORD_SIZE];
+    smt.insert(key_b, val_b);
+
+    // --- build Merkle store with equivalent data ----------------------------
+    let mut tree_root = get_init_root();
+    let index_a = NodeIndex::make(64, raw_a);
+    let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]);
+    tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
+
+    let index_b = NodeIndex::make(64, raw_b);
+    let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]);
+    tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
+
+    // --- verify that data is consistent between store and tree --------------
+
+    assert_eq!(smt.root(), tree_root.into());
+
+    assert_eq!(smt.get_value(key_a), val_a);
+    assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
+    let expected_path = store.get_path(tree_root, index_a).unwrap().path;
+    assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
+
+    assert_eq!(smt.get_value(key_b), val_b);
+    assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
+    let expected_path = store.get_path(tree_root, index_b).unwrap().path;
+    assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
+}
+
+// ERROR TESTS
+// ================================================================================================
+
+#[test]
+fn tsmt_node_not_available() {
+    let mut smt = TieredSmt::default();
+
+    let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
+    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
+    let value = [ONE; WORD_SIZE];
+
+    // build an index which is just below the inserted leaf node
+    let index = NodeIndex::make(17, raw >> 47);
+
+    // since we haven't inserted the node yet, we should be able to get node and path to this index
+    assert!(smt.get_node(index).is_ok());
+    assert!(smt.get_path(index).is_ok());
+
+    smt.insert(key, value);
+
+    // but once the node is inserted, everything under it should be unavailable
+    assert!(smt.get_node(index).is_err());
+    assert!(smt.get_path(index).is_err());
+
+    let index = NodeIndex::make(32, raw >> 32);
+    assert!(smt.get_node(index).is_err());
+    assert!(smt.get_path(index).is_err());
+
+    let index = NodeIndex::make(34, raw >> 30);
+    assert!(smt.get_node(index).is_err());
+    assert!(smt.get_path(index).is_err());
+
+    let index = NodeIndex::make(50, raw >> 14);
+    assert!(smt.get_node(index).is_err());
+    assert!(smt.get_path(index).is_err());
+
+    let index = NodeIndex::make(64, raw);
+    assert!(smt.get_node(index).is_err());
+    assert!(smt.get_path(index).is_err());
+}
+
+// HELPER FUNCTIONS
+// ================================================================================================
+
+fn get_init_root() -> Word {
+    EmptySubtreeRoots::empty_hashes(64)[0].into()
+}
+
+fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest {
+    let remaining_path = get_remaining_path(key, depth as u32);
+    Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into())
+}
+
+fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest {
+    assert_eq!(keys.len(), values.len());
+
+    let mut elements = Vec::with_capacity(keys.len());
+    for (key, val) in keys.iter().zip(values.iter()) {
+        let mut key = Word::from(key);
+        key[3] = ZERO;
+        elements.extend_from_slice(&key);
+        elements.extend_from_slice(val);
+    }
+
+    Rpo256::hash_elements(&elements)
+}