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

75
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,38 +965,46 @@ 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.
// Pre-assertions. let (nodes, subtrees): (Vec<BTreeMap<_, _>>, Vec<Vec<SubtreeLeaf>>) = leaf_subtrees
assert!( .into_iter()
subtree.is_sorted(), .enumerate()
"subtree {i} at bottom-depth {current_depth} is not sorted", .map(|(i, subtree)| {
); // Pre-assertions.
assert!( assert!(
!subtree.is_empty(), subtree.is_sorted(),
"subtree {i} at bottom-depth {current_depth} is empty!", "subtree {i} at bottom-depth {current_depth} is not sorted",
); );
assert!(
// Do actual things. !subtree.is_empty(),
let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth); "subtree {i} at bottom-depth {current_depth} is empty!",
// Post-assertions.
assert!(next_leaves.is_sorted());
for (&index, test_node) in nodes.iter() {
let control_node = control.get_inner_node(index);
assert_eq!(
test_node, &control_node,
"depth {} subtree {}: test node does not match control at index {:?}",
current_depth, i, index,
); );
}
// Update state. // Do actual things.
accumulated_nodes.extend(nodes); let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth);
// Post-assertions.
assert!(next_leaves.is_sorted());
for subtree_leaf in next_leaves { for (&index, test_node) in nodes.iter() {
super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf); let control_node = control.get_inner_node(index);
} assert_eq!(
} test_node, &control_node,
"depth {} subtree {}: test node does not match control at index {:?}",
current_depth, i, index,
);
}
(nodes, next_leaves)
})
.unzip();
// 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}"); assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
} }