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)] #[cfg(test)]
mod test { mod test {
use core::mem;
use alloc::{collections::BTreeMap, vec::Vec}; use alloc::{collections::BTreeMap, vec::Vec};
use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf}; use super::{
InnerNode, LeafIndex, PairComputations, SmtLeaf, SparseMerkleTree, SubtreeLeaf,
SubtreeLeavesIter,
};
use crate::{ use crate::{
hash::rpo::RpoDigest, hash::rpo::RpoDigest,
merkle::{LeafIndex, NodeIndex, Smt, SmtLeaf, SMT_DEPTH}, merkle::{NodeIndex, Smt, SMT_DEPTH},
Felt, Word, ONE, Felt, Word, ONE,
}; };
@ -964,7 +965,11 @@ mod test {
} = Smt::sorted_pairs_to_leaves(entries); } = Smt::sorted_pairs_to_leaves(entries);
for current_depth in (8..=SMT_DEPTH).step_by(8).rev() { 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. // Pre-assertions.
assert!( assert!(
subtree.is_sorted(), subtree.is_sorted(),
@ -977,9 +982,9 @@ mod test {
// Do actual things. // Do actual things.
let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth); let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth);
// Post-assertions. // Post-assertions.
assert!(next_leaves.is_sorted()); assert!(next_leaves.is_sorted());
for (&index, test_node) in nodes.iter() { for (&index, test_node) in nodes.iter() {
let control_node = control.get_inner_node(index); let control_node = control.get_inner_node(index);
assert_eq!( assert_eq!(
@ -989,13 +994,17 @@ mod test {
); );
} }
// Update state. (nodes, next_leaves)
accumulated_nodes.extend(nodes); })
.unzip();
for subtree_leaf in next_leaves { // Update state between each depth iteration.
super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
} // 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}"); assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
} }