Add missing methods to `Smt` (#268)

src/merkle/smt/full/mod.rs

@@ -109,11 +109,21 @@ impl Smt {
         <Self as SparseMerkleTree<SMT_DEPTH>>::root(self)
     }
 
-    /// Returns the leaf at the specified index.
+    /// Returns the leaf to which `key` maps
     pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
         <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
     }
 
+    /// Returns the value associated with `key`
+    pub fn get_value(&self, key: &RpoDigest) -> Word {
+        let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
+
+        match self.leaves.get(&leaf_pos) {
+            Some(leaf) => leaf.get_value(key),
+            None => EMPTY_WORD,
+        }
+    }
+
     /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
     /// path to the leaf, as well as the leaf itself.
     pub fn open(&self, key: &RpoDigest) -> (MerklePath, SmtLeaf) {
@@ -130,6 +140,11 @@ impl Smt {
             .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
     }
 
+    /// Returns an iterator over the key-value pairs of this [Smt].
+    pub fn entries(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
+        self.leaves().flat_map(|(_, leaf)| leaf.entries())
+    }
+
     /// Returns an iterator over the inner nodes of this [Smt].
     pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
         self.inner_nodes.values().map(|e| InnerNodeInfo {
@@ -270,11 +285,11 @@ impl SmtLeaf {
 
     /// Converts a leaf to a list of field elements
     pub fn into_elements(self) -> Vec<Felt> {
-        self.into_kv_pairs().into_iter().flat_map(kv_to_elements).collect()
+        self.into_entries().into_iter().flat_map(kv_to_elements).collect()
     }
 
     /// Returns the key-value pairs in the leaf
-    pub fn kv_pairs(&self) -> Vec<&(RpoDigest, Word)> {
+    pub fn entries(&self) -> Vec<&(RpoDigest, Word)> {
         match self {
             SmtLeaf::Empty => Vec::new(),
             SmtLeaf::Single(kv_pair) => vec![kv_pair],
@@ -283,7 +298,7 @@ impl SmtLeaf {
     }
 
     /// Converts a leaf the key-value pairs in the leaf
-    pub fn into_kv_pairs(self) -> Vec<(RpoDigest, Word)> {
+    pub fn into_entries(self) -> Vec<(RpoDigest, Word)> {
         match self {
             SmtLeaf::Empty => Vec::new(),
             SmtLeaf::Single(kv_pair) => vec![kv_pair],
@@ -306,6 +321,29 @@ impl SmtLeaf {
     // HELPERS
     // ---------------------------------------------------------------------------------------------
 
+    /// Returns the value associated with `key` in the leaf
+    fn get_value(&self, key: &RpoDigest) -> Word {
+        match self {
+            SmtLeaf::Empty => EMPTY_WORD,
+            SmtLeaf::Single((key_in_leaf, value_in_leaf)) => {
+                if key == key_in_leaf {
+                    *value_in_leaf
+                } else {
+                    EMPTY_WORD
+                }
+            }
+            SmtLeaf::Multiple(kv_pairs) => {
+                for (key_in_leaf, value_in_leaf) in kv_pairs {
+                    if key == key_in_leaf {
+                        return *value_in_leaf;
+                    }
+                }
+
+                EMPTY_WORD
+            }
+        }
+    }
+
     /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if
     /// any.
     fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {

src/merkle/smt/full/tests.rs

@@ -261,6 +261,51 @@ fn test_empty_leaf_hash() {
     assert_eq!(leaf.hash(), EMPTY_WORD.into());
 }
 
+/// Tests that `get_value()` works as expected
+#[test]
+fn test_smt_get_value() {
+    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
+    let key_2: RpoDigest =
+        RpoDigest::from([2_u64.into(), 2_u64.into(), 2_u64.into(), 2_u64.into()]);
+
+    let value_1 = [ONE; WORD_SIZE];
+    let value_2 = [2_u64.into(); WORD_SIZE];
+
+    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
+
+    let returned_value_1 = smt.get_value(&key_1);
+    let returned_value_2 = smt.get_value(&key_2);
+
+    assert_eq!(value_1, returned_value_1);
+    assert_eq!(value_2, returned_value_2);
+
+    // Check that a key with no inserted value returns the empty word
+    let key_no_value =
+        RpoDigest::from([42_u64.into(), 42_u64.into(), 42_u64.into(), 42_u64.into()]);
+
+    assert_eq!(EMPTY_WORD, smt.get_value(&key_no_value));
+}
+
+/// Tests that `entries()` works as expected
+#[test]
+fn test_smt_entries() {
+    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
+    let key_2: RpoDigest =
+        RpoDigest::from([2_u64.into(), 2_u64.into(), 2_u64.into(), 2_u64.into()]);
+
+    let value_1 = [ONE; WORD_SIZE];
+    let value_2 = [2_u64.into(); WORD_SIZE];
+
+    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
+
+    let mut entries = smt.entries();
+
+    // Note: for simplicity, we assume the order `(k1,v1), (k2,v2)`. If a new implementation
+    // switches the order, it is OK to modify the order here as well.
+    assert_eq!(&(key_1, value_1), entries.next().unwrap());
+    assert_eq!(&(key_2, value_2), entries.next().unwrap());
+    assert!(entries.next().is_none());
+}
 // HELPERS
 // --------------------------------------------------------------------------------------------