From 1bc790586aeb93cd076336e0f22ed3eebd66ebe6 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Tue, 29 Oct 2024 12:29:20 -0600 Subject: [PATCH] refactor test_build_subtree_from_leaves() --- src/merkle/smt/mod.rs | 33 +++++++++++++++------------------ 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 0c87724..0c55bd0 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -683,32 +683,29 @@ mod test { assert_eq!(control_subtree_leaves, test_subtree_leaves); } - #[test] - fn test_build_subtree_from_leaves() { - const PAIR_COUNT: u64 = u64::pow(2, 8); - - let entries: Vec<(RpoDigest, Word)> = (0..PAIR_COUNT) + // Helper for the below tests. + fn generate_entries(pair_count: u64) -> Vec<(RpoDigest, Word)> { + (0..pair_count) .map(|i| { - let leaf_index = ((i as f64 / PAIR_COUNT as f64) * (PAIR_COUNT as f64)) as u64; + let leaf_index = ((i as f64 / pair_count as f64) * (pair_count as f64)) as u64; let key = RpoDigest::new([ONE, ONE, Felt::new(i), Felt::new(leaf_index)]); let value = [ONE, ONE, ONE, Felt::new(i)]; (key, value) }) - .collect(); + .collect() + } + + #[test] + fn test_build_subtree_from_leaves() { + // A single subtree's worth of leaves. + const PAIR_COUNT: u64 = 256; + + let entries = generate_entries(PAIR_COUNT); let control = Smt::with_entries(entries.clone()).unwrap(); - let mut leaves: Vec = entries - .iter() - .map(|(key, value)| { - let leaf = SmtLeaf::new_single(*key, *value); - let col = leaf.index().index.value(); - let hash = leaf.hash(); - SubtreeLeaf { col, hash } - }) - .collect(); - leaves.sort(); - leaves.dedup_by_key(|leaf| leaf.col); + // `entries` should already be sorted by nature of how we constructed it. + let leaves = Smt::sorted_pairs_to_leaves(entries).subtrees; let (first_subtree, _) = Smt::build_subtree(leaves, SMT_DEPTH); assert!(!first_subtree.is_empty());