diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index cf3b421d..86ebaaf7 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -109,11 +109,21 @@ impl Smt { >::root(self) } - /// Returns the leaf at the specified index. + /// Returns the leaf to which `key` maps pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf { >::get_leaf(self, key) } + /// Returns the value associated with `key` + pub fn get_value(&self, key: &RpoDigest) -> Word { + let leaf_pos = LeafIndex::::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 { + self.leaves().flat_map(|(_, leaf)| leaf.entries()) + } + /// Returns an iterator over the inner nodes of this [Smt]. pub fn inner_nodes(&self) -> impl Iterator + '_ { 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 { - 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 { diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index b526cbf0..4a6d3e31 100644 --- a/src/merkle/smt/full/tests.rs +++ b/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 // --------------------------------------------------------------------------------------------