make build_subtree also return the next leaf row

This commit is contained in:
Qyriad 2024-10-28 14:06:49 -06:00
parent 125aa9c96e
commit d23cce314e
3 changed files with 6 additions and 6 deletions

View file

@ -52,7 +52,7 @@ fn smt_subtree_even(c: &mut Criterion) {
}, },
|leaves| { |leaves| {
// Benchmarked function. // Benchmarked function.
let subtree = let (subtree, _) =
Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH)); Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH));
assert!(!subtree.is_empty()); assert!(!subtree.is_empty());
}, },
@ -100,7 +100,7 @@ fn smt_subtree_random(c: &mut Criterion) {
leaves leaves
}, },
|leaves| { |leaves| {
let subtree = let (subtree, _) =
Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH)); Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH));
assert!(!subtree.is_empty()); assert!(!subtree.is_empty());
}, },

View file

@ -253,7 +253,7 @@ impl Smt {
pub fn build_subtree( pub fn build_subtree(
leaves: Vec<(u64, RpoDigest)>, leaves: Vec<(u64, RpoDigest)>,
bottom_depth: u8, bottom_depth: u8,
) -> BTreeMap<NodeIndex, InnerNode> { ) -> (BTreeMap<NodeIndex, InnerNode>, Vec<(u64, RpoDigest)>) {
<Self as SparseMerkleTree<SMT_DEPTH>>::build_subtree(leaves, bottom_depth) <Self as SparseMerkleTree<SMT_DEPTH>>::build_subtree(leaves, bottom_depth)
} }
} }

View file

@ -361,7 +361,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
fn build_subtree( fn build_subtree(
mut leaves: Vec<(u64, RpoDigest)>, mut leaves: Vec<(u64, RpoDigest)>,
bottom_depth: u8, bottom_depth: u8,
) -> BTreeMap<NodeIndex, InnerNode> { ) -> (BTreeMap<NodeIndex, InnerNode>, Vec<(u64, RpoDigest)>) {
debug_assert!(bottom_depth <= DEPTH); debug_assert!(bottom_depth <= DEPTH);
debug_assert!(bottom_depth.is_multiple_of(&8)); debug_assert!(bottom_depth.is_multiple_of(&8));
debug_assert!(leaves.len() <= usize::pow(2, 8)); debug_assert!(leaves.len() <= usize::pow(2, 8));
@ -445,7 +445,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
mem::swap(&mut leaves, &mut next_leaves); mem::swap(&mut leaves, &mut next_leaves);
} }
inner_nodes (inner_nodes, leaves)
} }
} }
@ -618,7 +618,7 @@ mod test {
leaves.sort(); leaves.sort();
leaves.dedup_by_key(|leaf| leaf.0); leaves.dedup_by_key(|leaf| leaf.0);
let first_subtree = Smt::build_subtree(leaves, SMT_DEPTH); let (first_subtree, _) = Smt::build_subtree(leaves, SMT_DEPTH);
assert!(!first_subtree.is_empty()); assert!(!first_subtree.is_empty());
for (index, node) in first_subtree.into_iter() { for (index, node) in first_subtree.into_iter() {