feat: implement additional leaf traversal methods on MerkleStore

CHANGELOG.md

@@ -1,6 +1,11 @@
 ## 0.7.0 (TBD)
 
 * Replaced `MerklePathSet` with `PartialMerkleTree` (#165).
+* Implemented clearing of nodes in `TieredSmt` (#173).
+* Added ability to generate inclusion proofs for `TieredSmt` (#174).
+* Added conditional `serde`` support for various structs (#180).
+* Implemented benchmarking for `TieredSmt` (#182).
+* Added more leaf traversal methods for `MerkleStore` (#185).
 
 ## 0.6.0 (2023-06-25)
 

src/merkle/path.rs

@@ -160,6 +160,16 @@ pub struct ValuePath {
     pub path: MerklePath,
 }
 
+impl ValuePath {
+    /// Returns a new [ValuePath] instantiated from the specified value and path.
+    pub fn new(value: RpoDigest, path: Vec<RpoDigest>) -> Self {
+        Self {
+            value,
+            path: MerklePath::new(path),
+        }
+    }
+}
+
 /// A container for a [MerklePath] and its [Word] root.
 ///
 /// This structure does not provide any guarantees regarding the correctness of the path to the

src/merkle/store/mod.rs

@@ -173,27 +173,24 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
         // the path is computed from root to leaf, so it must be reversed
         path.reverse();
 
-        Ok(ValuePath {
-            value: hash,
-            path: MerklePath::new(path),
-        })
+        Ok(ValuePath::new(hash, path))
     }
 
-    /// Reconstructs a path from the root until a leaf or empty node and returns its depth.
+    // LEAF TRAVERSAL
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the depth of the first leaf or an empty node encountered while traversing the tree
+    /// from the specified root down according to the provided index.
     ///
-    /// The `tree_depth` parameter defines up to which depth the tree will be traversed, starting
-    /// from `root`. The maximum value the argument accepts is [u64::BITS].
-    ///
-    /// The traversed path from leaf to root will start at the least significant bit of `index`,
-    /// and will be executed for `tree_depth` bits.
+    /// The `tree_depth` parameter specifies the depth of the tree rooted at `root`. The
+    /// maximum value the argument accepts is [u64::BITS].
     ///
     /// # Errors
     /// Will return an error if:
     /// - The provided root is not found.
-    /// - The path from the root continues to a depth greater than `tree_depth`.
-    /// - The provided `tree_depth` is greater than `64.
-    /// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more
-    ///   information, check [NodeIndex::new].
+    /// - The provided `tree_depth` is greater than 64.
+    /// - The provided `index` is not valid for a depth equivalent to `tree_depth`.
+    /// - No leaf or an empty node was found while traversing the tree down to `tree_depth`.
     pub fn get_leaf_depth(
         &self,
         root: RpoDigest,
@@ -206,13 +203,6 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
         }
         NodeIndex::new(tree_depth, index)?;
 
-        // it's not illegal to have a maximum depth of `0`; we should just return the root in that
-        // case. this check will simplify the implementation as we could overflow bits for depth
-        // `0`.
-        if tree_depth == 0 {
-            return Ok(0);
-        }
-
         // check if the root exists, providing the proper error report if it doesn't
         let empty = EmptySubtreeRoots::empty_hashes(tree_depth);
         let mut hash = root;
@@ -224,7 +214,7 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
         let mut path = (index << (64 - tree_depth)).reverse_bits();
 
         // iterate every depth and reconstruct the path from root to leaf
-        for depth in 0..tree_depth {
+        for depth in 0..=tree_depth {
             // we short-circuit if an empty node has been found
             if hash == empty[depth as usize] {
                 return Ok(depth);
@@ -241,13 +231,77 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
             path >>= 1;
         }
 
-        // at max depth assert it doesn't have sub-trees
-        if self.nodes.contains_key(&hash) {
-            return Err(MerkleError::DepthTooBig(tree_depth as u64 + 1));
+        // return an error because we exhausted the index but didn't find either a leaf or an
+        // empty node
+        Err(MerkleError::DepthTooBig(tree_depth as u64 + 1))
     }
 
-        // depleted bits; return max depth
-        Ok(tree_depth)
+    /// Returns index and value of a leaf node which is the only leaf node in a subtree defined by
+    /// the provided root. If the subtree contains zero or more than one leaf nodes None is
+    /// returned.
+    ///
+    /// The `tree_depth` parameter specifies the depth of the parent tree such that `root` is
+    /// located in this tree at `root_index`. The maximum value the argument accepts is
+    /// [u64::BITS].
+    ///
+    /// # Errors
+    /// Will return an error if:
+    /// - The provided root is not found.
+    /// - The provided `tree_depth` is greater than 64.
+    /// - The provided `root_index` has depth greater than `tree_depth`.
+    /// - A lone node at depth `tree_depth` is not a leaf node.
+    pub fn find_lone_leaf(
+        &self,
+        root: RpoDigest,
+        root_index: NodeIndex,
+        tree_depth: u8,
+    ) -> Result<Option<(NodeIndex, RpoDigest)>, MerkleError> {
+        // we set max depth at u64::BITS as this is the largest meaningful value for a 64-bit index
+        const MAX_DEPTH: u8 = u64::BITS as u8;
+        if tree_depth > MAX_DEPTH {
+            return Err(MerkleError::DepthTooBig(tree_depth as u64));
+        }
+        let empty = EmptySubtreeRoots::empty_hashes(MAX_DEPTH);
+
+        let mut node = root;
+        if !self.nodes.contains_key(&node) {
+            return Err(MerkleError::RootNotInStore(node));
+        }
+
+        let mut index = root_index;
+        if index.depth() > tree_depth {
+            return Err(MerkleError::DepthTooBig(index.depth() as u64));
+        }
+
+        // traverse down following the path of single non-empty nodes; this works because if a
+        // node has two empty children it cannot contain a lone leaf. similarly if a node has
+        // two non-empty children it must contain at least two leaves.
+        for depth in index.depth()..tree_depth {
+            // if the node is a leaf, return; otherwise, examine the node's children
+            let children = match self.nodes.get(&node) {
+                Some(node) => node,
+                None => return Ok(Some((index, node))),
+            };
+
+            let empty_node = empty[depth as usize + 1];
+            node = if children.left != empty_node && children.right == empty_node {
+                index = index.left_child();
+                children.left
+            } else if children.left == empty_node && children.right != empty_node {
+                index = index.right_child();
+                children.right
+            } else {
+                return Ok(None);
+            };
+        }
+
+        // if we are here, we got to `tree_depth`; thus, either the current node is a leaf node,
+        // and so we return it, or it is an internal node, and then we return an error
+        if self.nodes.contains_key(&node) {
+            Err(MerkleError::DepthTooBig(tree_depth as u64 + 1))
+        } else {
+            Ok(Some((index, node)))
+        }
     }
 
     // DATA EXTRACTORS

src/merkle/store/tests.rs

@@ -637,6 +637,9 @@ fn node_path_should_be_truncated_by_midtier_insert() {
     assert!(store.get_node(root, index).is_err());
 }
 
+// LEAF TRAVERSAL
+// ================================================================================================
+
 #[test]
 fn get_leaf_depth_works_depth_64() {
     let mut store = MerkleStore::new();
@@ -747,6 +750,67 @@ fn get_leaf_depth_works_with_depth_8() {
     assert_eq!(Err(MerkleError::DepthTooBig(9)), store.get_leaf_depth(root, 8, a));
 }
 
+#[test]
+fn find_lone_leaf() {
+    let mut store = MerkleStore::new();
+    let empty = EmptySubtreeRoots::empty_hashes(64);
+    let mut root: RpoDigest = empty[0];
+
+    // insert a single leaf into the store at depth 64
+    let key_a = 0b01010101_10101010_00001111_01110100_00111011_10101101_00000100_01000001_u64;
+    let idx_a = NodeIndex::make(64, key_a);
+    let val_a = RpoDigest::from([ONE, ONE, ONE, ONE]);
+    root = store.set_node(root, idx_a, val_a).unwrap().root;
+
+    // for every ancestor of A, A should be a long leaf
+    for depth in 1..64 {
+        let parent_index = NodeIndex::make(depth, key_a >> (64 - depth));
+        let parent = store.get_node(root, parent_index).unwrap();
+
+        let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
+        assert_eq!(res, Some((idx_a, val_a)));
+    }
+
+    // insert another leaf into the store such that it has the same 8 bit prefix as A
+    let key_b = 0b01010101_01111010_00001111_01110100_00111011_10101101_00000100_01000001_u64;
+    let idx_b = NodeIndex::make(64, key_b);
+    let val_b = RpoDigest::from([ONE, ONE, ONE, ZERO]);
+    root = store.set_node(root, idx_b, val_b).unwrap().root;
+
+    // for any node which is common between A and B, find_lone_leaf() should return None as the
+    // node has two descendants
+    for depth in 1..9 {
+        let parent_index = NodeIndex::make(depth, key_a >> (64 - depth));
+        let parent = store.get_node(root, parent_index).unwrap();
+
+        let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
+        assert_eq!(res, None);
+    }
+
+    // for other ancestors of A and B, A and B should be lone leaves respectively
+    for depth in 9..64 {
+        let parent_index = NodeIndex::make(depth, key_a >> (64 - depth));
+        let parent = store.get_node(root, parent_index).unwrap();
+
+        let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
+        assert_eq!(res, Some((idx_a, val_a)));
+    }
+
+    for depth in 9..64 {
+        let parent_index = NodeIndex::make(depth, key_b >> (64 - depth));
+        let parent = store.get_node(root, parent_index).unwrap();
+
+        let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
+        assert_eq!(res, Some((idx_b, val_b)));
+    }
+
+    // for any other node, find_lone_leaf() should return None as they have no leaf nodes
+    let parent_index = NodeIndex::make(16, 0b01010101_11111111);
+    let parent = store.get_node(root, parent_index).unwrap();
+    let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
+    assert_eq!(res, None);
+}
+
 // SUBSET EXTRACTION
 // ================================================================================================
 

src/merkle/tiered_smt/mod.rs

@@ -55,13 +55,13 @@ impl TieredSmt {
     // --------------------------------------------------------------------------------------------
 
     /// The number of levels between tiers.
-    const TIER_SIZE: u8 = 16;
+    pub const TIER_SIZE: u8 = 16;
 
     /// Depths at which leaves can exist in a tiered SMT.
-    const TIER_DEPTHS: [u8; 4] = [16, 32, 48, 64];
+    pub 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;
+    pub const MAX_DEPTH: u8 = 64;
 
     /// Value of an empty leaf.
     pub const EMPTY_VALUE: Word = super::empty_roots::EMPTY_WORD;