From 4301bc85a3f27f85797f7f379f72a0c200a811af Mon Sep 17 00:00:00 2001 From: Qyriad Date: Thu, 14 Nov 2024 16:46:28 -0700 Subject: [PATCH] smt: implement test for basic parallelized subtree computation w/ rayon Building on the previous commit, this commit implements a test proving that `SparseMerkleTree::build_subtree()` can be composed into itself not just concurrently, but in parallel, without issue. --- Cargo.lock | 1 + Cargo.toml | 4 +- src/merkle/smt/tests.rs | 101 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 3e822fb..fca1e2c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -534,6 +534,7 @@ dependencies = [ "rand", "rand_chacha", "rand_core", + "rayon", "seq-macro", "serde", "sha3", diff --git a/Cargo.toml b/Cargo.toml index 74df3ab..347cca8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,9 +40,10 @@ name = "store" harness = false [features] -default = ["std"] +default = ["std", "concurrent"] executable = ["dep:clap", "dep:rand-utils", "std"] serde = ["dep:serde", "serde?/alloc", "winter-math/serde"] +concurrent = ["dep:rayon"] std = [ "blake3/std", "dep:cc", @@ -66,6 +67,7 @@ sha3 = { version = "0.10", default-features = false } winter-crypto = { version = "0.10", default-features = false } winter-math = { version = "0.10", default-features = false } winter-utils = { version = "0.10", default-features = false } +rayon = { version = "1.10.0", optional = true } [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } diff --git a/src/merkle/smt/tests.rs b/src/merkle/smt/tests.rs index 3795df3..b16dc2e 100644 --- a/src/merkle/smt/tests.rs +++ b/src/merkle/smt/tests.rs @@ -311,3 +311,104 @@ fn test_singlethreaded_subtrees() { // And of course the root we got from each place should match. assert_eq!(control.root(), root_leaf.hash); } + +/// The parallel version of `test_singlethreaded_subtree()`. +#[test] +#[cfg(feature = "concurrent")] +fn test_multithreaded_subtrees() { + use rayon::prelude::*; + + const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64; + + let entries = generate_entries(PAIR_COUNT); + + let control = Smt::with_entries(entries.clone()).unwrap(); + + let mut accumulated_nodes: BTreeMap = Default::default(); + + let PairComputations { + leaves: mut leaf_subtrees, + nodes: test_leaves, + } = Smt::sorted_pairs_to_leaves(entries); + + for current_depth in (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() { + let (nodes, subtrees): (Vec>, Vec>) = leaf_subtrees + .into_par_iter() + .enumerate() + .map(|(i, subtree)| { + // Pre-assertions. + assert!( + subtree.is_sorted(), + "subtree {i} at bottom-depth {current_depth} is not sorted", + ); + assert!( + !subtree.is_empty(), + "subtree {i} at bottom-depth {current_depth} is empty!", + ); + + 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!( + test_node, &control_node, + "depth {} subtree {}: test node does not match control at index {:?}", + current_depth, i, index, + ); + } + + (nodes, next_leaves) + }) + .unzip(); + + let mut all_leaves: Vec = 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}"); + } + + // Make sure the true leaves match, checking length first and then each individual leaf. + let control_leaves: BTreeMap<_, _> = control.leaves().collect(); + let control_leaves_len = control_leaves.len(); + let test_leaves_len = test_leaves.len(); + assert_eq!(test_leaves_len, control_leaves_len); + for (col, ref test_leaf) in test_leaves { + let index = LeafIndex::new_max_depth(col); + let &control_leaf = control_leaves.get(&index).unwrap(); + assert_eq!(test_leaf, control_leaf); + } + + // Make sure the inner nodes match, checking length first and then each individual leaf. + let control_nodes_len = control.inner_nodes().count(); + let test_nodes_len = accumulated_nodes.len(); + assert_eq!(test_nodes_len, control_nodes_len); + for (index, test_node) in accumulated_nodes.clone() { + let control_node = control.get_inner_node(index); + assert_eq!(test_node, control_node, "test node does not match control at {index:?}"); + } + + // After the last iteration of the above for loop, we should have the new root node actually + // in two places: one in `accumulated_nodes`, and the other as the "next leaves" return from + // `build_subtree()`. So let's check both! + + let control_root = control.get_inner_node(NodeIndex::root()); + + // That for loop should have left us with only one leaf subtree... + let [leaf_subtree]: [_; 1] = leaf_subtrees.try_into().unwrap(); + // which itself contains only one 'leaf'... + let [root_leaf]: [_; 1] = leaf_subtree.try_into().unwrap(); + // which matches the expected root. + assert_eq!(control.root(), root_leaf.hash); + + // Likewise `accumulated_nodes` should contain a node at the root index... + assert!(accumulated_nodes.contains_key(&NodeIndex::root())); + // and it should match our actual root. + let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap(); + assert_eq!(control_root, *test_root); + // And of course the root we got from each place should match. + assert_eq!(control.root(), root_leaf.hash); +}