feat: add leaves() iterator to SimpleSmt

This commit is contained in:
Bobbin Threadbare
2023-05-09 18:46:06 -07:00
parent 13aeda5a27
commit 629494b601
4 changed files with 57 additions and 44 deletions

View File

@@ -1,9 +1,9 @@
use super::mmr::Mmr;
use super::{
BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, MerkleTree,
NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word,
mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet,
MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word,
};
use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
use core::borrow::Borrow;
#[cfg(test)]
mod tests;
@@ -264,11 +264,15 @@ impl MerkleStore {
/// nodes which are descendants of the specified roots.
///
/// The roots for which no descendants exist in this Merkle store are ignored.
pub fn subset(&self, roots: &[Word]) -> MerkleStore {
pub fn subset<I, R>(&self, roots: I) -> MerkleStore
where
I: Iterator<Item = R>,
R: Borrow<Word>,
{
let mut store = MerkleStore::new();
for root in roots {
let root = RpoDigest::from(*root);
self.clone_tree_into(root, &mut store);
let root = RpoDigest::from(*root.borrow());
store.clone_tree_from(root, self);
}
store
}
@@ -393,17 +397,17 @@ impl MerkleStore {
// HELPER METHODS
// --------------------------------------------------------------------------------------------
/// Recursively clones a tree starting at the specified root into the specified target.
/// Recursively clones a tree with the specified root from the specified source into self.
///
/// If this Merkle store does not contain a tree with the specified root, this is a noop.
fn clone_tree_into(&self, root: RpoDigest, target: &mut Self) {
// process the node only if it is in this store
if let Some(node) = self.nodes.get(&root) {
/// If the source store does not contain a tree with the specified root, this is a noop.
fn clone_tree_from(&mut self, root: RpoDigest, source: &Self) {
// process the node only if it is in the source
if let Some(node) = source.nodes.get(&root) {
// if the node has already been inserted, no need to process it further as all of its
// descendants should be already in the target store
if matches!(target.nodes.insert(root, *node), None) {
self.clone_tree_into(node.left, target);
self.clone_tree_into(node.right, target);
// descendants should be already cloned from the source store
if matches!(self.nodes.insert(root, *node), None) {
self.clone_tree_from(node.left, source);
self.clone_tree_from(node.right, source);
}
}
}