Merge pull request #153 from 0xPolygonMiden/bobbin-tsmt-iter

Tiered SMT iterators

src/merkle/index.rs

@@ -74,12 +74,26 @@ impl NodeIndex {
         Self { depth: 0, value: 0 }
     }
 
-    /// Computes the value of the sibling of the current node.
-    pub fn sibling(mut self) -> Self {
+    /// Computes sibling index of the current node.
+    pub const fn sibling(mut self) -> Self {
         self.value ^= 1;
         self
     }
 
+    /// Returns left child index of the current node.
+    pub const fn left_child(mut self) -> Self {
+        self.depth += 1;
+        self.value <<= 1;
+        self
+    }
+
+    /// Returns right child index of the current node.
+    pub const fn right_child(mut self) -> Self {
+        self.depth += 1;
+        self.value = (self.value << 1) + 1;
+        self
+    }
+
     // PROVIDERS
     // --------------------------------------------------------------------------------------------
 

src/merkle/store/mod.rs

@@ -1,6 +1,6 @@
 use super::{
     mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet,
-    MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word,
+    MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, TieredSmt, ValuePath, Vec, Word,
 };
 use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 use core::borrow::Borrow;
@@ -152,7 +152,6 @@ impl MerkleStore {
     /// The path starts at the sibling of the target leaf.
     ///
     /// # Errors
-    ///
     /// This method can return the following errors:
     /// - `RootNotInStore` if the `root` is not present in the store.
     /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
@@ -440,6 +439,14 @@ impl From<&Mmr> for MerkleStore {
     }
 }
 
+impl From<&TieredSmt> for MerkleStore {
+    fn from(value: &TieredSmt) -> Self {
+        let mut store = MerkleStore::new();
+        store.extend(value.inner_nodes());
+        store
+    }
+}
+
 impl FromIterator<InnerNodeInfo> for MerkleStore {
     fn from_iter<T: IntoIterator<Item = InnerNodeInfo>>(iter: T) -> Self {
         let mut store = MerkleStore::new();

src/merkle/tiered_smt/mod.rs

@@ -1,6 +1,6 @@
 use super::{
-    BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, MerkleError, MerklePath, NodeIndex, Rpo256,
-    RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO,
+    BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex,
+    Rpo256, RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO,
 };
 use core::cmp;
 
@@ -189,6 +189,52 @@ impl TieredSmt {
         old_value
     }
 
+    // ITERATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16
+    /// 32, 48, or 64).
+    ///
+    /// The iterator order is unspecified.
+    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
+        self.nodes.iter().filter_map(|(index, node)| {
+            if is_inner_node(index) {
+                Some(InnerNodeInfo {
+                    value: node.into(),
+                    left: self.get_node_unchecked(&index.left_child()).into(),
+                    right: self.get_node_unchecked(&index.right_child()).into(),
+                })
+            } else {
+                None
+            }
+        })
+    }
+
+    /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt].
+    ///
+    /// Each yielded item is a (node, key, value) tuple where key is a full un-truncated key (i.e.,
+    /// with key[3] element unmodified).
+    ///
+    /// The iterator order is unspecified.
+    pub fn upper_leaves(&self) -> impl Iterator<Item = (RpoDigest, RpoDigest, Word)> + '_ {
+        self.upper_leaves.iter().map(|(index, key)| {
+            let node = self.get_node_unchecked(index);
+            let value = self.get_value(*key);
+            (node, *key, value)
+        })
+    }
+
+    /// Returns an iterator over bottom leaves (i.e., depth = 64) of this [TieredSmt].
+    ///
+    /// Each yielded item consists of the hash of the leaf and its contents, where contents is
+    /// a vector containing key-value pairs of entries storied in this leaf. Note that keys are
+    /// un-truncated keys (i.e., with key[3] element unmodified).
+    ///
+    /// The iterator order is unspecified.
+    pub fn bottom_leaves(&self) -> impl Iterator<Item = (RpoDigest, Vec<(RpoDigest, Word)>)> + '_ {
+        self.bottom_leaves.values().map(|leaf| (leaf.hash(), leaf.contents()))
+    }
+
     // HELPER METHODS
     // --------------------------------------------------------------------------------------------
 
@@ -288,6 +334,7 @@ impl TieredSmt {
         }
 
         // update the root
+        self.nodes.insert(NodeIndex::root(), node);
         self.root = node;
     }
 }
@@ -367,6 +414,12 @@ const fn get_index_tier(index: &NodeIndex) -> usize {
     }
 }
 
+/// Returns true if the specified index is an index for an inner node (i.e., the depth is not 16,
+/// 32, 48, or 64).
+const fn is_inner_node(index: &NodeIndex) -> bool {
+    !matches!(index.depth(), 16 | 32 | 48 | 64)
+}
+
 // BOTTOM LEAF
 // ================================================================================================
 
@@ -378,16 +431,18 @@ const fn get_index_tier(index: &NodeIndex) -> usize {
 /// the same hash value.
 #[derive(Debug, Clone, PartialEq, Eq)]
 struct BottomLeaf {
+    prefix: u64,
     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 prefix = Word::from(key)[3].as_int();
         let mut values = BTreeMap::new();
         let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32);
         values.insert(key.into(), value);
-        Self { values }
+        Self { prefix, values }
     }
 
     /// Adds a new key-value pair to this leaf.
@@ -406,4 +461,22 @@ impl BottomLeaf {
         // TODO: hash in domain
         Rpo256::hash_elements(&elements)
     }
+
+    /// Returns contents of this leaf as a vector of (key, value) pairs.
+    ///
+    /// The keys are returned in their un-truncated form.
+    pub fn contents(&self) -> Vec<(RpoDigest, Word)> {
+        self.values
+            .iter()
+            .map(|(key, val)| {
+                let key = RpoDigest::from([
+                    Felt::new(key[0]),
+                    Felt::new(key[1]),
+                    Felt::new(key[2]),
+                    Felt::new(self.prefix),
+                ]);
+                (key, *val)
+            })
+            .collect()
+    }
 }

src/merkle/tiered_smt/tests.rs

@@ -1,9 +1,7 @@
 use super::{
-    super::{
-        super::{ONE, ZERO},
-        Felt, MerkleStore, WORD_SIZE,
-    },
-    get_remaining_path, EmptySubtreeRoots, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word,
+    super::{super::ONE, Felt, MerkleStore, WORD_SIZE, ZERO},
+    get_remaining_path, EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt,
+    Vec, Word,
 };
 
 #[test]
@@ -32,6 +30,17 @@ fn tsmt_insert_one() {
     // 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);
+
+    // make sure inner nodes match
+    let expected_nodes = get_non_empty_nodes(&store);
+    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
+    assert_eq!(actual_nodes.len(), expected_nodes.len());
+    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
+
+    // make sure leaves are returned correctly
+    let mut leaves = smt.upper_leaves();
+    assert_eq!(leaves.next(), Some((leaf_node, key, value)));
+    assert_eq!(leaves.next(), None);
 }
 
 #[test]
@@ -76,6 +85,18 @@ fn tsmt_insert_two_16() {
     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);
+
+    // make sure inner nodes match - the store contains more entries because it keeps track of
+    // all prior state - so, we don't check that the number of inner nodes is the same in both
+    let expected_nodes = get_non_empty_nodes(&store);
+    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
+    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
+
+    // make sure leaves are returned correctly
+    let mut leaves = smt.upper_leaves();
+    assert_eq!(leaves.next(), Some((leaf_node_a, key_a, val_a)));
+    assert_eq!(leaves.next(), Some((leaf_node_b, key_b, val_b)));
+    assert_eq!(leaves.next(), None);
 }
 
 #[test]
@@ -120,6 +141,12 @@ fn tsmt_insert_two_32() {
     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);
+
+    // make sure inner nodes match - the store contains more entries because it keeps track of
+    // all prior state - so, we don't check that the number of inner nodes is the same in both
+    let expected_nodes = get_non_empty_nodes(&store);
+    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
+    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
 }
 
 #[test]
@@ -182,6 +209,12 @@ fn tsmt_insert_three() {
     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);
+
+    // make sure inner nodes match - the store contains more entries because it keeps track of
+    // all prior state - so, we don't check that the number of inner nodes is the same in both
+    let expected_nodes = get_non_empty_nodes(&store);
+    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
+    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
 }
 
 #[test]
@@ -211,6 +244,12 @@ fn tsmt_update() {
     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);
+
+    // make sure inner nodes match - the store contains more entries because it keeps track of
+    // all prior state - so, we don't check that the number of inner nodes is the same in both
+    let expected_nodes = get_non_empty_nodes(&store);
+    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
+    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
 }
 
 // BOTTOM TIER TESTS
@@ -254,6 +293,17 @@ fn tsmt_bottom_tier() {
     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);
+
+    // make sure inner nodes match - the store contains more entries because it keeps track of
+    // all prior state - so, we don't check that the number of inner nodes is the same in both
+    let expected_nodes = get_non_empty_nodes(&store);
+    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
+    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
+
+    // make sure leaves are returned correctly
+    let mut leaves = smt.bottom_leaves();
+    assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a)])));
+    assert_eq!(leaves.next(), None);
 }
 
 #[test]
@@ -298,6 +348,18 @@ fn tsmt_bottom_tier_two() {
     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);
+
+    // make sure inner nodes match - the store contains more entries because it keeps track of
+    // all prior state - so, we don't check that the number of inner nodes is the same in both
+    let expected_nodes = get_non_empty_nodes(&store);
+    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
+    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
+
+    // make sure leaves are returned correctly
+    let mut leaves = smt.bottom_leaves();
+    assert_eq!(leaves.next(), Some((leaf_node_b, vec![(key_b, val_b)])));
+    assert_eq!(leaves.next(), Some((leaf_node_a, vec![(key_a, val_a)])));
+    assert_eq!(leaves.next(), None);
 }
 
 // ERROR TESTS
@@ -366,3 +428,14 @@ fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest {
 
     Rpo256::hash_elements(&elements)
 }
+
+fn get_non_empty_nodes(store: &MerkleStore) -> Vec<InnerNodeInfo> {
+    store
+        .inner_nodes()
+        .filter(|node| !is_empty_subtree(&RpoDigest::from(node.value)))
+        .collect::<Vec<_>>()
+}
+
+fn is_empty_subtree(node: &RpoDigest) -> bool {
+    EmptySubtreeRoots::empty_hashes(255).contains(&node)
+}