qyriad/miden-crypto
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

convert test_singlethreaded_subtree to use SubtreeLeavesIter

Browse source
This commit is contained in:
Qyriad 2024-11-05 13:04:24 -07:00
parent 3274990951
commit 5de20ade48
1 changed files with 42 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
src/merkle/smt/mod.rs
View file

@ -765,14 +765,15 @@ impl<'s> core::iter::Iterator for SubtreeLeavesIter<'s> {
// ================================================================================================
#[cfg(test)]
mod test {
use core::mem;
use alloc::{collections::BTreeMap, vec::Vec};
use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf};
use super::{
InnerNode, LeafIndex, PairComputations, SmtLeaf, SparseMerkleTree, SubtreeLeaf,
SubtreeLeavesIter,
};
use crate::{
hash::rpo::RpoDigest,
merkle::{LeafIndex, NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
merkle::{NodeIndex, Smt, SMT_DEPTH},
Felt, Word, ONE,
};
@ -964,7 +965,11 @@ mod test {
} = Smt::sorted_pairs_to_leaves(entries);
for current_depth in (8..=SMT_DEPTH).step_by(8).rev() {
for (i, subtree) in mem::take(&mut leaf_subtrees).into_iter().enumerate() {
// There's no flat_map_unzip(), so this is the best we can do.
let (nodes, subtrees): (Vec<BTreeMap<_, _>>, Vec<Vec<SubtreeLeaf>>) = leaf_subtrees
.into_iter()
.enumerate()
.map(|(i, subtree)| {
// Pre-assertions.
assert!(
subtree.is_sorted(),
@ -977,9 +982,9 @@ mod test {
// Do actual things.
let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth);
// Post-assertions.
assert!(next_leaves.is_sorted());
for (&index, test_node) in nodes.iter() {
let control_node = control.get_inner_node(index);
assert_eq!(
@ -989,13 +994,17 @@ mod test {
);
}
// Update state.
accumulated_nodes.extend(nodes);
(nodes, next_leaves)
})
.unzip();
for subtree_leaf in next_leaves {
super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
}
}
// Update state between each depth iteration.
// FIXME: is this flatten or Box<dyn Iterator> better?
let mut all_leaves: Vec<SubtreeLeaf> = subtrees.into_iter().flatten().collect();
leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut all_leaves).collect();
accumulated_nodes.extend(nodes.into_iter().flatten());
assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
}