feat: implement iterator over key-value pairs for TSMT

This commit is contained in:
Bobbin Threadbare 2023-08-01 11:02:29 -07:00
parent 1bb75e85dd
commit a03f2b5d5e
2 changed files with 11 additions and 0 deletions

View file

@ -214,6 +214,11 @@ impl TieredSmt {
// ITERATORS // ITERATORS
// -------------------------------------------------------------------------------------------- // --------------------------------------------------------------------------------------------
/// Returns an iterator over all key-value pairs in this [TieredSmt].
pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
self.values.iter()
}
/// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16 /// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16
/// 32, 48, or 64). /// 32, 48, or 64).
/// ///
@ -230,6 +235,7 @@ impl TieredSmt {
self.nodes.upper_leaves().map(|(index, node)| { self.nodes.upper_leaves().map(|(index, node)| {
let key_prefix = index_to_prefix(index); let key_prefix = index_to_prefix(index);
let (key, value) = self.values.get_first(key_prefix).expect("upper leaf not found"); let (key, value) = self.values.get_first(key_prefix).expect("upper leaf not found");
debug_assert_eq!(key_to_index(key, index.depth()), *index);
(*node, *key, *value) (*node, *key, *value)
}) })
} }

View file

@ -108,6 +108,11 @@ impl ValueStore {
None None
} }
/// Returns an iterator over all key-value pairs in this store.
pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
self.values.iter().flat_map(|(_, entry)| entry.iter())
}
// STATE MUTATORS // STATE MUTATORS
// -------------------------------------------------------------------------------------------- // --------------------------------------------------------------------------------------------