feat: added handling of bottom tier to TieredSmt

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

@@ -47,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),
 }
 
@@ -65,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}"
@@ -79,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

@@ -1,7 +1,8 @@
 use super::{
-    BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, StarkField,
-    Vec, Word, EMPTY_WORD,
+    BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, MerkleError, MerklePath, NodeIndex, Rpo256,
+    RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO,
 };
+use core::cmp;
 
 #[cfg(test)]
 mod tests;
@@ -9,12 +10,34 @@ 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>,
-    bottom_leaves: BTreeMap<u64, Vec<RpoDigest>>,
+    upper_leaves: BTreeMap<NodeIndex, RpoDigest>, // node_index |-> key map
+    bottom_leaves: BTreeMap<u64, BottomLeaf>,     // leaves of depth 64
     values: BTreeMap<RpoDigest, Word>,
 }
 
@@ -22,192 +45,365 @@ 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
     // --------------------------------------------------------------------------------------------
 
-    pub fn new() -> Self {
-        Self {
-            root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0],
-            nodes: BTreeMap::new(),
-            upper_leaves: BTreeMap::new(),
-            bottom_leaves: BTreeMap::new(),
-            values: BTreeMap::new(),
+    /// 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> {
-        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 if !self.is_node_available(index) {
-            todo!()
-        }
-
-        Ok(self.get_branch_node(&index))
+        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> {
-        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 if !self.is_node_available(index) {
-            todo!()
-        }
+        self.validate_node_access(index)?;
 
         let mut path = Vec::with_capacity(index.depth() as usize);
         for _ in 0..index.depth() {
-            let node = self.get_branch_node(&index.sibling());
+            let node = self.get_node_unchecked(&index.sibling());
             path.push(node.into());
             index.move_up();
         }
+
         Ok(path.into())
     }
 
-    pub fn get_value(&self, key: RpoDigest) -> Result<Word, MerkleError> {
+    /// 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) => Ok(*value),
-            None => Ok(EMPTY_WORD),
+            Some(value) => *value,
+            None => EMPTY_WORD,
         }
     }
 
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
-    pub fn insert(&mut self, key: RpoDigest, value: Word) -> Result<Word, MerkleError> {
+    /// 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 {
-                let common_prefix_len = get_common_prefix_length(&key, &other_key);
-                let depth = common_prefix_len + 16;
+                // 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);
-                self.move_leaf_node(other_key, index, other_index);
+                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);
             }
         }
 
-        let old_value = self.values.insert(key, value).unwrap_or(EMPTY_WORD);
-        if value != old_value {
-            self.upper_leaves.insert(index, key);
-            let new_node = build_leaf_node(key, value, index.depth().into());
-            self.root = self.update_path(index, new_node);
-        }
-
-        Ok(old_value)
+        // insert the node and return the old value
+        self.insert_node(index, key, value);
+        old_value
     }
 
     // HELPER METHODS
     // --------------------------------------------------------------------------------------------
 
-    fn is_node_available(&self, index: NodeIndex) -> bool {
-        match index.depth() {
-            32 => true,
-            48 => true,
-            _ => true,
+    /// 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(())
     }
 
-    fn get_branch_node(&self, index: &NodeIndex) -> RpoDigest {
+    /// 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 (16..64).step_by(16) {
-            let index = NodeIndex::new(depth, mse >> (Self::MAX_DEPTH - depth)).unwrap();
+        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) {
-                continue;
-            } else {
+            } else if !self.nodes.contains_key(&index) {
                 return (index, None);
             }
         }
 
-        // TODO: handle bottom tier
-        unimplemented!()
+        // 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)
     }
 
-    fn move_leaf_node(&mut self, key: RpoDigest, old_index: NodeIndex, new_index: NodeIndex) {
-        self.upper_leaves.remove(&old_index).unwrap();
-        self.upper_leaves.insert(new_index, key);
-        let value = *self.values.get(&key).unwrap();
-        let new_node = build_leaf_node(key, value, new_index.depth().into());
-        self.update_path(new_index, new_node);
-    }
+    /// 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())
+        };
 
-    fn update_path(&mut self, mut index: NodeIndex, mut node: RpoDigest) -> RpoDigest {
+        // 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_branch_node(&index.sibling());
+            let sibling = self.get_node_unchecked(&index.sibling());
             node = Rpo256::merge(&index.build_node(node, sibling));
             index.move_up();
         }
-        node
+
+        // update the root
+        self.root = node;
     }
 }
 
 impl Default for TieredSmt {
     fn default() -> Self {
-        Self::new()
+        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);
-    let remaining = (key[3].as_int() << depth) >> depth;
-    key[3] = remaining.into();
+    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()
 }
 
-fn build_leaf_node(key: RpoDigest, value: Word, depth: u32) -> RpoDigest {
-    let remaining_path = get_remaining_path(key, depth);
-    Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.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)
 }
 
-fn get_common_prefix_length(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
+/// 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
+}
 
-    if e1 == e2 {
-        64
-    } else if e1 >> 16 == e2 >> 16 {
-        48
-    } else if e1 >> 32 == e2 >> 32 {
-        32
-    } else if e1 >> 48 == e2 >> 48 {
-        16
-    } else {
-        0
+/// 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,
     }
 }
 
-fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex {
-    let mse = Word::from(key)[3].as_int();
-    let value = match depth {
-        16 | 32 | 48 => mse >> (depth as u32),
-        _ => unreachable!("invalid depth: {depth}"),
-    };
+// BOTTOM LEAF
+// ================================================================================================
 
-    // TODO: use unchecked version?
-    NodeIndex::new(depth, value).unwrap()
+/// 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

@@ -1,11 +1,14 @@
 use super::{
-    super::{super::ONE, Felt, MerkleStore, WORD_SIZE},
-    get_remaining_path, EmptySubtreeRoots, NodeIndex, Rpo256, RpoDigest, TieredSmt, Word,
+    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::new();
+    let mut smt = TieredSmt::default();
     let mut store = MerkleStore::default();
 
     let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
@@ -15,15 +18,15 @@ fn tsmt_insert_one() {
     // 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 = compute_leaf_node(key, value, 16);
+    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).unwrap();
+    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).unwrap(), value);
+    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
@@ -32,15 +35,15 @@ fn tsmt_insert_one() {
 }
 
 #[test]
-fn tsmt_insert_two() {
-    let mut smt = TieredSmt::new();
+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).unwrap();
+    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,
@@ -48,28 +51,72 @@ fn tsmt_insert_two() {
     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).unwrap();
+    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 = compute_leaf_node(key_a, val_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 = compute_leaf_node(key_b, val_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).unwrap(), val_a);
+    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).unwrap(), val_b);
+    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);
@@ -77,14 +124,14 @@ fn tsmt_insert_two() {
 
 #[test]
 fn tsmt_insert_three() {
-    let mut smt = TieredSmt::new();
+    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).unwrap();
+    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,
@@ -92,7 +139,7 @@ fn tsmt_insert_three() {
     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).unwrap();
+    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,
@@ -101,37 +148,37 @@ fn tsmt_insert_three() {
     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).unwrap();
+    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 = compute_leaf_node(key_a, val_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 = compute_leaf_node(key_b, val_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 = compute_leaf_node(key_c, val_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).unwrap(), val_a);
+    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).unwrap(), val_b);
+    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).unwrap(), val_c);
+    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);
@@ -139,33 +186,161 @@ fn tsmt_insert_three() {
 
 #[test]
 fn tsmt_update() {
-    let mut smt = TieredSmt::new();
+    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).unwrap();
+    smt.insert(key, value_a);
 
-    // --- update value ---------------------------------------
+    // --- update the value ---------------------------------------------------
     let value_b = [Felt::new(2); WORD_SIZE];
-    smt.insert(key, value_b).unwrap();
+    smt.insert(key, value_b);
 
     // --- verify consistency -------------------------------------------------
     let mut tree_root = get_init_root();
     let index = NodeIndex::make(16, raw >> 48);
-    let leaf_node = compute_leaf_node(key, value_b, 16);
+    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).unwrap(), value_b);
+    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
 // ================================================================================================
 
@@ -173,7 +348,21 @@ fn get_init_root() -> Word {
     EmptySubtreeRoots::empty_hashes(64)[0].into()
 }
 
-fn compute_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest {
+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)
+}