Skip to content

Commit

Permalink
Add missing methods to Smt (#268)
Browse files Browse the repository at this point in the history
  • Loading branch information
plafer authored and bobbinth committed Feb 14, 2024
1 parent 5fcf986 commit da12fd2
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 4 deletions.
46 changes: 42 additions & 4 deletions src/merkle/smt/full/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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 {
Expand Down Expand Up @@ -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],
Expand All @@ -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],
Expand All @@ -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> {
Expand Down
45 changes: 45 additions & 0 deletions src/merkle/smt/full/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
// --------------------------------------------------------------------------------------------

Expand Down

0 comments on commit da12fd2

Please sign in to comment.